From e8d89aa035de0a31e6897a65236db61e59c89a41 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 4 May 2026 13:31:15 +0300 Subject: [PATCH] dedup picture-only chunks at the search-result layer --- haiku_rag_slim/haiku/rag/client/search.py | 27 ++++++++++ haiku_rag_slim/haiku/rag/config/loader.py | 12 ++--- tests/test_search.py | 60 +++++++++++++++++++++++ 3 files changed, 92 insertions(+), 7 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 5dedc596..9009ec38 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -68,6 +68,7 @@ async def search( ) results = [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results] + results = _dedup_picture_chunks(results) if include_images: await _populate_image_data(client, results) @@ -75,6 +76,32 @@ async def search( return results +def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]: + """Collapse duplicate picture-only chunks to one result per ``self_ref``. + + A single picture can produce two chunks for the same self_ref: one whose + vector is the text embedding of the picture's description, and one whose + vector is the image embedding of the picture's bytes. Both can rank for + the same query. When two results share a single picture self_ref as + their only ref, keep the higher-scoring one. Wider chunks that span the + picture plus surrounding items pass through untouched. + """ + seen: dict[tuple[str | None, str], int] = {} + keep: list[bool] = [True] * len(results) + for i, r in enumerate(results): + if len(r.doc_item_refs) == 1 and r.doc_item_refs[0].startswith("#/pictures/"): + key = (r.document_id, r.doc_item_refs[0]) + prior = seen.get(key) + if prior is None: + seen[key] = i + elif r.score > results[prior].score: + keep[prior] = False + seen[key] = i + else: + keep[i] = False + return [r for r, k in zip(results, keep) if k] + + async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) -> None: """Attach base64 picture bytes to ``SearchResult.image_data`` in-place. diff --git a/haiku_rag_slim/haiku/rag/config/loader.py b/haiku_rag_slim/haiku/rag/config/loader.py index b1caa7c1..5205bfde 100644 --- a/haiku_rag_slim/haiku/rag/config/loader.py +++ b/haiku_rag_slim/haiku/rag/config/loader.py @@ -53,21 +53,19 @@ def load_yaml_config(path: Path) -> dict: def _translate_legacy_picture_fields(data: dict) -> None: - """Map pre-A4 picture knobs onto ``processing.pictures``. + """Map legacy picture knobs onto ``processing.pictures``. - Pre-A4 the same intent was expressed by two booleans on + Older configs expressed the same intent through two booleans on ``conversion_options``: ``generate_picture_images`` and ``picture_description.enabled``. Translation, in priority order: - ``picture_description.enabled = true`` (regardless of the image flag) - → ``pictures = "description"``. Mirrors the original behavior where - enabling the VLM implicitly forced docling to produce picture bytes. + → ``pictures = "description"``. - ``generate_picture_images = true`` (and no description) → ``"image"``. - both false / missing → no translation; default ``"none"`` applies. - If ``pictures`` is already set on the loaded YAML it wins — users who - have migrated keep their explicit choice. Mutates ``data`` in-place - and emits one warning per legacy field encountered. + If ``pictures`` is already set on the loaded YAML it wins. Mutates + ``data`` in-place and emits one warning per legacy field encountered. """ processing = data.get("processing") if not isinstance(processing, dict): diff --git a/tests/test_search.py b/tests/test_search.py index 5ced315e..8ab46e14 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -440,3 +440,63 @@ async def test_search_with_bytes_query_raises_for_text_only_embedder( async with HaikuRAG(temp_db_path, create=True) as rag: with pytest.raises(ValueError, match="multimodal embedder"): await rag.search(b"\x89PNG\r\n\x1a\n") + + +def _picture_only_result( + self_ref: str, score: float, document_id: str = "doc-1" +) -> SearchResult: + return SearchResult( + content="x", + score=score, + chunk_id=f"chunk-{self_ref}-{score}", + document_id=document_id, + doc_item_refs=[self_ref], + labels=["picture"], + ) + + +def test_dedup_keeps_higher_scoring_picture_chunk(): + """Two results referencing the same single picture self_ref collapse + to the one with the higher score.""" + from haiku.rag.client.search import _dedup_picture_chunks + + text_chunk = _picture_only_result("#/pictures/0", score=0.7) + pic_chunk = _picture_only_result("#/pictures/0", score=0.9) + other = _picture_only_result("#/pictures/1", score=0.6) + + deduped = _dedup_picture_chunks([text_chunk, pic_chunk, other]) + + assert len(deduped) == 2 + chosen = next(r for r in deduped if r.doc_item_refs == ["#/pictures/0"]) + assert chosen.score == 0.9 + assert any(r.doc_item_refs == ["#/pictures/1"] for r in deduped) + + +def test_dedup_preserves_wider_chunks_referencing_same_picture(): + """A wider chunk that contains the picture plus surrounding items + is independent signal — keep it alongside a picture-only chunk.""" + from haiku.rag.client.search import _dedup_picture_chunks + + pic_only = _picture_only_result("#/pictures/0", score=0.9) + wider = SearchResult( + content="surrounding paragraph text and a figure", + score=0.7, + chunk_id="wider", + document_id="doc-1", + doc_item_refs=["#/texts/3", "#/pictures/0", "#/texts/4"], + labels=["text", "picture", "text"], + ) + + deduped = _dedup_picture_chunks([pic_only, wider]) + assert len(deduped) == 2 + + +def test_dedup_does_not_collapse_across_documents(): + """Same self_ref in different documents is different content.""" + from haiku.rag.client.search import _dedup_picture_chunks + + a = _picture_only_result("#/pictures/0", score=0.5, document_id="doc-1") + b = _picture_only_result("#/pictures/0", score=0.9, document_id="doc-2") + + deduped = _dedup_picture_chunks([a, b]) + assert len(deduped) == 2