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:
parent
1d09b4e31b
commit
e94623ec37
7 changed files with 128 additions and 6 deletions
|
|
@ -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.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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(
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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):
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue