fdb5710491
Chunk 2 gave search a configured set to fan out over. ask and analyze
covered one database still: the RAG capability had no way to be told which
databases a question spanned, and the analysis sandbox mounted one
document tree.
The selection travels as sources on EvidenceState, beside the filter it
scopes with, so both capabilities read it the same way. clients_covering
is the one rule that turns a selection into clients, used by search, the
sandbox mount and the cite fallback, so a question scoped to some
databases cannot search, mount or cite another. Citations carry the
database they came from, and format_for_agent names it, so the model can
attribute evidence while it answers rather than only afterwards.
The sandbox keeps one flat /documents/{id}/ namespace and resolves each id
to the client holding it, which rests on ids being UUID4. A database
copied from another breaks that, so an id held twice is refused rather
than resolved to whichever arrived last.
On the CLI, search, ask and analyze cover the configured set and label
each result with its database. Every other command works on one, named
with --database NAME (a name reaches a database behind a URI, which --db
cannot) or --db PATH, and refuses a set it cannot choose from instead of
silently reading the default database. Cold databases open together, so a
first query costs the slowest open rather than their sum.
91 lines
3.3 KiB
Python
91 lines
3.3 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 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 configured database the cited chunk came from: the name
|
|
from ``lancedb.databases``, never a path or URI. It is None only where no
|
|
database is named, as with the single ``lancedb.uri``.
|
|
|
|
``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 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 "",
|
|
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
|