Tell the analysis and RAG capabilities about the databases

A capability covering several databases received `source` on every document
and search result and never used it: asked how many documents were in each
database, the model read the titles and answered that there was one corpus
of 67,581. The instruction files enumerate what a result carries, and both
enumerations had gone stale.

The note follows what the capability opens rather than what the
configuration names, through `covers_several_databases`: an explicit
`db_path` or a lent client covering one database is instructed as before,
as is every `uri` or path deployment and every eval dataset. The analysis
note separates the three interfaces, since they differ: an
`analysis_search` result carries a `Database:` line, in-code `search` and
`list_documents` return `source`, and the mounted files carry neither.

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -117,6 +117,20 @@ class EvidenceState(BaseModel):
         self.searches.clear()
 
 
+def covers_several_databases(
+    db_path: Path | None, config: AppConfig, rag: "HaikuRAG | None"
+) -> bool:
+    """Whether the capability will read from more than one database.
+
+    What the configuration names is not what a capability opens: an explicit
+    `db_path` opens that one database, and a lent client already knows what it
+    covers. Instructions follow coverage, not configuration.
+    """
+    if rag is not None:
+        return bool(rag._federated)
+    return db_path is None and len(config.lancedb.databases) > 1
+
+
 def _awaits_the_model(messages: list[ModelMessage]) -> bool:
     """Whether the history unmistakably leaves the model something to answer.
 
@@ -596,5 +610,6 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
 __all__ = [
     "CodeExecutionEntry",
     "RAGCapabilityBase",
+    "covers_several_databases",
     "resolve_db_path",
 ]

haiku_rag_slim/haiku/rag/capabilities/analysis.py

@@ -15,6 +15,7 @@ from haiku.rag.capabilities._base import (
     CodeExecutionEntry,
     EvidenceState,
     RAGCapabilityBase,
+    covers_several_databases,
     resolve_db_path,
 )
 from haiku.rag.capabilities._tools import merge_results
@@ -25,6 +26,9 @@ STATE_NAMESPACE = "analysis"
 _CAPABILITY_ID = "haiku-rag-analysis"
 _TOOL_NAMES = frozenset({"analysis_search", "analysis_execute_code", "analysis_cite"})
 _instructions_path = Path(__file__).parent / "instructions" / "analysis.md"
+_several_databases_path = (
+    Path(__file__).parent / "instructions" / "analysis_several_databases.md"
+)
 
 
 class AnalysisState(EvidenceState):
@@ -40,6 +44,13 @@ def instructions() -> str:
     return _instructions_path.read_text().strip()
 
 
+@cache
+def several_databases_instructions() -> str:
+    """Appended only where the capability covers several databases, so a single
+    database is instructed exactly as it was before they could be named."""
+    return _several_databases_path.read_text().rstrip()
+
+
 def _recovery_hint(stderr: str) -> str:
     """Name the workaround for sandbox limits models trip over repeatedly.
 
@@ -202,13 +213,17 @@ def create_capability(
 
         config = get_config()
     analysis_model = config.analysis.model or config.qa.model
+    resolved_db_path = resolve_db_path(db_path, config)
+    instruction_text = instructions()
+    if covers_several_databases(resolved_db_path, config, rag):
+        instruction_text += several_databases_instructions()
     return AnalysisCapability(
-        db_path=resolve_db_path(db_path, config),
+        db_path=resolved_db_path,
         config=config,
         borrowed_rag=rag,
         state_type=AnalysisState,
         state_namespace=STATE_NAMESPACE,
-        instruction_text=instructions(),
+        instruction_text=instruction_text,
         vision=analysis_model.vision if vision is None else vision,
         tool_names=_TOOL_NAMES,
         request_limit=request_limit,

haiku_rag_slim/haiku/rag/capabilities/instructions/analysis_several_databases.md

@@ -0,0 +1,14 @@
+
+## Several databases
+
+The corpus spans several databases, and each interface names them differently:
+
+- `analysis_search` results carry a `Database:` line.
+- In code, `await search(...)` and `await list_documents()` return `source`, the
+  configured database an item came from.
+- The mounted files do not. `/documents/{id}/metadata.json` has no `source`, so
+  map ids to databases with `await list_documents()` before reading the
+  filesystem per database.
+
+Group, count and compare by `source` when the question is about databases rather
+than about documents.

haiku_rag_slim/haiku/rag/capabilities/instructions/rag_several_databases.md

@@ -0,0 +1,6 @@
+
+## Several databases
+
+The corpus spans several databases. Each result carries a `Database:` line naming
+the configured database it came from. Say which database an answer draws on when
+the question compares them.

haiku_rag_slim/haiku/rag/capabilities/rag.py

@@ -13,6 +13,7 @@ if TYPE_CHECKING:
 from haiku.rag.capabilities._base import (
     EvidenceState,
     RAGCapabilityBase,
+    covers_several_databases,
     resolve_db_path,
 )
 from haiku.rag.config.models import AppConfig
@@ -29,6 +30,9 @@ STATE_NAMESPACE = "rag"
 _CAPABILITY_ID = "haiku-rag"
 _TOOL_NAMES = frozenset({"rag_search", "rag_cite"})
 _instructions_path = Path(__file__).parent / "instructions" / "rag.md"
+_several_databases_path = (
+    Path(__file__).parent / "instructions" / "rag_several_databases.md"
+)
 
 
 class RAGState(EvidenceState):
@@ -40,6 +44,13 @@ def instructions() -> str:
     return _instructions_path.read_text().strip()
 
 
+@cache
+def several_databases_instructions() -> str:
+    """Appended only where the capability covers several databases, so a single
+    database is instructed exactly as it was before they could be named."""
+    return _several_databases_path.read_text().rstrip()
+
+
 @dataclass
 class RAGCapability(RAGCapabilityBase[RAGState]):
     """Deferred, native Pydantic AI capability for grounded RAG queries."""
@@ -105,13 +116,17 @@ def create_capability(
         from haiku.rag.config import get_config
 
         config = get_config()
+    resolved_db_path = resolve_db_path(db_path, config)
+    instruction_text = instructions()
+    if covers_several_databases(resolved_db_path, config, rag):
+        instruction_text += several_databases_instructions()
     return RAGCapability(
-        db_path=resolve_db_path(db_path, config),
+        db_path=resolved_db_path,
         config=config,
         borrowed_rag=rag,
         state_type=RAGState,
         state_namespace=STATE_NAMESPACE,
-        instruction_text=instructions(),
+        instruction_text=instruction_text,
         vision=config.qa.model.vision if vision is None else vision,
         tool_names=_TOOL_NAMES,
         request_limit=request_limit,

tests/capabilities/test_capabilities.py

@@ -1521,3 +1521,87 @@ async def test_an_answered_question_is_no_longer_in_progress(temp_db_path):
     record = _record(deps, "rag")
     assert record.question == 0
     assert record.in_progress is False
+
+
+class TestSeveralDatabasesInstructions:
+    """The note follows what a capability opens, not what the configuration
+    names, so a single database is instructed exactly as it was before databases
+    could be named."""
+
+    @staticmethod
+    def _config(**databases):
+        from haiku.rag.config.models import LanceDBConfig
+
+        return AppConfig(lancedb=LanceDBConfig(databases=databases))
+
+    @staticmethod
+    def _client(federated):
+        client = AsyncMock()
+        client._federated = federated
+        return client
+
+    def test_one_database_is_instructed_as_before(self):
+        from haiku.rag.capabilities.analysis import instructions as analysis_text
+        from haiku.rag.capabilities.rag import instructions as rag_text
+
+        for factory, baseline in (
+            (create_rag, rag_text),
+            (create_analysis, analysis_text),
+        ):
+            for config in (AppConfig(), self._config(alpha="/a.lancedb")):
+                capability = factory(db_path=Path("/tmp/x.lancedb"), config=config)
+                assert capability.instruction_text == baseline()
+
+    def test_an_explicit_path_opens_one_database(self):
+        """A path names one database, whatever the configuration names."""
+        from haiku.rag.capabilities.analysis import instructions as analysis_text
+        from haiku.rag.capabilities.rag import instructions as rag_text
+
+        config = self._config(alpha="/a.lancedb", beta="/b.lancedb")
+        for factory, baseline in (
+            (create_rag, rag_text),
+            (create_analysis, analysis_text),
+        ):
+            capability = factory(db_path=Path("/tmp/one.lancedb"), config=config)
+            assert capability.instruction_text == baseline()
+
+    def test_a_lent_client_covering_one_database_is_instructed_as_before(self):
+        from haiku.rag.capabilities.analysis import instructions as analysis_text
+        from haiku.rag.capabilities.rag import instructions as rag_text
+
+        config = self._config(alpha="/a.lancedb", beta="/b.lancedb")
+        for factory, baseline in (
+            (create_rag, rag_text),
+            (create_analysis, analysis_text),
+        ):
+            capability = factory(config=config, rag=self._client({}))
+            assert capability.instruction_text == baseline()
+
+    def test_a_lent_client_covering_a_set_is_told_about_it(self):
+        config = self._config(alpha="/a.lancedb", beta="/b.lancedb")
+        covering = self._client({"alpha": "/a.lancedb", "beta": "/b.lancedb"})
+
+        for factory in (create_rag, create_analysis):
+            capability = factory(config=config, rag=covering)
+            assert "source" in capability.instruction_text or (
+                "Database:" in capability.instruction_text
+            )
+
+    def test_the_rag_note_names_the_line_a_result_carries(self):
+        config = self._config(alpha="/a.lancedb", beta="/b.lancedb")
+
+        text = create_rag(config=config).instruction_text
+
+        assert "Database:" in text
+
+    def test_the_analysis_note_separates_the_interfaces(self):
+        """The three interfaces name a database differently, and the mounted
+        files do not name it at all."""
+        config = self._config(alpha="/a.lancedb", beta="/b.lancedb")
+
+        text = create_analysis(config=config).instruction_text
+
+        assert "Database:" in text  # analysis_search results
+        assert "source" in text  # in-code search / list_documents
+        assert "metadata.json" in text  # the mounted files, which lack it
+        assert "list_documents" in text  # how to map ids to databases