diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index ad0d4329..6a809ad2 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/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", ] diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index 4f171f86..fdc3e9a1 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis_several_databases.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis_several_databases.md new file mode 100644 index 00000000..f3fd6d76 --- /dev/null +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/capabilities/instructions/rag_several_databases.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/rag_several_databases.md new file mode 100644 index 00000000..2a7b5987 --- /dev/null +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 46401483..db3e4515 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/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, diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index d8e17b86..de42dae3 100644 --- a/tests/capabilities/test_capabilities.py +++ b/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