import lancedb import lancedb.query import pytest from haiku.rag.client import HaikuRAG from haiku.rag.client.search import _attach_picture_data, _populate_image_data from haiku.rag.store.models import Chunk, 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_projections(monkeypatch): """Columns each document_items query projects.""" projections: list[list[str]] = [] select = lancedb.query.AsyncQuery.select def recording(self, columns): if isinstance(columns, list): projections.append([str(c) for c in columns]) return select(self, columns) monkeypatch.setattr(lancedb.query.AsyncQuery, "select", recording) return projections @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 async def _seed_expandable(rag: HaikuRAG, document_ids: list[str]) -> None: """A section header and two text items, so expansion has something to widen into. Positions and self_refs repeat across documents.""" for document_id in document_ids: await rag.document_item_repository.create_items( document_id, [ DocumentItem( document_id=document_id, position=0, self_ref="#/texts/0", label="section_header", text=f"Section of {document_id}", ), DocumentItem( document_id=document_id, position=1, self_ref="#/texts/1", label="text", text=f"anchor body of {document_id}", ), DocumentItem( document_id=document_id, position=2, self_ref="#/texts/2", label="text", text=f"neighbouring body of {document_id}", ), ], ) def _text_result(document_id: str) -> SearchResult: return SearchResult( chunk_id=f"{document_id}-anchor", document_id=document_id, content=f"anchor body of {document_id}", score=0.9, doc_item_refs=["#/texts/1"], ) @pytest.mark.asyncio async def test_expansion_query_count_is_flat_in_document_count( temp_db_path, item_queries ): async with HaikuRAG(temp_db_path, create=True) as rag: await _seed_expandable(rag, [f"doc-{i}" for i in range(5)]) counts = [] for n in (1, 5): results = [_text_result(f"doc-{i}") for i in range(n)] item_queries["n"] = 0 expanded = await rag.expand_context(results) counts.append(item_queries["n"]) assert len(expanded) == n assert counts == [2, 2], counts @pytest.mark.asyncio async def test_expansion_widens_each_document_with_its_own_items(temp_db_path): """Positions repeat across documents, so a batched window fetch keyed on position alone would splice one document's text into another's context.""" async with HaikuRAG(temp_db_path, create=True) as rag: await _seed_expandable(rag, ["doc-a", "doc-b"]) expanded = await rag.expand_context( [_text_result("doc-a"), _text_result("doc-b")] ) by_doc = {r.document_id: r.content for r in expanded} assert "neighbouring body of doc-a" in by_doc["doc-a"] assert "doc-b" not in by_doc["doc-a"] assert "neighbouring body of doc-b" in by_doc["doc-b"] assert "doc-a" not in by_doc["doc-b"] def _picture_chunk(document_id: str) -> Chunk: return Chunk( id=f"{document_id}-pic", document_id=document_id, content="a figure", metadata={"doc_item_refs": ["#/pictures/0"], "labels": ["picture"]}, ) @pytest.mark.asyncio async def test_reranker_blob_fetch_is_one_query_for_any_document_count( temp_db_path, item_queries, item_projections ): """This path runs over `limit * 10` candidates, so per-document fetching costs the most here.""" async with HaikuRAG(temp_db_path, create=True) as rag: await _seed(rag, [f"doc-{i}" for i in range(10)]) counts = [] for n in (1, 10): chunks = [_picture_chunk(f"doc-{i}") for i in range(n)] item_queries["n"] = 0 await _attach_picture_data(rag, chunks) counts.append(item_queries["n"]) assert all(c._picture_data for c in chunks) assert counts == [1, 1], counts # The reranker scores pixels, so `text` has no business in the projection. picture_projections = [p for p in item_projections if "picture_data" in p] assert picture_projections, "no picture query observed" assert all("text" not in p for p in picture_projections), picture_projections @pytest.mark.asyncio async def test_reranker_gives_each_chunk_its_own_document_picture(temp_db_path): async with HaikuRAG(temp_db_path, create=True) as rag: await _seed(rag, ["doc-a", "doc-b"]) chunks = [_picture_chunk("doc-a"), _picture_chunk("doc-b")] await _attach_picture_data(rag, chunks) 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]