78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
|
|
|
|
if TYPE_CHECKING:
|
|
from haiku.rag.store.models import SearchResult
|
|
|
|
|
|
class Citation(BaseModel):
|
|
"""Resolved citation with full metadata for display/visual grounding.
|
|
|
|
Used by the rag and analysis skills and rendered by the CLI / chat
|
|
application. The optional index field supports UI display ordering.
|
|
|
|
``picture_refs`` lists the ``self_ref`` values of picture items in the
|
|
cited chunk. Empty for text-only citations. UIs can fetch the picture
|
|
bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
|
|
and render them alongside the text content.
|
|
|
|
``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
|
|
into the cited result (always includes ``chunk_id``).
|
|
|
|
``doc_item_refs`` are the ``self_ref`` values of every item in the cited
|
|
content — the exact items the model saw. Visual grounding resolves bounding
|
|
boxes from them so the rendered pages match the citation precisely.
|
|
``picture_refs`` is the picture-labeled subset.
|
|
|
|
``document_meta`` carries the cited document's metadata for UIs.
|
|
"""
|
|
|
|
index: int | None = None
|
|
document_id: str
|
|
chunk_id: str
|
|
chunk_ids: list[str] = Field(default_factory=list)
|
|
document_uri: str
|
|
document_title: str | None = None
|
|
document_meta: dict = Field(default_factory=dict)
|
|
page_numbers: list[int] = Field(default_factory=list)
|
|
headings: list[str] | None = None
|
|
content: str
|
|
doc_item_refs: list[str] = Field(default_factory=list)
|
|
picture_refs: list[str] = Field(default_factory=list)
|
|
|
|
|
|
def resolve_citations(
|
|
cited_chunk_ids: list[str],
|
|
search_results: "list[SearchResult]",
|
|
) -> list[Citation]:
|
|
"""Resolve chunk IDs to full Citation objects with metadata."""
|
|
by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
|
|
|
|
citations = []
|
|
for raw_id in cited_chunk_ids:
|
|
chunk_id = raw_id.strip("[]")
|
|
r = by_id.get(chunk_id)
|
|
if not r:
|
|
continue
|
|
picture_refs = [
|
|
ref for ref in r.doc_item_refs if ref.startswith(PICTURE_REF_PREFIX)
|
|
]
|
|
citations.append(
|
|
Citation(
|
|
document_id=r.document_id or "",
|
|
chunk_id=chunk_id,
|
|
chunk_ids=r.chunk_ids or [chunk_id],
|
|
document_uri=r.document_uri or "",
|
|
document_title=r.document_title,
|
|
document_meta=r.document_meta,
|
|
page_numbers=r.page_numbers,
|
|
headings=r.headings,
|
|
content=r.content,
|
|
doc_item_refs=list(r.doc_item_refs),
|
|
picture_refs=picture_refs,
|
|
)
|
|
)
|
|
return citations
|