Merge pull request #509 from ggozad/feat/pass-metadata-citations

Search results and citations carry document metadata

CHANGELOG.md

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

docs/python.md

@@ -187,6 +187,8 @@ for result in results:
     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:
 ```python
 # 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)).
 

haiku_rag_slim/haiku/rag/context.py

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

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -376,6 +376,7 @@ def create_skill_tools(
                             doc_cache[did] = doc
                         chunk.document_uri = doc.uri 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))
                     if synthetic:
                         citations.extend(resolve_citations(missing, synthetic))

haiku_rag_slim/haiku/rag/store/models/chunk.py

@@ -129,6 +129,9 @@ class SearchResult(BaseModel):
     into this result; empty means just ``chunk_id``. It lets citation
     consumers (visual grounding) reproduce a merged expansion and is never
     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
@@ -138,6 +141,7 @@ class SearchResult(BaseModel):
     document_id: str | None = None
     document_uri: str | None = None
     document_title: str | None = None
+    document_meta: dict = {}
     order: int = 0
     doc_item_refs: list[str] = []
     page_numbers: list[int] = []
@@ -162,6 +166,7 @@ class SearchResult(BaseModel):
             document_id=chunk.document_id,
             document_uri=chunk.document_uri,
             document_title=chunk.document_title,
+            document_meta=chunk.document_meta,
             order=chunk.order,
             doc_item_refs=meta.doc_item_refs,
             page_numbers=meta.page_numbers,

haiku_rag_slim/haiku/rag/store/models/citation.py

@@ -26,6 +26,8 @@ class Citation(BaseModel):
     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
@@ -34,6 +36,7 @@ class Citation(BaseModel):
     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
@@ -64,6 +67,7 @@ def resolve_citations(
                 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,

tests/skills/conftest.py

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

tests/skills/test_rag.py

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

tests/store/test_citation.py

@@ -2,7 +2,11 @@ from haiku.rag.store.models.chunk import SearchResult
 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(
         content="content",
         score=0.9,
@@ -10,6 +14,7 @@ def _result(chunk_id: str, chunk_ids: list[str] | None = None) -> SearchResult:
         chunk_ids=chunk_ids or [],
         document_id="doc-1",
         document_uri="test://doc",
+        document_meta=document_meta or {},
     )
 
 
@@ -39,3 +44,13 @@ def test_resolve_citations_skips_unknown_ids():
     result = _result("c1")
     citations = resolve_citations(["c1", "missing"], [result])
     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"
+    }

tests/test_chunk.py

@@ -255,6 +255,37 @@ def test_chunk_metadata_resolve_empty_refs():
     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():
     """Test format_for_agent with rank and total parameters."""
     result = SearchResult(

tests/test_context.py

@@ -767,6 +767,38 @@ class TestExpandWithItems:
             # to the model.
             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):
         """A merged result's chunk_id anchors on the best-scoring constituent,
         not whichever chunk sits earliest in the document."""