expose citations in ag-ui-example
This commit is contained in:
parent
d808c6c425
commit
9382a6b132
4 changed files with 149 additions and 9 deletions
|
|
@ -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) {
|
|||
<Markdown content={qaResponse.answer} />
|
||||
</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>
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
Loading…
Reference in a new issue