everruns_core/capabilities/
citation_retrieval.rs

1//! `citation_retrieval` capability — claim-level citations from retrieval tools.
2//!
3//! Turns the sources surfaced by `search_index` / `search_knowledge` into
4//! claim-level provenance on the assistant's answer. It contributes no tools of
5//! its own; instead it registers a post-generation annotation hook (see
6//! [`crate::annotation_hook`]) that, once the model has answered, scans the
7//! turn's retrieval tool results, aligns each retrieved passage to the sentence
8//! it best supports, and attaches a [`TextAnnotation`] there. Alignment is
9//! deterministic token overlap — no extra model call, no text rewrite — so it is
10//! model- and provider-agnostic and never changes the streamed answer.
11//!
12//! See `specs/citations.md`. This is the retrieval *feed*; verification is a
13//! separate `citation_verification` capability.
14
15use std::collections::HashMap;
16use std::sync::Arc;
17
18use async_trait::async_trait;
19use serde::{Deserialize, Serialize};
20use serde_json::json;
21
22use crate::annotation_hook::{
23    AnnotationContext, AnnotationResult, PostGenerationAnnotationHook, citation_tokens,
24    token_overlap_ratio,
25};
26use crate::capabilities::Capability;
27use crate::capability_types::CapabilityStatus;
28use crate::message::{AnnotationSource, ContentPart, Message, TextAnnotation};
29
30/// Canonical capability id.
31pub const CITATION_RETRIEVAL_CAPABILITY_ID: &str = "citation_retrieval";
32
33/// Tool names whose results carry retrieval citations.
34const RETRIEVAL_TOOLS: &[&str] = &["search_index", "search_knowledge"];
35
36/// Default minimum token-overlap ratio for a passage to be attached to a
37/// sentence. Overlap is `|shared tokens| / |passage tokens|`, so 0.5 means at
38/// least half the passage's distinctive words appear in the sentence.
39const DEFAULT_MIN_OVERLAP: f32 = 0.5;
40
41/// Per-agent config for `citation_retrieval`.
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct CitationRetrievalConfig {
44    /// Minimum overlap ratio in `[0, 1]` for a citation to attach.
45    #[serde(default = "default_min_overlap")]
46    pub min_overlap: f32,
47}
48
49fn default_min_overlap() -> f32 {
50    DEFAULT_MIN_OVERLAP
51}
52
53impl Default for CitationRetrievalConfig {
54    fn default() -> Self {
55        Self {
56            min_overlap: DEFAULT_MIN_OVERLAP,
57        }
58    }
59}
60
61impl CitationRetrievalConfig {
62    fn from_value(config: &serde_json::Value) -> Self {
63        if config.is_null() {
64            return Self::default();
65        }
66        serde_json::from_value(config.clone()).unwrap_or_default()
67    }
68}
69
70/// The retrieval citation feed.
71pub struct CitationRetrievalCapability;
72
73impl Capability for CitationRetrievalCapability {
74    fn id(&self) -> &str {
75        CITATION_RETRIEVAL_CAPABILITY_ID
76    }
77
78    fn name(&self) -> &str {
79        "Retrieval citations"
80    }
81
82    fn description(&self) -> &str {
83        "Attach claim-level citations to the answer from search_index / \
84         search_knowledge results, so each grounded sentence links to its source."
85    }
86
87    fn status(&self) -> CapabilityStatus {
88        CapabilityStatus::Available
89    }
90
91    fn icon(&self) -> Option<&str> {
92        Some("quote")
93    }
94
95    fn category(&self) -> Option<&str> {
96        Some("Knowledge")
97    }
98
99    /// Surfaces the citation UI when any citation feed is active.
100    fn features(&self) -> Vec<&'static str> {
101        vec!["citations"]
102    }
103
104    fn config_schema(&self) -> Option<serde_json::Value> {
105        Some(json!({
106            "type": "object",
107            "additionalProperties": false,
108            "properties": {
109                "min_overlap": {
110                    "type": "number",
111                    "minimum": 0.0,
112                    "maximum": 1.0,
113                    "default": DEFAULT_MIN_OVERLAP,
114                    "description": "Minimum token-overlap ratio for a retrieved passage to be cited on a sentence."
115                }
116            }
117        }))
118    }
119
120    fn validate_config(&self, config: &serde_json::Value) -> Result<(), String> {
121        if config.is_null() {
122            return Ok(());
123        }
124        let cfg: CitationRetrievalConfig = serde_json::from_value(config.clone())
125            .map_err(|e| format!("invalid citation_retrieval config: {e}"))?;
126        if !(0.0..=1.0).contains(&cfg.min_overlap) {
127            return Err("min_overlap must be between 0.0 and 1.0".to_string());
128        }
129        Ok(())
130    }
131
132    fn post_output_annotation_hooks_with_config(
133        &self,
134        config: &serde_json::Value,
135    ) -> Vec<Arc<dyn PostGenerationAnnotationHook>> {
136        let cfg = CitationRetrievalConfig::from_value(config);
137        vec![Arc::new(RetrievalAnnotationHook {
138            min_overlap: cfg.min_overlap,
139        })]
140    }
141}
142
143/// The annotation hook: extract retrieval citations from the turn's tool
144/// results, then align each to the sentence it best supports.
145struct RetrievalAnnotationHook {
146    min_overlap: f32,
147}
148
149#[async_trait]
150impl PostGenerationAnnotationHook for RetrievalAnnotationHook {
151    fn id(&self) -> &str {
152        CITATION_RETRIEVAL_CAPABILITY_ID
153    }
154
155    async fn annotate(&self, ctx: &AnnotationContext<'_>) -> AnnotationResult {
156        let citations = extract_retrieval_citations(ctx.messages);
157        if citations.is_empty() {
158            return AnnotationResult::none();
159        }
160        let annotations = align_citations(ctx.message_text, &citations, self.min_overlap);
161        AnnotationResult {
162            annotations,
163            rewritten_text: None,
164        }
165    }
166}
167
168/// A retrieval citation extracted from a tool result, normalized across the
169/// `search_index` and `search_knowledge` result shapes.
170#[derive(Debug, Clone)]
171struct RetrievalCitation {
172    external_id: Option<String>,
173    uri: String,
174    title: Option<String>,
175    snippet: String,
176    location: Option<serde_json::Value>,
177}
178
179/// Scan the assembled context for retrieval tool results and normalize their
180/// entries. Tool results carry only a `tool_call_id`, so we first map call ids
181/// to tool names from the assistant `ToolCall` parts.
182fn extract_retrieval_citations(messages: &[Message]) -> Vec<RetrievalCitation> {
183    let mut tool_names: HashMap<&str, &str> = HashMap::new();
184    for msg in messages {
185        for part in &msg.content {
186            if let ContentPart::ToolCall(call) = part {
187                tool_names.insert(call.id.as_str(), call.name.as_str());
188            }
189        }
190    }
191
192    let mut out = Vec::new();
193    for msg in messages {
194        for part in &msg.content {
195            let ContentPart::ToolResult(result) = part else {
196                continue;
197            };
198            let Some(name) = tool_names.get(result.tool_call_id.as_str()) else {
199                continue;
200            };
201            if !RETRIEVAL_TOOLS.contains(name) {
202                continue;
203            }
204            let Some(value) = &result.result else {
205                continue;
206            };
207            let Some(results) = value.get("results").and_then(|r| r.as_array()) else {
208                continue;
209            };
210            for entry in results {
211                if let Some(citation) = normalize_entry(entry) {
212                    out.push(citation);
213                }
214            }
215        }
216    }
217    out
218}
219
220/// Normalize one result object, tolerating either the `KnowledgeIndexCitation`
221/// (`source_uri` / `document_title`) or `KnowledgeSearchHit` (`resource` /
222/// `title`) shape. Entries without a usable snippet are skipped — there is
223/// nothing to align.
224fn normalize_entry(entry: &serde_json::Value) -> Option<RetrievalCitation> {
225    let snippet = entry.get("snippet").and_then(|s| s.as_str())?.trim();
226    if snippet.is_empty() {
227        return None;
228    }
229    let external_id = entry.get("id").and_then(|v| v.as_str()).map(str::to_string);
230    let title = entry
231        .get("document_title")
232        .or_else(|| entry.get("title"))
233        .and_then(|v| v.as_str())
234        .map(str::to_string);
235    let uri = entry
236        .get("source_uri")
237        .or_else(|| entry.get("resource"))
238        .or_else(|| entry.get("uri"))
239        .and_then(|v| v.as_str())
240        .map(str::to_string)
241        .or_else(|| {
242            external_id
243                .as_ref()
244                .map(|id| format!("everruns://citation/{id}"))
245        })
246        .unwrap_or_else(|| "everruns://citation".to_string());
247    Some(RetrievalCitation {
248        external_id,
249        uri,
250        title,
251        snippet: snippet.to_string(),
252        location: entry.get("location").cloned(),
253    })
254}
255
256/// Attach each citation to the single sentence in `text` with the highest token
257/// overlap against the citation's snippet, when that overlap clears
258/// `min_overlap`.
259fn align_citations(
260    text: &str,
261    citations: &[RetrievalCitation],
262    min_overlap: f32,
263) -> Vec<TextAnnotation> {
264    let sentences = split_sentences(text);
265    if sentences.is_empty() {
266        return Vec::new();
267    }
268    let sentence_tokens: Vec<Vec<String>> =
269        sentences.iter().map(|s| citation_tokens(&s.text)).collect();
270
271    let mut annotations = Vec::new();
272    for citation in citations {
273        let needle = citation_tokens(&citation.snippet);
274        if needle.is_empty() {
275            continue;
276        }
277        let mut best: Option<(usize, f32)> = None;
278        for (idx, tokens) in sentence_tokens.iter().enumerate() {
279            let ratio = token_overlap_ratio(&needle, tokens);
280            if best.map(|(_, b)| ratio > b).unwrap_or(true) {
281                best = Some((idx, ratio));
282            }
283        }
284        if let Some((idx, ratio)) = best
285            && ratio >= min_overlap
286        {
287            let sentence = &sentences[idx];
288            annotations.push(TextAnnotation {
289                start: sentence.start,
290                end: sentence.end,
291                origin: CITATION_RETRIEVAL_CAPABILITY_ID.to_string(),
292                source: AnnotationSource {
293                    uri: citation.uri.clone(),
294                    title: citation.title.clone(),
295                    snippet: Some(citation.snippet.clone()),
296                    location: citation.location.clone(),
297                },
298                external_id: citation.external_id.clone(),
299                verified: None,
300            });
301        }
302    }
303    annotations
304}
305
306/// A sentence and its char span (0-indexed, exclusive end) into the source text.
307struct Sentence {
308    start: usize,
309    end: usize,
310    text: String,
311}
312
313/// Split into sentences on `.`, `!`, `?`, and newlines, tracking char offsets.
314/// Approximate (markdown-agnostic) but sufficient for span attachment.
315fn split_sentences(text: &str) -> Vec<Sentence> {
316    let chars: Vec<char> = text.chars().collect();
317    let mut out = Vec::new();
318    let mut start = 0usize;
319    let mut i = 0usize;
320    while i < chars.len() {
321        let c = chars[i];
322        let is_break = matches!(c, '.' | '!' | '?' | '\n');
323        let at_end = i + 1 == chars.len();
324        if is_break || at_end {
325            let end = i + 1;
326            let segment: String = chars[start..end].iter().collect();
327            if !segment.trim().is_empty() {
328                // Trim leading whitespace from the span so chips sit on words.
329                let lead_ws = segment.chars().take_while(|c| c.is_whitespace()).count();
330                out.push(Sentence {
331                    start: start + lead_ws,
332                    end,
333                    text: segment,
334                });
335            }
336            start = end;
337        }
338        i += 1;
339    }
340    out
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346    use crate::message::{MessageRole, ToolCallContentPart, ToolResultContentPart};
347
348    fn tool_call_msg(call_id: &str, tool: &str) -> Message {
349        let mut m = Message::assistant("");
350        m.role = MessageRole::Agent;
351        m.content = vec![ContentPart::ToolCall(ToolCallContentPart::new(
352            call_id,
353            tool,
354            json!({}),
355        ))];
356        m
357    }
358
359    fn tool_result_msg(call_id: &str, result: serde_json::Value) -> Message {
360        let mut m = Message::assistant("");
361        m.role = MessageRole::ToolResult;
362        m.content = vec![ContentPart::ToolResult(ToolResultContentPart::new(
363            call_id,
364            Some(result),
365            None,
366        ))];
367        m
368    }
369
370    #[test]
371    fn extracts_index_and_knowledge_shapes() {
372        let messages = vec![
373            tool_call_msg("c1", "search_index"),
374            tool_result_msg(
375                "c1",
376                json!({"results": [{
377                    "id": "kchk_1",
378                    "source_uri": "github://o/r@main/a.md",
379                    "document_title": "A",
380                    "snippet": "Photosynthesis converts sunlight into chemical energy.",
381                    "location": {"lines": [1, 4]},
382                    "score": 0.9
383                }]}),
384            ),
385            tool_call_msg("c2", "search_knowledge"),
386            tool_result_msg(
387                "c2",
388                json!({"count": 1, "results": [{
389                    "id": "kbe_2",
390                    "kb_id": "kb_x",
391                    "title": "B",
392                    "kind": "note",
393                    "tags": [],
394                    "snippet": "The mitochondria is the powerhouse of the cell."
395                }]}),
396            ),
397        ];
398        let cites = extract_retrieval_citations(&messages);
399        assert_eq!(cites.len(), 2);
400        assert_eq!(cites[0].external_id.as_deref(), Some("kchk_1"));
401        assert_eq!(cites[0].uri, "github://o/r@main/a.md");
402        assert_eq!(cites[1].external_id.as_deref(), Some("kbe_2"));
403        // KnowledgeSearchHit with no `resource` falls back to a synthetic uri.
404        assert_eq!(cites[1].uri, "everruns://citation/kbe_2");
405    }
406
407    #[test]
408    fn ignores_non_retrieval_tools() {
409        let messages = vec![
410            tool_call_msg("c1", "bash"),
411            tool_result_msg(
412                "c1",
413                json!({"results": [{"id": "x", "snippet": "hi there"}]}),
414            ),
415        ];
416        assert!(extract_retrieval_citations(&messages).is_empty());
417    }
418
419    #[test]
420    fn aligns_citation_to_best_sentence() {
421        let text = "Plants are green. Photosynthesis converts sunlight into chemical energy. \
422                    Cells divide.";
423        let cites = vec![RetrievalCitation {
424            external_id: Some("kchk_1".to_string()),
425            uri: "github://o/r@main/a.md".to_string(),
426            title: Some("A".to_string()),
427            snippet: "Photosynthesis converts sunlight into chemical energy.".to_string(),
428            location: None,
429        }];
430        let anns = align_citations(text, &cites, 0.5);
431        assert_eq!(anns.len(), 1);
432        let ann = &anns[0];
433        let cited: String = text
434            .chars()
435            .skip(ann.start)
436            .take(ann.end - ann.start)
437            .collect();
438        assert!(cited.contains("Photosynthesis converts sunlight"));
439        assert_eq!(ann.external_id.as_deref(), Some("kchk_1"));
440        assert_eq!(ann.origin, CITATION_RETRIEVAL_CAPABILITY_ID);
441    }
442
443    #[test]
444    fn drops_low_overlap_citations() {
445        let text = "The weather today is sunny and warm.";
446        let cites = vec![RetrievalCitation {
447            external_id: None,
448            uri: "u".to_string(),
449            title: None,
450            snippet: "Quantum chromodynamics describes the strong interaction.".to_string(),
451            location: None,
452        }];
453        assert!(align_citations(text, &cites, 0.5).is_empty());
454    }
455
456    #[test]
457    fn config_validation_rejects_out_of_range() {
458        let cap = CitationRetrievalCapability;
459        assert!(cap.validate_config(&json!({"min_overlap": 1.5})).is_err());
460        assert!(cap.validate_config(&json!({"min_overlap": 0.3})).is_ok());
461        assert!(cap.validate_config(&serde_json::Value::Null).is_ok());
462    }
463}
everruns_core/capabilities/citation_retrieval.rs

everruns_core/capabilities/
citation_retrieval.rs
