a04a16c717
HAIKU_RAG_DB and DB_PATH are gone: a capability covers what the configuration places or the db_path it is given, and the app backend and the AG-UI example load their configuration as the CLI does. The compose files point HAIKU_RAG_CONFIG_PATH at the mounted haiku.rag.yaml, which places the database at /data where DB_VOLUME is mounted; the backend refuses a configured set since it serves one database. The chat scopes a selection by source only over a set and names databases on filter rows only across several. Docstrings, docs and test fixtures stop describing an unnamed database; every database a search, listing or citation reports carries a name.
122 lines
4.7 KiB
Python
122 lines
4.7 KiB
Python
from collections.abc import Iterable
|
|
from typing import TYPE_CHECKING
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from haiku.rag.store.exceptions import AmbiguousCitationError
|
|
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 capabilities 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 ``HaikuRAG.get_picture_bytes(document_id, ref, source)`` 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``).
|
|
|
|
``source`` names the database the cited chunk came from: the name from
|
|
``lancedb.databases`` or a path's stem, never a path or URI. None only on a
|
|
citation resolved from a hand-built result.
|
|
|
|
``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.
|
|
|
|
``chunk_meta`` is the cited chunk's raw, unparsed ``Chunk.metadata``
|
|
dict — lossless and independent of the typed fields above, so a
|
|
third-party chunker's own fields survive here even as this schema
|
|
evolves.
|
|
"""
|
|
|
|
index: int | None = None
|
|
document_id: str
|
|
source: str | None = None
|
|
chunk_id: str
|
|
chunk_ids: list[str] = Field(default_factory=list)
|
|
chunk_meta: dict = Field(default_factory=dict)
|
|
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 ambiguous_citation(
|
|
chunk_id: str, sources: Iterable[str | None]
|
|
) -> AmbiguousCitationError:
|
|
"""The refusal for an id that names a chunk in more than one database."""
|
|
named = ", ".join(sorted(s or "unnamed" for s in sources))
|
|
return AmbiguousCitationError(
|
|
f"chunk id {chunk_id} names a chunk in more than one database "
|
|
f"({named}); a citation records the id alone and cannot say which"
|
|
)
|
|
|
|
|
|
def resolve_citations(
|
|
cited_chunk_ids: list[str],
|
|
search_results: "list[SearchResult]",
|
|
) -> list[Citation]:
|
|
"""Resolve chunk IDs to full Citation objects with metadata.
|
|
|
|
A chunk returned by more than one search resolves to its last occurrence.
|
|
Raises ``AmbiguousCitationError`` instead when a cited id names a chunk in
|
|
more than one of the databases searched: a citation records the id alone
|
|
and cannot say which.
|
|
"""
|
|
by_id: dict[str, SearchResult] = {}
|
|
ambiguous: dict[str, set[str | None]] = {}
|
|
for r in search_results:
|
|
# A result built by hand carries no id and nothing can cite it.
|
|
if cid := r.chunk_id:
|
|
if (held := by_id.get(cid)) is not None and held.source != r.source:
|
|
ambiguous.setdefault(cid, {held.source}).add(r.source)
|
|
# A chunk found by several searches is expanded once per search, so
|
|
# the copies differ in content, window and figures. The later entry
|
|
# wins.
|
|
by_id[cid] = r
|
|
|
|
citations = []
|
|
for raw_id in cited_chunk_ids:
|
|
chunk_id = raw_id.strip("[]")
|
|
if chunk_id in ambiguous:
|
|
raise ambiguous_citation(chunk_id, ambiguous[chunk_id])
|
|
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 "",
|
|
source=r.source,
|
|
chunk_id=chunk_id,
|
|
chunk_ids=r.chunk_ids or [chunk_id],
|
|
chunk_meta=r.chunk_meta,
|
|
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
|