af6a6b0bbe
`_populate_image_data` ran its stages once per result document, so a result set spanning N documents cost 4N `document_items` queries. Measured on a remote object-store corpus, a limit=5 search with expansion was 18 queries, 16 of them against `document_items`. The stages now run once each across every document, and flat in document count: two queries for the dependent caption-to-picture mapping when results ranked on a caption, one for the picture bytes. Two queries for a picture-ref result set, three at most. Picture text comes back with the bytes rather than from a second query, since it is on the same rows. Predicates are per document, `(document_id = 'a' AND self_ref IN (…)) OR (…)`, rather than `self_ref IN (union)`. self_ref and position values repeat across documents, so a union predicate would return other documents' rows: for picture_data that fetches blobs nobody asked for, and it can hand one document another document's picture.
134 lines
4.4 KiB
Python
134 lines
4.4 KiB
Python
import lancedb
|
|
import pytest
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.client.search import _populate_image_data
|
|
from haiku.rag.store.models import DocumentItem, SearchResult
|
|
|
|
|
|
def _picture_result(document_id: str, ref: str) -> SearchResult:
|
|
return SearchResult(
|
|
chunk_id=f"{document_id}-{ref}",
|
|
document_id=document_id,
|
|
content="body",
|
|
score=0.9,
|
|
doc_item_refs=[ref],
|
|
)
|
|
|
|
|
|
async def _seed(rag: HaikuRAG, document_ids: list[str]) -> None:
|
|
"""Each document gets the same self_refs, which is what real documents do:
|
|
`#/pictures/0` exists in every one of them."""
|
|
for document_id in document_ids:
|
|
await rag.document_item_repository.create_items(
|
|
document_id,
|
|
[
|
|
DocumentItem(
|
|
document_id=document_id,
|
|
position=0,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
text=f"caption for {document_id}",
|
|
picture_data=f"bytes-{document_id}".encode(),
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def item_queries(monkeypatch):
|
|
tally = {"n": 0}
|
|
query = lancedb.AsyncTable.query
|
|
|
|
def counted(self):
|
|
if self.name == "document_items":
|
|
tally["n"] += 1
|
|
return query(self)
|
|
|
|
monkeypatch.setattr(lancedb.AsyncTable, "query", counted)
|
|
return tally
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enrichment_query_count_does_not_grow_with_documents(
|
|
temp_db_path, item_queries
|
|
):
|
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
await _seed(rag, [f"doc-{i}" for i in range(6)])
|
|
|
|
results = [_picture_result("doc-0", "#/pictures/0")]
|
|
item_queries["n"] = 0
|
|
await _populate_image_data(rag, results)
|
|
one_document = item_queries["n"]
|
|
|
|
results = [_picture_result(f"doc-{i}", "#/pictures/0") for i in range(6)]
|
|
item_queries["n"] = 0
|
|
await _populate_image_data(rag, results)
|
|
six_documents = item_queries["n"]
|
|
|
|
assert (one_document, six_documents) == (2, 2), (
|
|
f"one document took {one_document} queries, six took {six_documents}"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_each_document_gets_its_own_pictures(temp_db_path):
|
|
"""self_refs collide across documents, so a batched fetch keyed on self_ref
|
|
alone would hand one document another's picture."""
|
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
await _seed(rag, ["doc-a", "doc-b"])
|
|
|
|
results = [
|
|
_picture_result("doc-a", "#/pictures/0"),
|
|
_picture_result("doc-b", "#/pictures/0"),
|
|
]
|
|
await _populate_image_data(rag, results)
|
|
|
|
import base64
|
|
|
|
for result, document_id in zip(results, ["doc-a", "doc-b"]):
|
|
assert result.image_data is not None
|
|
blob = base64.b64decode(result.image_data["#/pictures/0"])
|
|
assert blob == f"bytes-{document_id}".encode()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_caption_ranked_results_take_at_most_four_queries(
|
|
temp_db_path, item_queries
|
|
):
|
|
"""The worst case: results ranked on a caption, so the dependent
|
|
caption-to-picture mapping runs too. Two for that, one for the blobs and
|
|
their text. Still flat in document count."""
|
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
for i in range(4):
|
|
document_id = f"doc-{i}"
|
|
await rag.document_item_repository.create_items(
|
|
document_id,
|
|
[
|
|
DocumentItem(
|
|
document_id=document_id,
|
|
position=0,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
text=f"caption for {document_id}",
|
|
picture_data=f"bytes-{document_id}".encode(),
|
|
),
|
|
DocumentItem(
|
|
document_id=document_id,
|
|
position=1,
|
|
self_ref="#/texts/1",
|
|
label="caption",
|
|
text=f"figure 1 of {document_id}",
|
|
),
|
|
],
|
|
)
|
|
|
|
counts = []
|
|
for n in (1, 4):
|
|
results = [_picture_result(f"doc-{i}", "#/texts/1") for i in range(n)]
|
|
item_queries["n"] = 0
|
|
await _populate_image_data(rag, results)
|
|
counts.append(item_queries["n"])
|
|
assert all(r.image_data for r in results)
|
|
|
|
assert counts == [3, 3], counts
|