Let the classifier choose collections, and measure what that is worth

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.

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
 

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"]),

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")