diff --git a/CHANGELOG.md b/CHANGELOG.md index b165753a..989d444f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,6 +39,7 @@ - The `haiku.rag` package declares the `jina` extra, so `provider: jina-local` is supported by declaration rather than through `cross-encoder`'s transitive `transformers` and `torch`. Raises the full package's torch floor to 2.0. - `providers.docling_serve.timeout` (default 300 seconds), forwarded to the docling-serve client's per-request timeout. - `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. +- `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. ### Changed diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 30d8865c..68216a06 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -35,7 +35,7 @@ from haiku.rag.capabilities._tools import ( from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, EvidenceRef from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig -from haiku.rag.store.models.chunk import SearchResult +from haiku.rag.store.models.chunk import Chunk, SearchResult from haiku.rag.store.models.citation import Citation, resolve_citations from haiku.rag.tools.search import build_image_content_from_results @@ -91,6 +91,7 @@ class EvidenceState(BaseModel): citations: list[str] = Field(default_factory=list) evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord) document_filter: str | None = None + sources: list[str] | None = None searches: dict[str, list[SearchResult]] = Field(default_factory=dict) def begin_invocation(self) -> None: @@ -101,8 +102,8 @@ class EvidenceState(BaseModel): leaves a later citation unable to resolve against the expanded result the model saw, recording no provenance for it. - `document_filter` scopes the conversation, and `evidence` carries question - identity, so neither is working evidence. + `document_filter` and `sources` scope the conversation, and `evidence` + carries question identity, so none of them is working evidence. """ self.citations.clear() self.searches.clear() @@ -136,6 +137,21 @@ def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) - return False +async def _first_holding( + clients: "list[HaikuRAG]", chunk_id: str +) -> "tuple[HaikuRAG, Chunk] | None": + """The first client holding this chunk, and the chunk. + + A chunk id says nothing about which database holds it, so the only way to + place one is to ask. Returns None when none of them has it. + """ + for client in clients: + chunk = await client.get_chunk_by_id(chunk_id) + if chunk is not None: + return client, chunk + return None + + @dataclass class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): db_path: Path @@ -468,6 +484,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): query, limit=limit, document_filter=self.state.document_filter, + sources=self.state.sources, ) state = self.state # A model can search the same query twice with different limits, and the @@ -504,20 +521,35 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): if missing: async with self.rag_lock: rag = await self._ensure_rag() + # A chunk id says nothing about which database holds it, so a + # federating client looks through the ones it covers. + if rag._federated: + lookups = await rag.clients_for( + getattr(self.state, "sources", None) or list(rag._federated) + ) + else: + lookups = [rag] synthetic: list[SearchResult] = [] - documents: dict[str, Any] = {} + documents: dict[tuple[str | None, str], Any] = {} for chunk_id in missing: - chunk = await rag.get_chunk_by_id(chunk_id) - if chunk is None or not chunk.document_id: + found = await _first_holding(lookups, chunk_id) + if found is None: continue - document = documents.get(chunk.document_id) - if chunk.document_id not in documents: - document = await rag.get_document_by_id(chunk.document_id) - documents[chunk.document_id] = document + owner, chunk = found + if not chunk.document_id: + continue + key = (owner._source, chunk.document_id) + if key not in documents: + documents[key] = await owner.get_document_by_id( + chunk.document_id + ) + document = documents[key] chunk.document_uri = document.uri if document else None chunk.document_title = document.title if document else None chunk.document_meta = document.metadata if document else {} - synthetic.append(SearchResult.from_chunk(chunk, score=1.0)) + result = SearchResult.from_chunk(chunk, score=1.0) + result.source = owner._source + synthetic.append(result) citations.extend(resolve_citations(missing, synthetic)) if not citations: diff --git a/haiku_rag_slim/haiku/rag/capabilities/_tools.py b/haiku_rag_slim/haiku/rag/capabilities/_tools.py index da3e825d..a732157c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_tools.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_tools.py @@ -18,9 +18,12 @@ async def search_corpus( query: str, limit: int | None = None, document_filter: str | None = None, + sources: list[str] | None = None, ) -> tuple[str, list[SearchResult]]: """Search and context-expand results for a capability tool.""" - results = await rag.search(query, limit=limit, filter=document_filter) + results = await rag.search( + query, limit=limit, filter=document_filter, sources=sources + ) results = await rag.expand_context(results) formatted = "\n\n---\n\n".join( result.format_for_agent(rank=index + 1, total=len(results)) diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 2df90d77..7fa05334 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -635,10 +635,11 @@ class HaikuRAG: question: str, filter: str | None = None, images: Sequence[bytes] | None = None, + sources: list[str] | None = None, ) -> "tuple[str, list[Citation]]": from haiku.rag.client.agents import ask - return await ask(self, question, filter, images) + return await ask(self, question, filter, images, sources) async def analyze( self, diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index c85fa1ab..a18df5e5 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/haiku_rag_slim/haiku/rag/client/agents.py @@ -40,6 +40,7 @@ async def ask( question: str, filter: str | None = None, images: Sequence[bytes] | None = None, + sources: list[str] | None = None, ) -> "tuple[str, list[Citation]]": """Ask a question against the knowledge base via the RAG capability. @@ -61,13 +62,17 @@ async def ask( from haiku.rag.utils import get_model capability = create_capability( - db_path=client.store.db_path, + db_path=None if client._federated else client.store.db_path, config=client._config, rag=client, defer_loading=False, ) deps = _AgentDeps( - state={"rag": RAGState(document_filter=filter).model_dump(mode="json")} + state={ + "rag": RAGState(document_filter=filter, sources=sources).model_dump( + mode="json" + ) + } ) user_prompt = _build_user_prompt(question, images, client._config.qa.model) model = get_model(client._config.qa.model, client._config) diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index d8336de1..e50ee3a2 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -378,6 +378,28 @@ async def expand_context( chunks were created without docling metadata (e.g., custom chunks passed to import_document). """ + # A federating client has no repositories of its own, so each result expands + # through the database it came from. + if client._federated: + by_source: dict[str, list[SearchResult]] = {} + unsourced: list[SearchResult] = [] + for result in search_results: + if result.source: + by_source.setdefault(result.source, []).append(result) + else: + unsourced.append(result) + owners = await client.clients_for(list(by_source)) + expanded_groups = await asyncio.gather( + *( + expand_context(owner, by_source[owner._source]) + for owner in owners + if owner._source + ) + ) + merged = unsourced + [r for group in expanded_groups for r in group] + merged.sort(key=lambda r: r.score, reverse=True) + return merged + from haiku.rag.context import expand_with_items, window_for max_chars = client._config.search.max_context_chars diff --git a/haiku_rag_slim/haiku/rag/store/models/citation.py b/haiku_rag_slim/haiku/rag/store/models/citation.py index 161d8575..36b39937 100644 --- a/haiku_rag_slim/haiku/rag/store/models/citation.py +++ b/haiku_rag_slim/haiku/rag/store/models/citation.py @@ -22,6 +22,10 @@ class Citation(BaseModel): ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged into the cited result (always includes ``chunk_id``). + ``source`` names the configured database the cited chunk came from, and is + None when only one is configured. It is the name from ``lancedb.databases``, + never a path or URI. + ``doc_item_refs`` are the ``self_ref`` values of every item in the cited content — the exact items the model saw. Visual grounding resolves bounding boxes from them so the rendered pages match the citation precisely. @@ -37,6 +41,7 @@ class Citation(BaseModel): index: int | None = None document_id: str + source: str | None = None chunk_id: str chunk_ids: list[str] = Field(default_factory=list) chunk_meta: dict = Field(default_factory=dict) @@ -69,6 +74,7 @@ def resolve_citations( citations.append( Citation( document_id=r.document_id or "", + source=r.source, chunk_id=chunk_id, chunk_ids=r.chunk_ids or [chunk_id], chunk_meta=r.chunk_meta, diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 25727289..39396396 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -316,6 +316,10 @@ async def test_capability_isolated_per_run_and_round_trips_state(temp_db_path): async def test_run_error_closes_resources_and_propagates(temp_db_path): capability = create_rag(db_path=temp_db_path, config=AppConfig()) client = AsyncMock() + # Stands in for a single-database client: a bare AsyncMock's auto + # attributes are truthy Mocks, and `_source` reaches a validated field. + client._federated = {} + client._source = None capability.rag = client error = RuntimeError("model failed") @@ -385,6 +389,10 @@ async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db capability = create_rag(db_path=temp_db_path, config=AppConfig()) capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0)) client = AsyncMock() + # Stands in for a single-database client: a bare AsyncMock's auto + # attributes are truthy Mocks, and `_source` reaches a validated field. + client._federated = {} + client._source = None client.get_chunk_by_id.side_effect = [ Chunk(id="chunk-1", document_id="doc-1", content="first"), Chunk(id="chunk-2", document_id="doc-1", content="second"), @@ -411,6 +419,10 @@ async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path): capability = create_rag(db_path=temp_db_path, config=AppConfig()) capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0)) client = AsyncMock() + # Stands in for a single-database client: a bare AsyncMock's auto + # attributes are truthy Mocks, and `_source` reaches a validated field. + client._federated = {} + client._source = None client.get_chunk_by_id.side_effect = [ Chunk(id="chunk-1", document_id="doc-1", content="first"), None, @@ -455,6 +467,10 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path): }, ) client = AsyncMock() + # Stands in for a single-database client: a bare AsyncMock's auto + # attributes are truthy Mocks, and `_source` reaches a validated field. + client._federated = {} + client._source = None client.get_chunk_by_id.return_value = None capability.rag = client @@ -1214,6 +1230,10 @@ async def test_citing_without_searching_grounds_the_question(temp_db_path): return ModelResponse(parts=next(calls)) client = AsyncMock() + # Stands in for a single-database client: a bare AsyncMock's auto + # attributes are truthy Mocks, and `_source` reaches a validated field. + client._federated = {} + client._source = None client.get_chunk_by_id.return_value = Chunk( id="chunk-1", document_id="doc-1", content="evidence" ) @@ -1315,6 +1335,10 @@ async def test_a_capability_fetches_its_own_evidences_pictures(temp_db_path): """Compaction rehydrates through the owner, which already holds the connection.""" capability = create_rag(db_path=temp_db_path, config=AppConfig()) client = AsyncMock() + # Stands in for a single-database client: a bare AsyncMock's auto + # attributes are truthy Mocks, and `_source` reaches a validated field. + client._federated = {} + client._source = None client.document_item_repository.get_picture_bytes.return_value = b"picture-bytes" capability.rag = client @@ -1354,6 +1378,10 @@ async def test_citing_nothing_after_citing_something_keeps_it_grounded(temp_db_p capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0)) capability.epoch = 5 client = AsyncMock() + # Stands in for a single-database client: a bare AsyncMock's auto + # attributes are truthy Mocks, and `_source` reaches a validated field. + client._federated = {} + client._source = None client.get_chunk_by_id.return_value = Chunk( id="chunk-1", document_id="doc-1", content="evidence" ) diff --git a/tests/test_multi_db_ask.py b/tests/test_multi_db_ask.py new file mode 100644 index 00000000..ee9a932e --- /dev/null +++ b/tests/test_multi_db_ask.py @@ -0,0 +1,201 @@ +import pytest +from pydantic_ai import ModelRetry + +from haiku.rag.capabilities.rag import RAGState, create_capability +from haiku.rag.client import HaikuRAG +from haiku.rag.store.models import SearchResult +from haiku.rag.store.models.citation import resolve_citations +from tests.test_multi_db import _config, _seed + + +class TestExpansionRouting: + @pytest.mark.asyncio + async def test_expansion_routes_each_result_to_its_database(self, tmp_path): + """A federating client has no repositories of its own, so expansion has + to go through the database each result came from.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + results = await rag.search("cats", search_type="fts", limit=10) + expanded = await rag.expand_context(results) + + assert {r.source for r in expanded} == {"alpha", "beta"} + for r in expanded: + assert r.source is not None + assert r.source in r.content + + +class TestCitationSource: + def test_a_citation_carries_the_result_source(self): + result = SearchResult( + content="body", + score=0.9, + source="alpha", + chunk_id="c1", + document_id="d1", + document_uri="test://alpha/one", + ) + + [citation] = resolve_citations(["c1"], [result]) + + assert citation.source == "alpha" + + def test_a_single_database_citation_has_no_source(self): + result = SearchResult( + content="body", + score=0.9, + chunk_id="c1", + document_id="d1", + document_uri="test://one", + ) + + [citation] = resolve_citations(["c1"], [result]) + + assert citation.source is None + + +class TestAskAcrossDatabases: + @pytest.mark.asyncio + async def test_the_capability_searches_the_selected_databases(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + capability = create_capability(config=config, rag=rag, defer_loading=False) + capability.state = RAGState(sources=["alpha"]) + + formatted = await capability._search("cats", limit=10) + + assert isinstance(formatted, str) + assert "alpha" in formatted + assert "beta document" not in formatted + + @pytest.mark.asyncio + async def test_searching_all_databases_reaches_both(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + capability = create_capability(config=config, rag=rag, defer_loading=False) + capability.state = RAGState() + + formatted = await capability._search("cats", limit=10) + + assert isinstance(formatted, str) + assert "alpha document" in formatted + assert "beta document" in formatted + + +class TestCiteFallback: + @pytest.mark.asyncio + async def test_an_id_from_a_selected_database_resolves_with_its_source( + self, tmp_path + ): + """The fallback exists for a real id this run's searches did not return. + Across databases it looks through the selected ones and records which + held it.""" + from tests.capabilities.test_capabilities import Deps, make_context + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed( + config, "alpha", ["alpha document about cats", "alpha on aardvarks"] + ) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + alpha = (await rag.clients_for(["alpha"]))[0] + chunks = await alpha.chunk_repository.list_all() + [aardvark] = [c for c in chunks if "aardvark" in c.content] + assert aardvark.id is not None + + capability = create_capability(config=config, rag=rag, defer_loading=False) + deps = Deps( + state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")} + ) + run = await capability.for_run(make_context(deps)) + # The search returns the cats chunk, never the aardvark one. + await run._search("cats", limit=10) + + await run._cite([aardvark.id]) + + assert run.state is not None + [citation] = list(run.state.citation_index.values()) + assert citation.chunk_id == aardvark.id + assert citation.source == "alpha" + + @pytest.mark.asyncio + async def test_an_id_outside_the_selected_databases_does_not_resolve( + self, tmp_path + ): + """A question scoped to one database must not produce a citation from + another: the fallback looks only where the question looked.""" + from tests.capabilities.test_capabilities import Deps, make_context + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about dogs"]) + + async with HaikuRAG(config=config) as rag: + beta = (await rag.clients_for(["beta"]))[0] + [outside] = await beta.chunk_repository.list_all(limit=1) + assert outside.id is not None + + capability = create_capability(config=config, rag=rag, defer_loading=False) + deps = Deps( + state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")} + ) + run = await capability.for_run(make_context(deps)) + await run._search("cats", limit=10) + + with pytest.raises(ModelRetry): + await run._cite([outside.id]) + + +class TestFederatedEdges: + @pytest.mark.asyncio + async def test_expansion_passes_through_results_without_a_source(self, tmp_path): + """A caller can hand `expand_context` results it built itself. Those name + no database, so there is nowhere to expand them from.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + handmade = SearchResult(content="handmade", score=0.4, doc_item_refs=[]) + + async with HaikuRAG(config=config) as rag: + found = await rag.search("cats", search_type="fts", limit=10) + expanded = await rag.expand_context([*found, handmade]) + + assert "handmade" in [r.content for r in expanded] + scores = [r.score for r in expanded] + assert scores == sorted(scores, reverse=True), "merged in score order" + + @pytest.mark.asyncio + async def test_a_chunk_without_a_document_is_not_cited(self, tmp_path): + """`Chunk.document_id` is optional, and a citation without a document has + nothing to point at.""" + from unittest.mock import AsyncMock, patch + + from haiku.rag.capabilities.rag import RAGCapability + from haiku.rag.store.models import Chunk + from tests.capabilities.test_capabilities import Deps, make_context + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + orphan = AsyncMock() + orphan._federated = {} + orphan._source = None + orphan.get_chunk_by_id.return_value = Chunk( + id="orphan", document_id=None, content="no document" + ) + + capability = create_capability(config=config, defer_loading=False) + run = await capability.for_run(make_context(Deps())) + with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=orphan)): + with pytest.raises(ModelRetry): + await run._cite(["orphan"])