Match capability instructions to the run's collection scope

The collection block was chosen when the capability was built, so a run narrowed
through `state.sources` to one collection was still told how to attribute across
collections it could not reach, while its results correctly carried no
`Collection:` line.

`get_instructions` composes it per run instead, from `state.sources` where the
question narrowed the conversation and from the lent client or the scope
otherwise. Order is preamble, base instructions, collection block.

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}"
+            parts.insert(0, self.config.prompts.domain_preamble)
-        return self.instruction_text
+        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

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 = (
+_multiple_collections_path = (
-    Path(__file__).parent / "instructions" / "analysis_multiple_databases.md"
+    Path(__file__).parent / "instructions" / "analysis_multiple_collections.md"
 )
 
 
@@ -44,9 +44,9 @@ def instructions() -> str:
 
 
 @cache
-def multiple_databases_instructions() -> str:
+def multiple_collections_instructions() -> str:
-    """Appended only where the capability covers multiple databases."""
+    """Appended for a run that spans more than one collection."""
-    return _multiple_databases_path.read_text().rstrip()
+    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,

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.

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

@@ -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.

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.

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

@@ -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.

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 = (
+_multiple_collections_path = (
-    Path(__file__).parent / "instructions" / "rag_multiple_databases.md"
+    Path(__file__).parent / "instructions" / "rag_multiple_collections.md"
 )
 
 
@@ -44,9 +44,9 @@ def instructions() -> str:
 
 
 @cache
-def multiple_databases_instructions() -> str:
+def multiple_collections_instructions() -> str:
-    """Appended only where the capability covers multiple databases."""
+    """Appended for a run that spans more than one collection."""
-    return _multiple_databases_path.read_text().rstrip()
+    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,

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:
+class TestMultipleCollectionsInstructions:
-    """The note follows what a capability opens, not what the configuration
+    """The note follows what a run reads, not what the configuration names, so a
-    names, so a single database is instructed exactly as it was before databases
+    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 (
+            assert "Collection:" in capability.get_instructions()
-                "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
+        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:")