diff --git a/examples/ag-ui-research/frontend/components/StateDisplay.tsx b/examples/ag-ui-research/frontend/components/StateDisplay.tsx index 03e710b9..69a28b20 100644 --- a/examples/ag-ui-research/frontend/components/StateDisplay.tsx +++ b/examples/ag-ui-research/frontend/components/StateDisplay.tsx @@ -23,11 +23,29 @@ interface GapRecord { resolved_by: string[]; } +interface BoundingBox { + page_no: number; + left: number; + top: number; + right: number; + bottom: number; +} + +interface Citation { + document_uri: string; + document_title?: string; + page_numbers: number[]; + headings?: string[]; + content: string; + bounding_boxes?: BoundingBox[]; +} + interface SearchAnswer { - sub_question: string; + query: string; answer: string; confidence: number; - chunks_used: number; + cited_chunks: string[]; + citations: Citation[]; } interface ResearchContext { @@ -508,6 +526,125 @@ export default function StateDisplay({ state }: StateDisplayProps) { + + {/* Citations with visual grounding info */} + {qaResponse.citations && + qaResponse.citations.length > 0 && ( +
+
+ Citations ({qaResponse.citations.length}) +
+ {qaResponse.citations.map((citation, citIdx) => ( +
+
+
+ {citation.document_title || + citation.document_uri} +
+ {citation.page_numbers && + citation.page_numbers.length > 0 && ( +
+ {citation.page_numbers.length === 1 + ? `p. ${citation.page_numbers[0]}` + : `pp. ${citation.page_numbers[0]}-${citation.page_numbers[citation.page_numbers.length - 1]}`} +
+ )} +
+ {citation.headings && + citation.headings.length > 0 && ( +
+ {citation.headings.join(" › ")} +
+ )} +
+ {citation.content.slice(0, 200)} + {citation.content.length > 200 && "…"} +
+ {citation.bounding_boxes && + citation.bounding_boxes.length > 0 && ( +
+ {citation.bounding_boxes.map( + (bbox, bboxIdx) => ( + + 📍 p.{bbox.page_no} ({bbox.left.toFixed(0)},{bbox.top.toFixed(0)}) + + ), + )} +
+ )} +
+ ))} +
+ )} )} diff --git a/haiku_rag_slim/haiku/rag/graph/common/models.py b/haiku_rag_slim/haiku/rag/graph/common/models.py index 0cff7057..98f1eebf 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/models.py +++ b/haiku_rag_slim/haiku/rag/graph/common/models.py @@ -54,6 +54,10 @@ class SearchAnswer(BaseModel): ge=0.0, le=1.0, ) + citations: list[Citation] = Field( + default_factory=list, + description="Resolved citations with full metadata (populated after search)", + ) def resolve_citations( diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py index 0238b927..85b3fed7 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py @@ -1,7 +1,7 @@ from pydantic import BaseModel, Field from haiku.rag.client import HaikuRAG -from haiku.rag.graph.common.models import SearchAnswer +from haiku.rag.graph.common.models import SearchAnswer, resolve_citations from haiku.rag.store.models import SearchResult @@ -17,8 +17,8 @@ class DeepQAContext(BaseModel): def add_qa_response( self, qa: SearchAnswer, search_results: list[SearchResult] ) -> None: - """Add a QA response.""" - del search_results # Not needed with chunk ID-based citations + """Add a QA response with resolved citations.""" + qa.citations = resolve_citations(qa.cited_chunks, search_results) self.qa_responses.append(qa) diff --git a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py index d1b5ddaa..aba28df7 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py @@ -3,7 +3,7 @@ from collections.abc import Iterable from pydantic import BaseModel, Field, PrivateAttr from haiku.rag.client import HaikuRAG -from haiku.rag.graph.common.models import SearchAnswer +from haiku.rag.graph.common.models import SearchAnswer, resolve_citations from haiku.rag.graph.research.models import ( GapRecord, InsightAnalysis, @@ -41,9 +41,8 @@ class ResearchContext(BaseModel): def add_qa_response( self, qa: SearchAnswer, search_results: list[SearchResult] ) -> None: - """Add a structured QA response.""" - # Research doesn't accumulate search_results for citation resolution - del search_results + """Add a structured QA response with resolved citations.""" + qa.citations = resolve_citations(qa.cited_chunks, search_results) self.qa_responses.append(qa) def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]: