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/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/pyproject.toml b/pyproject.toml index 1a29ced..fb57fc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -217,9 +217,46 @@ 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 +# 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/specs/008-disease-variant-embeddings/spec.md b/specs/008-disease-variant-embeddings/spec.md new file mode 100644 index 0000000..e1f54d0 --- /dev/null +++ b/specs/008-disease-variant-embeddings/spec.md @@ -0,0 +1,168 @@ +# Feature Specification: Disease and Variant Embeddings + +**Feature Branch**: `spec/disease-variant-embeddings` + +**Created**: 2026-09-16 + +**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 +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 + +Both were delegated: *"I think you should decide on the columns ... make your best +design and go ahead."* + +### D1 -- what is embedded, and what is metadata: decided + +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. + +Embedded: `gene`, `variant`, `protein`, `residue_change`, `mutation_type`, `disease`, +`reaction`, `functional_status`, `disease_pathway`, `normal_reaction`, +`normal_pathway`, `normal_process`. A document reads: + +``` +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... +``` + +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 + +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. + +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 + +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. + +## Result, measured + +Built into a scratch bundle and asked through the real retriever: + +| | before | after | +|---|---|---| +| 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 + +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/specs/009-collection-routing/spec.md b/specs/009-collection-routing/spec.md new file mode 100644 index 0000000..c3f5a30 --- /dev/null +++ b/specs/009-collection-routing/spec.md @@ -0,0 +1,134 @@ +# 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."* + +## 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` +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 +- 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. + +## 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/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/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/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 new file mode 100644 index 0000000..74e8bc2 --- /dev/null +++ b/src/data_generation/disease_variant/__init__.py @@ -0,0 +1,213 @@ +"""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 json +import os +import shutil +from datetime import UTC, datetime +from pathlib import Path + +from chromadb.api.shared_system_client import SharedSystemClient +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. + # `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": "cross_references", + "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", + "cross_references", + "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 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 + 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, + 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)") + 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, + # 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, + 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/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 fc66d95..690716d 100644 --- a/src/data_generation/reactome/__init__.py +++ b/src/data_generation/reactome/__init__.py @@ -1,8 +1,14 @@ import os +import shutil from pathlib import Path +from chromadb.api.shared_system_client import SharedSystemClient from langchain_community.vectorstores import Chroma +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 @@ -16,13 +22,28 @@ def upload_to_chromadb( hf_model: str | None = None, device: str | None = None, ) -> Chroma: - metadata_columns: dict[str, list] = { + # `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[str]] = { "reactions": [ "st_id", "display_name", "pathway_id", "pathway_name", - "species", "input_id", "input_name", "output_id", @@ -30,7 +51,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", @@ -47,12 +68,28 @@ 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], 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( @@ -70,6 +107,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 +145,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/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/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..93c8e1e 100644 --- a/src/retrievers/reactome/metadata_info.py +++ b/src/retrievers/reactome/metadata_info.py @@ -11,11 +11,54 @@ "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. 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( + 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/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/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/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) 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} diff --git a/tests/data_generation/test_disease_variant.py b/tests/data_generation/test_disease_variant.py new file mode 100644 index 0000000..d8e4e96 --- /dev/null +++ b/tests/data_generation/test_disease_variant.py @@ -0,0 +1,175 @@ +"""The disease_variants collection, built from a release file rather than Neo4j.""" + +import csv +import json +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() + + +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" + + +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 diff --git a/tests/evaluation/test_answer_sweep.py b/tests/evaluation/test_answer_sweep.py index a230abb..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): @@ -162,3 +164,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 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"}