From 3d4a7f083c55e47ad8cd8b5d05084fab5c9ff244 Mon Sep 17 00:00:00 2001 From: Adam Wright Date: Wed, 9 Sep 2026 13:53:09 +0000 Subject: [PATCH] Stage 3: the caller supplies the bundle and the context budget Two module-level facts become arguments, closing out the retriever rewrite (FR-005, FR-006, SC-006). The bundle. Each create_*_rag had embeddings_directory: Path = EmbeddingEnvironment.get_dir("reactome") as a default argument, so it was evaluated at import time: importing a rag module required an installed bundle, which is why they could not be unit tested. get_dir returns `Path | None`, and None travelled into a parameter typed `Path` -- four mypy baseline entries existed solely to suppress that, and a ruff B008 suppression sat on each line. All five suppressions are now deleted rather than annotated. EmbeddingEnvironment.require_dir replaces the lookup at the three call sites that cannot proceed without a bundle, and also checks the directory exists. That second check is not tidiness: get_dir built a path without touching the disk, and Chroma CREATES a missing persist_directory rather than complaining, so an `embeddings/current` left pointing at a deleted bundle produced a chatbot that answered every question from an empty collection. Verified against main, the missing-bundle failure was "AttributeError: 'NoneType' object has no attribute 'glob'"; it now names the database, lists what is installed, and gives the command to fix it. The userguide bundle is deliberately optional, and react_to_me.py already resolved it itself and degraded on absence -- that path keeps get_dir, and require_dir raises FileNotFoundError precisely so the existing `except FileNotFoundError` there still means what it meant. The budget. MAX_DOCUMENTS_PER_COLLECTION is now DEFAULT_MAX_DOCUMENTS_PER_COLLECTION, the default of a constructor argument. As a module constant, comparing two budgets meant editing the file and restarting -- so the answer-quality evaluation that is supposed to settle the number could not be run at all. Verified end to end against the Release95 bundle: the assembled chain answers a real question identically through invoke and ainvoke (40 docs each), budget=3 over 4 collections yields 12 documents and budget=10 yields 40, and all three profiles construct or refuse correctly. --- pyproject.toml | 12 --- src/agent/profiles/cross_database.py | 9 ++- src/agent/profiles/plantreactome.py | 6 +- src/agent/profiles/react_to_me.py | 7 +- src/retrievers/csv_chroma.py | 22 +++++- src/retrievers/plantreactome/rag.py | 5 +- src/retrievers/reactome/rag.py | 5 +- src/retrievers/uniprot/rag.py | 5 +- src/retrievers/userguide/rag.py | 5 +- src/util/embedding_environment.py | 31 ++++++++ tests/retrievers/test_hybrid_retriever.py | 93 ++++++++++++++++++++++- tests/util/test_embedding_environment.py | 46 ++++++++++- 12 files changed, 207 insertions(+), 39 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9992ab2a..02865229 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -176,15 +176,3 @@ disable_error_code = ["typeddict-item", "no-any-return"] module = ["agent.profiles.react_to_me"] disable_error_code = ["typeddict-item", "override"] -[[tool.mypy.overrides]] -# TODO(phase-2): `EmbeddingEnvironment.get_dir()` returns `Path | None` but the -# parameter is typed `Path`, so an uninstalled bundle passes None straight through. -# Fixed properly by moving the lookup out of the default argument (see ruff B008). -module = [ - "retrievers.reactome.rag", - "retrievers.uniprot.rag", - "retrievers.plantreactome.rag", - "retrievers.userguide.rag", -] -disable_error_code = ["assignment"] - diff --git a/src/agent/profiles/cross_database.py b/src/agent/profiles/cross_database.py index 79763f92..eff99408 100644 --- a/src/agent/profiles/cross_database.py +++ b/src/agent/profiles/cross_database.py @@ -22,6 +22,7 @@ ) from retrievers.reactome.rag import create_reactome_rag from retrievers.uniprot.rag import create_uniprot_rag +from util.embedding_environment import EmbeddingEnvironment class CrossDatabaseState(BaseState): @@ -43,8 +44,12 @@ def __init__( super().__init__(llm, embedding) # Create runnables (tasks & tools) - self.reactome_rag: Runnable = create_reactome_rag(llm, embedding) - self.uniprot_rag: Runnable = create_uniprot_rag(llm, embedding) + self.reactome_rag: Runnable = create_reactome_rag( + llm, embedding, EmbeddingEnvironment.require_dir("reactome") + ) + self.uniprot_rag: Runnable = create_uniprot_rag( + llm, embedding, EmbeddingEnvironment.require_dir("uniprot") + ) self.completeness_checker = create_completeness_grader(llm) self.write_reactome_query = create_reactome_rewriter_w_uniprot(llm) diff --git a/src/agent/profiles/plantreactome.py b/src/agent/profiles/plantreactome.py index 5475214d..368b337e 100644 --- a/src/agent/profiles/plantreactome.py +++ b/src/agent/profiles/plantreactome.py @@ -9,6 +9,7 @@ from agent.profiles.base import BaseGraphBuilder, BaseState from agent.tasks.unsafe_question import create_unsafe_answer_generator from retrievers.plantreactome.rag import create_plantreactome_rag +from util.embedding_environment import EmbeddingEnvironment class PlantReactomeState(BaseState): @@ -28,7 +29,10 @@ def __init__( llm, streaming=True ) self.plantreactome_rag: Runnable = create_plantreactome_rag( - llm, embedding, streaming=True + llm, + embedding, + EmbeddingEnvironment.require_dir("plantreactome"), + streaming=True, ) # Create graph diff --git a/src/agent/profiles/react_to_me.py b/src/agent/profiles/react_to_me.py index 9ebb0cbf..bad9ef89 100644 --- a/src/agent/profiles/react_to_me.py +++ b/src/agent/profiles/react_to_me.py @@ -41,7 +41,12 @@ def __init__( ) self.rags: dict[SourceName, Runnable] = { - "reactome": create_reactome_rag(llm, embedding, streaming=True), + "reactome": create_reactome_rag( + llm, + embedding, + EmbeddingEnvironment.require_dir("reactome"), + streaming=True, + ), } self._available_sources: frozenset[SourceName] = frozenset({"reactome"}) self._register_userguide_rag(llm, embedding) diff --git a/src/retrievers/csv_chroma.py b/src/retrievers/csv_chroma.py index e07ca537..69f84965 100644 --- a/src/retrievers/csv_chroma.py +++ b/src/retrievers/csv_chroma.py @@ -83,7 +83,12 @@ # retriever returns. Changing it trades recall against the model's difficulty # attending to the middle of a long context; the right number should come from an # answer-quality evaluation rather than from taste. -MAX_DOCUMENTS_PER_COLLECTION = RESULTS_PER_RETRIEVER +# +# It is the DEFAULT, not the budget: a caller passes its own to +# HybridRetriever.from_subdirectory. That is what makes the evaluation above +# possible -- sweeping this value used to mean editing the module and restarting, +# so two budgets could not be compared within one process. +DEFAULT_MAX_DOCUMENTS_PER_COLLECTION = RESULTS_PER_RETRIEVER def dedupe_by_entity(docs: list[Document], limit: int) -> list[Document]: @@ -194,12 +199,15 @@ def create_bm25_chroma_ensemble_retriever( llm: BaseChatModel, embedding: Embeddings, embeddings_directory: Path, + *, + max_documents_per_collection: int = DEFAULT_MAX_DOCUMENTS_PER_COLLECTION, ) -> "HybridRetriever": return HybridRetriever.from_subdirectory( llm, embedding, embeddings_directory, include_original=True, + max_documents_per_collection=max_documents_per_collection, ) @@ -228,6 +236,8 @@ class HybridRetriever(BaseRetriever): query_expander: Runnable[dict[str, str], list[str]] include_original: bool = False collection_retrievers: dict[str, RetrieverDict] + # How many fused documents this instance contributes per collection. + max_documents_per_collection: int = DEFAULT_MAX_DOCUMENTS_PER_COLLECTION # BM25Retriever and the Chroma retriever are not pydantic models. model_config = ConfigDict(arbitrary_types_allowed=True) @@ -240,6 +250,7 @@ def from_subdirectory( embeddings_directory: Path, *, include_original: bool = False, + max_documents_per_collection: int = DEFAULT_MAX_DOCUMENTS_PER_COLLECTION, ) -> "HybridRetriever": _retrievers: dict[str, RetrieverDict] = {} for subdirectory in list_chroma_subdirectories(embeddings_directory): @@ -279,6 +290,7 @@ def from_subdirectory( query_expander=multi_query_prompt | llm | LineListOutputParser(), include_original=include_original, collection_retrievers=_retrievers, + max_documents_per_collection=max_documents_per_collection, ) def _get_relevant_documents( @@ -348,7 +360,9 @@ def retrieve_documents( doc_lists.append(dedupe_by_entity(bm25_docs, RESULTS_PER_RETRIEVER)) doc_lists.append(dedupe_by_entity(vector_docs, RESULTS_PER_RETRIEVER)) subdirectory_docs.extend( - self.weighted_reciprocal_rank(doc_lists)[:MAX_DOCUMENTS_PER_COLLECTION] + self.weighted_reciprocal_rank(doc_lists)[ + : self.max_documents_per_collection + ] ) return subdirectory_docs @@ -393,6 +407,8 @@ async def aretrieve_documents( for docs in await asyncio.gather(*subdir_results) ] subdirectory_docs.extend( - self.weighted_reciprocal_rank(doc_lists)[:MAX_DOCUMENTS_PER_COLLECTION] + self.weighted_reciprocal_rank(doc_lists)[ + : self.max_documents_per_collection + ] ) return subdirectory_docs diff --git a/src/retrievers/plantreactome/rag.py b/src/retrievers/plantreactome/rag.py index 17eafe1e..59b03312 100644 --- a/src/retrievers/plantreactome/rag.py +++ b/src/retrievers/plantreactome/rag.py @@ -7,15 +7,12 @@ from retrievers.csv_chroma import create_bm25_chroma_ensemble_retriever from retrievers.plantreactome.prompt import plantreactome_qa_prompt from retrievers.rag_chain import create_rag_chain -from util.embedding_environment import EmbeddingEnvironment def create_plantreactome_rag( llm: BaseChatModel, embedding: Embeddings, - # TODO(phase-2): resolved at import time, so importing this module requires an - # installed embeddings bundle. Blocks unit-testing; fix with the agent-API refactor. - embeddings_directory: Path = EmbeddingEnvironment.get_dir("plantreactome"), # noqa: B008 + embeddings_directory: Path, *, streaming: bool = False, ) -> Runnable: diff --git a/src/retrievers/reactome/rag.py b/src/retrievers/reactome/rag.py index 5ce5a3b3..b56abb74 100644 --- a/src/retrievers/reactome/rag.py +++ b/src/retrievers/reactome/rag.py @@ -7,15 +7,12 @@ from retrievers.csv_chroma import create_bm25_chroma_ensemble_retriever from retrievers.rag_chain import create_rag_chain from retrievers.reactome.prompt import reactome_qa_prompt -from util.embedding_environment import EmbeddingEnvironment def create_reactome_rag( llm: BaseChatModel, embedding: Embeddings, - # TODO(phase-2): resolved at import time, so importing this module requires an - # installed embeddings bundle. Blocks unit-testing; fix with the agent-API refactor. - embeddings_directory: Path = EmbeddingEnvironment.get_dir("reactome"), # noqa: B008 + embeddings_directory: Path, *, streaming: bool = False, ) -> Runnable: diff --git a/src/retrievers/uniprot/rag.py b/src/retrievers/uniprot/rag.py index dc2d9aa5..324083b1 100644 --- a/src/retrievers/uniprot/rag.py +++ b/src/retrievers/uniprot/rag.py @@ -7,15 +7,12 @@ from retrievers.csv_chroma import create_bm25_chroma_ensemble_retriever from retrievers.rag_chain import create_rag_chain from retrievers.uniprot.prompt import uniprot_qa_prompt -from util.embedding_environment import EmbeddingEnvironment def create_uniprot_rag( llm: BaseChatModel, embedding: Embeddings, - # TODO(phase-2): resolved at import time, so importing this module requires an - # installed embeddings bundle. Blocks unit-testing; fix with the agent-API refactor. - embeddings_directory: Path = EmbeddingEnvironment.get_dir("uniprot"), # noqa: B008 + embeddings_directory: Path, *, streaming: bool = False, ) -> Runnable: diff --git a/src/retrievers/userguide/rag.py b/src/retrievers/userguide/rag.py index 153910bd..31dc48a6 100644 --- a/src/retrievers/userguide/rag.py +++ b/src/retrievers/userguide/rag.py @@ -7,15 +7,12 @@ from retrievers.rag_chain import create_rag_chain from retrievers.userguide.prompt import userguide_qa_prompt from retrievers.userguide.retriever import create_userguide_retriever -from util.embedding_environment import EmbeddingEnvironment def create_userguide_rag( llm: BaseChatModel, embedding: Embeddings, - # TODO(phase-2): resolved at import time, so importing this module requires an - # installed embeddings bundle. Blocks unit-testing; fix with the agent-API refactor. - embeddings_directory: Path = EmbeddingEnvironment.get_dir("userguide"), # noqa: B008 + embeddings_directory: Path, *, streaming: bool = False, ) -> Runnable: diff --git a/src/util/embedding_environment.py b/src/util/embedding_environment.py index 140de601..9560672b 100644 --- a/src/util/embedding_environment.py +++ b/src/util/embedding_environment.py @@ -32,6 +32,37 @@ def get_dir(cls, key: str) -> Path | None: return EM_ARCHIVE / cls._get().embeddings[key] return None + @classmethod + def require_dir(cls, key: str) -> Path: + """`get_dir`, but for bundles the caller cannot run without. + + `get_dir` returns None for an unknown key, and returns a path for a + known one without checking that the path exists. Both failures used to + travel: None was passed into a parameter typed `Path` (four mypy + baseline entries suppressed the error), and a stale `embeddings/current` + pointing at a deleted bundle handed Chroma a missing directory, which + Chroma creates -- so the chatbot answered every question from an empty + collection instead of refusing to start. + + Raising FileNotFoundError rather than SystemExit: it is the accurate + exception, nothing catches it for a required bundle so the process still + stops, and the optional userguide path in react_to_me.py already catches + exactly this to degrade deliberately. + """ + directory = cls.get_dir(key) + if directory is None: + available = ", ".join(sorted(cls.get_dict())) or "none" + raise FileNotFoundError( + f"No embeddings bundle installed for {key!r} (installed: {available}). " + f"Run ./bin/embeddings_manager install ." + ) + if not directory.is_dir(): + raise FileNotFoundError( + f"{key!r} points at {directory}, which does not exist. " + f"{EM_CURRENT} is stale; re-run ./bin/embeddings_manager install." + ) + return directory + @classmethod def get_model(cls, key: str) -> str: return str(cls._get().embeddings[key].parent.parent) diff --git a/tests/retrievers/test_hybrid_retriever.py b/tests/retrievers/test_hybrid_retriever.py index 990b2287..e5586aa9 100644 --- a/tests/retrievers/test_hybrid_retriever.py +++ b/tests/retrievers/test_hybrid_retriever.py @@ -23,7 +23,7 @@ from retrievers import csv_chroma # noqa: E402 from retrievers.csv_chroma import ( # noqa: E402 - MAX_DOCUMENTS_PER_COLLECTION, + DEFAULT_MAX_DOCUMENTS_PER_COLLECTION, HybridRetriever, RetrieverDict, dedupe_by_entity, @@ -201,7 +201,7 @@ def test_fused_results_are_capped_per_collection() -> None: many = [[_doc(f"doc {i}") for i in range(50)]] assert len(_fuse(many)) == 50, "the fusion itself still returns everything" assert ( - MAX_DOCUMENTS_PER_COLLECTION < 50 + DEFAULT_MAX_DOCUMENTS_PER_COLLECTION < 50 ), "the cap, applied by the caller, is what bounds the prompt" @@ -232,3 +232,92 @@ def test_retriever_dict_accepts_any_base_retriever() -> None: """ hints = get_type_hints(RetrieverDict) assert hints["vector"] is BaseRetriever + + +class _StubVectorRetriever(BaseRetriever): + """Returns a fixed list, so a budget test needs no Chroma store and no LLM.""" + + docs: list[Document] + + def _get_relevant_documents(self, query: str, **kwargs: Any) -> list[Document]: + return self.docs + + +def _retriever_with_budget(budget: int | None) -> HybridRetriever: + from langchain_community.retrievers import BM25Retriever + from langchain_core.runnables import RunnableLambda + + corpus = [_doc(f"doc {i}") for i in range(30)] + bm25 = BM25Retriever.from_documents(corpus) # default whitespace tokenizer, no nltk + bm25.k = 30 + collection: RetrieverDict = { + "bm25": bm25, + "vector": _StubVectorRetriever(docs=corpus), + } + kwargs: dict[str, Any] = ( + {} if budget is None else {"max_documents_per_collection": budget} + ) + return HybridRetriever( + # No expansion: one query in, one query out, so the only thing varying + # between the two instances below is the budget. + query_expander=RunnableLambda(lambda inputs: [inputs["question"]]), + include_original=False, + collection_retrievers={"reactions": collection}, + **kwargs, + ) + + +def test_two_budgets_can_coexist_in_one_process() -> None: + """The point of Stage 3: the budget is an argument, not a module constant. + + While it was a module constant, comparing two budgets meant editing + csv_chroma.py and restarting -- so the answer-quality evaluation that is + supposed to settle the number could not be run. Asserting both instances in + the same test is the whole claim: not that the value can be changed, but that + two values can be live at once. + """ + small = _retriever_with_budget(3) + large = _retriever_with_budget(7) + + assert len(small.invoke("anything")) == 3 + assert len(large.invoke("anything")) == 7 + # and the smaller is a prefix of the larger: same ranking, less of it + assert [d.page_content for d in small.invoke("anything")] == [ + d.page_content for d in large.invoke("anything") + ][:3] + + +def test_omitting_the_budget_behaves_as_the_old_constant_did() -> None: + """Making it an argument must not quietly change production's context size. + + Production passes no budget, so the default is what actually ships. + """ + assert ( + len(_retriever_with_budget(None).invoke("x")) + == DEFAULT_MAX_DOCUMENTS_PER_COLLECTION + ) + + +@pytest.mark.parametrize( + "module", + ["reactome", "uniprot", "plantreactome", "userguide"], +) +def test_rag_factories_take_the_bundle_rather_than_resolving_it(module: str) -> None: + """FR-006: no default argument may call EmbeddingEnvironment at import time. + + The default used to be `EmbeddingEnvironment.get_dir(...)`, evaluated once + when the module was imported. That made importing a rag module require an + installed bundle, and it fed `Path | None` into a parameter typed `Path` -- + suppressed by four mypy baseline entries, now deleted. + """ + import importlib + import inspect + + mod = importlib.import_module(f"retrievers.{module}.rag") + fn = getattr(mod, f"create_{module}_rag") + parameter = inspect.signature(fn).parameters["embeddings_directory"] + + assert ( + parameter.default is inspect.Parameter.empty + ), "a default here is evaluated at import time; the caller must pass the bundle" + assert parameter.annotation is Path, "and it is a Path, never Path | None" diff --git a/tests/util/test_embedding_environment.py b/tests/util/test_embedding_environment.py index cebc5d86..ba4d95b4 100644 --- a/tests/util/test_embedding_environment.py +++ b/tests/util/test_embedding_environment.py @@ -1,8 +1,9 @@ """Characterization of embeddings path resolution. `embeddings/current` is a single colon-separated line mapping each database to the -bundle in use. `bin/embeddings_manager use` writes it; the retriever modules read it -at import time. +bundle in use. `bin/embeddings_manager use` writes it; the agent profiles read it +when they construct a RAG chain. They used to read it at import time, via a default +argument, which is why `require_dir` exists. """ from pathlib import Path @@ -61,3 +62,44 @@ def test_set_one_replaces_the_bundle_for_the_same_database(archive: Path) -> Non def test_empty_current_file_is_not_an_error(archive: Path) -> None: (archive / "current").write_text("") assert EmbeddingEnvironment.get_dict() == {} + + +def test_require_dir_names_what_is_installed_when_the_key_is_unknown( + archive: Path, +) -> None: + """The old path returned None here and passed it into a parameter typed Path.""" + (archive / "current").write_text(REACTOME) + (archive / REACTOME).mkdir(parents=True) + + with pytest.raises(FileNotFoundError) as exc: + EmbeddingEnvironment.require_dir("uniprot") + + message = str(exc.value) + assert "uniprot" in message + assert ( + "reactome" in message + ), "say which bundles ARE installed, not just which is not" + assert "embeddings_manager install" in message, "say what to do about it" + + +def test_require_dir_rejects_a_bundle_that_current_names_but_disk_lacks( + archive: Path, +) -> None: + """The failure that would otherwise be silent, and the reason for the is_dir check. + + `get_dir` builds a path without checking it exists, and Chroma CREATES a + missing persist_directory rather than complaining. So a `current` left + pointing at a deleted bundle produced a chatbot that answered every question + from an empty collection -- confidently, and with no error anywhere. + """ + (archive / "current").write_text(REACTOME) # note: directory never created + + assert EmbeddingEnvironment.get_dir("reactome") is not None, "the old, quiet path" + with pytest.raises(FileNotFoundError, match="does not exist"): + EmbeddingEnvironment.require_dir("reactome") + + +def test_require_dir_returns_the_directory_when_it_is_there(archive: Path) -> None: + (archive / "current").write_text(REACTOME) + (archive / REACTOME).mkdir(parents=True) + assert EmbeddingEnvironment.require_dir("reactome") == archive / REACTOME