diff --git a/examples/ag-ui-research/frontend/components/StateDisplay.tsx b/examples/ag-ui-research/frontend/components/StateDisplay.tsx index 69a28b20..a98a1442 100644 --- a/examples/ag-ui-research/frontend/components/StateDisplay.tsx +++ b/examples/ag-ui-research/frontend/components/StateDisplay.tsx @@ -23,21 +23,14 @@ interface GapRecord { resolved_by: string[]; } -interface BoundingBox { - page_no: number; - left: number; - top: number; - right: number; - bottom: number; -} - interface Citation { + document_id: string; + chunk_id: string; document_uri: string; document_title?: string; page_numbers: number[]; headings?: string[]; content: string; - bounding_boxes?: BoundingBox[]; } interface SearchAnswer { @@ -612,35 +605,6 @@ export default function StateDisplay({ state }: StateDisplayProps) { {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 8d1f7145..322b7d32 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/models.py +++ b/haiku_rag_slim/haiku/rag/graph/common/models.py @@ -4,8 +4,6 @@ from typing import TYPE_CHECKING from pydantic import BaseModel, Field, field_validator -from haiku.rag.store.models import BoundingBox - if TYPE_CHECKING: from haiku.rag.store.models import SearchResult @@ -38,7 +36,6 @@ class Citation(BaseModel): page_numbers: list[int] = Field(default_factory=list) headings: list[str] | None = None content: str - bounding_boxes: list[BoundingBox] | None = None class SearchAnswer(BaseModel): @@ -84,7 +81,6 @@ def resolve_citations( page_numbers=r.page_numbers, headings=r.headings, content=r.content, - bounding_boxes=r.bounding_boxes, ) ) return citations diff --git a/haiku_rag_slim/haiku/rag/graph/research/common.py b/haiku_rag_slim/haiku/rag/graph/research/common.py index a9d96c86..fcdb0820 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/common.py +++ b/haiku_rag_slim/haiku/rag/graph/research/common.py @@ -15,6 +15,15 @@ def format_context_for_prompt(context: ResearchContext) -> str: "question": qa.query, "answer": qa.answer, "confidence": qa.confidence, + "sources": [ + { + "document_uri": c.document_uri, + "document_title": c.document_title, + "page_numbers": c.page_numbers, + "headings": c.headings, + } + for c in qa.citations + ], } for qa in context.qa_responses ],