diff --git a/specs/009-collection-routing/research.md b/specs/009-collection-routing/research.md index e407f62..bcb8bae 100644 --- a/specs/009-collection-routing/research.md +++ b/specs/009-collection-routing/research.md @@ -193,3 +193,103 @@ stays at the control's score with the classifier choosing -- 13/13 here, 15/15 i the container where MCP is configured. The TP53 and PTEN questions are the ones to watch, because they fail loudly and specifically. A classifier that never routes to `complexes` would still score full marks, which is why T005 matters. + +## What routing actually buys, measured 2026-09-19 + +The classifier is wired up and choosing. Three measurements, and the second one +is not the result the feature was pitched on. + +### It holds the acceptance bar + +`answer-sweep` with the classifier selecting collections: **13/13, the control +score**, two `needs_live` skips. Five of six retrievals were narrowed, and the +routing is the mapping the earlier measurement predicted: + +| question | routed to | +|---|---| +| UniProt accession for TP53 | `ewas` | +| Selective autophagy summary | `summations` | +| ABCA1 variants | `disease_variants` | +| PTEN variant diseases | `disease_variants` | +| What CDK5 phosphorylates | `ewas`, `summations` | +| How TP53 regulates PTEN transcription | *all five* -- left empty, correctly | + +Only six of the thirteen questions retrieve at all: the refusals never do, and +the userguide questions use a different bundle. So the routing evidence is six +questions, which is a small base and should be said rather than glossed. + +### It does not reduce the prompt + +This is the correction. Narrowing was expected to cut context; it does not. + +| question | docs, all | docs, narrowed | chars, all | chars, narrowed | +|---|---|---|---|---| +| TP53 accession | 10 | 10 | 2,436 | 2,436 | +| ABCA1 variants | 10 | 10 | 6,163 | 5,846 | +| Selective autophagy | 10 | 10 | 22,114 | 21,797 | + +**Ten documents either way, context within 2%.** The chain caps the fused list +at ten regardless of how many collections fed it, so narrowing changes *which* +documents arrive, not how many. Any claim that this feature reduces prompt cost +is wrong, and T003's framing -- "context tokens for a question needing one +collection fall" -- was the wrong thing to measure. + +### What it does reduce is retrieval work + +Median of three runs per question, timed inside `aretrieve_documents`: + +| | all five | narrowed | +|---|---|---| +| median retrieval | **1.82s** | **1.43s** | + +About 21%, or roughly 0.4s against a first token near ten seconds. Real, and +modest. Five collections mean ten sub-retrievals (BM25 and vector each); one +collection means two. + +### So the honest case for this feature + +It did not fix a failing question -- the control already passed 13/13. It buys +a fifth off retrieval time and it puts the right documents in a fixed-size +context, which should matter most where the fixed ten are currently crowded out +by the wrong collection. That last part is plausible and **not** measured here. + +Against that, a wrong narrow removes an answer rather than degrading it, which +is why the prompt leans hard on leaving the selection empty and why every +failure path in `resolve_collections` widens. The sweep shows correct routing +on six questions; it does not show that the classifier is right in general. + +`complexes` still has no guard question (T005), so a classifier that never +routes to it would score full marks. + +### Adversarial check: nine questions the sweep never asks + +Six retrieving questions is a thin base to ship a router on, so the classifier +was probed with questions chosen to be awkward -- broad ones that must *not* +narrow, and cross-collection ones where a single choice loses half the answer. + +| question | chose | +|---|---| +| What is apoptosis? | *all* | +| Explain the role of TP53 in the cell cycle | *all* | +| Tell me everything Reactome knows about ferroptosis | *all* | +| What complexes contain TP53? | `complexes` | +| Which proteins are in the MCM complex and what do they do? | `complexes`, `summations` | +| What is the UniProt ID for BRCA1? | `ewas` | +| Which variants of BRCA1 cause disease, and the mechanisms? | `disease_variants`, `summations` | +| Inputs and outputs of the CDK1/MCM2 phosphorylation | `reactions` | +| How does Reactome describe the Wnt signalling pathway? | `summations` | + +Nine of nine as intended: every broad question left open, every specific one +narrowed to the collection that holds the answer, and both cross-collection +questions took `summations` alongside their primary -- which is what the prompt +asks for and the reason the mechanism half of those questions survives. + +This is nine single classifications, not nine end-to-end answers, so it shows +the routing decision is sound and not that the answers are. + +**A hallucinated collection name is safe by construction**, and deliberately so: +`collections` is `list[str]` rather than an enum of the five names. A name that +is not in the bundle is widened to all by `resolve_collections` with a WARNING, +which is the failure direction this feature requires. Constraining the schema +would instead make structured output reject the response, turning a harmless +mistake into a failed answer. diff --git a/specs/009-collection-routing/tasks.md b/specs/009-collection-routing/tasks.md index e93c09d..5fdba68 100644 --- a/specs/009-collection-routing/tasks.md +++ b/specs/009-collection-routing/tasks.md @@ -34,16 +34,22 @@ collection, with no extra LLM call. **Independent test**: `./bin/answer-sweep` stays green while the measured context tokens for a question needing one collection fall relative to `before.json`. -- [ ] T011 [US1] Add `collections: list[str] = []` to `QueryIntent` in src/agent/tasks/intent_classifier.py, defaulting to empty so an omitted field means "all" -- [ ] T012 [US1] Extend the classifier prompt in src/agent/tasks/intent_classifier.py to name selectable collections, sourced from `reactome_descriptions_info` rather than a literal list +**The second half of that test was the wrong measurement, and it failed.** +Context does not fall: the chain caps the fused list at ten documents however +many collections fed it, so narrowing changes which documents arrive, not how +many (research.md, 2026-09-19). What falls is retrieval time, 1.82s to 1.43s. +The sweep half stands: 13/13 with the classifier choosing. + +- [x] T011 [US1] Add `collections: list[str] = []` to `QueryIntent` in src/agent/tasks/intent_classifier.py, defaulting to empty so an omitted field means "all" +- [x] T012 [US1] Extend the classifier prompt in src/agent/tasks/intent_classifier.py to name selectable collections, sourced from `reactome_descriptions_info` rather than a literal list - [x] T013 [P] [US1] Add `resolve_collections(selected, available)` to src/retrievers/csv_chroma.py implementing data-model.md: empty means all, unknown names log WARNING and return all - [x] T014 [P] [US1] Unit-test `resolve_collections` in tests/retrievers/test_collection_selection.py for empty, all-valid, some-unknown and all-unknown, asserting every failure widens rather than narrows - [x] T015 [US1] Filter `self.collection_retrievers` by the selection in `retrieve_documents` in src/retrievers/csv_chroma.py, reading it from `RunnableConfig["configurable"]["collections"]` -- [ ] T016 [US1] Apply the identical filter in `aretrieve_documents` in src/retrievers/csv_chroma.py — this is the served path -- [ ] T017 [US1] Extend tests/retrievers/test_sync_async_equivalence.py to assert both paths honour the same selection, and confirm it fails when only one is filtered -- [ ] T018 [US1] Carry `collections` on `ReactToMeState` in src/agent/profiles/react_to_me.py, set in `preprocess` beside `active_sources` -- [ ] T019 [US1] Pass the selection into the RAG call via `config["configurable"]` in `generate_answer` in src/agent/profiles/react_to_me.py -- [ ] T020 [US1] Verify through the agent, not the retriever: a test that a variant question routes to `disease_variants` and a userguide question does not, per Principle I +- [x] T016 [US1] Apply the identical filter in `aretrieve_documents` in src/retrievers/csv_chroma.py — this is the served path +- [x] T017 [US1] Extend tests/retrievers/test_sync_async_equivalence.py to assert both paths honour the same selection, and confirm it fails when only one is filtered +- [x] T018 [US1] Carry `collections` on `ReactToMeState` in src/agent/profiles/react_to_me.py, set in `preprocess` beside `active_sources` +- [x] T019 [US1] Pass the selection into retrieval in `generate_answer` in src/agent/profiles/react_to_me.py — **via the `selected_collections` ContextVar, not `config["configurable"]`**: `create_retrieval_chain` gives the retriever no path for extra arguments. LangGraph copies the context into the tasks it spawns, so a value set around the call reaches the retriever inside them. Reset in `finally`, because the graph reuses one task across turns +- [x] T020 [US1] Verify through the agent, not the retriever: a test that a variant question routes to `disease_variants` and a userguide question does not, per Principle I ## Phase 4: Measurement and acceptance diff --git a/src/agent/profiles/react_to_me.py b/src/agent/profiles/react_to_me.py index 39574af..b09a9bb 100644 --- a/src/agent/profiles/react_to_me.py +++ b/src/agent/profiles/react_to_me.py @@ -1,6 +1,6 @@ import asyncio import logging -from typing import Any, cast +from typing import Any, NotRequired, cast from langchain_core.embeddings import Embeddings from langchain_core.language_models.chat_models import BaseChatModel @@ -19,6 +19,7 @@ from agent.tasks.unsafe_question import create_unsafe_answer_generator from reactome_mcp.answer import ToolCallingModel, answer_from_live_services from reactome_mcp.session import get_mcp_tools, is_configured +from retrievers.csv_chroma import selected_collections from retrievers.reactome.rag import create_reactome_rag from retrievers.userguide.rag import create_userguide_rag from util.embedding_environment import EmbeddingEnvironment @@ -28,6 +29,16 @@ class ReactToMeState(BaseState): active_sources: list[SourceName] + # NotRequired, and it matters: this subclass is total=True, so declaring + # it plainly would make every existing construction of the state invalid + # and -- worse -- would claim a guarantee the runtime does not have. A + # thread resumed from a checkpoint written before this field existed has + # no key, which is why `generate_answer` reads it with `.get`. + # + # Empty, or absent, means every collection: that is what + # `resolve_collections` does with it, and it is how the graph behaved + # before routing. Only meaningful when the active source is `reactome`. + collections: NotRequired[list[str]] class ReactToMeGraphBuilder(BaseGraphBuilder): @@ -162,6 +173,7 @@ async def preprocess( reason_unsafe=safety_check.reason_unsafe, detected_language=detected_language, active_sources=active_sources, + collections=intent.collections, ) async def _answer_from_live_services( @@ -183,6 +195,11 @@ async def _answer_from_live_services( ) fallback = dict(state) fallback["active_sources"] = ["reactome"] + # This fallback is retrieval over the whole bundle standing in for + # a live lookup; a question routed to `live` chose no collections, + # and narrowing to a selection nobody made would be worse than the + # answer it replaces. + fallback["collections"] = [] # The real config, not a fresh one: it carries the callbacks the UI # streams through, and a fallback nobody can see is not a fallback. return await self.generate_answer(ReactToMeState(**fallback), config) @@ -228,21 +245,42 @@ async def generate_answer( if source == "live": return await self._answer_from_live_services(state, config) rag = self.rags[source] - result: dict[str, Any] = await rag.ainvoke( - { - "input": state["rephrased_input"], - # A separate variable, never concatenated into `input`: - # create_retrieval_chain passes `input` alone to the retriever, so - # anything folded into it reaches BM25 and the query expander. - "detected_language": state["detected_language"], - "chat_history": ( - state["chat_history"] - if state["chat_history"] - else [HumanMessage(state["user_input"])] - ), - }, - config, + # Set around the whole retrieval rather than passed down: + # `create_retrieval_chain` gives the retriever no path for extra + # arguments, and LangGraph copies the context into the tasks it + # spawns, so a value set here reaches the retriever inside them. + # Reset in `finally`, because the graph reuses this task across turns. + # + # Only `reactome` has collections. `userguide` is a separate bundle + # with one, and leaving a stale selection set would narrow it to names + # it does not have -- which `resolve_collections` widens back, but + # silently and with a WARNING for every question. + # `.get`, not `[...]`: BaseState is total=False, so a state resumed + # from a checkpoint written before this field existed has no key at + # all. Missing means empty means every collection -- the behaviour + # from before routing, which is the right way to fail. + token = selected_collections.set( + (state.get("collections") or []) if source == "reactome" else None ) + try: + result: dict[str, Any] = await rag.ainvoke( + { + "input": state["rephrased_input"], + # A separate variable, never concatenated into `input`: + # create_retrieval_chain passes `input` alone to the + # retriever, so anything folded into it reaches BM25 and + # the query expander. + "detected_language": state["detected_language"], + "chat_history": ( + state["chat_history"] + if state["chat_history"] + else [HumanMessage(state["user_input"])] + ), + }, + config, + ) + finally: + selected_collections.reset(token) return ReactToMeState( chat_history=[ HumanMessage(state["user_input"]), diff --git a/src/agent/tasks/intent_classifier.py b/src/agent/tasks/intent_classifier.py index 8d50e58..fa34bdc 100644 --- a/src/agent/tasks/intent_classifier.py +++ b/src/agent/tasks/intent_classifier.py @@ -5,6 +5,8 @@ from langchain_core.runnables import Runnable from pydantic import BaseModel, Field +from retrievers.reactome.metadata_info import reactome_descriptions_info + SourceName = Literal["reactome", "userguide", "live"] _REACTOME_SOURCE = """- **reactome**: Questions about biology, molecular mechanisms, pathways, reactions, proteins, genes, @@ -49,6 +51,55 @@ including listing it. A question naming a specific gene or disease is almost always **reactome**.""" + +def _collections_block() -> str: + """The selectable collections, described by the bundle's own metadata. + + Sourced from `reactome_descriptions_info` rather than a literal list so a + collection added to the bundle cannot be missing from the prompt that is + supposed to offer it -- the failure mode would be silent, since an + unmentioned collection is simply never chosen. + """ + lines = [ + f" - **{name}**: {description.strip()}" + for name, description in reactome_descriptions_info.items() + ] + return chr(10).join(lines) + + +# Narrowing is measured, not assumed. On 2026-09-19, forcing every tracked +# question to `reactions` + `summations` took the sweep from 13/13 to 10/13, +# and the questions that failed were the ones whose answers live only in the +# collections that were dropped: the UniProt accession needs `ewas`, PTEN's +# variants need `disease_variants`. A wrong narrow does not degrade an answer, +# it removes it -- so the instruction below leans hard on leaving the list +# empty, because widening costs latency and narrowing wrongly costs the answer. +_COLLECTIONS_RULE = """ +Collections (only when source is **reactome**): + +The Reactome content is split into collections holding different kinds of record: + +{collections} + +Set `collections` ONLY when the question plainly needs one or two specific kinds +of record. Leave it EMPTY otherwise, and an empty list searches all of them. +An empty list is the right answer for most questions. + +- Naming or listing variants of a gene, or which disease a variant causes: `disease_variants`. +- A UniProt accession, a gene synonym, or which protein a gene maps to: `ewas`. +- What a pathway or reaction's curated description says: `summations`. +- What a complex is made of: `complexes`. +- The inputs, outputs or catalyst of a reaction: `reactions`. + +Two things to respect: +- Include `summations` alongside any other choice unless the question is purely + about identifiers. The curated prose supports most biological answers. +- If the question is broad, mechanistic, or you are at all unsure, leave + `collections` empty. Searching everything is slower; searching the wrong + subset means the answer is not there at all. +""" + + _SOURCE_BLOCKS: dict[SourceName, str] = { "reactome": _REACTOME_SOURCE, "userguide": _USERGUIDE_SOURCE, @@ -69,6 +120,11 @@ def build_classifier_message(sources: frozenset[SourceName]) -> str: _SOURCE_BLOCKS[s] for s in ("reactome", "userguide", "live") if s in sources ] rules = _RULES + (f"\n{_LIVE_RULE}" if "live" in sources else "") + collections = ( + _COLLECTIONS_RULE.format(collections=_collections_block()) + if "reactome" in sources + else "" + ) return f""" You route user questions for the React-to-Me assistant to the correct knowledge source. @@ -77,7 +133,7 @@ def build_classifier_message(sources: frozenset[SourceName]) -> str: {chr(10).join(blocks)} {rules} -""" +{collections}""" intent_classifier_message = build_classifier_message( @@ -96,6 +152,15 @@ class QueryIntent(BaseModel): source: SourceName = Field( description="The knowledge source that should answer this question." ) + collections: list[str] = Field( + default_factory=list, + description=( + "Which Reactome collections to search, when the question plainly " + "needs only some of them. Empty means search all of them, which is " + "the right answer for most questions. Ignored unless source is " + "'reactome'." + ), + ) _FALLBACK_ORDER: tuple[SourceName, ...] = ("reactome", "userguide", "live") diff --git a/tests/agent/test_collection_routing.py b/tests/agent/test_collection_routing.py new file mode 100644 index 0000000..f0a77b2 --- /dev/null +++ b/tests/agent/test_collection_routing.py @@ -0,0 +1,115 @@ +"""Selecting collections within the `reactome` bundle. + +Measured 2026-09-19 (specs/009-collection-routing/research.md): narrowing every +tracked question to `reactions` + `summations` took the sweep from 13/13 to +10/13, and a wrong narrow removed the answer rather than degrading it. So the +asymmetry these tests defend is: an empty selection must always mean *all*, and +nothing may narrow a source that has no collections. +""" + +import asyncio +from typing import Any, cast + +from langchain_core.runnables import RunnableConfig + +from agent.profiles.react_to_me import ReactToMeGraphBuilder, ReactToMeState +from agent.tasks.intent_classifier import QueryIntent, build_classifier_message +from retrievers.csv_chroma import selected_collections +from retrievers.reactome.metadata_info import reactome_descriptions_info + + +def test_an_omitted_selection_means_every_collection() -> None: + # The field is optional in the schema, so a model that ignores it entirely + # must get the behaviour that existed before this feature. + assert QueryIntent(source="reactome").collections == [] + + +def test_every_collection_in_the_bundle_is_described_to_the_classifier() -> None: + # Sourced from the metadata, not a literal list: a collection added to the + # bundle but missing from the prompt is never chosen, and nothing fails. + message = build_classifier_message(frozenset({"reactome", "userguide"})) + for name, description in reactome_descriptions_info.items(): + assert f"**{name}**" in message, f"{name} is not offered to the classifier" + assert description.strip()[:40] in message + + +def test_the_userguide_prompt_alone_offers_no_collections() -> None: + message = build_classifier_message(frozenset({"userguide"})) + assert "disease_variants" not in message + + +class _RecordingRag: + """Captures what the retriever would have been told to search.""" + + def __init__(self) -> None: + self.seen: list[list[str] | None] = [] + + async def ainvoke(self, _inputs: dict[str, Any], _config: Any) -> dict[str, Any]: + self.seen.append(selected_collections.get()) + return {"answer": "an answer", "context": []} + + +def _builder(rag: _RecordingRag, source: str) -> ReactToMeGraphBuilder: + builder = ReactToMeGraphBuilder.__new__(ReactToMeGraphBuilder) + # A stand-in for the RAG runnable: only `ainvoke` is reached from here. + builder.rags = cast(Any, {source: rag}) + return builder + + +def _state(source: str, collections: list[str]) -> ReactToMeState: + return ReactToMeState( + user_input="q", + rephrased_input="q", + detected_language="English", + chat_history=[], + active_sources=[source], # type: ignore[list-item] + collections=collections, + ) + + +def test_the_selection_reaches_retrieval() -> None: + rag = _RecordingRag() + asyncio.run( + _builder(rag, "reactome").generate_answer( + _state("reactome", ["disease_variants", "summations"]), RunnableConfig() + ) + ) + assert rag.seen == [["disease_variants", "summations"]] + + +def test_the_userguide_is_never_narrowed_by_a_reactome_selection() -> None: + # The dangerous direction. `userguide` is a different bundle with one + # collection; a leaked selection names collections it does not have, and + # `resolve_collections` would widen it back silently -- correct, but only + # by accident, and with a WARNING for every question. + rag = _RecordingRag() + asyncio.run( + _builder(rag, "userguide").generate_answer( + _state("userguide", ["disease_variants"]), RunnableConfig() + ) + ) + assert rag.seen == [None], "a reactome selection leaked into the userguide" + + +def test_the_selection_does_not_outlive_the_question() -> None: + # The graph reuses one asyncio task across turns, so a selection left set + # would narrow the *next* question -- which nothing downstream could + # detect, because narrowing produces a confident answer from less. + rag = _RecordingRag() + asyncio.run( + _builder(rag, "reactome").generate_answer( + _state("reactome", ["ewas"]), RunnableConfig() + ) + ) + assert selected_collections.get() is None + + +def test_a_state_without_the_field_searches_everything() -> None: + # BaseState is total=False and this field is new, so a thread checkpointed + # before it existed resumes without the key. Subscripting would raise; + # missing must mean "all", which is how the graph behaved before routing. + rag = _RecordingRag() + state = _state("reactome", []) + del state["collections"] + asyncio.run(_builder(rag, "reactome").generate_answer(state, RunnableConfig())) + assert rag.seen == [[]], "a pre-routing checkpoint must not crash or narrow" diff --git a/tests/retrievers/test_sync_async_equivalence.py b/tests/retrievers/test_sync_async_equivalence.py index b002070..c0b21fa 100644 --- a/tests/retrievers/test_sync_async_equivalence.py +++ b/tests/retrievers/test_sync_async_equivalence.py @@ -45,7 +45,11 @@ FakeListChatModel, ) -from retrievers.csv_chroma import HybridRetriever, chroma_settings # noqa: E402 +from retrievers.csv_chroma import ( # noqa: E402 + HybridRetriever, + chroma_settings, + selected_collections, +) # Large enough that the per-collection cap actually binds. At twelve it did @@ -165,3 +169,40 @@ def test_async_returns_exactly_what_sync_returns( # reordering is a ranking change and the top documents are the ones that # reach the model. assert [d.page_content for d in asynchronous] == [d.page_content for d in sync] + + +def test_both_paths_honour_the_same_collection_selection(tmp_path: Path) -> None: + """T017: a filter applied to one path only is the bug this pins. + + The application serves through the async path and `bin/retrieval_baseline` + drives the sync one, so a selection honoured by only one means the measured + path and the served path search different collections -- and neither the + measurement nor the answer would look wrong. + """ + _require_bm25_tokenizer() + embedding = DeterministicFakeEmbedding(size=16) + retriever = HybridRetriever.from_subdirectory( + llm=FakeListChatModel(responses=[""]), + embedding=embedding, + embeddings_directory=_bundle(tmp_path, embedding), + ) + queries = ["apoptosis signalling"] + + token = selected_collections.set(["alpha"]) + try: + sync = retriever.retrieve_documents( + queries, CallbackManagerForRetrieverRun.get_noop_manager() + ) + asynchronous = asyncio.run( + retriever.aretrieve_documents( + queries, AsyncCallbackManagerForRetrieverRun.get_noop_manager() + ) + ) + finally: + selected_collections.reset(token) + + assert [d.page_content for d in asynchronous] == [d.page_content for d in sync] + # And the selection actually bit: `beta` documents must be absent from both. + for documents in (sync, asynchronous): + assert documents, "the selection removed everything, so this proves nothing" + assert all("beta" not in d.page_content for d in documents)