diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index ee9e80ce..f8dbdf4a 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -273,6 +273,12 @@ one retrieved result has the ID, that result is cited. For an ID absent from search results, the fallback checks every selected database and rejects multiple holders. A shared ID that nothing cites is ignored. +`get_document_by_id`, `get_chunk_by_id` and `get_picture_bytes` take an optional +`source`, and ask that database alone. A name the client does not cover raises +`UnknownDatabaseError`. Without one, the document and chunk lookups ask every +covered database and answer from the first that holds the ID; `get_picture_bytes` +requires one whenever the client covers a set. + The analysis sandbox rejects shared document IDs because its mount path is `/documents/{id}/`. diff --git a/docs/python.md b/docs/python.md index 1e28f89a..c86f4c72 100644 --- a/docs/python.md +++ b/docs/python.md @@ -93,6 +93,8 @@ PDFs that carry attachments via the `/EmbeddedFiles` table are split into one Do By ID: ```python doc = await client.get_document_by_id("document-id-string") +doc = await client.get_document_by_id("document-id-string", "papers") +chunk = await client.get_chunk_by_id("chunk-id-string", "papers") ``` By URI: diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 160ec134..276cee6a 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -655,15 +655,28 @@ class HaikuRAG: uri, ) - async def get_document_by_id(self, document_id: str) -> Document | None: - """Get a document by its ID. + async def get_document_by_id( + self, document_id: str, source: str | None = None + ) -> Document | None: + """Get a document by its ID, from the database named by `source`. Args: document_id: The unique identifier of the document. + source: The database it came from, which this client must cover. + Without one every covered database is asked, and ids repeat + between copies of a database, so a caller holding a source must + pass it. Returns: The Document instance if found, None otherwise. + + Raises: + UnknownDatabaseError: If `source` names a database this client does + not cover. """ + if source is not None: + (owner,) = await self.clients_covering([source]) + return await owner.get_document_by_id(document_id) if self.covers_multiple: return await self._from_any_covered( lambda owner: owner.get_document_by_id(document_id) @@ -672,15 +685,28 @@ class HaikuRAG: document_id ) - async def get_chunk_by_id(self, chunk_id: str) -> Chunk | None: - """Get a chunk by its ID. + async def get_chunk_by_id( + self, chunk_id: str, source: str | None = None + ) -> Chunk | None: + """Get a chunk by its ID, from the database named by `source`. Args: chunk_id: The unique identifier of the chunk. + source: The database it came from, which this client must cover. + Without one every covered database is asked, and ids repeat + between copies of a database, so a caller holding a source must + pass it. Returns: The Chunk instance if found, None otherwise. + + Raises: + UnknownDatabaseError: If `source` names a database this client does + not cover. """ + if source is not None: + (owner,) = await self.clients_covering([source]) + return await owner.get_chunk_by_id(chunk_id) if self.covers_multiple: return await self._from_any_covered( lambda owner: owner.get_chunk_by_id(chunk_id) @@ -695,21 +721,26 @@ class HaikuRAG: Args: document_id: The document holding the picture. self_ref: The picture's `self_ref`. - source: The database it came from. Required when federating. + source: The database it came from, which this client must cover. + Required when covering a set. Returns: The picture bytes if found, None otherwise. + + Raises: + UnknownDatabaseError: If `source` names a database this client does + not cover. """ - if not self.covers_multiple: - return await self.document_item_repository.get_picture_bytes( + if source is not None: + (owner,) = await self.clients_covering([source]) + return await owner.document_item_repository.get_picture_bytes( document_id, self_ref ) - if source is None: + if self.covers_multiple: raise ValueError( "a picture lookup across databases needs the source it came from" ) - (owner,) = await self.clients_for([source]) - return await owner.document_item_repository.get_picture_bytes( + return await self.document_item_repository.get_picture_bytes( document_id, self_ref ) @@ -744,13 +775,10 @@ class HaikuRAG: return doc safe_input = escape_sql_string(id_or_title) - docs = await self.list_documents(filter=f"title = '{safe_input}'") - if docs and docs[0].id: - return await self.get_document_by_id(docs[0].id) - - docs = await self.list_documents(filter=f"uri = '{safe_input}'") - if docs and docs[0].id: - return await self.get_document_by_id(docs[0].id) + for column in ("title", "uri"): + docs = await self.list_documents(filter=f"{column} = '{safe_input}'") + if docs and docs[0].id: + return await self.get_document_by_id(docs[0].id, docs[0].source) return None diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index f12062d6..c4bb7ddb 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -53,14 +53,14 @@ async def find_document(client: HaikuRAG, query: str): filter=f"LOWER(uri) LIKE LOWER('%{escaped_query}%') OR LOWER(uri) LIKE LOWER('%{no_spaces}%')", ) if docs and docs[0].id: - return await client.get_document_by_id(docs[0].id) + return await client.get_document_by_id(docs[0].id, docs[0].source) docs = await client.list_documents( limit=1, filter=f"LOWER(title) LIKE LOWER('%{escaped_query}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')", ) if docs and docs[0].id: - return await client.get_document_by_id(docs[0].id) + return await client.get_document_by_id(docs[0].id, docs[0].source) return None diff --git a/tests/multi_db/test_documents.py b/tests/multi_db/test_documents.py index e3feb35c..72e94904 100644 --- a/tests/multi_db/test_documents.py +++ b/tests/multi_db/test_documents.py @@ -7,7 +7,9 @@ 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.exceptions import UnknownDatabaseError from haiku.rag.store.models import Chunk +from haiku.rag.store.models.document_item import DocumentItem from tests.multi_db.helpers import ( _config, _seed, @@ -118,6 +120,36 @@ class TestLookupByIdentifier: assert found is not None and found.content == "beta one" + @pytest.mark.asyncio + async def test_a_chunk_is_read_from_the_database_its_source_names(self, tmp_path): + """One chunk id in two databases, holding different content. A result + carries the database it came from, so a caller holding one must be able + to say which of the two it means.""" + import shutil + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["shared body"]) + shutil.copytree(tmp_path / "alpha.lancedb", tmp_path / "beta.lancedb") + + async with HaikuRAG(config=config, sources=["beta"]) as beta: + [held] = await beta.chunk_repository.list_all(limit=1) + assert held.id is not None + await beta.store.chunks_table.update( + {"content": "only in beta"}, where=f"id = '{held.id}'" + ) + + async with HaikuRAG(config=config) as rag: + from_alpha = await rag.get_chunk_by_id(held.id, "alpha") + from_beta = await rag.get_chunk_by_id(held.id, "beta") + unqualified = await rag.get_chunk_by_id(held.id) + with pytest.raises(UnknownDatabaseError): + await rag.get_chunk_by_id(held.id, "gamma") + + assert from_alpha is not None and from_alpha.content == "shared body" + assert from_beta is not None and from_beta.content == "only in beta" + # Configured order, as an unqualified lookup has always answered. + assert unqualified is not None and unqualified.content == "shared body" + @pytest.mark.asyncio async def test_a_document_held_by_two_databases_answers_from_the_first( self, tmp_path @@ -140,6 +172,118 @@ class TestLookupByIdentifier: assert found is not None and found.source == "alpha" + @staticmethod + async def _collided(tmp_path): + """Two databases holding one document id, where only beta's answers to + the title and URI asked for.""" + import shutil + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["shared body"]) + shutil.copytree(tmp_path / "alpha.lancedb", tmp_path / "beta.lancedb") + + async with HaikuRAG(config=config, sources=["beta"]) as beta: + [target] = await beta.document_repository.list_all(limit=1) + target.title = "only in beta" + target.uri = "test://beta/only" + await beta.document_repository.update_meta(target) + return config + + @pytest.mark.asyncio + async def test_a_title_match_is_read_from_the_database_that_matched(self, tmp_path): + config = await self._collided(tmp_path) + + async with HaikuRAG(config=config) as rag: + by_title = await rag.resolve_document("only in beta") + by_uri = await rag.resolve_document("test://beta/only") + + assert by_title is not None + assert (by_title.source, by_title.title) == ("beta", "only in beta") + assert by_uri is not None + assert (by_uri.source, by_uri.uri) == ("beta", "test://beta/only") + + @pytest.mark.asyncio + async def test_a_partial_match_is_read_from_the_database_that_matched( + self, tmp_path + ): + from haiku.rag.tools.document import find_document + + config = await self._collided(tmp_path) + + async with HaikuRAG(config=config) as rag: + by_uri = await find_document(rag, "beta/onl") + by_title = await find_document(rag, "only in bet") + + assert by_uri is not None and by_uri.source == "beta" + assert by_title is not None and by_title.source == "beta" + + @pytest.mark.asyncio + async def test_a_source_is_checked_against_what_the_client_covers(self, tmp_path): + """A lookup naming a database the client does not cover is wrong rather + than answerable from the one it does cover.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config, sources=["alpha"]) as alpha: + [target] = await alpha.document_repository.list_all(limit=1) + assert target.id is not None + + await alpha.document_item_repository.create_all( + [ + DocumentItem( + document_id=target.id, + self_ref="#/pictures/0", + position=0, + label="picture", + text="", + picture_data=b"alpha-picture", + ) + ] + ) + + [held] = await alpha.chunk_repository.list_all(limit=1) + assert held.id is not None + + found = await alpha.get_document_by_id(target.id, "alpha") + picture = await alpha.get_picture_bytes(target.id, "#/pictures/0", "alpha") + chunk = await alpha.get_chunk_by_id(held.id, "alpha") + with pytest.raises(UnknownDatabaseError): + await alpha.get_document_by_id(target.id, "beta") + with pytest.raises(UnknownDatabaseError): + await alpha.get_picture_bytes(target.id, "#/pictures/0", "beta") + with pytest.raises(UnknownDatabaseError): + await alpha.get_chunk_by_id(held.id, "beta") + + assert found is not None and found.uri == "test://alpha/alpha one" + assert picture == b"alpha-picture" + assert chunk is not None and chunk.content == "alpha one" + + @pytest.mark.asyncio + async def test_an_unnamed_database_answers_to_no_name(self, temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as rag: + docling = DoclingDocument(name="one") + docling.add_text(label=DocItemLabel.TEXT, text="body") + dim = get_config().embeddings.model.vector_dim + doc = await rag.import_document( + docling, + [Chunk(content="body", embedding=[0.1] * dim, order=0)], + uri="test://one", + ) + assert doc.id is not None + + [held] = await rag.chunk_repository.list_all(limit=1) + assert held.id is not None + + assert await rag.get_document_by_id(doc.id) is not None + assert await rag.get_chunk_by_id(held.id) is not None + with pytest.raises(UnknownDatabaseError): + await rag.get_document_by_id(doc.id, "alpha") + with pytest.raises(UnknownDatabaseError): + await rag.get_picture_bytes(doc.id, "#/pictures/0", "alpha") + with pytest.raises(UnknownDatabaseError): + await rag.get_chunk_by_id(held.id, "alpha") + @pytest.mark.asyncio async def test_an_unknown_identifier_is_absent_rather_than_an_error(self, tmp_path): config = _config(tmp_path, ["alpha", "beta"])