diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 31f5498b..b15f51fc 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -184,6 +184,7 @@ async def _store_document_with_chunks( if session.config.storage.auto_vacuum: session.schedule_vacuum() + session.name(stored_doc) return stored_doc @@ -237,6 +238,7 @@ async def _update_document_with_chunks( if session.config.storage.auto_vacuum: session.schedule_vacuum() + session.name(updated_doc) return updated_doc @@ -343,7 +345,7 @@ async def _store_documents_with_chunks( if session.config.storage.auto_vacuum: session.schedule_vacuum() - return created + return session.name_all(created) async def import_documents( @@ -402,7 +404,9 @@ async def _refresh_doc_metadata( # The vacuum is debounced, and document_meta is tiny, so this is cheap. if session.config.storage.auto_vacuum: session.schedule_vacuum() + session.name(result) return result + session.name(doc) return doc @@ -879,6 +883,7 @@ async def update_document( updated = await session.document_repository.update_meta(existing_doc) if session.config.storage.auto_vacuum: session.schedule_vacuum() + session.name(updated) return updated if chunks is not None: diff --git a/haiku_rag_slim/haiku/rag/client/session.py b/haiku_rag_slim/haiku/rag/client/session.py index 3db23cc4..f52af32f 100644 --- a/haiku_rag_slim/haiku/rag/client/session.py +++ b/haiku_rag_slim/haiku/rag/client/session.py @@ -174,6 +174,12 @@ class SingleDatabaseSession: document.source = self.source return document + def name_all(self, documents: "list[Document]") -> "list[Document]": + """`documents`, each told which database it came from.""" + for document in documents: + document.source = self.source + return documents + async def get_document_by_id(self, document_id: str) -> "Document | None": return self.name(await self.document_repository.get_by_id(document_id)) diff --git a/tests/multi_db/test_documents.py b/tests/multi_db/test_documents.py index a7333d9f..e3feb35c 100644 --- a/tests/multi_db/test_documents.py +++ b/tests/multi_db/test_documents.py @@ -5,6 +5,7 @@ from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.labels import DocItemLabel from haiku.rag.client import HaikuRAG +from haiku.rag.client.session import SingleDatabaseSession from haiku.rag.config import get_config from haiku.rag.store.models import Chunk from tests.multi_db.helpers import ( @@ -157,6 +158,107 @@ class TestLookupByIdentifier: assert await rag.get_document_by_uri("test://nowhere") is None +class TestWritesNameTheirDatabase: + """A write returns the document it wrote, and it came from a database. A + read of the same document names it, so the write has to as well.""" + + @staticmethod + def _doc(text: str): + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + doc = DoclingDocument(name=text) + doc.add_text(label=DocItemLabel.TEXT, text=text) + return doc + + @pytest.mark.asyncio + async def test_import_names_the_database(self, tmp_path): + config = _config(tmp_path, ["alpha"]) + dim = get_config().embeddings.model.vector_dim + + async with HaikuRAG(config=config, create=True, sources=["alpha"]) as rag: + written = await rag.import_document( + self._doc("cats"), + [Chunk(content="cats", embedding=[0.1] * dim, order=0)], + uri="test://alpha/cats", + ) + assert written.id is not None + read = await rag.get_document_by_id(written.id) + + assert written.source == "alpha" + assert read is not None and read.source == written.source + + @pytest.mark.asyncio + async def test_a_batch_import_names_every_document(self, tmp_path): + from haiku.rag.client.documents import DocumentImport + + config = _config(tmp_path, ["alpha"]) + dim = get_config().embeddings.model.vector_dim + + async with HaikuRAG(config=config, create=True, sources=["alpha"]) as rag: + written = await rag.import_documents( + [ + DocumentImport( + docling_document=self._doc(text), + chunks=[Chunk(content=text, embedding=[0.1] * dim, order=0)], + uri=f"test://alpha/{text}", + ) + for text in ("cats", "dogs") + ] + ) + + assert [d.source for d in written] == ["alpha", "alpha"] + + @pytest.mark.asyncio + async def test_a_metadata_only_update_names_the_database(self, tmp_path): + """Changing only metadata rewrites the row without re-chunking, so it + never reaches the paths that name a document on the way through.""" + config = _config(tmp_path, ["alpha"]) + dim = get_config().embeddings.model.vector_dim + + async with HaikuRAG(config=config, create=True, sources=["alpha"]) as rag: + stored = await rag.import_document( + self._doc("cats"), + [Chunk(content="cats", embedding=[0.1] * dim, order=0)], + uri="test://alpha/cats", + ) + assert stored.id is not None + + updated = await rag.update_document( + stored.id, title="Cats", metadata={"k": "v"} + ) + + assert updated is not None + assert updated.title == "Cats" + assert updated.source == "alpha" + + @pytest.mark.asyncio + async def test_the_revision_short_circuit_names_the_database(self, tmp_path): + """`create_document_from_source` refreshes metadata in place when the + revision is unchanged, returning the document it rewrote.""" + from haiku.rag.client.documents import _refresh_doc_metadata + + config = _config(tmp_path, ["alpha"]) + dim = get_config().embeddings.model.vector_dim + + async with HaikuRAG(config=config, create=True, sources=["alpha"]) as rag: + stored = await rag.import_document( + self._doc("cats"), + [Chunk(content="cats", embedding=[0.1] * dim, order=0)], + uri="test://alpha/cats", + ) + assert isinstance(rag._session, SingleDatabaseSession) + refreshed = await _refresh_doc_metadata( + rag._session, + stored, + title="Cats", + user_metadata={"k": "v"}, + source_metadata=None, + ) + + assert refreshed.source == "alpha" + + class TestDocumentsNameTheirDatabase: """A listing that spans databases is unreadable when the documents do not say which one they came from, the same reason a search result carries one."""