From 65f6d72cd5d373af91ca3c78cae96fb48d3b9ec6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 17:12:10 +0300 Subject: [PATCH] Remove the single-document item accessors the batching replaced `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs` and `get_all_items_grouped` have no callers left: the grouped equivalents serve every path that used them. `get_all_items_grouped` had none even before this branch. Tests whose subject was a removed method go with it. Tests that only used one to fetch a fixture now use the grouped call, so what they assert is unchanged. --- CHANGELOG.md | 2 +- .../rag/store/repositories/document_item.py | 109 +---------------- tests/store/test_document_items.py | 115 ++++-------------- tests/test_context.py | 4 +- tests/test_context_enhancement.py | 4 +- 5 files changed, 30 insertions(+), 204 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1dc425bd..e84e192d 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`. `get_text_for_refs` is removed; `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 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. - 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/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index b8902cbf..0426801b 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -108,69 +108,10 @@ class DocumentItemRepository: items.sort(key=lambda x: x.position) return items - async def get_all_items_grouped( - self, document_ids: list[str] | None = None - ) -> dict[str, list[DocumentItem]]: - """Get all items grouped by document_id in a single query. - - Args: - document_ids: If provided, only fetch items for these documents. - If None, fetches all items. - - Returns: - Dict mapping document_id to sorted list of DocumentItem. - """ - query = self.store.document_items_table.query().select(_METADATA_COLUMNS) - if document_ids is not None: - safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids) - query = query.where(f"document_id IN ({safe_ids})") - rows = await query.to_list() - - grouped: dict[str, list[DocumentItem]] = {} - for row in rows: - item = self._record_to_item(row) - grouped.setdefault(item.document_id, []).append(item) - for items in grouped.values(): - items.sort(key=lambda x: x.position) - return grouped - - async def get_items_in_range( - self, document_id: str, start: int, end: int - ) -> list[DocumentItem]: - """Get items for a document within a position range (inclusive).""" - safe_id = escape_sql_string(document_id) - rows = await ( - self.store.document_items_table.query() - .select(_METADATA_COLUMNS) - .where( - f"document_id = '{safe_id}' " - f"AND position >= {start} AND position <= {end}" - ) - .to_list() - ) - items = [self._record_to_item(row) for row in rows] - items.sort(key=lambda x: x.position) - return items - - async def resolve_refs(self, document_id: str, refs: list[str]) -> dict[str, int]: - """Resolve self_refs to positions. Returns {self_ref: position}.""" - if not refs: - return {} - - safe_id = escape_sql_string(document_id) - refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs) - rows = await ( - self.store.document_items_table.query() - .select(["self_ref", "position"]) - .where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})") - .to_list() - ) - return {row["self_ref"]: row["position"] for row in rows} - async def resolve_refs_grouped( self, refs_by_document: "Mapping[str, Sequence[str]]" ) -> dict[str, dict[str, int]]: - """`resolve_refs` across documents in one query.""" + """Resolve self_refs to positions, across documents, in one query.""" predicate = self._per_document_predicate(refs_by_document, "self_ref") if predicate is None: return {} @@ -190,7 +131,7 @@ class DocumentItemRepository: async def get_items_in_ranges( self, ranges_by_document: "Mapping[str, tuple[int, int]]" ) -> dict[str, list[DocumentItem]]: - """`get_items_in_range` across documents in one query. + """Items within a position range per document, in one query. Each document keeps its own inclusive range. Positions repeat across documents, so a shared range would splice one document's items into @@ -352,7 +293,7 @@ class DocumentItemRepository: async def get_caption_picture_refs_grouped( self, refs_by_document: "Mapping[str, list[str]]" ) -> dict[str, dict[str, str]]: - """`get_caption_picture_refs` across documents in two queries. + """Map caption refs to the picture preceding them, in two queries. Two rather than one because the stages are dependent: a caption's picture is the item at `position - 1`, which the first query is what @@ -393,47 +334,3 @@ class DocumentItemRepository: if caption: grouped.setdefault(row["document_id"], {})[caption] = row["self_ref"] return grouped - - async def get_caption_picture_refs( - self, document_id: str, refs: list[str] - ) -> dict[str, str]: - """Map caption refs to the picture item immediately preceding them. - - Docling emits a figure's caption at the position right after its - picture, so a caption's picture is the picture item at - ``position - 1``. Returns ``{caption_ref: picture_ref}`` for the - caption refs among ``refs`` that have a picture predecessor. Non-caption - refs, and captions whose predecessor is not a picture (table captions), - map to nothing. - """ - if not refs: - return {} - - safe_id = escape_sql_string(document_id) - refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs) - caption_rows = await ( - self.store.document_items_table.query() - .select(["self_ref", "position"]) - .where( - f"document_id = '{safe_id}' AND label = 'caption' " - f"AND self_ref IN ({refs_sql})" - ) - .to_list() - ) - if not caption_rows: - return {} - - prev_to_caption = {row["position"] - 1: row["self_ref"] for row in caption_rows} - positions_sql = ", ".join(str(p) for p in prev_to_caption) - picture_rows = await ( - self.store.document_items_table.query() - .select(["self_ref", "position"]) - .where( - f"document_id = '{safe_id}' AND label = 'picture' " - f"AND position IN ({positions_sql})" - ) - .to_list() - ) - return { - prev_to_caption[row["position"]]: row["self_ref"] for row in picture_rows - } diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index cfdb0f69..6104f2fc 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -265,33 +265,14 @@ class TestDocumentItemRepository: ] await repo.create_items("doc-1", items) - result = await repo.get_items_in_range("doc-1", 3, 7) + result = (await repo.get_items_in_ranges({"doc-1": (3, 7)})).get( + "doc-1", [] + ) assert len(result) == 5 assert result[0].position == 3 assert result[-1].position == 7 assert result[0].text == "Item 3" - async def test_resolve_refs(self, temp_db_path): - async with HaikuRAG(temp_db_path, create=True) as rag: - repo = DocumentItemRepository(rag.store) - - items = [ - DocumentItem( - document_id="doc-1", - position=i, - self_ref=f"#/texts/{i}", - label="paragraph", - text=f"Item {i}", - ) - for i in range(10) - ] - await repo.create_items("doc-1", items) - - refs = await repo.resolve_refs( - "doc-1", ["#/texts/2", "#/texts/7", "#/texts/999"] - ) - assert refs == {"#/texts/2": 2, "#/texts/7": 7} - async def test_get_item_count(self, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as rag: repo = DocumentItemRepository(rag.store) @@ -376,25 +357,24 @@ class TestDocumentItemRepository: (0, 2), ] - in_range = await repo.get_items_in_range("doc-1", 0, 2) + in_range = (await repo.get_items_in_ranges({"doc-1": (0, 2)})).get( + "doc-1", [] + ) assert [(i.heading_level, i.tree_depth) for i in in_range] == [ (1, 1), (2, 2), (0, 2), ] - grouped = await repo.get_all_items_grouped(["doc-1"]) - assert [(i.heading_level, i.tree_depth) for i in grouped["doc-1"]] == [ + assert [ + (i.heading_level, i.tree_depth) + for i in await repo.get_all_items("doc-1") + ] == [ (1, 1), (2, 2), (0, 2), ] - async def test_empty_refs_returns_empty(self, temp_db_path): - async with HaikuRAG(temp_db_path, create=True) as rag: - repo = DocumentItemRepository(rag.store) - assert await repo.resolve_refs("doc-1", []) == {} - async def test_items_sorted_by_position(self, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as rag: repo = DocumentItemRepository(rag.store) @@ -412,7 +392,9 @@ class TestDocumentItemRepository: ] await repo.create_items("doc-1", items) - result = await repo.get_items_in_range("doc-1", 0, 9) + result = (await repo.get_items_in_ranges({"doc-1": (0, 9)})).get( + "doc-1", [] + ) positions = [item.position for item in result] assert positions == sorted(positions) @@ -440,9 +422,11 @@ class TestDocumentItemPopulation: count = await rag.document_item_repository.get_item_count(created.id) assert count == 6 - items = await rag.document_item_repository.get_items_in_range( - created.id, 0, count - ) + items = ( + await rag.document_item_repository.get_items_in_ranges( + {created.id: (0, count)} + ) + ).get(created.id, []) assert items[0].label == "section_header" assert items[0].text == "Introduction" assert items[1].label == "paragraph" @@ -607,62 +591,6 @@ class TestPictureDataStorage: # Empty refs returns empty dict assert await repo.get_pictures_for_chunk("doc-1", []) == {} - async def test_get_caption_picture_refs(self, temp_db_path): - """A caption ref resolves to the picture at the immediately preceding - position; a table caption (no preceding picture) resolves to nothing.""" - async with HaikuRAG(temp_db_path, create=True) as rag: - repo = DocumentItemRepository(rag.store) - await repo.create_items( - "doc-1", - [ - DocumentItem( - document_id="doc-1", - position=0, - self_ref="#/pictures/0", - label="picture", - picture_data=b"\x89PNG\r\n\x1a\nfake", - ), - DocumentItem( - document_id="doc-1", - position=1, - self_ref="#/texts/0", - label="caption", - text="Figure 1. A figure caption.", - ), - DocumentItem( - document_id="doc-1", - position=2, - self_ref="#/texts/1", - label="paragraph", - text="Body prose.", - ), - DocumentItem( - document_id="doc-1", - position=3, - self_ref="#/tables/0", - label="table", - text="| a | b |", - ), - DocumentItem( - document_id="doc-1", - position=4, - self_ref="#/texts/2", - label="caption", - text="Table 1. A table caption.", - ), - ], - ) - - # Figure caption resolves to its picture; table caption does not. - got = await repo.get_caption_picture_refs( - "doc-1", ["#/texts/0", "#/texts/1", "#/texts/2"] - ) - assert got == {"#/texts/0": "#/pictures/0"} - - # A non-caption ref alone yields nothing. - assert await repo.get_caption_picture_refs("doc-1", ["#/texts/1"]) == {} - assert await repo.get_caption_picture_refs("doc-1", []) == {} - async def test_hot_paths_exclude_picture_data(self, temp_db_path): """Light read paths must NOT pull picture_data into memory.""" async with HaikuRAG(temp_db_path, create=True) as rag: @@ -685,10 +613,9 @@ class TestPictureDataStorage: for item in await repo.get_all_items("doc-1"): assert item.picture_data is None - for item in await repo.get_items_in_range("doc-1", 0, 10): - assert item.picture_data is None - grouped = await repo.get_all_items_grouped(["doc-1"]) - for item in grouped["doc-1"]: + for item in (await repo.get_items_in_ranges({"doc-1": (0, 10)})).get( + "doc-1", [] + ): assert item.picture_data is None # But the picture-byte accessors still work diff --git a/tests/test_context.py b/tests/test_context.py index 2d73c792..48130fdd 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -1462,10 +1462,10 @@ class TestExpandWithItemsWindowEdges: ) async def no_window(*_args, **_kwargs): - return [] + return {} monkeypatch.setattr( - rag.document_item_repository, "get_items_in_range", no_window + rag.document_item_repository, "get_items_in_ranges", no_window ) result = SearchResult( diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py index 6696edf7..9fa11e72 100644 --- a/tests/test_context_enhancement.py +++ b/tests/test_context_enhancement.py @@ -322,7 +322,9 @@ async def test_expand_context_single_item_document(temp_db_path): assert doc.id is not None # Create a search result with a doc_item_ref pointing to the item - items = await client.document_item_repository.get_items_in_range(doc.id, 0, 10) + items = ( + await client.document_item_repository.get_items_in_ranges({doc.id: (0, 10)}) + ).get(doc.id, []) assert len(items) > 0 search_results = [