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.
This commit is contained in:
Yiorgis Gozadinos
parent
460215158d
commit
65f6d72cd5
5 changed files with 30 additions and 204 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
Loading…
Reference in a new issue