Tell a document which database it came from

Document.source names the configured database, as SearchResult and
Citation already do. A listing spanning databases is unreadable
without it, and `--database NAME list` could not name the one it
opened.
This commit is contained in:
Yiorgis Gozadinos 2026-08-23 12:22:47 +03:00
parent 1d09b4e31b
commit e94623ec37
No known key found for this signature in database
7 changed files with 128 additions and 6 deletions

View file

@ -5,8 +5,7 @@
### Added ### 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. - `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. `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`.
- `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`.
- `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.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. - `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.

View file

@ -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 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 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 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 **Configure a reranker when searching several databases.** Reciprocal rank fusion
compares ranks, not scores, so every database contributes its own best matches compares ranks, not scores, so every database contributes its own best matches

View file

@ -850,10 +850,16 @@ class HaikuRAGApp:
if doc.title if doc.title
else "" else ""
) )
database_part = (
f" [repr.attrib_name]database[/repr.attrib_name]: {doc.source}"
if doc.source
else ""
)
self.console.print( self.console.print(
f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} " f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} "
f"[repr.attrib_name]uri[/repr.attrib_name]: {doc.uri}" f"[repr.attrib_name]uri[/repr.attrib_name]: {doc.uri}"
+ title_part + title_part
+ database_part
+ f" [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}" + f" [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}"
) )
self.console.print( self.console.print(

View file

@ -509,7 +509,7 @@ class HaikuRAG:
return await self._from_any_covered( return await self._from_any_covered(
lambda owner: owner.get_document_by_id(document_id) 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: async def get_chunk_by_id(self, chunk_id: str) -> Chunk | None:
"""Get a chunk by its ID. """Get a chunk by its ID.
@ -565,7 +565,7 @@ class HaikuRAG:
return await self._from_any_covered( return await self._from_any_covered(
lambda owner: owner.get_document_by_uri(uri) 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: async def resolve_document(self, id_or_title: str) -> Document | None:
"""Resolve a document by ID, title, or URI (in that order). """Resolve a document by ID, title, or URI (in that order).
@ -671,9 +671,12 @@ class HaikuRAG:
] ]
start = offset or 0 start = offset or 0
return merged[start:] if limit is None else merged[start : start + limit] 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 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: async def count_documents(self, filter: str | None = None) -> int:
"""Count documents with optional filtering. """Count documents with optional filtering.
@ -694,6 +697,15 @@ class HaikuRAG:
return sum(counts) return sum(counts)
return await self.document_repository.count(filter=filter) 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( async def _from_any_covered(
self, lookup: "Callable[[HaikuRAG], Coroutine[Any, Any, Any]]" self, lookup: "Callable[[HaikuRAG], Coroutine[Any, Any, Any]]"
) -> Any: ) -> Any:

View file

@ -13,12 +13,18 @@ if TYPE_CHECKING:
class Document(BaseModel): class Document(BaseModel):
""" """
Represents a document with an ID, content, and metadata. 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 id: str | None = None
content: str content: str
uri: str | None = None uri: str | None = None
title: str | None = None title: str | None = None
source: str | None = None
metadata: dict = {} metadata: dict = {}
docling_document: bytes | None = Field(default=None, exclude=True) docling_document: bytes | None = Field(default=None, exclude=True)
docling_pages: bytes | None = Field(default=None, exclude=True) docling_pages: bytes | None = Field(default=None, exclude=True)

View file

@ -64,6 +64,28 @@ async def test_list_documents_prints_each_document(app, client):
assert "second" in out(app) 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): async def test_add_document_from_text_reports_the_new_id(app, client):
client.create_document.return_value = _doc("added body") client.create_document.return_value = _doc("added body")

View file

@ -279,6 +279,81 @@ class TestLookupByIdentifier:
assert await rag.get_document_by_uri("test://nowhere") is None 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: class TestFederatedSearch:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_results_carry_their_source(self, tmp_path): async def test_results_carry_their_source(self, tmp_path):