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.
This commit is contained in:
Yiorgis Gozadinos
parent
0d7810c78a
commit
8a7cfb949f
8 changed files with 106 additions and 56 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -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.
|
||||
|
|
@ -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.
|
||||
|
|
@ -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.
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:")
|
||||
|
|
|
|||
Loading…
Reference in a new issue