From 0fe7d88181ddd96edd245b1a69e92915232306a8 Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Sat, 19 Sep 2026 05:09:13 +0000 Subject: [PATCH 1/3] Let the classifier choose collections, and measure what that is worth Completes spec 009. `QueryIntent` gains a `collections` field, the classifier prompt describes each collection from `reactome_descriptions_info` rather than a literal list, and `generate_answer` sets the selection around retrieval. The selection rides the `selected_collections` ContextVar rather than `config["configurable"]` as tasks.md assumed: `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. It is reset in `finally`, because the graph reuses one asyncio task across turns and a selection left set would narrow the next question -- which nothing downstream could detect, since narrowing yields a confident answer from less. Three measurements, and the middle one is not the result this feature was pitched on. It holds the bar. The sweep is 13/13 with the classifier choosing, which is the control score, and five of six retrievals were narrowed. The routing is the mapping the earlier measurement predicted: the accession question to `ewas`, both variant questions to `disease_variants`, the summary question to `summations`, and the broad mechanistic one left empty. It does not reduce the prompt. Ten documents reach the model either way and context characters are within 2%, because the chain caps the fused list at ten however many collections fed it. Narrowing changes which documents arrive, not how many. T003's framing -- that context tokens would fall -- was the wrong thing to measure, and it is corrected rather than quietly dropped. What it does reduce is retrieval work: 1.82s to 1.43s median, three runs per question, timed inside `aretrieve_documents`. Five collections mean ten sub-retrievals; one means two. About a fifth, or 0.4s against a first token near ten seconds. So the honest case is narrower than the spec's: it did not fix a failing question, the control already passed. It buys a fifth off retrieval and puts the right documents in a fixed-size context, and only the first of those is measured. Tests cover the dangerous directions rather than the happy path: that an omitted field still means all collections, that a reactome selection never leaks into the userguide bundle, that the selection does not outlive the question, and that the sync and async retrieval paths honour the same selection -- the served path and the measured path are different code. Co-Authored-By: Claude Opus 5 --- specs/009-collection-routing/research.md | 67 ++++++++++++ specs/009-collection-routing/tasks.md | 20 ++-- src/agent/profiles/react_to_me.py | 55 +++++++--- src/agent/tasks/intent_classifier.py | 67 +++++++++++- tests/agent/test_collection_routing.py | 103 ++++++++++++++++++ .../retrievers/test_sync_async_equivalence.py | 43 +++++++- 6 files changed, 332 insertions(+), 23 deletions(-) create mode 100644 tests/agent/test_collection_routing.py diff --git a/specs/009-collection-routing/research.md b/specs/009-collection-routing/research.md index e407f62..7d09d4d 100644 --- a/specs/009-collection-routing/research.md +++ b/specs/009-collection-routing/research.md @@ -193,3 +193,70 @@ 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. 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..bf6366b 100644 --- a/src/agent/profiles/react_to_me.py +++ b/src/agent/profiles/react_to_me.py @@ -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,9 @@ class ReactToMeState(BaseState): active_sources: list[SourceName] + # Empty means every collection, which is what `resolve_collections` does + # with it. Only meaningful when the active source is `reactome`. + collections: list[str] class ReactToMeGraphBuilder(BaseGraphBuilder): @@ -162,6 +166,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 +188,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 +238,38 @@ 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. + token = selected_collections.set( + state["collections"] 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..e4ad382 --- /dev/null +++ b/tests/agent/test_collection_routing.py @@ -0,0 +1,103 @@ +"""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 + +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) + builder.rags = {source: rag} # type: ignore[attr-defined] + 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 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) From ef4ea718b177f90a2206073b28ea4e96d57ccdba Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Sat, 19 Sep 2026 05:13:17 +0000 Subject: [PATCH 2/3] Adversarial review: a missing key crashed, and nine unseen questions route right Two things the first commit did not cover. `state["collections"]` would raise. BaseState is total=False and the field is new, so a thread resumed from a checkpoint written before this landed has no key at all -- a KeyError on the served path, in a state that is invisible until someone resumes an old conversation. Now `.get(...) or []`, which means every collection, which is how the graph behaved before routing. Pinned by a test that deletes the key and fails against the subscript. And the routing evidence was six questions, all of them ones the sweep already asks. Probed with nine it has never seen, chosen to be awkward: three broad questions that must not narrow, two cross-collection ones where a single choice loses half the answer, four specific ones. Nine of nine as intended -- every broad question left open, every specific one narrowed to the collection holding the answer, and both cross-collection questions took `summations` alongside their primary, which is what the prompt asks for. Recorded with its limit: that is nine classifications, not nine answers. It shows the decision is sound, not that the answers are. Also written down rather than left implicit: `collections` is `list[str]` and not an enum of the five names on purpose. A hallucinated name widens to all with a WARNING, which is the direction this feature requires. Constraining the schema would make structured output reject the response instead, turning a harmless mistake into a failed answer. Co-Authored-By: Claude Opus 5 --- specs/009-collection-routing/research.md | 33 ++++++++++++++++++++++++ src/agent/profiles/react_to_me.py | 6 ++++- tests/agent/test_collection_routing.py | 11 ++++++++ 3 files changed, 49 insertions(+), 1 deletion(-) diff --git a/specs/009-collection-routing/research.md b/specs/009-collection-routing/research.md index 7d09d4d..bcb8bae 100644 --- a/specs/009-collection-routing/research.md +++ b/specs/009-collection-routing/research.md @@ -260,3 +260,36 @@ 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/src/agent/profiles/react_to_me.py b/src/agent/profiles/react_to_me.py index bf6366b..8a4d88f 100644 --- a/src/agent/profiles/react_to_me.py +++ b/src/agent/profiles/react_to_me.py @@ -248,8 +248,12 @@ async def generate_answer( # 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["collections"] if source == "reactome" else None + (state.get("collections") or []) if source == "reactome" else None ) try: result: dict[str, Any] = await rag.ainvoke( diff --git a/tests/agent/test_collection_routing.py b/tests/agent/test_collection_routing.py index e4ad382..5d094a2 100644 --- a/tests/agent/test_collection_routing.py +++ b/tests/agent/test_collection_routing.py @@ -101,3 +101,14 @@ def test_the_selection_does_not_outlive_the_question() -> None: ) ) 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"] # type: ignore[misc] + asyncio.run(_builder(rag, "reactome").generate_answer(state, RunnableConfig())) + assert rag.seen == [[]], "a pre-routing checkpoint must not crash or narrow" From fbccfa57fec727e84a5bd7741aec9ed2c7438de1 Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Sat, 19 Sep 2026 05:27:04 +0000 Subject: [PATCH 3/3] Declare collections NotRequired, which is what the runtime already assumed CI's lint failed where mine passed, because CI runs `poetry run mypy` over everything and I ran `mypy src`. Verifying a narrower scope than the gate is the mistake; the finding underneath it is real. `ReactToMeState` subclasses a total=False TypedDict but is itself total=True, so declaring `collections: list[str]` made it a *required* key. That invalidated four existing state constructions, and it claimed a guarantee the runtime does not have -- the previous commit had already established that a thread resumed from an older checkpoint has no such key, which is why `generate_answer` reads it with `.get`. `NotRequired[list[str]]` says what is true: absent and empty both mean every collection, and the type no longer disagrees with the accessor. Also replaced a `type: ignore[attr-defined]` in the new test that was hiding a second error rather than describing one. Co-Authored-By: Claude Opus 5 --- src/agent/profiles/react_to_me.py | 15 +++++++++++---- tests/agent/test_collection_routing.py | 7 ++++--- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/src/agent/profiles/react_to_me.py b/src/agent/profiles/react_to_me.py index 8a4d88f..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 @@ -29,9 +29,16 @@ class ReactToMeState(BaseState): active_sources: list[SourceName] - # Empty means every collection, which is what `resolve_collections` does - # with it. Only meaningful when the active source is `reactome`. - collections: list[str] + # 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): diff --git a/tests/agent/test_collection_routing.py b/tests/agent/test_collection_routing.py index 5d094a2..f0a77b2 100644 --- a/tests/agent/test_collection_routing.py +++ b/tests/agent/test_collection_routing.py @@ -8,7 +8,7 @@ """ import asyncio -from typing import Any +from typing import Any, cast from langchain_core.runnables import RunnableConfig @@ -51,7 +51,8 @@ async def ainvoke(self, _inputs: dict[str, Any], _config: Any) -> dict[str, Any] def _builder(rag: _RecordingRag, source: str) -> ReactToMeGraphBuilder: builder = ReactToMeGraphBuilder.__new__(ReactToMeGraphBuilder) - builder.rags = {source: rag} # type: ignore[attr-defined] + # A stand-in for the RAG runnable: only `ainvoke` is reached from here. + builder.rags = cast(Any, {source: rag}) return builder @@ -109,6 +110,6 @@ def test_a_state_without_the_field_searches_everything() -> None: # missing must mean "all", which is how the graph behaved before routing. rag = _RecordingRag() state = _state("reactome", []) - del state["collections"] # type: ignore[misc] + del state["collections"] asyncio.run(_builder(rag, "reactome").generate_answer(state, RunnableConfig())) assert rag.seen == [[]], "a pre-routing checkpoint must not crash or narrow"