diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bcc7be9..c063d37e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ### Changed - `import_documents` embeds chunks across the whole batch in one pass instead of per document. +- Raw chunk metadata is now exposed to search and citation results, through `SearchResult.chunk_meta` and `Citation.chunk_meta`. For context-expanded results, the metadata is that of the anchor chunk. ### Removed diff --git a/docs/custom-pipelines.md b/docs/custom-pipelines.md index a523b4af..aa003a85 100644 --- a/docs/custom-pipelines.md +++ b/docs/custom-pipelines.md @@ -105,6 +105,8 @@ for chunk in chunks: print(f"Headings: {meta.headings}") print(f"Page numbers: {meta.page_numbers}") print(f"Labels: {meta.labels}") + # Access raw metadata (including headings, page_numbers and labels) + print(f"Raw metadata: {chunk.metadata}") ``` Chunks are returned with: @@ -115,6 +117,10 @@ Chunks are returned with: - `embedding` - `None` (not yet embedded) - `document_id` - `None` (not yet stored) +A custom `DocumentChunker` can provide other keys and values in `metadata`. They will +be stored with the chunk and are accessible when it is returned in a search result or citation, within +`SearchResult.chunk_meta` / `Citation.chunk_meta`. + ## Embed `embed_chunks()` generates embeddings for chunks using the client's embedder. It automatically contextualizes chunks (prepends section headings) before embedding for better semantic search, without modifying the stored content: diff --git a/docs/python.md b/docs/python.md index 48075426..e261d6cd 100644 --- a/docs/python.md +++ b/docs/python.md @@ -348,7 +348,7 @@ answer, citations = await client.ask( Images are passed to the model alongside the question. Retrieval stays text-based. The QA model must have `vision: true` in its configuration. -`client.ask` runs the [RAG capability](capabilities/rag.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. +`client.ask` runs the [RAG capability](capabilities/rag.md) and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, document references, the document's metadata (`document_meta`), and the cited chunk's raw, unparsed metadata (`chunk_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)). diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index 51e8375d..42073423 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -408,6 +408,7 @@ def _build_result( score=max(r.score for r in original_results), chunk_id=first.chunk_id, chunk_ids=chunk_ids, + chunk_meta=first.chunk_meta, document_id=first.document_id, document_uri=first.document_uri, document_title=first.document_title, diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 95666a39..9a6cbadb 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -132,12 +132,17 @@ class SearchResult(BaseModel): ``document_meta`` carries the parent document's metadata for citation consumers (UIs). Never part of ``format_for_agent`` output. + + ``chunk_meta``is the anchor chunk's unparsed ``Chunk.metadata`` and does not + include the metadata of any other chunks merged with it. Never part of + ``format_for_agent`` output. """ content: str score: float chunk_id: str | None = None chunk_ids: list[str] = [] + chunk_meta: dict = {} document_id: str | None = None document_uri: str | None = None document_title: str | None = None @@ -172,6 +177,7 @@ class SearchResult(BaseModel): page_numbers=meta.page_numbers, headings=meta.headings, labels=meta.labels, + chunk_meta=chunk.metadata, image_data=image_data, ) diff --git a/haiku_rag_slim/haiku/rag/store/models/citation.py b/haiku_rag_slim/haiku/rag/store/models/citation.py index b7b38eba..161d8575 100644 --- a/haiku_rag_slim/haiku/rag/store/models/citation.py +++ b/haiku_rag_slim/haiku/rag/store/models/citation.py @@ -28,12 +28,18 @@ class Citation(BaseModel): ``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 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) @@ -65,6 +71,7 @@ def resolve_citations( document_id=r.document_id or "", 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, diff --git a/tests/store/test_citation.py b/tests/store/test_citation.py index 2bf8ce38..2f765112 100644 --- a/tests/store/test_citation.py +++ b/tests/store/test_citation.py @@ -6,6 +6,7 @@ def _result( chunk_id: str, chunk_ids: list[str] | None = None, document_meta: dict | None = None, + chunk_meta: dict | None = None, ) -> SearchResult: return SearchResult( content="content", @@ -15,6 +16,7 @@ def _result( document_id="doc-1", document_uri="test://doc", document_meta=document_meta or {}, + chunk_meta=chunk_meta or {}, ) @@ -54,3 +56,9 @@ def test_resolve_citations_copies_document_meta(): assert citations[0].document_meta == { "source_url": "https://example.org/report/view" } + + +def test_resolve_citations_copies_chunk_meta(): + result = _result("c1", chunk_meta={"para_no": "12", "speaker": "MR SMITH"}) + citations = resolve_citations(["c1"], [result]) + assert citations[0].chunk_meta == {"para_no": "12", "speaker": "MR SMITH"} diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 5fe774f3..d052299f 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -222,6 +222,42 @@ def test_search_result_from_chunk_preserves_document_meta(): assert result.document_meta == {"source_url": "https://example.org/report/view"} +def test_search_result_from_chunk_preserves_chunk_meta(): + """Test flow through of unparsed chunk metadata from Chunk to SearchResult""" + chunk = Chunk( + id="chunk-1", + document_id="doc-1", + content="Some content.", + metadata={ + "headings": ["Chapter 1"], + "para_no": "12", + "speaker": "MR SMITH", + }, + ) + + result = SearchResult.from_chunk(chunk, score=0.9) + + assert result.chunk_meta == { + "headings": ["Chapter 1"], + "para_no": "12", + "speaker": "MR SMITH", + } + + +def test_search_result_format_for_agent_omits_chunk_meta(): + """Test that chunk_meta is never shown to the model""" + result = SearchResult( + content="Some content.", + score=0.9, + chunk_id="chunk-1", + chunk_meta={"para_no": "12"}, + ) + + formatted = result.format_for_agent(rank=1, total=1) + + assert "para_no" not in formatted + + def test_search_result_format_for_agent_omits_document_meta(): """Document metadata is UI plumbing, never shown to the model.""" result = SearchResult( diff --git a/tests/test_context.py b/tests/test_context.py index d99a70a9..4d307f60 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -799,6 +799,46 @@ class TestExpandWithItems: "source_url": "https://example.org/report/view" } + async def test_expanded_result_carries_anchor_chunk_meta(self, temp_db_path): + """chunk_meta belongs only to the anchor chunk (ie whichever constituent chunk earned the result its rank)""" + 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"], + chunk_meta={"para_no": "12"}, + ) + r2 = SearchResult( + content="Paragraph 3.", + score=0.85, + chunk_id="c2", + document_id="doc-1", + doc_item_refs=["#/texts/3"], + chunk_meta={"para_no": "14"}, + ) + expanded = await expand_with_items( + rag.document_item_repository, "doc-1", [r1, r2], 5000 + ) + assert len(expanded) == 1 + assert expanded[0].chunk_id == "c1" + assert expanded[0].chunk_meta == {"para_no": "12"} + 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."""