Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
194 changes: 14 additions & 180 deletions packages/elf-service/src/docs/source_capture.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
mod helpers;

use crate::docs::{
ByteChunk, DOC_SOURCE_CAPTURE_SCHEMA_V1, DOC_SOURCE_SPAN_SCHEMA_V1, DocChunk, DocType,
DocsSourceCaptureSummary, DocsSourceSpanRef, Error, Map, OffsetDateTime, Result, Rfc3339,
SourceCaptureSummaryInput, Uuid, Value, WritePolicyAudit,
DocsSourceCaptureSummary, DocsSourceSpanRef, Error, Map, OffsetDateTime, Result,
SourceCaptureSummaryInput, Uuid, Value,
};

pub(super) fn build_doc_chunk_rows(
Expand Down Expand Up @@ -47,7 +49,7 @@ pub(super) fn source_record_id_for(
agent_id,
scope,
doc_type.as_str(),
source_identity_value(source_ref, doc_type),
helpers::source_identity_value(source_ref, doc_type),
content_hash,
])
.to_string();
Expand Down Expand Up @@ -86,7 +88,7 @@ pub(super) fn build_source_capture_summary(
.get("captured_at")
.and_then(Value::as_str)
.map(ToString::to_string)
.unwrap_or(format_timestamp(now)?);
.unwrap_or(helpers::format_timestamp(now)?);
let source_spans = chunks
.iter()
.map(|chunk| DocsSourceSpanRef {
Expand All @@ -106,73 +108,22 @@ pub(super) fn build_source_capture_summary(
chunk_hash: Some(chunk.chunk_hash.clone()),
})
.collect();
let policy_spans = source_policy_spans(raw_content_hash, write_policy_audit);
let policy_spans = helpers::source_policy_spans(raw_content_hash, write_policy_audit);

Ok(DocsSourceCaptureSummary {
schema: DOC_SOURCE_CAPTURE_SCHEMA_V1.to_string(),
source_record_id: doc_id,
origin: source_origin(source_ref, doc_type),
origin: helpers::source_origin(source_ref, doc_type),
captured_at,
content_hash: content_hash.to_string(),
visibility_scope: scope.to_string(),
title: title.map(ToString::to_string),
source_type: source_type(source_ref, doc_type),
source_type: helpers::source_type(source_ref, doc_type),
source_spans,
policy_spans,
})
}

pub(super) fn source_policy_spans(
raw_content_hash: &str,
audit: Option<&WritePolicyAudit>,
) -> Vec<DocsSourceSpanRef> {
let Some(audit) = audit else {
return Vec::new();
};
let mut spans = Vec::with_capacity(audit.exclusions.len() + audit.redactions.len());

for span in &audit.exclusions {
spans.push(policy_span_ref(
raw_content_hash,
span.start,
span.end,
"excluded",
"WRITE_POLICY_EXCLUSION",
));
}
for redaction in &audit.redactions {
spans.push(policy_span_ref(
raw_content_hash,
redaction.span.start,
redaction.span.end,
"redacted",
"WRITE_POLICY_REDACTION",
));
}

spans
}

pub(super) fn policy_span_ref(
content_hash: &str,
start: usize,
end: usize,
status: &str,
reason_code: &str,
) -> DocsSourceSpanRef {
DocsSourceSpanRef {
schema: DOC_SOURCE_SPAN_SCHEMA_V1.to_string(),
span_id: source_span_id(content_hash, start, end, reason_code),
chunk_id: None,
status: status.to_string(),
reason_code: Some(reason_code.to_string()),
start_offset: start,
end_offset: end,
content_hash: content_hash.to_string(),
chunk_hash: None,
}
}

pub(super) fn normalize_source_ref_for_capture(
source_ref: Value,
source_capture: &DocsSourceCaptureSummary,
Expand All @@ -199,134 +150,17 @@ pub(super) fn normalize_source_ref_for_capture(
}

source_ref.insert("source_type".to_string(), Value::String(source_capture.source_type.clone()));
source_ref
.insert("source_spans".to_string(), source_spans_to_value(&source_capture.source_spans)?);
source_ref.insert(
"source_spans".to_string(),
helpers::source_spans_to_value(&source_capture.source_spans)?,
);

if !source_capture.policy_spans.is_empty() {
source_ref.insert(
"policy_spans".to_string(),
source_spans_to_value(&source_capture.policy_spans)?,
helpers::source_spans_to_value(&source_capture.policy_spans)?,
);
}

Ok(Value::Object(source_ref))
}

pub(super) fn source_spans_to_value(spans: &[DocsSourceSpanRef]) -> Result<Value> {
serde_json::to_value(spans).map_err(|err| Error::InvalidRequest {
message: format!("failed to encode source span metadata: {err}"),
})
}

pub(super) fn source_type(source_ref: &Map<String, Value>, doc_type: DocType) -> String {
source_ref
.get("source_kind")
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| doc_type.as_str())
.to_string()
}

pub(super) fn source_origin(source_ref: &Map<String, Value>, doc_type: DocType) -> String {
if let Some(origin) = source_ref_string(source_ref, "canonical_uri")
.or_else(|| source_ref_string(source_ref, "url"))
.or_else(|| source_ref_string(source_ref, "uri"))
{
return origin.to_string();
}

match doc_type {
DocType::Chat => source_ref_string(source_ref, "message_id")
.map(|message_id| {
format!(
"thread:{}#{}",
source_ref_string(source_ref, "thread_id").unwrap_or("unknown"),
message_id
)
})
.unwrap_or_else(|| {
format!(
"thread:{}",
source_ref_string(source_ref, "thread_id").unwrap_or("unknown")
)
}),
DocType::Search => source_ref_string(source_ref, "domain")
.map(|domain| format!("search:{domain}"))
.unwrap_or_else(|| "search:unknown".to_string()),
DocType::Dev => dev_origin(source_ref),
DocType::Knowledge => source_ref_string(source_ref, "ts")
.map(|ts| format!("knowledge:{ts}"))
.unwrap_or_else(|| "knowledge:unknown".to_string()),
}
}

pub(super) fn dev_origin(source_ref: &Map<String, Value>) -> String {
let repo = source_ref_string(source_ref, "repo").unwrap_or("unknown");
let path = source_ref_string(source_ref, "path").unwrap_or("");
let revision = source_ref_string(source_ref, "commit_sha")
.map(|commit| format!("@{commit}"))
.or_else(|| source_ref_i64(source_ref, "pr_number").map(|pr| format!("#pr-{pr}")))
.or_else(|| {
source_ref_i64(source_ref, "issue_number").map(|issue| format!("#issue-{issue}"))
})
.unwrap_or_default();

if path.is_empty() {
format!("repo:{repo}{revision}")
} else {
format!("repo:{repo}/{path}{revision}")
}
}

pub(super) fn source_identity_value(source_ref: &Map<String, Value>, doc_type: DocType) -> Value {
if let Some(canonical_uri) = source_ref_string(source_ref, "canonical_uri") {
return serde_json::json!(["canonical_uri", canonical_uri]);
}

match doc_type {
DocType::Chat => serde_json::json!([
"chat",
source_ref_string(source_ref, "thread_id"),
source_ref_string(source_ref, "message_id"),
source_ref_string(source_ref, "role"),
source_ref_string(source_ref, "ts"),
]),
DocType::Search => serde_json::json!([
"search",
source_ref_string(source_ref, "url"),
source_ref_string(source_ref, "domain"),
source_ref_string(source_ref, "query"),
source_ref_string(source_ref, "ts"),
]),
DocType::Dev => serde_json::json!([
"dev",
source_ref_string(source_ref, "repo"),
source_ref_string(source_ref, "path"),
source_ref_string(source_ref, "commit_sha"),
source_ref_i64(source_ref, "pr_number"),
source_ref_i64(source_ref, "issue_number"),
]),
DocType::Knowledge => serde_json::json!([
"knowledge",
source_ref_string(source_ref, "uri"),
source_ref_string(source_ref, "ts"),
]),
}
}

pub(super) fn source_ref_string<'a>(
source_ref: &'a Map<String, Value>,
key: &str,
) -> Option<&'a str> {
source_ref.get(key).and_then(Value::as_str).filter(|value| !value.trim().is_empty())
}

pub(super) fn source_ref_i64(source_ref: &Map<String, Value>, key: &str) -> Option<i64> {
source_ref.get(key).and_then(Value::as_i64)
}

pub(super) fn format_timestamp(ts: OffsetDateTime) -> Result<String> {
ts.format(&Rfc3339).map_err(|err| Error::InvalidRequest {
message: format!("failed to format RFC3339 timestamp: {err}"),
})
}
Loading