danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

expose citations in ag-ui-example

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-03 09:08:01 +02:00
parent d808c6c425
commit 9382a6b132
No known key found for this signature in database
4 changed files with 149 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

141
examples/ag-ui-research/frontend/components/StateDisplay.tsx
View file

@ -23,11 +23,29 @@ interface GapRecord {
resolved_by: string[]; 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 { interface SearchAnswer {
sub_question: string; query: string;
answer: string; answer: string;
confidence: number; confidence: number;
chunks_used: number; cited_chunks: string[];
citations: Citation[];
} }
interface ResearchContext { interface ResearchContext {
@ -508,6 +526,125 @@ export default function StateDisplay({ state }: StateDisplayProps) {
<Markdown content={qaResponse.answer} /> <Markdown content={qaResponse.answer} />
</div> </div>
</div> </div>
{/* Citations with visual grounding info */}
{qaResponse.citations &&
qaResponse.citations.length > 0 && (
<div style={{ marginTop: "1rem" }}>
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginBottom: "0.5rem",
fontWeight: "600",
}}
>
Citations ({qaResponse.citations.length})
</div>
{qaResponse.citations.map((citation, citIdx) => (
<div
key={citIdx}
style={{
padding: "0.75rem",
background: "white",
borderRadius: "4px",
border: "1px solid #e2e8f0",
marginBottom: "0.5rem",
}}
>
<div
style={{
display: "flex",
justifyContent: "space-between",
alignItems: "flex-start",
marginBottom: "0.5rem",
}}
>
<div
style={{
fontSize: "0.75rem",
fontWeight: "600",
color: "#2d3748",
}}
>
{citation.document_title ||
citation.document_uri}
</div>
{citation.page_numbers &&
citation.page_numbers.length > 0 && (
<div
style={{
fontSize: "0.7rem",
color: "#718096",
background: "#edf2f7",
padding: "0.125rem 0.375rem",
borderRadius: "4px",
}}
>
{citation.page_numbers.length === 1
? `p. ${citation.page_numbers[0]}`
: `pp. ${citation.page_numbers[0]}-${citation.page_numbers[citation.page_numbers.length - 1]}`}
</div>
)}
</div>
{citation.headings &&
citation.headings.length > 0 && (
<div
style={{
fontSize: "0.7rem",
color: "#718096",
marginBottom: "0.375rem",
}}
>
{citation.headings.join(" › ")}
</div>
)}
<div
style={{
fontSize: "0.8rem",
color: "#4a5568",
lineHeight: "1.4",
maxHeight: "4.5rem",
overflow: "hidden",
textOverflow: "ellipsis",
}}
>
{citation.content.slice(0, 200)}
{citation.content.length > 200 && "…"}
</div>
{citation.bounding_boxes &&
citation.bounding_boxes.length > 0 && (
<div
style={{
fontSize: "0.65rem",
color: "#a0aec0",
marginTop: "0.375rem",
display: "flex",
gap: "0.25rem",
flexWrap: "wrap",
}}
>
{citation.bounding_boxes.map(
(bbox, bboxIdx) => (
<span
key={bboxIdx}
style={{
background: "#f7fafc",
padding: "0.125rem 0.25rem",
borderRadius: "2px",
border: "1px solid #e2e8f0",
}}
>
📍 p.{bbox.page_no} ({bbox.left.toFixed(0)},{bbox.top.toFixed(0)})
</span>
),
)}
</div>
)}
</div>
))}
</div>
)}
</div> </div>
)} )}
</div> </div>

4
haiku_rag_slim/haiku/rag/graph/common/models.py
View file

@ -54,6 +54,10 @@ class SearchAnswer(BaseModel):
ge=0.0, ge=0.0,
le=1.0, le=1.0,
) )
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata (populated after search)",
)
def resolve_citations( def resolve_citations(

6
haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py
View file

@ -1,7 +1,7 @@
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG 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 from haiku.rag.store.models import SearchResult
@ -17,8 +17,8 @@ class DeepQAContext(BaseModel):
def add_qa_response( def add_qa_response(
self, qa: SearchAnswer, search_results: list[SearchResult] self, qa: SearchAnswer, search_results: list[SearchResult]
) -> None: ) -> None:
"""Add a QA response.""" """Add a QA response with resolved citations."""
del search_results # Not needed with chunk ID-based citations qa.citations = resolve_citations(qa.cited_chunks, search_results)
self.qa_responses.append(qa) self.qa_responses.append(qa)

7
haiku_rag_slim/haiku/rag/graph/research/dependencies.py
View file

@ -3,7 +3,7 @@ from collections.abc import Iterable
from pydantic import BaseModel, Field, PrivateAttr from pydantic import BaseModel, Field, PrivateAttr
from haiku.rag.client import HaikuRAG 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 ( from haiku.rag.graph.research.models import (
GapRecord, GapRecord,
InsightAnalysis, InsightAnalysis,
@ -41,9 +41,8 @@ class ResearchContext(BaseModel):
def add_qa_response( def add_qa_response(
self, qa: SearchAnswer, search_results: list[SearchResult] self, qa: SearchAnswer, search_results: list[SearchResult]
) -> None: ) -> None:
"""Add a structured QA response.""" """Add a structured QA response with resolved citations."""
# Research doesn't accumulate search_results for citation resolution qa.citations = resolve_citations(qa.cited_chunks, search_results)
del search_results
self.qa_responses.append(qa) self.qa_responses.append(qa)
def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]: def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]: