diff --git a/CHANGELOG.md b/CHANGELOG.md index 98d8e70a..6fa78a95 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,8 +5,7 @@ ### Added - `api_key` on model and embedding-model config, overriding the provider's environment variable. Honored on the `openai` and `ollama` providers, `vllm` embedders and rerankers, the picture-description VLM endpoint, and `doctor`'s endpoint probes; other providers raise. -- `lancedb.databases`: a name-to-location mapping for searching several databases at once, mutually exclusive with `lancedb.uri`. `client.search(..., sources=[...])` selects which to search, `sources=None` searches all of them, and `SearchResult.source` carries the configured name a result came from. Candidates are fused by the configured reranker over the union, or by reciprocal rank fusion when none is configured. `SearchResult.format_for_agent` names the database, so the model can attribute evidence to one while it answers. `haiku-rag search`, `ask` and `analyze` cover the configured set and label each result and citation with its database; every other command works on one, named with `--database NAME` or `--db PATH`. -- `lancedb.databases`: a name-to-location mapping for searching several databases at once, mutually exclusive with `lancedb.uri`. `client.search(..., sources=[...])` selects which to search, `sources=None` searches all of them, and `SearchResult.source` carries the configured name a result came from. Candidates are fused by the configured reranker over the union, or by reciprocal rank fusion when none is configured. `SearchResult.format_for_agent` names the database, so the model can attribute evidence to one while it answers. `haiku-rag search`, `ask`, `analyze` and `chat` cover the configured set and label each result and citation with its database; every other command works on one, named with `--database NAME` or `--db PATH`. +- `lancedb.databases`: a name-to-location mapping for searching several databases at once, mutually exclusive with `lancedb.uri`. `client.search(..., sources=[...])` selects which to search, `sources=None` searches all of them, and `SearchResult.source` carries the configured name a result came from. `Document.source` names it on a document from a listing or a lookup. Candidates are fused by the configured reranker over the union, or by reciprocal rank fusion when none is configured. `SearchResult.format_for_agent` names the database, so the model can attribute evidence to one while it answers. `haiku-rag search`, `ask`, `analyze` and `chat` cover the configured set and label each result and citation with its database; every other command works on one, named with `--database NAME` or `--db PATH`. - `client.ask(..., sources=[...])` asks across the selected databases, and `Citation.source` names the one a cited chunk came from. The cite fallback for an id absent from the run's results looks only in the selected databases, so a question scoped to some cannot cite another. - `client.analyze(..., sources=[...])` analyzes across the selected databases: the sandbox mounts their documents under one flat `/documents/{id}/` namespace, resolving each id to the database holding it, and in-code `search()` covers the same selection. diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index 6a8a52b7..3b7ce5eb 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -211,7 +211,9 @@ results = await client.search("query", sources=["medic"]) # one of them Candidates from each database are fused into one ranked list, by the configured reranker where there is one and by reciprocal rank fusion otherwise. Each result carries `source`, the name of the database it came from, and so does each -citation. +citation. A document from `list_documents`, `get_document_by_id`, +`get_document_by_uri` or `resolve_document` carries it too, whether it came from +a set or from one database named in `lancedb.databases`. **Configure a reranker when searching several databases.** Reciprocal rank fusion compares ranks, not scores, so every database contributes its own best matches diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 4a71ac9b..8725f4da 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -850,10 +850,16 @@ class HaikuRAGApp: if doc.title else "" ) + database_part = ( + f" [repr.attrib_name]database[/repr.attrib_name]: {doc.source}" + if doc.source + else "" + ) self.console.print( f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} " f"[repr.attrib_name]uri[/repr.attrib_name]: {doc.uri}" + title_part + + database_part + f" [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}" ) self.console.print( diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 44641ec4..70a7ad21 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -509,7 +509,7 @@ class HaikuRAG: return await self._from_any_covered( lambda owner: owner.get_document_by_id(document_id) ) - return await self.document_repository.get_by_id(document_id) + return self._name(await self.document_repository.get_by_id(document_id)) async def get_chunk_by_id(self, chunk_id: str) -> Chunk | None: """Get a chunk by its ID. @@ -565,7 +565,7 @@ class HaikuRAG: return await self._from_any_covered( lambda owner: owner.get_document_by_uri(uri) ) - return await self.document_repository.get_by_uri(uri) + return self._name(await self.document_repository.get_by_uri(uri)) async def resolve_document(self, id_or_title: str) -> Document | None: """Resolve a document by ID, title, or URI (in that order). @@ -671,9 +671,12 @@ class HaikuRAG: ] start = offset or 0 return merged[start:] if limit is None else merged[start : start + limit] - return await self.document_repository.list_all( + documents = await self.document_repository.list_all( limit=limit, offset=offset, filter=filter, include_content=include_content ) + for document in documents: + document.source = self._source + return documents async def count_documents(self, filter: str | None = None) -> int: """Count documents with optional filtering. @@ -694,6 +697,15 @@ class HaikuRAG: return sum(counts) return await self.document_repository.count(filter=filter) + def _name(self, document: Document | None) -> Document | None: + """`document`, told which configured database it came from. + + None where no database is named, as with a single ``lancedb.uri``. + """ + if document is not None: + document.source = self._source + return document + async def _from_any_covered( self, lookup: "Callable[[HaikuRAG], Coroutine[Any, Any, Any]]" ) -> Any: diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index f045b939..06407edf 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/haiku_rag_slim/haiku/rag/store/models/document.py @@ -13,12 +13,18 @@ if TYPE_CHECKING: class Document(BaseModel): """ Represents a document with an ID, content, and metadata. + + ``source`` names the configured database a document came from: the name + from ``lancedb.databases``, never a path or URI. It is None where no + database is named, as with the single ``lancedb.uri``, and is never + persisted. """ id: str | None = None content: str uri: str | None = None title: str | None = None + source: str | None = None metadata: dict = {} docling_document: bytes | None = Field(default=None, exclude=True) docling_pages: bytes | None = Field(default=None, exclude=True) diff --git a/tests/test_app.py b/tests/test_app.py index bca18636..ed13346a 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -64,6 +64,28 @@ async def test_list_documents_prints_each_document(app, client): assert "second" in out(app) +async def test_list_documents_names_the_database_of_each(app, client): + """A listing spanning databases has to say which one each document is from.""" + client.list_documents.return_value = [ + _doc("first", source="arxiv"), + _doc("second", source="wiki"), + ] + + await app.list_documents() + + printed = out(app) + assert "database: arxiv" in printed + assert "database: wiki" in printed + + +async def test_list_documents_omits_the_database_when_unnamed(app, client): + client.list_documents.return_value = [_doc("only")] + + await app.list_documents() + + assert "database:" not in out(app) + + async def test_add_document_from_text_reports_the_new_id(app, client): client.create_document.return_value = _doc("added body") diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index 23b752bc..b0a7d135 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -279,6 +279,81 @@ class TestLookupByIdentifier: assert await rag.get_document_by_uri("test://nowhere") is None +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.""" + + @pytest.mark.asyncio + async def test_a_listing_names_each_document_s_database(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one", "alpha two"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + docs = await rag.list_documents() + + assert {d.uri: d.source for d in docs} == { + "test://alpha/alpha one": "alpha", + "test://alpha/alpha two": "alpha", + "test://beta/beta one": "beta", + } + + @pytest.mark.asyncio + async def test_a_looked_up_document_names_its_database(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + beta = (await rag.clients_for(["beta"]))[0] + [target] = await beta.document_repository.list_all(limit=1) + assert target.id is not None + + by_id = await rag.get_document_by_id(target.id) + by_uri = await rag.get_document_by_uri("test://alpha/alpha one") + resolved = await rag.resolve_document(target.id) + + assert by_id is not None and by_id.source == "beta" + assert by_uri is not None and by_uri.source == "alpha" + assert resolved is not None and resolved.source == "beta" + + @pytest.mark.asyncio + async def test_one_named_database_still_names_itself(self, tmp_path): + """`haiku-rag --database alpha list` opens one database, and its name is + the whole reason the option exists.""" + 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 rag: + [listed] = await rag.list_documents() + assert listed.id is not None + by_id = await rag.get_document_by_id(listed.id) + by_uri = await rag.get_document_by_uri("test://alpha/alpha one") + + assert listed.source == "alpha" + assert by_id is not None and by_id.source == "alpha" + assert by_uri is not None and by_uri.source == "alpha" + + @pytest.mark.asyncio + async def test_one_database_leaves_the_source_unset(self, tmp_path, temp_db_path): + """Nothing names the database when there is only one to name.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + dim = get_config().embeddings.model.vector_dim + doc = DoclingDocument(name="solo") + doc.add_text(label=DocItemLabel.TEXT, text="solo") + await rag.import_document( + doc, + [Chunk(content="solo", embedding=[0.1] * dim, order=0)], + uri="test://solo", + ) + + [listed] = await rag.list_documents() + assert listed.source is None + assert listed.id is not None + assert (await rag.get_document_by_id(listed.id)).source is None + + class TestFederatedSearch: @pytest.mark.asyncio async def test_results_carry_their_source(self, tmp_path):