diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9ef33fd3..c443ce34 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -49,6 +49,13 @@ jobs: - name: Set up Python and Poetry uses: ./.github/actions/install_python_poetry + # The same data the Dockerfile installs. BM25 tokenises with + # word_tokenize(..., language="english"), so any test that drives + # retrieval needs punkt_tab -- without it the suite passes locally, + # where a developer has downloaded it, and fails here. + - name: Fetch NLTK data used by BM25 + run: poetry run python -m nltk.downloader punkt_tab + - name: Run tests run: poetry run pytest diff --git a/specs/009-collection-routing/contracts/intent_classifier.md b/specs/009-collection-routing/contracts/intent_classifier.md new file mode 100644 index 00000000..c0f07169 --- /dev/null +++ b/specs/009-collection-routing/contracts/intent_classifier.md @@ -0,0 +1,43 @@ +# Contract: Intent Classifier Output + +The classifier is the only interface this feature changes. One LLM call per question, +already made. + +## Before + +```json +{"source": "reactome"} +``` + +## After + +```json +{"source": "reactome", "collections": ["disease_variants"]} +``` + +`collections` is optional and defaults to `[]`. Every consumer treats `[]` as "search +everything", so an older prompt, a model that omits the field, or a failed parse all +produce today's behaviour rather than a narrower search. + +## Constraints + +- Names must be collection directories present in the installed bundle. Validity is + decided against `list_chroma_subdirectories()`, not a literal list +- Unknown names do not narrow the search. They are logged at WARNING and the whole + selection falls back to all collections +- `collections` applies only when `source == "reactome"`. `userguide` has one + collection; `live` does not use the vector store at all + +## Prompt input + +The per-collection descriptions already written in +`src/retrievers/reactome/metadata_info.py` (`reactome_descriptions_info`), which +exist for routing and are currently read only by `bin/retrieval_baseline`. They +become a serving input, so they must stay accurate as collections are added -- +already true of `disease_variants`, described when it was added. + +## Compatibility + +A deployment running an older prompt against newer code returns no `collections`, +gets `[]`, and searches everything. The feature degrades to the current system rather +than to a broken one. diff --git a/specs/009-collection-routing/data-model.md b/specs/009-collection-routing/data-model.md new file mode 100644 index 00000000..fa526b3d --- /dev/null +++ b/specs/009-collection-routing/data-model.md @@ -0,0 +1,64 @@ +# Phase 1 Data Model: Collection Routing + +## Selection + +The unit this feature adds: which collections a single question should search. + +| Field | Type | Notes | +|---|---|---| +| `collections` | `list[str]` | Collection directory names, e.g. `["disease_variants", "summations"]`. Empty means "all", which is also the fallback for every failure | + +Not a new class. It is a field on the classifier's existing structured output and a +key in `RunnableConfig["configurable"]`, because adding a type for it would mean +threading that type through three layers that do not otherwise know about each other. + +### Validity + +Decided against the **live bundle**, never a literal: + +```python +valid = set(list_chroma_subdirectories(embeddings_directory)) +``` + +Principle V: the same list that decides what is searched decides what is nameable, so +the two cannot drift. A collection added to a bundle is immediately selectable; one +removed cannot be selected. + +### Resolution rules + +In order. Every failure path widens the search, never narrows it. + +| Input | Result | Why | +|---|---|---| +| Empty or absent | all collections | The default, and what any failure degrades to | +| All names valid | those collections | The feature | +| Some names unknown | **all** collections, WARNING naming the unknown ones | The prompt and bundle disagree; narrowing on a misunderstanding is worse than not routing. Principle IV | +| All names unknown | all collections, WARNING | As above | +| Names valid but none match the question well | those collections | Not detectable here. This is what the measurement is for, and what the sweep catches | + +## QueryIntent (existing, extended) + +```python +class QueryIntent(BaseModel): + source: SourceName + collections: list[str] = [] # new; empty means all +``` + +The default matters: a model that omits the field, an older prompt, or a parse +failure all produce `[]`, which searches everything. The feature cannot fail closed. + +## Where it travels + +``` +intent_classifier ──> QueryIntent.collections + │ +ReactToMeState["collections"] (preprocess, beside active_sources) + │ +RunnableConfig["configurable"]["collections"] (generate_answer) + │ +HybridRetriever.{retrieve,aretrieve}_documents (filters collection_retrievers) +``` + +The retriever is constructed once at startup and reads the selection per call. See +research.md R1 for why it is not a constructor argument: `BM25Retriever` is built over +17,004 documents for `reactions` alone, which is startup work. diff --git a/specs/009-collection-routing/plan.md b/specs/009-collection-routing/plan.md new file mode 100644 index 00000000..e06a35de --- /dev/null +++ b/specs/009-collection-routing/plan.md @@ -0,0 +1,116 @@ +# Implementation Plan: Searching Only the Collections a Question Needs + +**Branch**: `009-collection-routing` | **Date**: 2026-09-17 | **Spec**: [spec.md](./spec.md) + +**Input**: Feature specification from `/specs/009-collection-routing/spec.md` + +## Summary + +Every collection in the bundle is searched for every question, and each contributes +a fixed ten documents whether or not it had anything to say. Adding a fifth cost +**+2,376 context tokens and +3.4s** on a question about CDK5 that has nothing to do +with variants. + +The intent classifier already makes one LLM call per question and returns a source. +It will also return which collections to search. That adds no call and no latency, +and the per-collection descriptions it needs already exist in +`reactome_descriptions_info` -- written for routing, currently read only by +`bin/retrieval_baseline`. + +## Technical Context + +**Language/Version**: Python 3.12 + +**Primary Dependencies**: langchain 1.x, langchain-chroma, chromadb 0.6.3, pydantic 2 + +**Storage**: Chroma collections on disk under `embeddings//reactome//` + +**Testing**: pytest. `bin/retrieval_baseline` for retrieval diffs, `bin/answer-sweep` +for end-to-end answers + +**Target Platform**: Linux container behind FastAPI/Chainlit + +**Project Type**: single project + +**Performance Goals**: reduce context tokens and retrieval latency on questions that +do not need every collection. Today's baseline, measured: 5 collections = 50 docs, +9,437 tokens, 14.9s; 4 collections = 40 docs, 7,061 tokens, 11.5s + +**Constraints**: no additional LLM call; a classifier failure must degrade to +today's behaviour, never to a narrower search + +**Scale/Scope**: 5 collections today, 112,000 documents; the feature exists because +that number will grow + +## Constitution Check + +| Principle | Gate | How this plan satisfies it | +|---|---|---| +| I. Verify the path a user takes | Tests must go through the agent, not only the retriever | The sync retriever is not the served path (`aretrieve_documents` is). Acceptance runs `bin/answer-sweep`, which drives the whole graph | +| II. Measure retrieval changes | `bin/retrieval_baseline` capture before and after, diff reported | Phase 1 captures the baseline **before** any code changes, so the comparison exists to be made | +| III. Characterization tests | Current behaviour pinned before changing it | A test asserting all collections are searched today, so switching to selection is a deliberate edit of test and code together | +| IV. Fail loudly | Config that cannot be honoured stops rather than substitutes | A classifier naming a collection that is not in the bundle is a bug in the prompt or the bundle. It is logged at WARNING and the selection falls back to all collections -- never silently dropped, never a narrower search | +| V. Derive from the source of truth | No hand-synchronised lists | The selectable collections are derived from `list_chroma_subdirectories` on the live bundle, not a literal. `reactome_descriptions_info` is the one place a collection is described | + +**Gate result**: pass. No violations to justify. + +### The gate that is not yet met + +The spec records it: thirteen sweep questions cannot cover five collections, and +**four of the five have no question that fails if routing stops searching them**. +Principle I is only satisfied once they do. That is task work in Phase 1, before the +routing change lands, not after. + +## Project Structure + +### Documentation (this feature) + +``` +specs/009-collection-routing/ +├── spec.md +├── plan.md # this file +├── research.md # Phase 0 +├── data-model.md # Phase 1 +├── contracts/ +│ └── intent_classifier.md +├── quickstart.md +└── tasks.md # /speckit-tasks +``` + +### Source Code (repository root) + +``` +src/ +├── agent/tasks/intent_classifier.py # gains the collection selection +├── agent/profiles/react_to_me.py # threads the selection to the RAG +├── retrievers/ +│ ├── csv_chroma.py # HybridRetriever filters by selection +│ └── reactome/metadata_info.py # descriptions become a serving input +└── evaluation/answer_sweep.py # per-collection questions + +tests/ +├── agent/test_intent_classifier_sources.py +├── retrievers/test_collection_selection.py # new +└── retrievers/test_sync_async_equivalence.py # both paths must filter alike +``` + +## Complexity Tracking + +| Addition | Current need | Why the simpler option is insufficient | +|---|---|---| +| Collection names in the classifier's structured output | Selection must cost no extra latency | A second LLM call doubles the routing cost, which is what the feature exists to reduce | +| Selection passed through `ReactToMeState` | The retriever is built once at startup, per profile | Rebuilding a retriever per question would re-read every BM25 index on every message | + +## Phase 0: Research + +See [research.md](./research.md). + +## Phase 1: Design + +See [data-model.md](./data-model.md), [contracts/](./contracts/), +[quickstart.md](./quickstart.md). + +### Post-design constitution re-check + +Unchanged: pass. The design adds no new LLM call, derives the collection list from +the bundle rather than a literal, and fails toward the wider search. diff --git a/specs/009-collection-routing/quickstart.md b/specs/009-collection-routing/quickstart.md new file mode 100644 index 00000000..ce4f0492 --- /dev/null +++ b/specs/009-collection-routing/quickstart.md @@ -0,0 +1,65 @@ +# Quickstart: Validating Collection Routing + +## Prerequisites + +- An installed Release97 bundle (`./bin/embeddings_manager which`) +- `OPENAI_API_KEY` +- A reachable MCP server for the two live questions, or accept them as skipped + +## 1. Capture the baseline BEFORE changing anything + +Principle II: the comparison has to exist before the change. + +```bash +./bin/retrieval_baseline capture --out before.json +``` + +## 2. Record the cost the feature exists to reduce + +On a question that needs no variant data: + +```bash +./bin/answer-sweep --only "CDK5" +``` + +Known baseline, 2026-09-17: 5 collections, 50 documents, 9,437 context tokens, +14.9s retrieval. Four collections was 40, 7,061 and 11.5s. + +## 3. After the change + +```bash +./bin/retrieval_baseline capture --out after.json +./bin/retrieval_baseline compare before.json after.json +``` + +The diff is **reported, not gated**: dropping documents is what this change does. It +exits non-zero when anything changed, which is information here rather than failure. + +## 4. The gate + +```bash +./bin/answer-sweep +``` + +Must stay green. Two questions fail if variant questions stop reaching +`disease_variants`; the species and release questions fail if `live` routing breaks. + +**This is the pass/fail.** See spec.md "Clarifications" for why no numeric +document-loss threshold is set. + +## 5. The part that does not exist yet + +Four of five collections have **no** question that fails if routing stops searching +them. Adding those is task work before the routing change lands -- until then the +gate has a hole exactly where this feature could break things. + +## 6. Verify through the path a user takes + +Principle I. The served path is async; `retrieval_baseline` drives the sync one. + +```bash +python -m pytest tests/retrievers/test_sync_async_equivalence.py +``` + +Both paths must filter by the same selection. A change that routes only the sync path +would pass every measurement above and serve unrouted results. diff --git a/specs/009-collection-routing/research.md b/specs/009-collection-routing/research.md new file mode 100644 index 00000000..8488d9cb --- /dev/null +++ b/specs/009-collection-routing/research.md @@ -0,0 +1,91 @@ +# Phase 0 Research: Collection Routing + +## R1. How does a per-question selection reach a retriever built once? + +**The problem, from the code rather than assumed.** `ReactToMeProfile.__init__` builds +the RAG chain once and holds it in `self.rags["reactome"]`. Inside, +`create_retrieval_chain(retriever=..., combine_docs_chain=...)` closes over a single +`HybridRetriever` instance. The retriever's entry point is +`_get_relevant_documents(query: str, *, run_manager)` -- a string, with no room for +"and search only these collections". + +Rebuilding per question is not an option: `HybridRetriever.from_subdirectory` reads +every collection's CSV and constructs a `BM25Retriever` over it. For `reactions` alone +that is 17,004 documents tokenised with `word_tokenize`. That is startup work, not +per-message work. + +**Options considered** + +| Option | How | Why not / why | +|---|---|---| +| A. `RunnableConfig["configurable"]` | The profile calls `rag.ainvoke(..., config)`; the retriever reads `config["configurable"]["collections"]` | `config` already flows to the retriever -- `postprocess` reads `config["configurable"].get("enable_postprocess")` today, so the mechanism is in use in this codebase already | +| B. `ConfigurableField` | Declare the field configurable and bind per call | Works, but pydantic-configurable retrievers re-validate on every bind; A gets the same result with machinery already present | +| C. contextvar | Set around the call | Invisible coupling, and wrong under concurrency if it ever leaks across tasks | +| D. Encode in the query string | Prefix the question | Reaches BM25 and the query expander as text. The codebase already carries a comment warning against exactly this for `detected_language` | + +**Decision: A.** It reuses a path this repo already relies on and adds no new +LangChain surface. + +**Rationale.** `src/agent/profiles/base.py` reads `config["configurable"]` in +`postprocess`, so configurable state is established practice here, and the failure +mode is a missing key -- which defaults to "all collections" and is therefore safe. + +## R2. Can the classifier name collections reliably? + +Unknown, and the plan does not depend on it being perfect. The measurement decides. + +What is known: the same call already routes `reactome` / `userguide` / `live` +correctly for the thirteen sweep questions, and sharpening one rule on 2026-09-17 +fixed the variant questions without breaking the live ones. A per-collection +description already exists for each collection in `reactome_descriptions_info`. + +**Decision**: ask for collections in the same structured output, and treat an empty +or unparseable list as "all". Measure with `bin/retrieval_baseline` before deciding +whether it is good enough. + +## R3. What happens when the classifier names a collection that does not exist? + +It means the prompt and the bundle disagree -- a new collection added without a +description, or a renamed one. Principle IV says configuration that cannot be +honoured must not substitute something plausible. + +**Decision**: log at WARNING naming the unknown collection, and fall back to searching +**all** collections for that question. + +**Rationale.** The alternatives are worse. Dropping the unknown name silently narrows +the search for a reason nobody can see. Raising takes the chatbot down for what is a +recoverable prompt/bundle mismatch. Falling back to all is the only option whose +failure mode is today's behaviour. + +The selectable set is derived from `list_chroma_subdirectories()` on the live bundle, +so "does not exist" is decided against what is actually installed, not a literal list +that would itself drift (Principle V). + +## R4. What is the baseline, and when is it captured? + +Captured **before** any code change, or there is nothing to compare against. + +Already measured on 2026-09-17, on *"What does CDK5 phosphorylate in Alzheimer +disease?"*, a question needing no variant data: + +| | docs | context tokens | retrieval | +|---|---|---|---| +| 4 collections | 40 | 7,061 | 11.5s | +| 5 collections | 50 | 9,437 | 14.9s | + +And the inversion: `disease_variants` contributed **15%** of context to that question +and **5%** to the ABCA1 question it exists for, because every collection gets ten +documents regardless of quality. + +**Decision**: `bin/retrieval_baseline capture` over the fixed question set on the +Release97 bundle is task 1, before anything else. + +## R5. Which path must be measured? + +The served path is `aretrieve_documents`; `bin/retrieval_baseline` drives +`retrieve_documents`. They are separate implementations and were verified equivalent +on 2026-09-17 (PR #227), so measuring the sync path is currently valid. + +**Decision**: the equivalence test is a prerequisite of this work, not a nice-to-have. +Both paths must filter by the same selection, and the existing test must be extended +to assert that -- otherwise the measurement stops describing what users get. diff --git a/specs/009-collection-routing/spec.md b/specs/009-collection-routing/spec.md index c3f5a30d..33e83ba3 100644 --- a/specs/009-collection-routing/spec.md +++ b/specs/009-collection-routing/spec.md @@ -4,7 +4,8 @@ **Created**: 2026-09-17 -**Status**: Draft. Two decisions (D1, D2) for the team. +**Status**: Planned. D1, D2 and the acceptance bar are all decided and recorded +below; plan.md, data-model.md and tasks.md exist and depend on them. **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 @@ -79,9 +80,14 @@ instead of all of them -- a filter on `self.collection_retrievers.items()`. ## Decisions +*Both were left as recommendations after the Clarifications session, which settled +only the acceptance bar. `/speckit-analyze` caught that plan.md and data-model.md had +already built on them as though decided -- so they are recorded as decisions here, +on 2026-09-17, rather than left to be inferred from downstream documents.* + ### D1 -- what happens when the classifier is unsure -Recommendation: **select all collections**. A wrong selection costs recall silently, +**Decided: 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. @@ -100,7 +106,7 @@ Option (a): leave the cap alone. Simple, and routing alone removes most of the w 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, +**Decided: (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 diff --git a/specs/009-collection-routing/tasks.md b/specs/009-collection-routing/tasks.md new file mode 100644 index 00000000..5478aa78 --- /dev/null +++ b/specs/009-collection-routing/tasks.md @@ -0,0 +1,89 @@ +# Tasks: Searching Only the Collections a Question Needs + +**Feature**: 009-collection-routing | **Plan**: [plan.md](./plan.md) | **Spec**: [spec.md](./spec.md) + +Test tasks are included. Principle II requires a before-and-after measurement on any +change to what reaches the LLM, and Principle III requires current behaviour pinned +before it is changed. + +## Phase 1: Setup + +- [ ] T001 Create branch `009-collection-routing` from main and set `.specify/feature.json` to `specs/009-collection-routing` +- [ ] T002 Capture the retrieval baseline before any code change: `./bin/retrieval_baseline capture --out specs/009-collection-routing/before.json` against the installed Release97 bundle +- [ ] T003 [P] Record the current cost in `specs/009-collection-routing/quickstart.md`: documents, context tokens and retrieval seconds for one question needing no variant data + +## Phase 2: Foundational (blocks every user story) + +**The gate has a hole.** Four of the five collections have no question that fails if +routing stops searching them. Routing must not land before this closes, or the +acceptance criterion cannot detect the failure the feature can cause. + +- [ ] T004 [P] Add a `summations`-dependent question to `EXPECTATIONS` in src/evaluation/answer_sweep.py, with a `must`/`must_match` that fails if that collection is not searched +- [ ] T005 [P] Add a `complexes`-dependent question to `EXPECTATIONS` in src/evaluation/answer_sweep.py +- [ ] T006 [P] Add an `ewas`-dependent question to `EXPECTATIONS` in src/evaluation/answer_sweep.py +- [ ] T007 [P] Add a `reactions`-dependent question to `EXPECTATIONS` in src/evaluation/answer_sweep.py +- [ ] T008 Verify each new question FAILS when its collection is removed from the bundle copy, and passes with it present; record the evidence in the PR +- [ ] T009 Pin current behaviour: a characterization test in tests/retrievers/test_collection_selection.py asserting that with no selection every collection in the bundle is searched +- [ ] T010 Run `./bin/answer-sweep` against Release97 and confirm green before any behaviour change + +## Phase 3: User Story 1 — a question searches only the collections it needs (P1) + +**Goal**: cut context tokens and retrieval latency on questions that do not need every +collection, with no extra LLM call. + +**Independent test**: `./bin/answer-sweep` stays green while the measured context +tokens for a question needing one collection fall relative to `before.json`. + +- [ ] T011 [US1] Add `collections: list[str] = []` to `QueryIntent` in src/agent/tasks/intent_classifier.py, defaulting to empty so an omitted field means "all" +- [ ] T012 [US1] Extend the classifier prompt in src/agent/tasks/intent_classifier.py to name selectable collections, sourced from `reactome_descriptions_info` rather than a literal list +- [ ] T013 [P] [US1] Add `resolve_collections(selected, available)` to src/retrievers/csv_chroma.py implementing data-model.md: empty means all, unknown names log WARNING and return all +- [ ] T014 [P] [US1] Unit-test `resolve_collections` in tests/retrievers/test_collection_selection.py for empty, all-valid, some-unknown and all-unknown, asserting every failure widens rather than narrows +- [ ] T015 [US1] Filter `self.collection_retrievers` by the selection in `retrieve_documents` in src/retrievers/csv_chroma.py, reading it from `RunnableConfig["configurable"]["collections"]` +- [ ] T016 [US1] Apply the identical filter in `aretrieve_documents` in src/retrievers/csv_chroma.py — this is the served path +- [ ] T017 [US1] Extend tests/retrievers/test_sync_async_equivalence.py to assert both paths honour the same selection, and confirm it fails when only one is filtered +- [ ] T018 [US1] Carry `collections` on `ReactToMeState` in src/agent/profiles/react_to_me.py, set in `preprocess` beside `active_sources` +- [ ] T019 [US1] Pass the selection into the RAG call via `config["configurable"]` in `generate_answer` in src/agent/profiles/react_to_me.py +- [ ] T020 [US1] Verify through the agent, not the retriever: a test that a variant question routes to `disease_variants` and a userguide question does not, per Principle I + +## Phase 4: Measurement and acceptance + +- [ ] T021 Capture `./bin/retrieval_baseline capture --out specs/009-collection-routing/after.json` and `compare` it with before.json; report the diff in the PR as information, not as a gate +- [ ] T022 Re-measure context tokens and retrieval seconds for the same question as T003 and state the change as a number +- [ ] T023 Run `./bin/answer-sweep` against Release97 with live MCP; it must be green including all nine collection-dependent questions +- [ ] T024 Record in the PR how often the classifier selected a subset, and how often it named an unknown collection + +## Phase 5: Polish + +- [ ] T025 [P] Update the Status line and add a Results section to specs/009-collection-routing/spec.md with the measured before/after +- [ ] T026 [P] Note in src/retrievers/reactome/metadata_info.py that `reactome_descriptions_info` is now a serving input, so an inaccurate description degrades routing +- [ ] T027 Run `/speckit-analyze` across spec, plan and tasks and resolve anything CRITICAL or HIGH + +## Dependencies + +``` +Phase 1 (T001-T003) + ↓ +Phase 2 (T004-T010) ← blocking: the gate must detect routing mistakes first + ↓ +Phase 3 (T011-T020) ← the feature + ↓ +Phase 4 (T021-T024) ← acceptance + ↓ +Phase 5 (T025-T027) +``` + +Within Phase 3: T011 → T012; T013 → T014, T015, T016; T015 and T016 → T017; T018 → T019 → T020. + +## Parallel opportunities + +- T004-T007: four sweep questions, four independent edits to the same list — write together, run T008 once +- T013 and T014 alongside T011 and T012: `resolve_collections` is pure and does not depend on the prompt +- T025 and T026: different files + +## Implementation strategy + +**MVP is Phase 2 plus Phase 3.** Phase 2 alone is worth landing on its own: it closes +a hole in the gate that exists today, independently of whether routing is ever built. + +Ship Phase 2 as its own PR. If routing then turns out to cost more recall than it +saves, that PR still stands on its own merit. diff --git a/tests/retrievers/test_sync_async_equivalence.py b/tests/retrievers/test_sync_async_equivalence.py new file mode 100644 index 00000000..b0020705 --- /dev/null +++ b/tests/retrievers/test_sync_async_equivalence.py @@ -0,0 +1,167 @@ +"""The async retrieval path must return what the sync one does. + +`HybridRetriever` implements retrieval twice: `retrieve_documents` and +`aretrieve_documents`, the second gathering coroutines so the collections are +queried concurrently. They are separate implementations of the same fusion, and +nothing exercised the async one at all. + +That asymmetry matters beyond the usual duplication argument: the application +serves through the async path, while `bin/retrieval_baseline` -- the tool the +constitution names for measuring retrieval changes -- drives the sync one. If +they drift, every measurement is of a path no user takes, and Principle II's +before-and-after comparison measures the wrong thing. + +Uses DeterministicFakeEmbedding rather than FakeEmbeddings: the latter returns a +fresh random vector per call, so the same query embeds differently on the second +run and the two paths appear to disagree when they do not. That false result is +what prompted this test. + +What this does NOT cover, measured rather than assumed: the per-collection cap. +Deterministic fake vectors carry no relation to the text, so every query retrieves +much the same set, fusion lands on exactly `max_documents_per_collection` however +many queries are given, and changing the cap on one path alone stays invisible. +Tried at 12 and 60 documents per collection and with one, three and seven +queries; fusion produced ten per collection every time. Covering the cap needs +real embeddings and therefore an installed bundle, which is what +`bin/retrieval_baseline` is for. +""" + +import asyncio +import csv +from pathlib import Path + +import pytest + +pytest.importorskip("langchain_chroma") + +from langchain_chroma import Chroma # noqa: E402 +from langchain_core.callbacks import ( # noqa: E402 + AsyncCallbackManagerForRetrieverRun, + CallbackManagerForRetrieverRun, +) +from langchain_core.documents import Document # noqa: E402 +from langchain_core.embeddings import DeterministicFakeEmbedding # noqa: E402 +from langchain_core.language_models.fake_chat_models import ( # noqa: E402 + FakeListChatModel, +) + +from retrievers.csv_chroma import HybridRetriever, chroma_settings # noqa: E402 + + +# Large enough that the per-collection cap actually binds. At twelve it did +# not: fusion produced fewer documents than the cap, so a test that changed +# the cap on one path only still passed. +def _require_bm25_tokenizer() -> None: + """BM25 tokenises with nltk's word_tokenize, which needs punkt_tab. + + CI installs it, matching the Dockerfile. A developer who has not downloaded + it should get a skip saying so rather than a LookupError from inside nltk. + """ + import nltk + + try: + nltk.data.find("tokenizers/punkt_tab") + except LookupError: + pytest.skip( + "nltk punkt_tab not downloaded: python -m nltk.downloader punkt_tab" + ) + + +COLLECTIONS = {"alpha": 60, "beta": 60} + +# Varied on purpose. With near-identical text BM25 and vector search return the +# same ten documents, fusion never exceeds the per-collection cap, and a test +# that changes the cap on one path only still passes -- which this one did. +WORDS = [ + "kinase phosphorylation cascade", + "cholesterol transport vesicle", + "ubiquitin ligase complex", + "mitochondrial respiratory chain", + "DNA mismatch repair", + "interferon signalling", + "collagen assembly", +] + + +def _bundle(tmp_path: Path, embedding: DeterministicFakeEmbedding) -> Path: + csv_dir = tmp_path / "csv_files" + csv_dir.mkdir(parents=True, exist_ok=True) + for collection, count in COLLECTIONS.items(): + with open(csv_dir / f"{collection}.csv", "w", newline="") as handle: + writer = csv.DictWriter( + handle, fieldnames=["st_id", "display_name", "text"] + ) + writer.writeheader() + for i in range(count): + writer.writerow( + { + "st_id": f"{collection}-{i}", + "display_name": f"{collection} item {i}", + "text": WORDS[i % len(WORDS)] + f" {collection} {i}", + } + ) + Chroma.from_documents( + documents=[ + Document( + page_content=( + f"st_id: {collection}-{i}\n" + f"text: {WORDS[i % len(WORDS)]} {collection} {i}" + ), + metadata={"st_id": f"{collection}-{i}"}, + ) + for i in range(count) + ], + embedding=embedding, + persist_directory=str(tmp_path / collection), + client_settings=chroma_settings(), + ) + return tmp_path + + +@pytest.mark.requires_retrieval_stack +@pytest.mark.parametrize( + "queries", + [ + ["kinase phosphorylation"], + ["nothing matches this at all"], + # Enough distinct queries that fusion overflows the per-collection cap, + # so a change to the cap on one path only is visible. With one query the + # two retrievers return largely the same documents, fusion stays under + # the cap, and such a change passes unnoticed -- as it did here. + [ + "kinase phosphorylation cascade", + "cholesterol transport vesicle", + "ubiquitin ligase complex", + "mitochondrial respiratory chain", + "DNA mismatch repair", + "interferon signalling", + "collagen assembly", + ], + ], +) +def test_async_returns_exactly_what_sync_returns( + tmp_path: Path, queries: list[str] +) -> None: + _require_bm25_tokenizer() + embedding = DeterministicFakeEmbedding(size=16) + retriever = HybridRetriever.from_subdirectory( + # Never called: these tests drive retrieve_documents directly, below the + # query-expansion step. It is here because the constructor requires one. + llm=FakeListChatModel(responses=[""]), + embedding=embedding, + embeddings_directory=_bundle(tmp_path, embedding), + ) + + sync = retriever.retrieve_documents( + queries, CallbackManagerForRetrieverRun.get_noop_manager() + ) + asynchronous = asyncio.run( + retriever.aretrieve_documents( + queries, AsyncCallbackManagerForRetrieverRun.get_noop_manager() + ) + ) + + # Order, not just membership: RRF resolves ties by first appearance, so a + # reordering is a ranking change and the top documents are the ones that + # reach the model. + assert [d.page_content for d in asynchronous] == [d.page_content for d in sync]