Coverage for agentos/rag/citation.py: 0%
118 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 19:15 +0800
1"""Citation tracing for RAG pipeline.
3Tracks which source documents contributed to generated text,
4enabling answer provenance and fact-checking.
5"""
7from __future__ import annotations
9import hashlib
10import re
11from dataclasses import dataclass, field
12from typing import Any
15@dataclass
16class Citation:
17 """A single citation linking generated text to a source chunk."""
19 chunk_id: str
20 chunk_text: str
21 source_doc: str = "" # source document identifier
22 relevance_score: float = 0.0
23 context: str = "" # surrounding context window
24 start_char: int = 0 # position in generated answer
25 end_char: int = 0
28@dataclass
29class CitationReport:
30 """Complete citation analysis for a generated response."""
32 answer: str = ""
33 citations: list[Citation] = field(default_factory=list)
34 source_count: int = 0
35 coverage: float = 0.0 # fraction of answer covered by citations
36 unused_sources: list[str] = field(default_factory=list)
38 def to_dict(self) -> dict[str, Any]:
39 return {
40 "answer": self.answer,
41 "num_citations": len(self.citations),
42 "source_count": self.source_count,
43 "coverage": self.coverage,
44 "citations": [
45 {
46 "chunk_id": c.chunk_id,
47 "source_doc": c.source_doc,
48 "relevance": c.relevance_score,
49 "span": f"{c.start_char}-{c.end_char}",
50 "text_preview": c.chunk_text[:200],
51 }
52 for c in self.citations
53 ],
54 }
57class CitationTracer:
58 """Track which retrieved chunks contributed to an answer.
60 Two modes:
61 - token_overlap: Match answer spans to chunk texts by token overlap.
62 - explicit: Parse answer for explicit citation markers like [1], [doc1].
63 """
65 def __init__(
66 self,
67 mode: str = "token_overlap",
68 min_overlap: int = 20, # minimum characters of overlap
69 overlap_ratio: float = 0.3, # minimum overlap ratio
70 ):
71 self.mode = mode
72 self.min_overlap = min_overlap
73 self.overlap_ratio = overlap_ratio
75 def trace(
76 self,
77 answer: str,
78 sources: list[dict[str, Any]],
79 ) -> CitationReport:
80 """Trace answer back to source chunks.
82 Args:
83 answer: Generated text response.
84 sources: Retrieved chunks with 'text', 'score', 'index' keys.
86 Returns:
87 CitationReport with matched citations.
88 """
89 if self.mode == "explicit":
90 citations = self._trace_explicit(answer, sources)
91 else:
92 citations = self._trace_overlap(answer, sources)
94 # Compute coverage
95 if answer and citations:
96 covered_chars = self._compute_covered_chars(answer, citations)
97 coverage = covered_chars / len(answer)
98 else:
99 coverage = 0.0
101 # Find unused sources
102 used_ids = {c.chunk_id for c in citations}
103 unused = [
104 f"chunk_{s.get('index', i)}"
105 for i, s in enumerate(sources)
106 if f"chunk_{s.get('index', i)}" not in used_ids
107 ]
109 return CitationReport(
110 answer=answer,
111 citations=citations,
112 source_count=len(sources),
113 coverage=round(coverage, 3),
114 unused_sources=unused,
115 )
117 def _trace_overlap(
118 self,
119 answer: str,
120 sources: list[dict[str, Any]],
121 ) -> list[Citation]:
122 """Find answer spans that overlap with source chunks."""
123 citations = []
125 for i, src in enumerate(sources):
126 chunk_text = src.get("text", "")
127 if not chunk_text:
128 continue
130 chunk_id = f"chunk_{src.get('index', i)}"
132 # Find longest common substrings
133 matches = self._find_substring_matches(answer, chunk_text)
134 for start, end in matches:
135 citations.append(
136 Citation(
137 chunk_id=chunk_id,
138 chunk_text=chunk_text,
139 source_doc=src.get("source", src.get("document", "")),
140 relevance_score=src.get("score", 0.0),
141 context=self._get_context(chunk_text, start, end),
142 start_char=start,
143 end_char=end,
144 )
145 )
147 return citations
149 def _trace_explicit(
150 self,
151 answer: str,
152 sources: list[dict[str, Any]],
153 ) -> list[Citation]:
154 """Parse explicit citation markers like [1], [source1], [doc:1]."""
155 citations = []
157 # Match [N], [docN], [source N]
158 pattern = r"\[(?:doc|source|ref)?\s*(\d+)\]"
159 matches = re.finditer(pattern, answer, re.IGNORECASE)
161 for m in matches:
162 ref_num = int(m.group(1))
163 if 1 <= ref_num <= len(sources):
164 src = sources[ref_num - 1]
165 citations.append(
166 Citation(
167 chunk_id=f"chunk_{src.get('index', ref_num - 1)}",
168 chunk_text=src.get("text", ""),
169 source_doc=src.get("source", ""),
170 relevance_score=src.get("score", 0.0),
171 context=src.get("text", "")[:500],
172 start_char=m.start(),
173 end_char=m.end(),
174 )
175 )
177 return citations
179 def _find_substring_matches(
180 self,
181 answer: str,
182 chunk: str,
183 ) -> list[tuple[int, int]]:
184 """Find spans in answer that match substrings from chunk."""
185 matches = []
186 min_len = min(self.min_overlap, len(chunk) // 4)
188 # Use sliding window of sentences/phrases from chunk
189 sentences = re.split(r"(?<=[.!?。！？])\s+", chunk)
190 for sent in sentences:
191 sent = sent.strip()
192 if len(sent) < min_len:
193 continue
195 pos = answer.find(sent)
196 if pos >= 0:
197 matches.append((pos, pos + len(sent)))
198 else:
199 # Try shorter windows
200 window = max(min_len, len(sent) // 2)
201 step = window // 2
202 for start in range(0, len(sent) - window + 1, step):
203 sub = sent[start : start + window]
204 pos = answer.find(sub)
205 if pos >= 0:
206 matches.append((pos, pos + len(sub)))
207 break
209 return self._merge_overlapping(matches)
211 def _merge_overlapping(
212 self,
213 spans: list[tuple[int, int]],
214 ) -> list[tuple[int, int]]:
215 """Merge overlapping citation spans."""
216 if not spans:
217 return []
219 sorted_spans = sorted(spans)
220 merged = [sorted_spans[0]]
222 for span in sorted_spans[1:]:
223 last = merged[-1]
224 if span[0] <= last[1]:
225 merged[-1] = (last[0], max(last[1], span[1]))
226 else:
227 merged.append(span)
229 return merged
231 def _get_context(self, chunk_text: str, start: int, end: int) -> str:
232 """Get surrounding context around a match."""
233 ctx_start = max(0, start - 100)
234 ctx_end = min(len(chunk_text), end + 100)
235 return chunk_text[ctx_start:ctx_end]
237 def _compute_covered_chars(
238 self,
239 answer: str,
240 citations: list[Citation],
241 ) -> int:
242 """Compute total characters covered by citations."""
243 if not citations:
244 return 0
246 coverage = [False] * len(answer)
247 for c in citations:
248 for i in range(max(0, c.start_char), min(c.end_char, len(answer))):
249 coverage[i] = True
251 return sum(coverage)
253 def build_attribution_map(
254 self,
255 answer: str,
256 sources: list[dict[str, Any]],
257 ) -> str:
258 """Build HTML attribution map for the answer.
260 Wraps cited spans in <cite> tags with source references.
261 """
262 citations = self._trace_overlap(answer, sources)
264 # Sort citations by start position and apply in reverse to preserve indices
265 citations.sort(key=lambda c: c.start_char, reverse=True)
267 result = answer
268 for c in citations:
269 prefix = result[: c.start_char]
270 cited = result[c.start_char : c.end_char]
271 suffix = result[c.end_char :]
273 source_id = c.chunk_id.replace("chunk_", "")
274 cited_wrapped = (
275 f'<cite data-source="{c.source_doc}" '
276 f'data-chunk="{source_id}" '
277 f'data-score="{c.relevance_score:.2f}">'
278 f"{cited}</cite>"
279 )
280 result = prefix + cited_wrapped + suffix
282 return result
285def hash_chunk_id(text: str, index: int) -> str:
286 """Generate a stable chunk ID from text content."""
287 digest = hashlib.sha256(text.encode("utf-8")).hexdigest()[:12]
288 return f"chunk_{index}_{digest}"