diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 2bb04dd0..b92f4b6d 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -166,6 +166,8 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): state_type: type[StateT] state_namespace: str instruction_text: str + collection_instructions: str + """Appended for a run that spans more than one collection.""" vision: bool tool_names: frozenset[str] request_limit: int | None = None @@ -235,10 +237,27 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): run_capability._sync_state() return run_capability + @property + def spans_collections(self) -> bool: + """Whether this run reads more than one collection. + + A question narrows the conversation through `sources`, so a capability + built over a set can still run against one collection, and telling it + how to attribute across collections it cannot reach is noise. + """ + if self.state is not None and self.state.sources is not None: + return len(set(self.state.sources)) > 1 + if self.borrowed_rag is not None: + return self.borrowed_rag.covers_multiple + return self.scope.covers_multiple + def get_instructions(self) -> str: + parts = [self.instruction_text] if self.config.prompts.domain_preamble: - return f"{self.config.prompts.domain_preamble}\n\n{self.instruction_text}" - return self.instruction_text + parts.insert(0, self.config.prompts.domain_preamble) + if self.spans_collections: + parts.append(self.collection_instructions) + return "\n\n".join(parts) async def before_model_request( self, ctx: RunContext[Any], request_context: ModelRequestContext diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index c9a59da6..fb856ccc 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -25,8 +25,8 @@ 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" -_multiple_databases_path = ( - Path(__file__).parent / "instructions" / "analysis_multiple_databases.md" +_multiple_collections_path = ( + Path(__file__).parent / "instructions" / "analysis_multiple_collections.md" ) @@ -44,9 +44,9 @@ def instructions() -> str: @cache -def multiple_databases_instructions() -> str: - """Appended only where the capability covers multiple databases.""" - return _multiple_databases_path.read_text().rstrip() +def multiple_collections_instructions() -> str: + """Appended for a run that spans more than one collection.""" + return _multiple_collections_path.read_text().rstrip() def _recovery_hint(stderr: str) -> str: @@ -212,18 +212,14 @@ def create_capability( config = get_config() analysis_model = config.analysis.model or config.qa.model scope = resolve_scope(db_path, config) - instruction_text = instructions() - # A lent client covers what it covers; otherwise the scope says. - covers_multiple = rag.covers_multiple if rag is not None else scope.covers_multiple - if covers_multiple: - instruction_text += multiple_databases_instructions() return AnalysisCapability( scope=scope, config=config, borrowed_rag=rag, state_type=AnalysisState, state_namespace=STATE_NAMESPACE, - instruction_text=instruction_text, + instruction_text=instructions(), + collection_instructions=multiple_collections_instructions(), 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_multiple_collections.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis_multiple_collections.md new file mode 100644 index 00000000..e07942e2 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis_multiple_collections.md @@ -0,0 +1,14 @@ + +## Multiple collections + +The corpus spans multiple collections, and each interface names them differently: + +- Results from searches spanning multiple collections carry a `Collection:` line. +- In code, `await search(...)` and `await list_documents()` return `source`, the + collection an item came from. +- The mounted files do not. `/documents/{id}/metadata.json` has no `source`, so + map ids to collections with `await list_documents()` before reading the + filesystem per collection. + +Group, count and compare by `source` when the question is about collections +rather than about documents. diff --git a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis_multiple_databases.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis_multiple_databases.md deleted file mode 100644 index c85d2217..00000000 --- a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis_multiple_databases.md +++ /dev/null @@ -1,14 +0,0 @@ - -## Multiple 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_multiple_collections.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/rag_multiple_collections.md new file mode 100644 index 00000000..a26dc98f --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/instructions/rag_multiple_collections.md @@ -0,0 +1,6 @@ + +## Multiple collections + +The corpus spans multiple collections. Results from searches spanning multiple +collections carry a `Collection:` line naming the collection they came from. Say +which collection an answer draws on when the question compares them. diff --git a/haiku_rag_slim/haiku/rag/capabilities/instructions/rag_multiple_databases.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/rag_multiple_databases.md deleted file mode 100644 index 97270437..00000000 --- a/haiku_rag_slim/haiku/rag/capabilities/instructions/rag_multiple_databases.md +++ /dev/null @@ -1,6 +0,0 @@ - -## Multiple 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 03b536a5..e4dfb138 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -29,8 +29,8 @@ STATE_NAMESPACE = "rag" _CAPABILITY_ID = "haiku-rag" _TOOL_NAMES = frozenset({"rag_search", "rag_cite"}) _instructions_path = Path(__file__).parent / "instructions" / "rag.md" -_multiple_databases_path = ( - Path(__file__).parent / "instructions" / "rag_multiple_databases.md" +_multiple_collections_path = ( + Path(__file__).parent / "instructions" / "rag_multiple_collections.md" ) @@ -44,9 +44,9 @@ def instructions() -> str: @cache -def multiple_databases_instructions() -> str: - """Appended only where the capability covers multiple databases.""" - return _multiple_databases_path.read_text().rstrip() +def multiple_collections_instructions() -> str: + """Appended for a run that spans more than one collection.""" + return _multiple_collections_path.read_text().rstrip() @dataclass @@ -115,18 +115,14 @@ def create_capability( config = get_config() scope = resolve_scope(db_path, config) - instruction_text = instructions() - # A lent client covers what it covers; otherwise the scope says. - covers_multiple = rag.covers_multiple if rag is not None else scope.covers_multiple - if covers_multiple: - instruction_text += multiple_databases_instructions() return RAGCapability( scope=scope, config=config, borrowed_rag=rag, state_type=RAGState, state_namespace=STATE_NAMESPACE, - instruction_text=instruction_text, + instruction_text=instructions(), + collection_instructions=multiple_collections_instructions(), 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 f62810a3..4b500494 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -1613,9 +1613,9 @@ async def test_an_answered_question_is_no_longer_in_progress(temp_db_path): 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 +class TestMultipleCollectionsInstructions: + """The note follows what a run reads, not what the configuration names, so a + run over one collection is instructed exactly as it was before collections could be named.""" @staticmethod @@ -1674,25 +1674,64 @@ class TestSeveralDatabasesInstructions: 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 - ) + assert "Collection:" in capability.get_instructions() 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 + text = create_rag(config=config).get_instructions() - assert "Database:" in text + assert "Collection:" in text def test_the_analysis_note_separates_the_interfaces(self): - """The three interfaces name a database differently, and the mounted + """The three interfaces name a collection 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 + text = create_analysis(config=config).get_instructions() - assert "Database:" in text # analysis_search results + assert "Collection:" 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 + assert "list_documents" in text # how to map ids to collections + + def test_a_run_narrowed_to_one_collection_drops_the_note(self): + """A question narrows the conversation, so a capability over a set can + still read one collection.""" + config = self._config(alpha="/a.lancedb", beta="/b.lancedb") + + rag = create_rag(config=config) + rag.state = RAGState(sources=["alpha"]) + analysis = create_analysis(config=config) + analysis.state = AnalysisState(sources=["alpha"]) + + assert "Collection:" not in rag.get_instructions() + assert "Collection:" not in analysis.get_instructions() + + def test_a_run_narrowed_to_two_collections_keeps_the_note(self): + config = self._config(alpha="/a.lancedb", beta="/b.lancedb", gamma="/c.lancedb") + + capability = create_rag(config=config) + capability.state = RAGState(sources=["alpha", "beta"]) + + assert "Collection:" in capability.get_instructions() + + def test_an_unnarrowed_run_follows_the_lent_client(self): + config = self._config(alpha="/a.lancedb", beta="/b.lancedb") + one = create_rag(config=config, rag=self._client({})) + one.state = RAGState() + covering = create_rag( + config=config, rag=self._client({"alpha": "/a", "beta": "/b"}) + ) + covering.state = RAGState() + + assert "Collection:" not in one.get_instructions() + assert "Collection:" in covering.get_instructions() + + def test_the_note_follows_the_preamble_and_the_base(self): + config = self._config(alpha="/a.lancedb", beta="/b.lancedb") + config.prompts.domain_preamble = "PREAMBLE" + + text = create_rag(config=config).get_instructions() + + assert text.index("PREAMBLE") < text.index("# RAG") < text.index("Collection:")