Search results and citations carry document metadata

This commit is contained in:
Yiorgis Gozadinos 2026-07-23 10:13:29 +03:00
parent 4f0fb7322e
commit 75a21c82c2
No known key found for this signature in database
11 changed files with 102 additions and 5 deletions

View file

@ -1,6 +1,10 @@
# Changelog # Changelog
## [Unreleased] ## [Unreleased]
### Added
- `SearchResult.document_meta` and `Citation.document_meta` carry the parent document's metadata.
## [0.67.1] - 2026-07-22 ## [0.67.1] - 2026-07-22
### Added ### Added

View file

@ -187,6 +187,8 @@ for result in results:
print(f"Document ID: {result.document_id}") print(f"Document ID: {result.document_id}")
``` ```
Each result carries the parent document's metadata in `result.document_meta`. It is not shown to the model during QA.
Search with different search types: Search with different search types:
```python ```python
# Vector search only # Vector search only
@ -326,7 +328,7 @@ answer, citations = await client.ask(
) )
``` ```
`client.ask` runs the [rag skill](skills/index.md) and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, and document references. `client.ask` runs the [rag skill](skills/index.md) and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, document references, and the document's metadata (`document_meta`), so UIs can render metadata keys such as a public source URL alongside the citation.
The QA provider and model are configured in `haiku.rag.yaml` or can be passed directly to the client (see [Configuration](configuration/index.md)). The QA provider and model are configured in `haiku.rag.yaml` or can be passed directly to the client (see [Configuration](configuration/index.md)).

View file

@ -415,6 +415,7 @@ def _build_result(
document_id=first.document_id, document_id=first.document_id,
document_uri=first.document_uri, document_uri=first.document_uri,
document_title=first.document_title, document_title=first.document_title,
document_meta=first.document_meta,
doc_item_refs=refs or first.doc_item_refs, doc_item_refs=refs or first.doc_item_refs,
page_numbers=sorted(pages) or first.page_numbers, page_numbers=sorted(pages) or first.page_numbers,
headings=all_headings or None, headings=all_headings or None,

View file

@ -376,6 +376,7 @@ def create_skill_tools(
doc_cache[did] = doc doc_cache[did] = doc
chunk.document_uri = doc.uri if doc else None chunk.document_uri = doc.uri if doc else None
chunk.document_title = doc.title if doc else None chunk.document_title = doc.title if doc else None
chunk.document_meta = doc.metadata if doc else {}
synthetic.append(SearchResult.from_chunk(chunk, score=1.0)) synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
if synthetic: if synthetic:
citations.extend(resolve_citations(missing, synthetic)) citations.extend(resolve_citations(missing, synthetic))

View file

@ -129,6 +129,9 @@ class SearchResult(BaseModel):
into this result; empty means just ``chunk_id``. It lets citation into this result; empty means just ``chunk_id``. It lets citation
consumers (visual grounding) reproduce a merged expansion and is never consumers (visual grounding) reproduce a merged expansion and is never
part of ``format_for_agent`` output. part of ``format_for_agent`` output.
``document_meta`` carries the parent document's metadata for citation
consumers (UIs). Never part of ``format_for_agent`` output.
""" """
content: str content: str
@ -138,6 +141,7 @@ class SearchResult(BaseModel):
document_id: str | None = None document_id: str | None = None
document_uri: str | None = None document_uri: str | None = None
document_title: str | None = None document_title: str | None = None
document_meta: dict = {}
order: int = 0 order: int = 0
doc_item_refs: list[str] = [] doc_item_refs: list[str] = []
page_numbers: list[int] = [] page_numbers: list[int] = []
@ -162,6 +166,7 @@ class SearchResult(BaseModel):
document_id=chunk.document_id, document_id=chunk.document_id,
document_uri=chunk.document_uri, document_uri=chunk.document_uri,
document_title=chunk.document_title, document_title=chunk.document_title,
document_meta=chunk.document_meta,
order=chunk.order, order=chunk.order,
doc_item_refs=meta.doc_item_refs, doc_item_refs=meta.doc_item_refs,
page_numbers=meta.page_numbers, page_numbers=meta.page_numbers,

View file

@ -26,6 +26,8 @@ class Citation(BaseModel):
content the exact items the model saw. Visual grounding resolves bounding content the exact items the model saw. Visual grounding resolves bounding
boxes from them so the rendered pages match the citation precisely. boxes from them so the rendered pages match the citation precisely.
``picture_refs`` is the picture-labeled subset. ``picture_refs`` is the picture-labeled subset.
``document_meta`` carries the cited document's metadata for UIs.
""" """
index: int | None = None index: int | None = None
@ -34,6 +36,7 @@ class Citation(BaseModel):
chunk_ids: list[str] = Field(default_factory=list) chunk_ids: list[str] = Field(default_factory=list)
document_uri: str document_uri: str
document_title: str | None = None document_title: str | None = None
document_meta: dict = Field(default_factory=dict)
page_numbers: list[int] = Field(default_factory=list) page_numbers: list[int] = Field(default_factory=list)
headings: list[str] | None = None headings: list[str] | None = None
content: str content: str
@ -64,6 +67,7 @@ def resolve_citations(
chunk_ids=r.chunk_ids or [chunk_id], chunk_ids=r.chunk_ids or [chunk_id],
document_uri=r.document_uri or "", document_uri=r.document_uri or "",
document_title=r.document_title, document_title=r.document_title,
document_meta=r.document_meta,
page_numbers=r.page_numbers, page_numbers=r.page_numbers,
headings=r.headings, headings=r.headings,
content=r.content, content=r.content,

View file

@ -77,6 +77,7 @@ async def rag_db(tmp_path_factory):
"Deep learning models are used in healthcare, finance, and transportation.", "Deep learning models are used in healthcare, finance, and transportation.",
title="AI Overview", title="AI Overview",
uri="test://ai-overview", uri="test://ai-overview",
metadata={"topic": "ai"},
) )
await rag.create_document( await rag.create_document(
"Machine learning is a subset of artificial intelligence. " "Machine learning is a subset of artificial intelligence. "

View file

@ -406,9 +406,9 @@ class TestCiteTool:
""" """
from haiku.rag.skills.rag import RAGState, create_skill from haiku.rag.skills.rag import RAGState, create_skill
docs = await rag_client.list_documents(limit=1) doc = await rag_client.get_document_by_uri("test://ai-overview")
assert docs, "fixture should have at least one document" assert doc, "fixture should have the ai-overview document"
doc_id = docs[0].id doc_id = doc.id
chunks = await rag_client.chunk_repository.get_by_document_id(doc_id) chunks = await rag_client.chunk_repository.get_by_document_id(doc_id)
assert chunks, "fixture document should have chunks" assert chunks, "fixture document should have chunks"
chunk_id = chunks[0].id chunk_id = chunks[0].id
@ -425,6 +425,7 @@ class TestCiteTool:
registered = state.citation_index[chunk_id] registered = state.citation_index[chunk_id]
assert registered.document_id == doc_id assert registered.document_id == doc_id
assert registered.document_uri # uri must be populated from doc lookup assert registered.document_uri # uri must be populated from doc lookup
assert registered.document_meta.get("topic") == "ai"
class TestLifespan: class TestLifespan:

View file

@ -2,7 +2,11 @@ from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import resolve_citations from haiku.rag.store.models.citation import resolve_citations
def _result(chunk_id: str, chunk_ids: list[str] | None = None) -> SearchResult: def _result(
chunk_id: str,
chunk_ids: list[str] | None = None,
document_meta: dict | None = None,
) -> SearchResult:
return SearchResult( return SearchResult(
content="content", content="content",
score=0.9, score=0.9,
@ -10,6 +14,7 @@ def _result(chunk_id: str, chunk_ids: list[str] | None = None) -> SearchResult:
chunk_ids=chunk_ids or [], chunk_ids=chunk_ids or [],
document_id="doc-1", document_id="doc-1",
document_uri="test://doc", document_uri="test://doc",
document_meta=document_meta or {},
) )
@ -39,3 +44,13 @@ def test_resolve_citations_skips_unknown_ids():
result = _result("c1") result = _result("c1")
citations = resolve_citations(["c1", "missing"], [result]) citations = resolve_citations(["c1", "missing"], [result])
assert len(citations) == 1 assert len(citations) == 1
def test_resolve_citations_copies_document_meta():
result = _result(
"c1", document_meta={"source_url": "https://example.org/report/view"}
)
citations = resolve_citations(["c1"], [result])
assert citations[0].document_meta == {
"source_url": "https://example.org/report/view"
}

View file

@ -255,6 +255,37 @@ def test_chunk_metadata_resolve_empty_refs():
assert doc_items == [] assert doc_items == []
def test_search_result_from_chunk_preserves_document_meta():
"""Document metadata flows from Chunk to SearchResult for citation
consumers (UIs)."""
chunk = Chunk(
id="chunk-1",
document_id="doc-1",
content="Some content.",
document_uri="file:///docs/report.pdf",
document_meta={"source_url": "https://example.org/report/view"},
)
result = SearchResult.from_chunk(chunk, score=0.9)
assert result.document_meta == {"source_url": "https://example.org/report/view"}
def test_search_result_format_for_agent_omits_document_meta():
"""Document metadata is UI plumbing, never shown to the model."""
result = SearchResult(
content="Some content.",
score=0.9,
chunk_id="chunk-1",
document_meta={"source_url": "https://example.org/report/view"},
)
formatted = result.format_for_agent(rank=1, total=1)
assert "source_url" not in formatted
assert "https://example.org/report/view" not in formatted
def test_search_result_format_for_agent_with_rank(): def test_search_result_format_for_agent_with_rank():
"""Test format_for_agent with rank and total parameters.""" """Test format_for_agent with rank and total parameters."""
result = SearchResult( result = SearchResult(

View file

@ -767,6 +767,38 @@ class TestExpandWithItems:
# to the model. # to the model.
assert "c2" not in expanded[0].format_for_agent() assert "c2" not in expanded[0].format_for_agent()
async def test_expanded_result_carries_document_meta(self, temp_db_path):
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="text",
text=f"Paragraph {i}. " * 10,
)
for i in range(5)
]
await rag.document_item_repository.create_items("doc-1", items)
r1 = SearchResult(
content="Paragraph 1.",
score=0.9,
chunk_id="c1",
document_id="doc-1",
doc_item_refs=["#/texts/1"],
document_meta={"source_url": "https://example.org/report/view"},
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [r1], 5000
)
assert len(expanded) == 1
assert expanded[0].document_meta == {
"source_url": "https://example.org/report/view"
}
async def test_merged_anchor_is_highest_scoring_constituent(self, temp_db_path): async def test_merged_anchor_is_highest_scoring_constituent(self, temp_db_path):
"""A merged result's chunk_id anchors on the best-scoring constituent, """A merged result's chunk_id anchors on the best-scoring constituent,
not whichever chunk sits earliest in the document.""" not whichever chunk sits earliest in the document."""