From 7582bfa467a4dca0702ac893736f58630f9734af Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Wed, 16 Sep 2026 19:32:11 +0000 Subject: [PATCH 01/10] Spec the disease and variant collection The chatbot answers about disease one level above the question. Asked to list the ABCA1 variants in Reactome it names none, and says "Defective ABCA1 causes Tangier Disease" instead; asked which diseases involve PTEN variants it gives the PTEN Loss of Function pathway. Reactome curates six ABCA1 variants and 108 PTEN variants across 86 diseases, and disease_variant_ewas_mapping.tsv has all of them with the residue change, the disease identifiers, and the normal reaction each defective one replaces. It is also the only embeddings work not blocked. The four existing collections are generated from Neo4j and we have no credentials for it; this one is a flat TSV, so it follows the alliance path -- MetaDataCSVLoader and build_embeddings, no database -- and can be built and tested while Release 97 waits on access. Two decisions are left for the team: which columns are embedded rather than filterable metadata, and what to do about the disease field, where a third of rows pack up to nineteen diseases into one pipe-delimited string and turn 443 real diseases into 968 apparent ones. Co-Authored-By: Claude Opus 5 --- specs/008-disease-variant-embeddings/spec.md | 130 +++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 specs/008-disease-variant-embeddings/spec.md diff --git a/specs/008-disease-variant-embeddings/spec.md b/specs/008-disease-variant-embeddings/spec.md new file mode 100644 index 0000000..8334ecf --- /dev/null +++ b/specs/008-disease-variant-embeddings/spec.md @@ -0,0 +1,130 @@ +# Feature Specification: Disease and Variant Embeddings + +**Feature Branch**: `spec/disease-variant-embeddings` + +**Created**: 2026-09-16 + +**Status**: Draft. Two decisions (D1, D2) for the team. + +**Input**: Add disease and variant data to the embeddings, from +`disease_variant_ewas_mapping.tsv` in the release download directory. Do it as a new +collection alongside regenerating the existing ones for Release 97. + +## Why: the chatbot answers about diseases, but cannot name a variant + +Asked on 2026-09-16 against beta running `a4de878`, with live MCP: + +> **List the ABCA1 variants in Reactome and the disease each one causes.** +> +> "The Reactome database lists the following disease associated with ABCA1 variants: +> **Defective ABCA1 causes Tangier Disease (TGD)** ... This is the primary disease +> explicitly mentioned in the context of ABCA1 variants. Other diseases may be +> associated with ABCA1, bu[t] ..." + +Zero variants named. Reactome curates **six** ABCA1 variants, and the file has them: +`ABCA1 W590S`, `N935S`, `C1417R` and three more, each with its stable ID and the +residue change in prose ("L-tryptophan 590 replaced with L-serine"). + +> **Which diseases involve variants of the PTEN gene in Reactome?** +> +> "... **PTEN Loss of Function in Cancer**: This pathway highlights ..." + +Answered at pathway level. Reactome curates **108 PTEN variants across 86 diseases**. + +The pattern is consistent: the four existing collections (`complexes`, `ewas`, +`reactions`, `summations`) hold pathway- and reaction-level prose about disease, so +the chatbot says *"defective ABCA1 causes Tangier disease"* fluently. What it has no +document for is the individual variant entity. It therefore answers a level up from +the question and hedges, which is the failure mode this project keeps finding: a +fluent answer that does not contain the fact asked for. + +## Why this one is worth doing now + +It is the only part of the embeddings work that is **not blocked**. + +| | source | needs | +|---|---|---| +| `complexes`, `ewas`, `reactions`, `summations` | Neo4j at `bolt://localhost:7687` | credentials we do not have | +| **disease/variant** | **a flat TSV in the download directory** | **nothing** | + +`generate_reactome_embeddings` reads Neo4j. `generate_alliance_embeddings` does not -- +it reads flat files through `MetaDataCSVLoader` and `build_embeddings`. That is the +precedent to follow, and it means this collection can be built, installed and tested +today while the Release 97 regeneration waits on access. + +## The data + +`.../static/download/97/disease_variant_ewas_mapping.tsv`, 5.3 MB, **6,294 rows**, +26 columns, one row per variant entity (6,294 distinct `stable_id`, so no duplication). + +- **400** genes, **462** disease pathways +- **443** distinct diseases -- *not* the 968 distinct `disease` strings, see D2 +- Best covered: cancer (1,700 variants), Kabuki syndrome (564), acute myeloid + leukaemia (126), ornithine carbamoyltransferase deficiency (104) + +Each row already carries the chain a user actually asks about: gene, variant display +name, the residue change in prose, the disease with Mondo/DOID identifiers, the +reaction the variant takes part in, its functional status, **and the normal reaction +and pathway it is the defective counterpart of**. + +Fill rates are high: 20 of 26 columns are ≥95% populated. Two are effectively empty +and should be dropped (`normal_reaction_like_event_go_biological_process_*`, 6.3%); +`entityWithAccessionedSequence_literatureReference_pubMedIdentifier` is 30%. + +## Decisions + +### D1 -- what goes in the embedded text, and what stays metadata + +`MetaDataCSVLoader` takes `content_columns` and `metadata_columns`. Getting this wrong +is how a collection retrieves badly: stable IDs embedded as text add tokens and match +nothing a person types. + +Recommendation -- content: `Genename`, `displayName`, `hasModifiedResidue_displayName`, +`disease`, the reaction and pathway `displayName`s, and the *normal* reaction and +pathway `displayName`s. Metadata (filterable, not embedded): every `stable_id`, +`referenceEntity_id`, `disease_identifier`, `cross_reference`, `modifiedResidue_class`, +PubMed identifiers. + +The normal-counterpart names belong in content because "what is the healthy version of +this reaction" is a question the chain uniquely answers. + +### D2 -- the pipe-delimited `disease` field + +33% of rows (2,105) carry more than one disease in one field, up to nineteen: + +``` +Barrett's esophagus|esophagus squamous cell carcinoma|brain meningioma|...|astrocytoma +``` + +Left as-is, "968 diseases" is really 443, a search for `melanoma` competes with a +wall of unrelated text in the same document, and the metadata value is unfilterable. + +Options: (a) leave as-is, simplest, retrieval suffers on the long ones; (b) split into +a list for metadata, keep the joined string in content; (c) one document per +variant-disease pair, which fixes retrieval but inflates 6,294 rows to ~10,000 +documents and repeats the variant text. + +Recommendation: **(b)**. It costs nothing at generation time and makes the identifiers +usable, without multiplying near-duplicate documents. + +## Scope + +In: one new collection built from this one file; wiring it into the retriever +alongside the existing four; answer-sweep expectations that fail if it regresses. + +Out: `HumanDiseasePathways.txt` and `Reactome2OMIM.txt` -- not needed for this +(confirmed 2026-09-16); regenerating the four Neo4j-backed collections for Release 97; +publishing to S3, which is blocked separately -- this host's instance profile is +`EC2CloudwatchAgentRole` and `head_bucket` on `download.reactome.org` returns 403. + +## How we will know it worked + +`src/evaluation/answer_sweep.py` gains expectations that fail today and pass after: + +| question | must contain | why | +|---|---|---| +| List the ABCA1 variants in Reactome | `W590S` | today it names none | +| Which diseases involve PTEN variants in Reactome? | `Cowden` | today it answers at pathway level | + +Cost is not a blocker: 6,294 short documents on `text-embedding-3-large` is cents. +Disk is not either -- the collection is a fraction of the 3.4 GB the existing four take. From 2855de4c6b454342cca5d256da82ecd83a1ff429 Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Wed, 16 Sep 2026 20:08:21 +0000 Subject: [PATCH 02/10] Build the disease_variants collection from the release file Asked to list the ABCA1 variants, the chatbot answered "Defective ABCA1 causes Tangier Disease" and named none of the six. The four existing collections hold pathway- and reaction-level prose about disease, so it talks about disease fluently and has no document for the variant itself. It now answers with C1417R, Q537R, N935S, R587W and S1446L, each with its disease and whether it is a loss of function; PTEN moves from "PTEN Loss of Function in Cancer" to naming Q17*, Q97* and Q171* against endometrial cancer. Both open decisions were delegated back, and one of the spec's own recommendations turned out to be impossible. Chroma accepts only str, int, float or bool as a metadata value, so the disease field cannot be split into a list. Fanning out to one document per variant-disease pair was measured rather than guessed: 6,294 documents become 10,500 and p16INK4A R80* repeats 45 times, which is enough near-identical documents to fill a result set. One document per variant it is, with the pipes rewritten as commas. The other decision was which columns to embed. The loader renders each field as "name: value", so the column names are embedded too -- and the release names are the query paths that produced them, up to 104 characters of entityWithAccessionedSequence_reactionLikeEvent_entityFunctionalStatus_... They are renamed to what a person would call them, identifiers are metadata rather than content because nobody types R-HSA-5682201 at a chatbot, and the variant's own identifier is named st_id because that is the key csv_chroma de-duplicates on. The two sweep expectations match a pattern for a named variant rather than a specific one, since which of the six come back depends on retrieval order, and they skip until the collection is installed rather than turning the deploy gate red for a reason nobody can act on. Co-Authored-By: Claude Opus 5 --- bin/embeddings_manager | 6 + specs/008-disease-variant-embeddings/spec.md | 102 +++++++---- .../disease_variant/__init__.py | 164 ++++++++++++++++++ src/data_generation/reactome/__init__.py | 19 ++ src/evaluation/answer_sweep.py | 53 +++++- src/retrievers/reactome/metadata_info.py | 41 +++++ tests/data_generation/test_disease_variant.py | 129 ++++++++++++++ tests/evaluation/test_answer_sweep.py | 34 ++++ 8 files changed, 511 insertions(+), 37 deletions(-) create mode 100644 src/data_generation/disease_variant/__init__.py create mode 100644 tests/data_generation/test_disease_variant.py diff --git a/bin/embeddings_manager b/bin/embeddings_manager index 088ad82..abcd564 100755 --- a/bin/embeddings_manager +++ b/bin/embeddings_manager @@ -276,6 +276,12 @@ if __name__ == "__main__": action="store_true", help="Force regeneration of CSV files" ) + make_parser.add_argument( + "--disease-variant-tsv", + help="disease_variant_ewas_mapping.tsv from the release download " + "directory. Adds the disease_variants collection, which is built " + "from that file rather than from Neo4j.", + ) args = parser.parse_args() func = args.func diff --git a/specs/008-disease-variant-embeddings/spec.md b/specs/008-disease-variant-embeddings/spec.md index 8334ecf..e1f54d0 100644 --- a/specs/008-disease-variant-embeddings/spec.md +++ b/specs/008-disease-variant-embeddings/spec.md @@ -4,7 +4,8 @@ **Created**: 2026-09-16 -**Status**: Draft. Two decisions (D1, D2) for the team. +**Status**: Implemented. Both decisions were delegated back and are recorded below, +with one of the original recommendations withdrawn as impossible. **Input**: Add disease and variant data to the embeddings, from `disease_variant_ewas_mapping.tsv` in the release download directory. Do it as a new @@ -73,39 +74,58 @@ and should be dropped (`normal_reaction_like_event_go_biological_process_*`, 6.3 ## Decisions -### D1 -- what goes in the embedded text, and what stays metadata +Both were delegated: *"I think you should decide on the columns ... make your best +design and go ahead."* -`MetaDataCSVLoader` takes `content_columns` and `metadata_columns`. Getting this wrong -is how a collection retrieves badly: stable IDs embedded as text add tokens and match -nothing a person types. +### D1 -- what is embedded, and what is metadata: decided -Recommendation -- content: `Genename`, `displayName`, `hasModifiedResidue_displayName`, -`disease`, the reaction and pathway `displayName`s, and the *normal* reaction and -pathway `displayName`s. Metadata (filterable, not embedded): every `stable_id`, -`referenceEntity_id`, `disease_identifier`, `cross_reference`, `modifiedResidue_class`, -PubMed identifiers. +The loader renders each field as `name: value`, so the **column names are embedded +too**. The release names are the query paths that produced them, up to 104 characters +of `entityWithAccessionedSequence_reactionLikeEvent_entityFunctionalStatus_...`. Left +alone they would contribute more tokens than the values, identically in every +document. They are renamed to what a person would call them. -The normal-counterpart names belong in content because "what is the healthy version of -this reaction" is a question the chain uniquely answers. - -### D2 -- the pipe-delimited `disease` field - -33% of rows (2,105) carry more than one disease in one field, up to nineteen: +Embedded: `gene`, `variant`, `protein`, `residue_change`, `mutation_type`, `disease`, +`reaction`, `functional_status`, `disease_pathway`, `normal_reaction`, +`normal_pathway`, `normal_process`. A document reads: ``` -Barrett's esophagus|esophagus squamous cell carcinoma|brain meningioma|...|astrocytoma +gene: ABCA1 +variant: ABCA1 W590S [plasma membrane] +residue_change: L-tryptophan 590 replaced with L-serine +mutation_type: ReplacedResidue +disease: Tangier disease +reaction: Defective ABCA1 does not transport CHOL from transport vesicle membrane... +functional_status: loss_of_function +normal_reaction: 4xPALM-C-p-2S-ABCA1 tetramer transports CHOL from transport vesicle... ``` -Left as-is, "968 diseases" is really 443, a search for `melanoma` competes with a -wall of unrelated text in the same document, and the metadata value is unfilterable. +Metadata, not embedded: every identifier -- `st_id`, `uniprot_id`, `disease_id`, +`disease_cross_reference`, and the four Reactome stable IDs. Nobody types +`R-HSA-5682201` at a chatbot, and embedding it costs tokens in every document. The +variant's own identifier is named **`st_id`** because that is the key `csv_chroma` +de-duplicates on; a different name would silently disable de-duplication here. + +Dropped: the two `go_biological_process` columns at 6.3% fill, and +`first_entitySet`. A column empty in 94% of documents earns nothing and costs a line +of `name:` in every one. + +Median document: 556 characters, about 140 tokens. + +### D2 -- the pipe-delimited `disease` field: decided, and the earlier recommendation withdrawn -Options: (a) leave as-is, simplest, retrieval suffers on the long ones; (b) split into -a list for metadata, keep the joined string in content; (c) one document per -variant-disease pair, which fixes retrieval but inflates 6,294 rows to ~10,000 -documents and repeats the variant text. +This spec first recommended splitting the field into a list for metadata. **That is +not possible**: Chroma accepts only `str`, `int`, `float` or `bool` as a metadata +value and rejects a list outright. -Recommendation: **(b)**. It costs nothing at generation time and makes the identifiers -usable, without multiplying near-duplicate documents. +Fanning out to one document per variant-disease pair was measured rather than +estimated: 6,294 documents become **10,500** (+67%), and `p16INK4A R80*` would be +repeated **45 times**. Forty-five near-identical documents can fill an entire result +set, which is a worse failure than a long disease string. + +Decided: **one document per variant**, with `|` rewritten to `, ` so the field reads +as a list rather than a path. All disease names stay searchable in the content, and +the metadata value stays a string Chroma will accept. ## Scope @@ -117,14 +137,32 @@ Out: `HumanDiseasePathways.txt` and `Reactome2OMIM.txt` -- not needed for this publishing to S3, which is blocked separately -- this host's instance profile is `EC2CloudwatchAgentRole` and `head_bucket` on `download.reactome.org` returns 403. -## How we will know it worked +## Result, measured -`src/evaluation/answer_sweep.py` gains expectations that fail today and pass after: +Built into a scratch bundle and asked through the real retriever: -| question | must contain | why | +| | before | after | |---|---|---| -| List the ABCA1 variants in Reactome | `W590S` | today it names none | -| Which diseases involve PTEN variants in Reactome? | `Cowden` | today it answers at pathway level | +| ABCA1 | *"Defective ABCA1 causes Tangier Disease ... Other diseases may be associated with ABCA1, bu[t]"* -- no variant named | **C1417R, Q537R, N935S, R587W, S1446L**, each with its disease and loss-of-function status | +| PTEN | *"PTEN Loss of Function in Cancer"* -- a pathway | **Q17\*, Q97\*, Q171\***, named against endometrial cancer | + +Two answer-sweep expectations pin this. They match a *pattern* for a named variant +(`\bABCA1 [A-Z]\d{2,4}[A-Z*]`) rather than a specific one, because which of the six +come back depends on retrieval order and pinning one would fail a good answer -- +the mistake made in `007` and fixed there. Checked against the recorded answers: the +pattern does not match the old answer and does match the new one. + +They carry `needs_collection="disease_variants"` and skip, loudly, until the bundle +ships -- otherwise they would turn the deploy gate red for a reason nobody can act +on. A test covers the direction that matters: that they run once it is installed. + +Cost was cents. The collection is a fraction of the 3.4 GB the existing four take. + +## Installing it -Cost is not a blocker: 6,294 short documents on `text-embedding-3-large` is cents. -Disk is not either -- the collection is a fraction of the 3.4 GB the existing four take. +The embeddings tree is root-owned, so installing needs sudo -- and using sudo is what +keeps it root-owned. It is not the container's doing: the image has run as `appuser` +(uid 3001) since 2025-04-17 and the bundle was created 2026-09-02, so something was +run under sudo on the host. `~/fix-embeddings-ownership.sh` sets `awright:reactome` +with world-read, which suits both: the owner can manage bundles without sudo, and the +container, whose uid is not a host user and which only reads at runtime, still can. diff --git a/src/data_generation/disease_variant/__init__.py b/src/data_generation/disease_variant/__init__.py new file mode 100644 index 0000000..79e5418 --- /dev/null +++ b/src/data_generation/disease_variant/__init__.py @@ -0,0 +1,164 @@ +"""Disease variants: the collection built from a file rather than from Neo4j. + +Asked to list the ABCA1 variants in Reactome, the chatbot answered "Defective +ABCA1 causes Tangier Disease" and named none of the six. The four Neo4j-backed +collections hold pathway- and reaction-level prose about disease, so it talks +about disease fluently and has no document for the variant itself. + +`disease_variant_ewas_mapping.tsv` in the release download directory has all of +them -- 6,294 variants over 400 genes -- each with the residue change already in +prose, the disease with its Mondo and DOID identifiers, and the normal reaction +the defective one replaces. It needs no database, which is also why this +collection can be built when the others cannot. +""" + +import csv +import os +from pathlib import Path + +from langchain_chroma import Chroma + +from data_generation.embeddings import build_embeddings +from data_generation.metadata_csv_loader import MetaDataCSVLoader + +COLLECTION = "disease_variants" + +# The source column names are the query paths that produced them, up to 104 +# characters of `entityWithAccessionedSequence_reactionLikeEvent_...`. That +# matters because the loader embeds each field as "name: value", so the column +# names are themselves embedded: left alone they would contribute more tokens +# than the values, identically in every document. These are the names a person +# would use. +COLUMNS: dict[str, str] = { + "Genename": "gene", + "displayName": "variant", + "referenceEntity_name": "protein", + "hasModifiedResidue_displayName": "residue_change", + "modifiedResidue_class": "mutation_type", + "disease": "disease", + "entityWithAccessionedSequence_reactionLikeEvent_displayName": "reaction", + "entityWithAccessionedSequence_reactionLikeEvent_entityFunctionalStatus_functionalStatus_functionalStatusType_displayName": "functional_status", + "entityWithAccessionedSequence_pathway_displayName": "disease_pathway", + "entityWithAccessionedSequence_reactionLikeEvent_normalReaction_displayName": "normal_reaction", + "entityWithAccessionedSequence_pathway_normalPathway_displayName": "normal_pathway", + "entityWithAccessionedSequence_pathway_normalPathway_goBiologicalProcess_displayName": "normal_process", + # Identifiers. Kept out of the embedded text -- nobody types R-HSA-5682201 + # at a chatbot, and embedding it costs tokens in every document. + "stable_id": "st_id", + "referenceEntity_id": "uniprot_id", + "cross_reference": "disease_cross_reference", + "disease_identifier": "disease_id", + "entityWithAccessionedSequence_reactionLikeEvent_stable_id": "reaction_id", + "entityWithAccessionedSequence_pathway_stable_id": "disease_pathway_id", + "entityWithAccessionedSequence_reactionLikeEvent_normalReaction_stable_id": "normal_reaction_id", + "entityWithAccessionedSequence_pathway_normalPathway_stable_id": "normal_pathway_id", + "reactionLikeEvent_literatureReference_pubMedIdentifier": "reaction_pubmed_ids", +} + +# What a question is actually about: names, the change in prose, the disease, +# and what the variant breaks. +CONTENT_COLUMNS: list[str] = [ + "gene", + "variant", + "protein", + "residue_change", + "mutation_type", + "disease", + "reaction", + "functional_status", + "disease_pathway", + "normal_reaction", + "normal_pathway", + "normal_process", +] + +# Filterable, and carried into the answer so a citation can be made. `st_id` is +# named for the key `csv_chroma` de-duplicates on. +METADATA_COLUMNS: list[str] = [ + "st_id", + "gene", + "protein", + "uniprot_id", + "disease", + "disease_id", + "disease_cross_reference", + "mutation_type", + "reaction_id", + "disease_pathway_id", + "normal_reaction_id", + "normal_pathway_id", + "reaction_pubmed_ids", +] + +# 6.3% populated. A column that is empty in 94% of documents earns nothing and +# costs a line of "name: " in every one of them. +DROPPED = ( + "normal_reaction_like_event_go_biological_process_accession", + "normal_reaction_like_event_go_biological_process_displayName", + "entityWithAccessionedSequence_pathway_normalPathway_goBiologicalProcess_accession", + "entityWithAccessionedSequence_literatureReference_pubMedIdentifier", + "first_entitySet", +) + + +def _tidy(value: str) -> str: + """One row's field, made readable. + + A third of rows pack several diseases into one pipe-delimited string, up to + nineteen of them. They stay in one document -- fanning out to one document + per variant-disease pair would repeat `p16INK4A R80*` forty-five times, and + forty-five near-identical documents can fill a whole result set. The + separator becomes a comma so it reads as a list rather than a path. + """ + return ", ".join(part.strip() for part in value.split("|") if part.strip()) + + +def write_csv(tsv_path: Path, csv_path: Path) -> int: + """Rewrite the release TSV as the CSV both retrievers read. Returns rows.""" + csv_path.parent.mkdir(parents=True, exist_ok=True) + with open(tsv_path, newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle, delimiter="\t")) + + missing = set(COLUMNS) - set(rows[0]) if rows else set(COLUMNS) + if missing: + raise ValueError( + f"{tsv_path} is missing expected columns: {sorted(missing)}. " + "The release file's shape has changed; update COLUMNS." + ) + + with open(csv_path, "w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=list(COLUMNS.values())) + writer.writeheader() + for row in rows: + writer.writerow( + {new: _tidy(row[old] or "") for old, new in COLUMNS.items()} + ) + return len(rows) + + +def generate_disease_variant_embeddings( + embeddings_dir: str, + tsv_path: str | Path, + hf_model: str | None = None, + device: str | None = None, +) -> Chroma: + """Build the disease_variants collection inside an existing bundle.""" + bundle = Path(embeddings_dir) + csv_path = bundle / "csv_files" / f"{COLLECTION}.csv" + count = write_csv(Path(tsv_path), csv_path) + print(f" wrote {csv_path} ({count} variants)") + + loader = MetaDataCSVLoader( + file_path=str(csv_path), + content_columns=CONTENT_COLUMNS, + metadata_columns=METADATA_COLUMNS, + encoding="utf-8", + ) + docs = loader.load() + print(f" loaded {len(docs)} documents") + + return Chroma.from_documents( + documents=docs, + embedding=build_embeddings(hf_model, device, chunk_size=400), + persist_directory=os.path.join(embeddings_dir, COLLECTION), + ) diff --git a/src/data_generation/reactome/__init__.py b/src/data_generation/reactome/__init__.py index fc66d95..9f159b5 100644 --- a/src/data_generation/reactome/__init__.py +++ b/src/data_generation/reactome/__init__.py @@ -3,6 +3,7 @@ from langchain_community.vectorstores import Chroma +from data_generation.disease_variant import generate_disease_variant_embeddings from data_generation.embeddings import build_embeddings from data_generation.metadata_csv_loader import MetaDataCSVLoader from data_generation.reactome.csv_generator import generate_all_csvs @@ -70,6 +71,7 @@ def generate_reactome_embeddings( force: bool = False, hf_model: str | None = None, device: str | None = None, + disease_variant_tsv: str | Path | None = None, ) -> None: csv_dir = Path(embeddings_dir) / "csv_files" reactions_csv = str(csv_dir / "reactions.csv") @@ -107,3 +109,20 @@ def generate_reactome_embeddings( print(db._collection.count()) db = upload_to_chromadb(embeddings_dir, ewas_csv, "ewas", hf_model, device) print(db._collection.count()) + + # Built from a release file, not from Neo4j, so it is the one collection + # that can be generated without database access. Skipped rather than + # failed when no file is given: the other four are still a valid bundle. + if disease_variant_tsv: + # Its own name: this returns langchain_chroma's Chroma, while the four + # above still come back as the deprecated langchain_community one. + variants_db = generate_disease_variant_embeddings( + embeddings_dir, disease_variant_tsv, hf_model, device + ) + print(variants_db._collection.count()) + else: + print( + "No --disease-variant-tsv given; skipping the disease_variants " + "collection. The chatbot will answer about disease at the level of " + "a pathway and will not be able to name a variant." + ) diff --git a/src/evaluation/answer_sweep.py b/src/evaluation/answer_sweep.py index cb4b11b..91e3271 100644 --- a/src/evaluation/answer_sweep.py +++ b/src/evaluation/answer_sweep.py @@ -33,6 +33,7 @@ from agent.graph import AgentGraph from agent.profile_names import ProfileName from reactome_mcp.session import is_configured +from util.embedding_environment import EmbeddingEnvironment @dataclass(frozen=True) @@ -55,6 +56,9 @@ class Expectation: # a plain local checkout -- these cannot pass, and reporting them as # regressions is how a gate teaches people to ignore it. needs_live: bool = False + # A collection that may not be in the installed bundle yet. Same reasoning + # as needs_live: a question that cannot pass is not a regression. + needs_collection: str = "" EXPECTATIONS: tuple[Expectation, ...] = ( @@ -87,6 +91,24 @@ class Expectation: why="Preferring the web tool must not bury the plugin for someone who wants it.", must=("FIViz",), ), + # --- disease variants, which only the new collection can name ----------- + Expectation( + question="List the ABCA1 variants in Reactome and the disease each one causes.", + why="Answered 'Defective ABCA1 causes Tangier Disease' and named none of " + "the six curated variants. Pathway-level prose instead of the variant.", + must=("Tangier",), + # A named variant, not a specific one: which of the six come back + # depends on retrieval order, and pinning one would fail a good answer. + must_match=(r"\bABCA1 [A-Z]\d{2,4}[A-Z*]",), + needs_collection="disease_variants", + ), + Expectation( + question="Which diseases involve variants of the PTEN gene in Reactome?", + why="Answered with the PTEN Loss of Function pathway. Reactome curates " + "108 PTEN variants across 86 diseases.", + must_match=(r"\bPTEN [A-Z]\d{2,4}[A-Z*]",), + needs_collection="disease_variants", + ), # --- facts about the database, which retrieval cannot answer ------------ Expectation( question="what species are in reactome", @@ -194,6 +216,11 @@ def _contains(haystack: str, needle: str) -> bool: ) +def _has_collection(name: str) -> bool: + bundle = EmbeddingEnvironment.get_dir("reactome") + return bool(bundle and (bundle / name / "chroma.sqlite3").exists()) + + def _looks_transient(result: "Result") -> bool: return any(_contains(result.answer, marker) for marker in TRANSIENT) @@ -204,6 +231,17 @@ async def run(expectations: tuple[Expectation, ...], retries: int = 1) -> list[R live = is_configured() try: for index, expectation in enumerate(expectations, start=1): + if expectation.needs_collection and not _has_collection( + expectation.needs_collection + ): + results.append( + Result( + expectation=expectation, + skipped=f"the {expectation.needs_collection} collection " + "is not in the installed bundle", + ) + ) + continue if expectation.needs_live and not live: results.append( Result( @@ -296,11 +334,16 @@ def report(results: list[Result]) -> int: ran = len(results) - len(skipped) print(f"\n {ran - len(failures)}/{ran} passed, {total:.0f}s total") if skipped: - print( - f" {len(skipped)} skipped: they need the live service, so a local run" - " cannot check them." - ) - print(" The deploy runs this inside the container, where MCP is configured.") + # Grouped by reason: "they need the live service" was printed for + # every skip, including one whose collection was simply not installed. + reasons: dict[str, int] = {} + for r in skipped: + reasons[r.skipped] = reasons.get(r.skipped, 0) + 1 + print(f" {len(skipped)} skipped, neither a pass nor a failure:") + for reason, count in reasons.items(): + print(f" {count}x {reason}") + print(" The deploy runs this inside the container, where the bundle is") + print(" installed and MCP is configured.") if failures: print( " A failure here is a question the chatbot used to get wrong and does again." diff --git a/src/retrievers/reactome/metadata_info.py b/src/retrievers/reactome/metadata_info.py index fff4783..31c810a 100644 --- a/src/retrievers/reactome/metadata_info.py +++ b/src/retrievers/reactome/metadata_info.py @@ -11,11 +11,52 @@ "ewas": "Contains data on proteins and nucleic acids with known sequences. Includes entity names, IDs, canonical and synonymous gene names, and functions.", "complexes": "Catalogs biological complexes, listing complex names and IDs along with the names and IDs of their components. ", "reactions": "Documents biological pathways and their constituent reactions, detailing pathway and reaction names and IDs. It includes information on the inputs, outputs, and catalysts for each reaction, emphasizing the interconnected nature of cellular processes. Inputs and outputs, critical to the initiation and conclusion of reactions, along with catalysts that facilitate these processes, are cataloged to highlight their roles across various reactions and pathways", + "disease_variants": "Individual disease-causing variants of proteins: the gene, the variant name, the amino acid change in prose, the disease it causes with its Mondo and DOID identifiers, the reaction the variant takes part in and whether it is a loss or gain of function, and the normal reaction and pathway the defective one replaces. This is the only collection holding the variants themselves; the others describe disease at the level of a pathway or reaction.", "summations": "Enumerates biological reactions, accompanied by concise summaries ('summations') of each reaction. These summations encapsulate the essence and biochemical significance of the reactions, offering insights into their roles within cellular processes and pathways.", } reactome_field_info: dict[str, list[AttributeInfo]] = { + "disease_variants": [ + AttributeInfo( + name="st_id", + description="The Reactome Identifier for the variant entity itself, " + "e.g. R-HSA-5682201 for ABCA1 W590S.", + type="string", + ), + AttributeInfo( + name="gene", + description="The gene the variant belongs to, e.g. ABCA1 or PTEN. " + "One gene has many variants: Reactome curates 108 for PTEN.", + type="string", + ), + AttributeInfo( + name="disease", + description="The disease or diseases the variant causes, comma " + "separated. A third of variants list more than one, and the " + "well-covered ones list many.", + type="string", + ), + AttributeInfo( + name="disease_id", + description="Disease Ontology identifier, e.g. DOID:1388 for " + "Tangier disease. Matches `disease` position for position.", + type="string", + ), + AttributeInfo( + name="mutation_type", + description="How the sequence differs: ReplacedResidue, " + "NonsenseMutation, FragmentDeletionModification, " + "FragmentInsertionModification or FragmentReplacedModification.", + type="string", + ), + AttributeInfo( + name="normal_reaction_id", + description="The Reactome Identifier of the normal reaction this " + "defective one replaces, for comparing against healthy biology.", + type="string", + ), + ], "summations": [ AttributeInfo( name="st_id", diff --git a/tests/data_generation/test_disease_variant.py b/tests/data_generation/test_disease_variant.py new file mode 100644 index 0000000..509c88f --- /dev/null +++ b/tests/data_generation/test_disease_variant.py @@ -0,0 +1,129 @@ +"""The disease_variants collection, built from a release file rather than Neo4j.""" + +import csv +from pathlib import Path + +import pytest + +from data_generation.disease_variant import ( + COLUMNS, + CONTENT_COLUMNS, + METADATA_COLUMNS, + _tidy, + write_csv, +) +from data_generation.metadata_csv_loader import MetaDataCSVLoader + +ROW = { + "Genename": "ABCA1", + "displayName": "ABCA1 W590S [plasma membrane]", + "stable_id": "R-HSA-5682201", + "referenceEntity_name": "ABCA1", + "referenceEntity_id": "UniProt:O95477", + "hasModifiedResidue_displayName": "L-tryptophan 590 replaced with L-serine", + "modifiedResidue_class": "ReplacedResidue", + "cross_reference": "Mondo:0008783", + "disease": "Tangier disease", + "disease_identifier": "DOID:1388", + "entityWithAccessionedSequence_literatureReference_pubMedIdentifier": "", + "first_entitySet": "", + "entityWithAccessionedSequence_reactionLikeEvent_stable_id": "R-HSA-5682111", + "entityWithAccessionedSequence_reactionLikeEvent_displayName": "Defective ABCA1 does not transport CHOL", + "entityWithAccessionedSequence_reactionLikeEvent_entityFunctionalStatus_functionalStatus_functionalStatusType_displayName": "loss_of_function", + "reactionLikeEvent_literatureReference_pubMedIdentifier": "pubmed:12509412", + "entityWithAccessionedSequence_pathway_stable_id": "R-HSA-5682113", + "entityWithAccessionedSequence_pathway_displayName": "Defective ABCA1 causes TGD", + "entityWithAccessionedSequence_reactionLikeEvent_normalReaction_stable_id": "R-HSA-216723", + "entityWithAccessionedSequence_reactionLikeEvent_normalReaction_displayName": "ABCA1 tetramer transports CHOL", + "entityWithAccessionedSequence_pathway_normalPathway_displayName": "Plasma lipoprotein assembly", + "entityWithAccessionedSequence_pathway_normalPathway_stable_id": "R-HSA-174824", + "normal_reaction_like_event_go_biological_process_accession": "", + "normal_reaction_like_event_go_biological_process_displayName": "", + "entityWithAccessionedSequence_pathway_normalPathway_goBiologicalProcess_accession": "GO:0071827", + "entityWithAccessionedSequence_pathway_normalPathway_goBiologicalProcess_displayName": "plasma lipoprotein particle organization", +} + + +def _tsv(tmp_path: Path, rows: list[dict[str, str]]) -> Path: + path = tmp_path / "mapping.tsv" + with open(path, "w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=list(ROW), delimiter="\t") + writer.writeheader() + writer.writerows(rows) + return path + + +def test_the_release_columns_are_renamed_to_readable_ones(tmp_path: Path) -> None: + # The loader embeds each field as "name: value", so the column names are + # themselves embedded. The release names are query paths up to 104 chars. + out = tmp_path / "csv_files" / "disease_variants.csv" + assert write_csv(_tsv(tmp_path, [ROW]), out) == 1 + header = out.read_text().splitlines()[0] + assert "entityWithAccessionedSequence" not in header + assert "residue_change" in header + assert "normal_reaction" in header + + +def test_identifiers_are_metadata_and_never_embedded(tmp_path: Path) -> None: + # Nobody types R-HSA-5682201 at a chatbot, and embedding it costs tokens in + # every document. + out = tmp_path / "csv_files" / "disease_variants.csv" + write_csv(_tsv(tmp_path, [ROW]), out) + (doc,) = MetaDataCSVLoader( + file_path=str(out), + content_columns=CONTENT_COLUMNS, + metadata_columns=METADATA_COLUMNS, + encoding="utf-8", + ).load() + + assert "R-HSA-5682201" not in doc.page_content + assert "UniProt:O95477" not in doc.page_content + assert doc.metadata["st_id"] == "R-HSA-5682201" + # csv_chroma de-duplicates on metadata["st_id"]; a different key there + # would silently stop that working for this collection. + assert "st_id" in METADATA_COLUMNS + + assert "ABCA1 W590S [plasma membrane]" in doc.page_content + assert "L-tryptophan 590 replaced with L-serine" in doc.page_content + assert "Tangier disease" in doc.page_content + # The normal counterpart is content, because "what is the healthy version + # of this" is a question only this chain answers. + assert "ABCA1 tetramer transports CHOL" in doc.page_content + + +def test_multiple_diseases_stay_in_one_document(tmp_path: Path) -> None: + # Fanning out to one document per variant-disease pair would repeat + # p16INK4A R80* forty-five times, and that many near-identical documents + # can fill a whole result set. + row = dict(ROW, disease="Cowden syndrome|breast cancer|melanoma") + out = tmp_path / "csv_files" / "disease_variants.csv" + assert write_csv(_tsv(tmp_path, [row]), out) == 1 + (doc,) = MetaDataCSVLoader( + file_path=str(out), + content_columns=CONTENT_COLUMNS, + metadata_columns=METADATA_COLUMNS, + encoding="utf-8", + ).load() + assert "Cowden syndrome, breast cancer, melanoma" in doc.page_content + + +def test_tidy_normalises_the_pipe_separator() -> None: + assert _tidy("a|b|c") == "a, b, c" + assert _tidy("") == "" + assert _tidy("only") == "only" + assert _tidy("a||b") == "a, b" + + +def test_a_changed_release_file_fails_loudly(tmp_path: Path) -> None: + # Silently writing empty columns would produce a bundle that embeds + # nothing useful and reports no error. + path = tmp_path / "mapping.tsv" + path.write_text("Genename\tdisease\nABCA1\tTangier disease\n") + with pytest.raises(ValueError, match="missing expected columns"): + write_csv(path, tmp_path / "out.csv") + + +def test_every_renamed_column_is_used_somewhere() -> None: + # A column renamed but left out of both lists is silently dropped. + used = set(CONTENT_COLUMNS) | set(METADATA_COLUMNS) + assert set(COLUMNS.values()) - used == set() diff --git a/tests/evaluation/test_answer_sweep.py b/tests/evaluation/test_answer_sweep.py index a230abb..8feb973 100644 --- a/tests/evaluation/test_answer_sweep.py +++ b/tests/evaluation/test_answer_sweep.py @@ -162,3 +162,37 @@ def test_a_must_not_guard_still_catches_inflections() -> None: # A number stays closed at both ends: a longer one is a different number. assert not _contains("released in 1996", "96") assert not _contains("there are 965 of them", "96") + + +COLLECTION = ( + Expectation( + question="List the ABCA1 variants in Reactome.", + why="Pathway-level prose instead of the variant.", + must=("W590S",), + needs_collection="disease_variants", + ), +) + + +def test_a_question_is_skipped_when_its_collection_is_not_installed( + stub: Install, monkeypatch: pytest.MonkeyPatch +) -> None: + graph = stub(["Defective ABCA1 causes Tangier Disease."]) + monkeypatch.setattr("evaluation.answer_sweep._has_collection", lambda _n: False) + (result,) = asyncio.run(run(COLLECTION)) + assert result.skipped + assert result.ok + assert graph.asked == [] + + +def test_it_runs_once_the_collection_is_installed( + stub: Install, monkeypatch: pytest.MonkeyPatch +) -> None: + # The direction that matters: once the bundle ships, this must actually be + # checked rather than skipped forever. + graph = stub(["Defective ABCA1 causes Tangier Disease."]) + monkeypatch.setattr("evaluation.answer_sweep._has_collection", lambda _n: True) + (result,) = asyncio.run(run(COLLECTION)) + assert not result.skipped + assert not result.ok + assert len(graph.asked) == 1 From 8fa1280a4559e1e98233605105132d26e16829c4 Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Wed, 16 Sep 2026 21:06:40 +0000 Subject: [PATCH 03/10] Stop the router sending variant questions past the collection that has them With the collection installed the sweep still failed. Retrieval was not the problem: the ABCA1 question retrieves three named variants, and the RAG chain answers it with five in 2,546 characters. The agent answered in 393 and named none. The difference was the router. `live` is described as answering "whether some specific thing exists in it at all", and "List the ABCA1 variants in Reactome" reads exactly like a question about what the database contains -- so with MCP configured it went to the live services, which answer at the level of the pathway and never see the new documents. Beta always has MCP configured, so the collection would have been dead there while passing every local test. The rule now draws the line at scope versus content: live answers how many, which species, which release and whether a thing exists at all; reactome answers what is curated about a named gene, disease or pathway, including listing it. Checked in both directions -- the two variant questions now pass and the species and release questions still route live and still pass. 13/13 against the running MCP sibling. The prompt is unchanged where `live` is not offered, which FR-007 requires byte for byte, and a test covers that. Co-Authored-By: Claude Opus 5 --- src/agent/tasks/intent_classifier.py | 9 ++++++- src/retrievers/reactome/prompt.py | 3 +++ tests/agent/test_intent_classifier_sources.py | 25 +++++++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/src/agent/tasks/intent_classifier.py b/src/agent/tasks/intent_classifier.py index 07670e7..8d50e58 100644 --- a/src/agent/tasks/intent_classifier.py +++ b/src/agent/tasks/intent_classifier.py @@ -40,7 +40,14 @@ _LIVE_RULE = """- If the user asks what the database *contains* or *covers*, rather than asking about the biology in it, choose **live**. "What does CDK5 do?" is **reactome**; "does Reactome have - CDK5?" is **live**.""" + CDK5?" is **live**. +- Naming or listing curated entities is **reactome**, not **live**. "Which ABCA1 variants + are there?" and "which diseases involve PTEN variants?" are answered from stored + documents, which hold the variants themselves. The line is scope versus content: + **live** answers how many, which species, which release, and whether a thing exists at + all; **reactome** answers what is curated about a given gene, disease or pathway -- + including listing it. A question naming a specific gene or disease is almost always + **reactome**.""" _SOURCE_BLOCKS: dict[SourceName, str] = { "reactome": _REACTOME_SOURCE, diff --git a/src/retrievers/reactome/prompt.py b/src/retrievers/reactome/prompt.py index d348939..4dc788e 100644 --- a/src/retrievers/reactome/prompt.py +++ b/src/retrievers/reactome/prompt.py @@ -18,6 +18,9 @@ 2. Inline citations required: Every factual statement must include ≥1 inline anchor citation in the format: display_name - If multiple entries support the same fact, cite them together (space-separated). 3. Comprehensiveness: Capture all mechanistically relevant details available in Reactome, focusing on processes, complexes, regulations, and interactions. + - When the question asks **which**, or asks you to **list** or **name** specific entities -- variants, complexes, participants, reactions -- name each one in the context individually. Do not answer at the level of the pathway that groups them. + - "Defective ABCA1 causes Tangier Disease" does not answer "which ABCA1 variants are there". If the context contains `ABCA1 W590S` and `ABCA1 C1417R`, those are the answer, and a pathway describing them collectively is the background to it. + - The narrative style above is for questions about mechanism. A question asking which things exist wants the things. 4. Tone & Style: - Write in a clear, engaging, and conversational tone. - Use accessible language while maintaining technical precision. diff --git a/tests/agent/test_intent_classifier_sources.py b/tests/agent/test_intent_classifier_sources.py index cfa6aa6..40f4973 100644 --- a/tests/agent/test_intent_classifier_sources.py +++ b/tests/agent/test_intent_classifier_sources.py @@ -51,3 +51,28 @@ def test_an_unavailable_source_falls_back_rather_than_failing() -> None: assert resolve_active_sources("live", TWO) == ["reactome"] assert resolve_active_sources("live", THREE) == ["live"] assert resolve_active_sources("userguide", ONE) == ["reactome"] + + +def test_listing_curated_entities_is_steered_away_from_live() -> None: + """With `live` offered, "List the ABCA1 variants" routed there. + + The live prompt says `live` answers "whether some specific thing exists in + it at all", and a request to list variants reads like a question about what + the database contains. It is not: the variants are documents in the + `disease_variants` collection, and the live services answer that question + at the level of the pathway -- 393 characters naming no variant, where + retrieval names five. + """ + message = build_classifier_message(THREE) + assert "Naming or listing curated entities is **reactome**" in message + assert "scope versus content" in message + # The examples are the two questions that actually failed. + assert "Which ABCA1 variants" in message + assert "which diseases involve PTEN variants" in message + + +def test_that_steer_is_absent_when_live_is_not_offered() -> None: + # It only makes sense next to `live`; without it the rule is noise, and + # FR-007 requires the no-MCP prompt to stay byte-for-byte unchanged. + assert "Naming or listing curated entities" not in build_classifier_message(TWO) + assert "Naming or listing curated entities" not in build_classifier_message(ONE) From 7b733ead612b407156a1423dccb09af68f113996 Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Wed, 16 Sep 2026 21:25:21 +0000 Subject: [PATCH 04/10] Adversarial review of the disease_variants work: three fixes Regenerating appended instead of replacing. Chroma.from_documents adds to whatever is already in the directory, so a second `make` would have produced 12,588 documents -- every variant twice, retrieved twice in a result set, with no error anywhere. It now removes the collection first, and clears chromadb's cached system client, which otherwise keeps pointing at the deleted sqlite file and fails the next write with "attempt to write a readonly database". `cross_reference` was named `disease_cross_reference`, and it is not all disease: every row carries a Mondo disease id, but 4,239 also carry a COSMIC variant id, plus ClinVar, ClinGen and LOVD. Anyone filtering on that name would have got variant identifiers back. It is `cross_references` now, and the description says what is really in it. The disease identifier proper is `disease_id`, which is DOID throughout -- and that one I did check: it lines up with `disease` position for position on all 6,294 rows. The enumeration rule I added to the reactome prompt was dead weight. I added it on the theory that the model was summarising instead of listing; the real cause was the router sending the question to the live services. Reverted and verified: both variant questions pass without it, so it was costing tokens in every Reactome answer for nothing. Co-Authored-By: Claude Opus 5 --- .../disease_variant/__init__.py | 23 +++++++++++++++++-- src/retrievers/reactome/metadata_info.py | 4 +++- src/retrievers/reactome/prompt.py | 3 --- tests/data_generation/test_disease_variant.py | 23 +++++++++++++++++++ 4 files changed, 47 insertions(+), 6 deletions(-) diff --git a/src/data_generation/disease_variant/__init__.py b/src/data_generation/disease_variant/__init__.py index 79e5418..109cdcf 100644 --- a/src/data_generation/disease_variant/__init__.py +++ b/src/data_generation/disease_variant/__init__.py @@ -14,8 +14,10 @@ import csv import os +import shutil from pathlib import Path +from chromadb.api.client import SharedSystemClient from langchain_chroma import Chroma from data_generation.embeddings import build_embeddings @@ -44,9 +46,14 @@ "entityWithAccessionedSequence_pathway_normalPathway_goBiologicalProcess_displayName": "normal_process", # Identifiers. Kept out of the embedded text -- nobody types R-HSA-5682201 # at a chatbot, and embedding it costs tokens in every document. + # `cross_references` is deliberately not called a disease identifier: it + # carries a Mondo disease id on every row AND COSMIC, ClinVar, ClinGen or + # LOVD *variant* ids on thousands of them (4,239 rows have a COSMIC id). + # The disease identifier proper is `disease_id`, which is DOID throughout + # and lines up with `disease` position for position on all 6,294 rows. "stable_id": "st_id", "referenceEntity_id": "uniprot_id", - "cross_reference": "disease_cross_reference", + "cross_reference": "cross_references", "disease_identifier": "disease_id", "entityWithAccessionedSequence_reactionLikeEvent_stable_id": "reaction_id", "entityWithAccessionedSequence_pathway_stable_id": "disease_pathway_id", @@ -81,7 +88,7 @@ "uniprot_id", "disease", "disease_id", - "disease_cross_reference", + "cross_references", "mutation_type", "reaction_id", "disease_pathway_id", @@ -148,6 +155,18 @@ def generate_disease_variant_embeddings( count = write_csv(Path(tsv_path), csv_path) print(f" wrote {csv_path} ({count} variants)") + # Chroma.from_documents appends to whatever is already in the directory, so + # a second run would silently double the collection -- 12,588 documents, + # every variant twice, and no error anywhere. Regeneration replaces. + persist = bundle / COLLECTION + if persist.exists(): + shutil.rmtree(persist) + # chromadb caches one system client per path. Removing the directory + # underneath it leaves that client pointing at a deleted sqlite file, + # and the next write fails with "attempt to write a readonly database". + SharedSystemClient.clear_system_cache() + print(f" removed the previous {COLLECTION} collection") + loader = MetaDataCSVLoader( file_path=str(csv_path), content_columns=CONTENT_COLUMNS, diff --git a/src/retrievers/reactome/metadata_info.py b/src/retrievers/reactome/metadata_info.py index 31c810a..93c8e1e 100644 --- a/src/retrievers/reactome/metadata_info.py +++ b/src/retrievers/reactome/metadata_info.py @@ -40,7 +40,9 @@ AttributeInfo( name="disease_id", description="Disease Ontology identifier, e.g. DOID:1388 for " - "Tangier disease. Matches `disease` position for position.", + "Tangier disease. DOID throughout, and lines up with `disease` " + "position for position. Not to be confused with `cross_references`, " + "which mixes a Mondo disease id with COSMIC and ClinVar variant ids.", type="string", ), AttributeInfo( diff --git a/src/retrievers/reactome/prompt.py b/src/retrievers/reactome/prompt.py index 4dc788e..d348939 100644 --- a/src/retrievers/reactome/prompt.py +++ b/src/retrievers/reactome/prompt.py @@ -18,9 +18,6 @@ 2. Inline citations required: Every factual statement must include ≥1 inline anchor citation in the format: display_name - If multiple entries support the same fact, cite them together (space-separated). 3. Comprehensiveness: Capture all mechanistically relevant details available in Reactome, focusing on processes, complexes, regulations, and interactions. - - When the question asks **which**, or asks you to **list** or **name** specific entities -- variants, complexes, participants, reactions -- name each one in the context individually. Do not answer at the level of the pathway that groups them. - - "Defective ABCA1 causes Tangier Disease" does not answer "which ABCA1 variants are there". If the context contains `ABCA1 W590S` and `ABCA1 C1417R`, those are the answer, and a pathway describing them collectively is the background to it. - - The narrative style above is for questions about mechanism. A question asking which things exist wants the things. 4. Tone & Style: - Write in a clear, engaging, and conversational tone. - Use accessible language while maintaining technical precision. diff --git a/tests/data_generation/test_disease_variant.py b/tests/data_generation/test_disease_variant.py index 509c88f..8f211e1 100644 --- a/tests/data_generation/test_disease_variant.py +++ b/tests/data_generation/test_disease_variant.py @@ -127,3 +127,26 @@ def test_every_renamed_column_is_used_somewhere() -> None: # A column renamed but left out of both lists is silently dropped. used = set(CONTENT_COLUMNS) | set(METADATA_COLUMNS) assert set(COLUMNS.values()) - used == set() + + +def test_regenerating_replaces_rather_than_appends( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Chroma.from_documents appends. Twice would double the collection. + + Every variant stored twice, with no error: the bundle would look fine and + retrieve the same document twice in a result set. + """ + from langchain_core.embeddings import FakeEmbeddings + + import data_generation.disease_variant as dv + + monkeypatch.setattr(dv, "build_embeddings", lambda *a, **k: FakeEmbeddings(size=8)) + tsv = _tsv( + tmp_path, [ROW, dict(ROW, stable_id="R-HSA-2", displayName="ABCA1 N935S")] + ) + + first = dv.generate_disease_variant_embeddings(str(tmp_path), tsv) + assert first._collection.count() == 2 + second = dv.generate_disease_variant_embeddings(str(tmp_path), tsv) + assert second._collection.count() == 2, "regenerating must not append" From ebe338ccf46457b250eb3c8701b562277ad65dd7 Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Thu, 17 Sep 2026 00:50:34 +0000 Subject: [PATCH 05/10] Record where a collection came from, and spec searching only what is needed The bundle directory is named Release95 and now holds one collection built from the Release 97 download directory. Nothing recorded that, which is the failure build_embeddings already warns about: a bundle whose contents disagree with its path is unusable in a way nothing reports. Generating now writes provenance.json beside the collections, merged so each keeps its own entry. Spec 009 is Adam's concern, measured. Every collection is searched for every question and contributes a fixed ten documents. Adding the fifth cost 2,376 tokens and 3.4 seconds on a question about CDK5 and Alzheimer disease, which has nothing to do with variants -- +34% and +30%, so three more tables is roughly +7K tokens and +10s on everything. The allocation is also inverted: disease_variants contributed 15% of the context to that CDK5 question and 5% to the ABCA1 question it exists for, because every collection gets ten documents whether or not they are any good and variant documents are short where pathway prose is long. The proposal costs nothing to run: the intent classifier already makes one LLM call per question, and the per-collection descriptions it would need already exist in metadata_info.py, read today only by bin/retrieval_baseline. Co-Authored-By: Claude Opus 5 --- specs/009-collection-routing/spec.md | 111 ++++++++++++++++++ .../disease_variant/__init__.py | 28 +++++ tests/data_generation/test_disease_variant.py | 23 ++++ 3 files changed, 162 insertions(+) create mode 100644 specs/009-collection-routing/spec.md diff --git a/specs/009-collection-routing/spec.md b/specs/009-collection-routing/spec.md new file mode 100644 index 0000000..4c81246 --- /dev/null +++ b/specs/009-collection-routing/spec.md @@ -0,0 +1,111 @@ +# Feature Specification: Searching Only the Collections a Question Needs + +**Feature Branch**: `spec/collection-routing` + +**Created**: 2026-09-17 + +**Status**: Draft. Two decisions (D1, D2) for the team. + +**Input**: *"the more tables we add the more tokens we use up and the longer the +responses take to return, it would be nice to search the disease variant table if +that is what they are asking about."* + +## The cost is real, and it was measured + +Every collection in the bundle is searched for every question. `retrieve_documents` +loops over all of them and each contributes up to `max_documents_per_collection` +(10), regardless of whether it had anything to say. + +Adding `disease_variants` as a fifth collection, asked a question with nothing to do +with variants -- *"What does CDK5 phosphorylate in Alzheimer disease?"*: + +| | docs | context tokens | retrieval | +|---|---|---|---| +| 4 collections | 40 | 7,061 | 11.5s | +| 5 collections | 50 | 9,437 | **14.9s** | + +**+2,376 tokens (+34%) and +3.4s (+30%)**, all of it noise for that question. The +cost is per collection and per question, so three more tables is roughly +7K tokens +and +10s on every question anyone asks. + +## The allocation is also inverted + +Worse than linear cost: `disease_variants` contributes **more** to questions it is +irrelevant to than to the one it exists for. + +| question | tokens from `disease_variants` | share of context | +|---|---|---| +| What does CDK5 phosphorylate in Alzheimer disease? | 1,427 | **15%** | +| How does TP53 regulate PTEN transcription? | 1,135 | 8% | +| List the ABCA1 variants in Reactome | 616 | **5%** | + +Each collection gets a fixed ten documents whether or not they are any good, and +variant documents are short (556 characters median) where pathway prose is long. So +the fixed allocation spends the most context on the collection with the least to +contribute, and the least on the one that answers the question. + +## The proposal: decide collections in the call that already happens + +The intent classifier already runs **one LLM call per question** and returns a +source (`reactome`, `userguide`, `live`). Having it also name collections adds **no +latency and no call**. + +The descriptions it would need already exist. `reactome_descriptions_info` in +`src/retrievers/reactome/metadata_info.py` holds a written description of every +collection and is currently read only by `bin/retrieval_baseline`. It was written +for exactly this kind of routing and is otherwise unused in the serving path. + +`HybridRetriever.retrieve_documents` then loops over the selected collections +instead of all of them -- a filter on `self.collection_retrievers.items()`. + +## Decisions + +### D1 -- what happens when the classifier is unsure + +Recommendation: **select all collections**. A wrong selection costs recall silently, +which is the failure this project keeps finding; a wrong *default* costs only what +we already pay today. So the change can only make things faster, never worse than +the current behaviour, unless the classifier actively picks a wrong subset. + +The alternative -- always include a core set and gate only specialised collections +-- does not scale: the fifth collection is specialised today, but the reason for +this spec is that there will be a tenth. + +### D2 -- fixed allocation, or relevance-weighted + +Routing fixes *which* collections are searched. It does not fix the inversion above, +which is the fixed ten-per-collection cap. + +Option (a): leave the cap alone. Simple, and routing alone removes most of the waste. + +Option (b): give the 50-document budget out by score rather than by collection, so a +collection with nothing relevant contributes nothing even when it is searched. + +Recommendation: **(a) first, measured, then (b) separately**. They are independent, +and doing both at once makes it impossible to say which one moved the numbers. + +## How we will know it worked + +`bin/retrieval_baseline` exists precisely for this: it captures what each retriever +returns for a fixed question set and diffs two captures. Measured across two +identical runs BM25 is byte-identical on 80/80 question-collections, so a diff after +this change is a real difference in behaviour rather than noise. + +- `capture` before and after; the diff names every question whose documents changed +- the answer sweep must stay at 13/13 -- including the two variant questions, which + fail if routing sends them past `disease_variants`, and the species and release + questions, which fail if it stops routing to `live` +- the token and latency numbers above, re-measured + +That last one matters: the point of this change is a number going down, and it +should be reported as one. + +## Scope + +In: collection selection for the `reactome` source, defaulting to all; the +measurement above. + +Out: the relevance-weighted allocation of D2, which is its own change; `userguide`, +which has one collection; anything about the four Neo4j-backed collections being +Release95 while `disease_variants` came from the Release 97 download directory -- +that is recorded in the bundle's `provenance.json` and is spec 008's problem. diff --git a/src/data_generation/disease_variant/__init__.py b/src/data_generation/disease_variant/__init__.py index 109cdcf..dc48c5f 100644 --- a/src/data_generation/disease_variant/__init__.py +++ b/src/data_generation/disease_variant/__init__.py @@ -13,8 +13,10 @@ """ import csv +import json import os import shutil +from datetime import UTC, datetime from pathlib import Path from chromadb.api.client import SharedSystemClient @@ -143,6 +145,31 @@ def write_csv(tsv_path: Path, csv_path: Path) -> int: return len(rows) +def record_provenance(bundle: Path, source: Path, rows: int) -> None: + """Note where this collection came from, beside the bundle. + + The bundle directory is named for a release, and this collection is built + from a release file that may not be the same one -- the four Neo4j-backed + collections are Release95 while `disease_variant_ewas_mapping.tsv` came + from the Release 97 download directory. Nothing else in the bundle records + that, and `build_embeddings` already warns that a bundle whose contents + disagree with its path is unusable in a way nothing reports. + + Merged rather than overwritten, so each collection keeps its own entry. + """ + path = bundle / "provenance.json" + try: + known = json.loads(path.read_text()) + except (OSError, ValueError): + known = {} + known[COLLECTION] = { + "source": str(source), + "rows": rows, + "generated": datetime.now(UTC).isoformat(timespec="seconds"), + } + path.write_text(json.dumps(known, indent=2, sort_keys=True) + "\n") + + def generate_disease_variant_embeddings( embeddings_dir: str, tsv_path: str | Path, @@ -154,6 +181,7 @@ def generate_disease_variant_embeddings( csv_path = bundle / "csv_files" / f"{COLLECTION}.csv" count = write_csv(Path(tsv_path), csv_path) print(f" wrote {csv_path} ({count} variants)") + record_provenance(bundle, Path(tsv_path), count) # Chroma.from_documents appends to whatever is already in the directory, so # a second run would silently double the collection -- 12,588 documents, diff --git a/tests/data_generation/test_disease_variant.py b/tests/data_generation/test_disease_variant.py index 8f211e1..d8e4e96 100644 --- a/tests/data_generation/test_disease_variant.py +++ b/tests/data_generation/test_disease_variant.py @@ -1,6 +1,7 @@ """The disease_variants collection, built from a release file rather than Neo4j.""" import csv +import json from pathlib import Path import pytest @@ -150,3 +151,25 @@ def test_regenerating_replaces_rather_than_appends( assert first._collection.count() == 2 second = dv.generate_disease_variant_embeddings(str(tmp_path), tsv) assert second._collection.count() == 2, "regenerating must not append" + + +def test_provenance_records_which_release_file_was_used(tmp_path: Path) -> None: + # The bundle directory is named for a release; this collection can come + # from a different one, and nothing else in the bundle says so. + from data_generation.disease_variant import record_provenance + + record_provenance( + tmp_path, Path("/downloads/97/disease_variant_ewas_mapping.tsv"), 6294 + ) + written = json.loads((tmp_path / "provenance.json").read_text()) + assert written["disease_variants"]["rows"] == 6294 + assert "97" in written["disease_variants"]["source"] + + # Another collection's entry must survive. + (tmp_path / "provenance.json").write_text( + json.dumps({"reactions": {"source": "neo4j"}}) + ) + record_provenance(tmp_path, Path("/downloads/97/x.tsv"), 1) + written = json.loads((tmp_path / "provenance.json").read_text()) + assert written["reactions"]["source"] == "neo4j" + assert "disease_variants" in written From 237f166c25d14e0a47a38dfe28566fbc712ae90b Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Thu, 17 Sep 2026 03:15:06 +0000 Subject: [PATCH 06/10] Fix the two bugs that made the reactome bundle unbuildable and doubled `make` has been broken for this bundle. `metadata_columns` asked for a `species` column that the reactions and summations queries never return -- both filter on speciesName = "Homo sapiens" and then do not select it -- so generation died on the first upload with "Metadata column 'species' not found in CSV file". The Release95 bundle predates it, which is why nobody hit it. Adding it to the queries would have been worse: these collections pass no `content_columns`, so every CSV column is embedded, and a line reading "species: Homo sapiens" in all 52,000 documents is noise in every one of them. The second bug shipped. Chroma.from_documents appends to whatever is already in the directory, and the live Release95 bundle holds 33,498 reaction documents for 16,749 CSV rows -- exactly 2.00x, every reaction stored twice, because generation ran twice. Nothing reported it, and since the vector retriever overfetches and then de-duplicates on st_id, half that overfetch was being spent on duplicates. Regeneration now replaces, and clears chromadb's cached system client so the next write does not fail on a deleted sqlite file. Release 97 builds 17,004 reactions for 17,004 rows. Co-Authored-By: Claude Opus 5 --- src/data_generation/reactome/__init__.py | 34 ++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/src/data_generation/reactome/__init__.py b/src/data_generation/reactome/__init__.py index 9f159b5..03fdf2d 100644 --- a/src/data_generation/reactome/__init__.py +++ b/src/data_generation/reactome/__init__.py @@ -1,6 +1,8 @@ import os +import shutil from pathlib import Path +from chromadb.api.client import SharedSystemClient from langchain_community.vectorstores import Chroma from data_generation.disease_variant import generate_disease_variant_embeddings @@ -17,13 +19,28 @@ def upload_to_chromadb( hf_model: str | None = None, device: str | None = None, ) -> Chroma: + # `species` is deliberately absent from reactions and summations. Both + # queries filter on speciesName = "Homo sapiens" and never return it, so + # asking for it here killed generation on the first upload with + # "Metadata column 'species' not found in CSV file" -- `make` has been + # broken for the reactome bundle, and the Release95 bundle predates it. + # + # Adding it to the queries would be worse than removing it here: these + # collections pass no `content_columns`, so every CSV column is embedded, + # and a column that reads "species: Homo sapiens" in all 52,000 documents + # is noise in every one of them. + # + # `complexes` does return it, and its value is constant too -- all 38,967 + # rows are Homo sapiens, because that query filters on it as well. That is + # existing noise in an existing bundle rather than noise this adds, so it + # is left alone here; removing it means changing the query and rebuilding, + # which is a separate change with its own before-and-after measurement. metadata_columns: dict[str, list] = { "reactions": [ "st_id", "display_name", "pathway_id", "pathway_name", - "species", "input_id", "input_name", "output_id", @@ -31,7 +48,7 @@ def upload_to_chromadb( "catalyst_id", "catalyst_name", ], - "summations": ["st_id", "display_name", "labels", "species", "summation"], + "summations": ["st_id", "display_name", "labels", "summation"], "complexes": [ "st_id", "display_name", @@ -48,6 +65,19 @@ def upload_to_chromadb( ], } + # Chroma.from_documents appends to whatever is already in the directory. + # The shipped Release95 bundle has 33,498 reaction documents for 16,749 CSV + # rows -- exactly 2.00x, every reaction stored twice, because generation ran + # twice. Nothing reported it, and the vector retriever spends half its + # overfetch on duplicates. Regeneration replaces. + persist = Path(embeddings_dir) / embedding_table + if persist.exists(): + shutil.rmtree(persist) + # chromadb caches one system client per path; without this the next + # write fails with "attempt to write a readonly database". + SharedSystemClient.clear_system_cache() + print(f" replaced the previous {embedding_table} collection") + loader = MetaDataCSVLoader( file_path=file, metadata_columns=metadata_columns[embedding_table], From ebc431808d6f3a2c6f06146c55251f64306746ff Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Thu, 17 Sep 2026 03:15:06 +0000 Subject: [PATCH 07/10] Take the strict type-checking settings that cost nothing Measured each flag rather than guessing: warn_unused_configs and disallow_subclassing_any both report zero errors, so they cost nothing to hold and stop that slippage arriving unnoticed. warn_unused_ignores reported five, all stale "type: ignore" comments in one test file -- a stale ignore is a claim about the code that is no longer true. Left for a per-module ratchet, with the numbers: no_implicit_reexport (10), disallow_untyped_calls (22), disallow_any_generics (52). Those three bite because ignore_missing_imports is on repo-wide, so untyped third-party surfaces leak Any inward; fixing them is about stubs, not annotations. Co-Authored-By: Claude Opus 5 --- pyproject.toml | 7 +++++++ tests/agent/test_preprocess_concurrency.py | 10 +++++----- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1a29ced..24db7e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -217,6 +217,13 @@ disallow_incomplete_defs = true warn_no_return = true warn_return_any = true extra_checks = true +# Free: both already pass with zero errors, so they cost nothing to hold and +# stop the corresponding slippage from arriving unnoticed. +warn_unused_configs = true +disallow_subclassing_any = true +# A stale ignore is a claim about the code that is no longer true. There were +# five, all in one test file, all obsolete. +warn_unused_ignores = true # disallow_untyped_decorators stays off: chainlit's @cl.* decorators are # untyped upstream, so it only ever fires on bin/chat-chainlit.py. diff --git a/tests/agent/test_preprocess_concurrency.py b/tests/agent/test_preprocess_concurrency.py index d53cc36..671f777 100644 --- a/tests/agent/test_preprocess_concurrency.py +++ b/tests/agent/test_preprocess_concurrency.py @@ -34,17 +34,17 @@ async def run(_: Any) -> Any: def _builder(log: list[tuple[str, float, float]]) -> BaseGraphBuilder: """Build without __init__, which would construct real LLM chains.""" b = BaseGraphBuilder.__new__(BaseGraphBuilder) - b.rephrase_chain = _slow("rephrased", log, "rephrase") # type: ignore[assignment] - b.safety_checker = _slow( # type: ignore[assignment] + b.rephrase_chain = _slow("rephrased", log, "rephrase") + b.safety_checker = _slow( SafetyCheck(safety="true", reason_unsafe=""), log, "safety" ) - b.language_detector = _slow("en", log, "language") # type: ignore[assignment] + b.language_detector = _slow("en", log, "language") return b def test_safety_and_language_overlap() -> None: log: list[tuple[str, float, float]] = [] - state = BaseState(user_input="what is TP53?") # type: ignore[typeddict-item] + state = BaseState(user_input="what is TP53?") result = asyncio.run(_builder(log).preprocess(state, RunnableConfig())) @@ -78,7 +78,7 @@ def test_safety_and_language_overlap() -> None: def test_rephrase_still_precedes_the_safety_check() -> None: """Ordering that must not be lost: the safety check reads the rephrased text.""" log: list[tuple[str, float, float]] = [] - state = BaseState(user_input="what is TP53?") # type: ignore[typeddict-item] + state = BaseState(user_input="what is TP53?") asyncio.run(_builder(log).preprocess(state, RunnableConfig())) spans = {name: (start, end) for name, start, end in log} From 1a73741f931bec5c170a7d05396ab5c748f7718e Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Thu, 17 Sep 2026 03:42:18 +0000 Subject: [PATCH 08/10] Restore the enumeration rule: I removed it on two runs of evidence Yesterday's review called this rule dead weight, because both variant questions passed without it. Two runs was not enough evidence for a case this marginal. Against the Release 97 bundle the PTEN question fails without it and passes with it: without the rule: 1 of 3 runs named a variant with the rule: 3 of 3 Retrieval is not the difference -- both bundles return the same seven disease_variants documents with the same PTEN variants in them. The model has the variants in context either way and chooses whether to enumerate them or summarise at the level of the pathway, and the rule is what decides that. The router fix from yesterday was necessary and not sufficient: routing gets the question to the collection, and this gets the answer to name what the collection returned. 13/13 against Release 97 with the live MCP. Co-Authored-By: Claude Opus 5 --- src/retrievers/reactome/prompt.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/retrievers/reactome/prompt.py b/src/retrievers/reactome/prompt.py index d348939..4dc788e 100644 --- a/src/retrievers/reactome/prompt.py +++ b/src/retrievers/reactome/prompt.py @@ -18,6 +18,9 @@ 2. Inline citations required: Every factual statement must include ≥1 inline anchor citation in the format: display_name - If multiple entries support the same fact, cite them together (space-separated). 3. Comprehensiveness: Capture all mechanistically relevant details available in Reactome, focusing on processes, complexes, regulations, and interactions. + - When the question asks **which**, or asks you to **list** or **name** specific entities -- variants, complexes, participants, reactions -- name each one in the context individually. Do not answer at the level of the pathway that groups them. + - "Defective ABCA1 causes Tangier Disease" does not answer "which ABCA1 variants are there". If the context contains `ABCA1 W590S` and `ABCA1 C1417R`, those are the answer, and a pathway describing them collectively is the background to it. + - The narrative style above is for questions about mechanism. A question asking which things exist wants the things. 4. Tone & Style: - Write in a clear, engaging, and conversational tone. - Use accessible language while maintaining technical precision. From 945f4716379d0e58d2bbffe2400b58bf3e290301 Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Thu, 17 Sep 2026 03:56:00 +0000 Subject: [PATCH 09/10] Take three more strict settings, and say why the fourth is still off no_implicit_reexport: names now come from where they are defined rather than from whichever module imported them first -- ProfileName from agent.profile_names, TriggerEvent from util.config_yml.messages, SharedSystemClient from chromadb.api.shared_system_client. The one deliberate re-export is `logging` from util.logging, now declared in that module's __all__: importing it is what runs dictConfig, so callers take their logger through it on purpose. disallow_untyped_calls: on repo-wide, with three modules exempted. All 22 violations were calls into chainlit's unannotated API or boto3's Session -- nothing fixable by annotating our code -- so exempting those three keeps the rule meaningful everywhere else. bin/chat-chainlit.py carries the exemption as an inline directive because a hyphenated filename cannot be named in an overrides section; warn_unused_configs, enabled earlier today, is what caught that the section was doing nothing. disallow_any_generics stays off, but every bare builtin it flagged is fixed -- nine of them, including `-> tuple` and three `dict[str, list]`. The 43 left are framework generics, and writing Runnable[X, Y] without checking each chain would encode a wrong type rather than an absent one. 317 tests, ruff and mypy clean. Co-Authored-By: Claude Opus 5 --- bin/chat-chainlit.py | 10 +++++-- pyproject.toml | 30 +++++++++++++++++++ src/agent/graph.py | 3 +- src/data_generation/alliance/__init__.py | 2 +- src/data_generation/alliance/csv_generator.py | 2 +- .../disease_variant/__init__.py | 2 +- src/data_generation/metadata_csv_loader.py | 2 +- src/data_generation/reactome/__init__.py | 4 +-- src/data_generation/uniprot/__init__.py | 2 +- src/util/chainlit_helpers.py | 7 +++-- src/util/logging.py | 6 ++++ src/util/orcid_provider.py | 2 +- tests/evaluation/test_answer_sweep.py | 4 ++- tests/reactome_mcp/test_live_answer.py | 4 ++- 14 files changed, 65 insertions(+), 15 deletions(-) diff --git a/bin/chat-chainlit.py b/bin/chat-chainlit.py index fbbbae8..db0df97 100644 --- a/bin/chat-chainlit.py +++ b/bin/chat-chainlit.py @@ -1,3 +1,7 @@ +# mypy: disallow-untyped-calls=False +# chainlit ships no annotations for cl.user_session.get/set, cl.Message.send +# or get_data_layer. Not fixable here; it needs stubs upstream. The file is +# named with a hyphen, so it cannot be listed in [[tool.mypy.overrides]]. import os import chainlit as cl @@ -9,7 +13,8 @@ from langchain_community.callbacks import OpenAICallbackHandler from agent.graph import AgentGraph -from agent.profiles import ProfileName, get_chat_profiles +from agent.profile_names import ProfileName +from agent.profiles import get_chat_profiles from agent.profiles.base import OutputState from util.chainlit_helpers import ( PrefixedS3StorageClient, @@ -19,7 +24,8 @@ static_messages, update_search_results, ) -from util.config_yml import Config, TriggerEvent +from util.config_yml import Config +from util.config_yml.messages import TriggerEvent from util.logging import logging from util.orcid_provider import ORCIDOAuthProvider from util.secrets import ( diff --git a/pyproject.toml b/pyproject.toml index 24db7e9..fb57fc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -224,9 +224,39 @@ disallow_subclassing_any = true # A stale ignore is a claim about the code that is no longer true. There were # five, all in one test file, all obsolete. warn_unused_ignores = true +# Names now come from where they are defined rather than from whichever module +# happened to import them first: ProfileName from agent.profile_names, +# TriggerEvent from util.config_yml.messages, SharedSystemClient from +# chromadb.api.shared_system_client. The one deliberate re-export, `logging` +# from util.logging, is declared in that module's __all__ -- importing it is +# what runs dictConfig, so callers take the logger through it on purpose. +no_implicit_reexport = true +# Enabled repo-wide, with three modules exempted below. Every violation is a +# call into an untyped third-party API rather than anything we could annotate, +# so exempting those three keeps the rule meaningful everywhere else: new +# untyped calls cannot appear in ordinary code. +disallow_untyped_calls = true +# disallow_any_generics is the one strict setting still off. Every bare builtin +# it flagged is now parameterised; the 43 left are all framework generics -- +# Runnable (29), StateGraph (9), CompiledStateGraph (3), RunnableLambda and +# BasePromptTemplate. Turning it on means writing the real input and output +# type of every chain and graph, and a guess there encodes a type that is +# wrong rather than absent, which is worse than the bare form. It needs doing +# per chain, with the chain in front of you. # disallow_untyped_decorators stays off: chainlit's @cl.* decorators are # untyped upstream, so it only ever fires on bin/chat-chainlit.py. +# Untyped third-party surfaces, not our debt. chainlit ships no annotations for +# `cl.user_session.get/set`, `cl.Message.send` or `get_data_layer`, and boto3's +# `Session` is untyped without boto3-stubs. Nothing here is fixable by +# annotating our own code; it needs stubs upstream or a typed wrapper of our +# own, which is a change with its own design rather than a lint fix. +# bin/chat-chainlit.py carries the same exemption as an inline directive: its +# hyphenated name cannot appear in an overrides section. +[[tool.mypy.overrides]] +module = ["util.chainlit_helpers", "util.secrets"] +disallow_untyped_calls = false + # --- mypy baseline ------------------------------------------------------------- # Modules that do not yet meet the floor above. Each entry is debt: fix the module, # then delete its block. Do not add to this list without a TODO. diff --git a/src/agent/graph.py b/src/agent/graph.py index 561e983..d1804ce 100644 --- a/src/agent/graph.py +++ b/src/agent/graph.py @@ -15,7 +15,8 @@ from psycopg_pool import AsyncConnectionPool from agent.models import get_embedding, get_llm -from agent.profiles import ProfileName, create_profile_graphs +from agent.profile_names import ProfileName +from agent.profiles import create_profile_graphs from agent.profiles.base import InputState, OutputState from util.config_yml.models import LLMConfig from util.embedding_environment import EmbeddingEnvironment diff --git a/src/data_generation/alliance/__init__.py b/src/data_generation/alliance/__init__.py index 0038e51..aed6b7d 100644 --- a/src/data_generation/alliance/__init__.py +++ b/src/data_generation/alliance/__init__.py @@ -43,7 +43,7 @@ def upload_to_chromadb( hf_model: str | None = None, device: str | None = None, ) -> Chroma | None: - metadata_columns: dict[str, list] = { + metadata_columns: dict[str, list[str]] = { "genes": [ "Your Input", "Gene ID", diff --git a/src/data_generation/alliance/csv_generator.py b/src/data_generation/alliance/csv_generator.py index f159b65..c92263a 100644 --- a/src/data_generation/alliance/csv_generator.py +++ b/src/data_generation/alliance/csv_generator.py @@ -106,7 +106,7 @@ def get_genes(version: str, force: bool) -> str: return gene_csv -def generate_all_csvs(version: str, force: bool) -> tuple: +def generate_all_csvs(version: str, force: bool) -> tuple[str, ...]: files = [] # Download gene file diff --git a/src/data_generation/disease_variant/__init__.py b/src/data_generation/disease_variant/__init__.py index dc48c5f..7162bb5 100644 --- a/src/data_generation/disease_variant/__init__.py +++ b/src/data_generation/disease_variant/__init__.py @@ -19,7 +19,7 @@ from datetime import UTC, datetime from pathlib import Path -from chromadb.api.client import SharedSystemClient +from chromadb.api.shared_system_client import SharedSystemClient from langchain_chroma import Chroma from data_generation.embeddings import build_embeddings diff --git a/src/data_generation/metadata_csv_loader.py b/src/data_generation/metadata_csv_loader.py index 093ef50..5e6e6ab 100644 --- a/src/data_generation/metadata_csv_loader.py +++ b/src/data_generation/metadata_csv_loader.py @@ -91,7 +91,7 @@ def __read_file(self, csvfile: TextIOWrapper) -> list[Document]: # Skip lines starting with '#' valid_lines = (line for line in csvfile if not line.startswith("#")) - csv_reader: csv.DictReader = csv.DictReader(valid_lines, **self.csv_args) + csv_reader: csv.DictReader[str] = csv.DictReader(valid_lines, **self.csv_args) for i, row in enumerate(csv_reader): try: source = ( diff --git a/src/data_generation/reactome/__init__.py b/src/data_generation/reactome/__init__.py index 03fdf2d..d9bea22 100644 --- a/src/data_generation/reactome/__init__.py +++ b/src/data_generation/reactome/__init__.py @@ -2,7 +2,7 @@ import shutil from pathlib import Path -from chromadb.api.client import SharedSystemClient +from chromadb.api.shared_system_client import SharedSystemClient from langchain_community.vectorstores import Chroma from data_generation.disease_variant import generate_disease_variant_embeddings @@ -35,7 +35,7 @@ def upload_to_chromadb( # existing noise in an existing bundle rather than noise this adds, so it # is left alone here; removing it means changing the query and rebuilding, # which is a separate change with its own before-and-after measurement. - metadata_columns: dict[str, list] = { + metadata_columns: dict[str, list[str]] = { "reactions": [ "st_id", "display_name", diff --git a/src/data_generation/uniprot/__init__.py b/src/data_generation/uniprot/__init__.py index 6c3d9ff..37de125 100644 --- a/src/data_generation/uniprot/__init__.py +++ b/src/data_generation/uniprot/__init__.py @@ -15,7 +15,7 @@ def upload_to_chromadb( hf_model: str | None = None, device: str | None = None, ) -> Chroma: - metadata_columns: dict[str, list] = { + metadata_columns: dict[str, list[str]] = { "uniprot_data": [ "gene_names", "short_protein_name", diff --git a/src/util/chainlit_helpers.py b/src/util/chainlit_helpers.py index e6b5e58..15b75d0 100644 --- a/src/util/chainlit_helpers.py +++ b/src/util/chainlit_helpers.py @@ -10,7 +10,8 @@ from langchain_community.callbacks import OpenAICallbackHandler from tools.external_search.state import WebSearchResult -from util.config_yml import Config, TriggerEvent +from util.config_yml import Config +from util.config_yml.messages import TriggerEvent from util.config_yml.usage_limits import MessageRate _GUEST_METADATA_KEY = "_guest_metadata" @@ -80,7 +81,9 @@ def is_feature_enabled(config: Config | None, feature_id: str) -> bool: def save_openai_metrics(message_id: str, openai_cb: OpenAICallbackHandler) -> None: - openai_metrics: dict[str, dict] = cl.user_session.get("openai_metrics", {}) + openai_metrics: dict[str, dict[str, Any]] = cl.user_session.get( + "openai_metrics", {} + ) openai_metrics[message_id] = { prop: openai_cb.__dict__.get(prop, None) for prop in [ diff --git a/src/util/logging.py b/src/util/logging.py index d531876..8a7bf13 100644 --- a/src/util/logging.py +++ b/src/util/logging.py @@ -23,3 +23,9 @@ }, } logging.config.dictConfig(LOGGING_CONFIG) + +# Importing this module configures logging as a side effect, and callers write +# `from util.logging import logging` so that the configuration is guaranteed to +# have run before they take a logger. That re-export is the point of the module, +# so declare it rather than leaving it implicit. +__all__ = ["DEFAULT_LOG_LEVEL", "LOGGING_CONFIG", "logging"] diff --git a/src/util/orcid_provider.py b/src/util/orcid_provider.py index c1402b7..9a8eea7 100644 --- a/src/util/orcid_provider.py +++ b/src/util/orcid_provider.py @@ -27,7 +27,7 @@ def __init__(self) -> None: if prompt := self.get_prompt(): self.authorize_params["prompt"] = prompt - async def get_raw_token_response(self, code: str, url: str) -> dict: + async def get_raw_token_response(self, code: str, url: str) -> dict[str, Any]: payload = { "client_id": self.client_id, "client_secret": self.client_secret, diff --git a/tests/evaluation/test_answer_sweep.py b/tests/evaluation/test_answer_sweep.py index 8feb973..e297168 100644 --- a/tests/evaluation/test_answer_sweep.py +++ b/tests/evaluation/test_answer_sweep.py @@ -21,7 +21,9 @@ def __init__(self, answers: list[str | Exception]) -> None: self.answers = answers self.asked: list[str] = [] - async def ainvoke(self, question: str, *_args: object, **_kwargs: object) -> dict: + async def ainvoke( + self, question: str, *_args: object, **_kwargs: object + ) -> dict[str, object]: self.asked.append(question) answer = self.answers[min(len(self.asked) - 1, len(self.answers) - 1)] if isinstance(answer, Exception): diff --git a/tests/reactome_mcp/test_live_answer.py b/tests/reactome_mcp/test_live_answer.py index 1713adb..5fd72dc 100644 --- a/tests/reactome_mcp/test_live_answer.py +++ b/tests/reactome_mcp/test_live_answer.py @@ -48,7 +48,9 @@ async def ainvoke(self, messages: Any, *args: Any, **kwargs: Any) -> AIMessage: return self._replies.pop(0) if self._replies else AIMessage("out of replies") -def _call(name: str, call_id: str = "1", args: dict | None = None) -> dict: +def _call( + name: str, call_id: str = "1", args: dict[str, Any] | None = None +) -> dict[str, Any]: return {"name": name, "args": args or {}, "id": call_id, "type": "tool_call"} From 431d9ee2e85b0494ece8c1a16748c791a083f52a Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Thu, 17 Sep 2026 04:04:45 +0000 Subject: [PATCH 10/10] Record provenance for every collection, and settle 009's acceptance bar provenance.json covered disease_variants alone, which implied the other four were unknown rather than simply unrecorded. upload_to_chromadb now records every collection it builds, and the Release97 bundle is backfilled so all five carry their source and row count. Spec 009 gains the decision it was missing: the answer sweep is the pass/fail, retrieval_baseline is captured and its diff reported because Principle II requires the measurement, and no numeric document-loss threshold is set until there is one real measurement to set it from. A dropped document is what this change is for; a wrong answer is the defect. Any threshold chosen now would be invented, and a number that looks rigorous and means nothing is the failure this project keeps finding. Its weakness is written down rather than hidden: thirteen questions cannot cover five collections, so growing the tracked set is part of the work. Four of the five collections currently have no question that fails if routing stops searching them. Co-Authored-By: Claude Opus 5 --- specs/009-collection-routing/spec.md | 23 +++++++++++++++++++ .../disease_variant/__init__.py | 6 +++-- src/data_generation/reactome/__init__.py | 8 ++++++- 3 files changed, 34 insertions(+), 3 deletions(-) diff --git a/specs/009-collection-routing/spec.md b/specs/009-collection-routing/spec.md index 4c81246..c3f5a30 100644 --- a/specs/009-collection-routing/spec.md +++ b/specs/009-collection-routing/spec.md @@ -10,6 +10,25 @@ responses take to return, it would be nice to search the disease variant table if that is what they are asking about."* +## Clarifications + +### Session 2026-09-17 + +- Q: When routing skips a collection that would have contributed a document, how much recall loss should reject the change? → A: Gate on answers, measure documents. No numeric document-loss threshold is set until there is one real measurement to set it from. + +The answer sweep staying green is the pass/fail. `bin/retrieval_baseline` is captured +before and after and its diff reported, because Principle II requires the measurement +-- but a changed document set does not by itself block, since dropping documents is +what this change is for. A wrong answer is the defect; a dropped document is data. + +Any threshold chosen today would be invented: there is no evidence yet on how often +the classifier picks wrong, and a number that looks rigorous and means nothing is the +failure this project keeps finding. + +The weakness in that, stated rather than hidden: thirteen questions is a thin net, and +routing could break something nobody tracks. So growing the tracked set is part of +this work rather than a follow-up -- see Success Criteria. + ## The cost is real, and it was measured Every collection in the bundle is searched for every question. `retrieve_documents` @@ -96,6 +115,10 @@ this change is a real difference in behaviour rather than noise. fail if routing sends them past `disease_variants`, and the species and release questions, which fail if it stops routing to `live` - the token and latency numbers above, re-measured +- the tracked question set grown before the change lands, not after. Thirteen + questions cannot cover five collections; each collection needs at least one + question that fails if routing stops searching it. `disease_variants` has two + already, and the other four have none. That last one matters: the point of this change is a number going down, and it should be reported as one. diff --git a/src/data_generation/disease_variant/__init__.py b/src/data_generation/disease_variant/__init__.py index 7162bb5..74e8bc2 100644 --- a/src/data_generation/disease_variant/__init__.py +++ b/src/data_generation/disease_variant/__init__.py @@ -145,7 +145,9 @@ def write_csv(tsv_path: Path, csv_path: Path) -> int: return len(rows) -def record_provenance(bundle: Path, source: Path, rows: int) -> None: +def record_provenance( + bundle: Path, source: Path, rows: int, collection: str = COLLECTION +) -> None: """Note where this collection came from, beside the bundle. The bundle directory is named for a release, and this collection is built @@ -162,7 +164,7 @@ def record_provenance(bundle: Path, source: Path, rows: int) -> None: known = json.loads(path.read_text()) except (OSError, ValueError): known = {} - known[COLLECTION] = { + known[collection] = { "source": str(source), "rows": rows, "generated": datetime.now(UTC).isoformat(timespec="seconds"), diff --git a/src/data_generation/reactome/__init__.py b/src/data_generation/reactome/__init__.py index d9bea22..690716d 100644 --- a/src/data_generation/reactome/__init__.py +++ b/src/data_generation/reactome/__init__.py @@ -5,7 +5,10 @@ from chromadb.api.shared_system_client import SharedSystemClient from langchain_community.vectorstores import Chroma -from data_generation.disease_variant import generate_disease_variant_embeddings +from data_generation.disease_variant import ( + generate_disease_variant_embeddings, + record_provenance, +) from data_generation.embeddings import build_embeddings from data_generation.metadata_csv_loader import MetaDataCSVLoader from data_generation.reactome.csv_generator import generate_all_csvs @@ -84,6 +87,9 @@ def upload_to_chromadb( encoding="utf-8", ) docs = loader.load() + # Recorded for every collection, not just disease_variants: a provenance + # file that covers one of five implies the other four are unknown. + record_provenance(Path(embeddings_dir), Path(file), len(docs), embedding_table) embeddings_instance = build_embeddings(hf_model, device, chunk_size=400) return Chroma.from_documents(