From 28217fcf824966dcfa6d1d9e87df3716da8c7d0a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 17:41:10 +0300 Subject: [PATCH] Keep result order when expansion is batched Splitting the assembly into a passthrough pass and an expandable pass reordered equal-scored results: the score sort that follows is stable, so the order results arrive in is the tiebreak. Results are assembled in document_groups order again, after the batched fetch rather than around it. Also ports the caption negative cases the removed single-document test carried: a table's caption and an ordinary text reference map to no picture. --- CHANGELOG.md | 2 +- haiku_rag_slim/haiku/rag/client/search.py | 12 ++++--- tests/store/test_document_item_grouped.py | 42 +++++++++++++++++++++++ tests/test_enrichment_batching.py | 22 ++++++++++++ 4 files changed, 72 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e84e192d..cdd0613c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ ### Changed -- Search enrichment, the multimodal reranker's picture fetch, and context expansion each issue a fixed number of `document_items` queries regardless of how many documents a result set spans, instead of one set per document. `expand_with_items` takes the items to expand from rather than fetching them, and `DocumentItemRepository` gains `resolve_refs_grouped`, `get_items_in_ranges`, `get_pictures_grouped` and `get_caption_picture_refs_grouped`. The single-document methods they replace are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns the text alongside the bytes. +- Search enrichment, the multimodal reranker's picture fetch, and context expansion each issue a fixed number of `document_items` queries regardless of how many documents a result set spans, instead of one set per document. `expand_with_items` takes the items to expand from rather than fetching them, and `DocumentItemRepository` gains `resolve_refs_grouped`, `get_items_in_ranges`, `get_pictures_grouped` and `get_caption_picture_refs_grouped`. The superseded methods are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns the text alongside the bytes. - Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged. - `create_capability(rag=...)` lends a capability an open client rather than having it open its own; `client.ask`/`client.analyze` now pass theirs. - The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift. diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 8f8e9e0c..4bd1dc0d 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -237,10 +237,6 @@ async def expand_context( for doc_id, doc_results in document_groups.items() if doc_id is not None and any(r.doc_item_refs for r in doc_results) } - for doc_id, doc_results in document_groups.items(): - if doc_id not in expandable: - expanded_results.extend(doc_results) - repo = client.document_item_repository positions_by_document = await repo.resolve_refs_grouped( { @@ -255,7 +251,13 @@ async def expand_context( } items_by_document = await repo.get_items_in_ranges(windows) - for doc_id, doc_results in expandable.items(): + # In document_groups order: the score sort below is stable, so assembling + # expandable and passthrough documents in separate passes would reorder + # equal-scored results. + for doc_id, doc_results in document_groups.items(): + if doc_id not in expandable: + expanded_results.extend(doc_results) + continue expanded_results.extend( expand_with_items( doc_results, diff --git a/tests/store/test_document_item_grouped.py b/tests/store/test_document_item_grouped.py index 2daf03c7..9a975bab 100644 --- a/tests/store/test_document_item_grouped.py +++ b/tests/store/test_document_item_grouped.py @@ -114,3 +114,45 @@ async def test_grouped_calls_with_nothing_asked_for_do_not_query( assert await repo.get_caption_picture_refs_grouped({}) == {} assert item_queries["n"] == 0 + + +@pytest.mark.asyncio +async def test_caption_picture_refs_grouped_ignores_non_picture_predecessors( + temp_db_path, +): + """A caption maps to a picture only. A table's caption, and an ordinary text + reference, map to nothing.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentItemRepository(store) + await repo.create_items( + "doc-1", + [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/tables/0", + label="table", + text="a table", + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/texts/table-caption", + label="caption", + text="Table 1", + ), + DocumentItem( + document_id="doc-1", + position=2, + self_ref="#/texts/plain", + label="text", + text="ordinary prose", + ), + ], + ) + + got = await repo.get_caption_picture_refs_grouped( + {"doc-1": ["#/texts/table-caption", "#/texts/plain"]} + ) + + assert got == {} diff --git a/tests/test_enrichment_batching.py b/tests/test_enrichment_batching.py index 34af4a7d..3c13abd4 100644 --- a/tests/test_enrichment_batching.py +++ b/tests/test_enrichment_batching.py @@ -251,3 +251,25 @@ async def test_reranker_gives_each_chunk_its_own_document_picture(temp_db_path): assert chunks[0]._picture_data == b"bytes-doc-a" assert chunks[1]._picture_data == b"bytes-doc-b" + + +@pytest.mark.asyncio +async def test_expansion_keeps_document_order_for_tied_scores(temp_db_path): + """The score sort is stable, so equal-scored results must come back in the + order they arrived, whether or not their document expands.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + await _seed_expandable(rag, ["doc-expandable"]) + + passthrough = SearchResult( + chunk_id="doc-plain-anchor", + document_id="doc-plain", + content="plain body", + score=0.5, + doc_item_refs=[], + ) + expandable = _text_result("doc-expandable") + expandable.score = 0.5 + + for order in ([passthrough, expandable], [expandable, passthrough]): + expanded = await rag.expand_context(list(order)) + assert [r.chunk_id for r in expanded] == [r.chunk_id for r in order]