diff --git a/deploy/beta/env.beta.template b/deploy/beta/env.beta.template index d3f2787..8a2f7e2 100644 --- a/deploy/beta/env.beta.template +++ b/deploy/beta/env.beta.template @@ -16,3 +16,10 @@ UVICORN_LOG_LEVEL=info # Chainlit is mounted here; the landing page lives at /chat/ CHAINLIT_URI=/chat/guest CHAINLIT_URL=https://beta.reactome.org + +# Embedding model. Leave unset: it defaults to whatever built the installed +# bundle, which is what the stored vectors were produced with. Setting it to +# anything else makes retrieval meaningless. +#EMBEDDING_MODEL= +# LLM used for answers; defaults to gpt-4o-mini. +#LLM_MODEL= diff --git a/env_template b/env_template index 48f0ea2..e096380 100644 --- a/env_template +++ b/env_template @@ -14,3 +14,10 @@ CHAINLIT_IMAGE=public.ecr.aws/reactome/reactome-chatbot:latest CHAINLIT_URI=/chat CHAINLIT_ROOT_PATH= TAVILY_API_KEY= + +# Embedding model. Leave unset: it defaults to whatever built the installed +# bundle, which is what the stored vectors were produced with. Setting it to +# anything else makes retrieval meaningless. +#EMBEDDING_MODEL= +# LLM used for answers; defaults to gpt-4o-mini. +#LLM_MODEL= diff --git a/src/agent/graph.py b/src/agent/graph.py index 7b50cf7..4fafd9d 100644 --- a/src/agent/graph.py +++ b/src/agent/graph.py @@ -16,6 +16,7 @@ from agent.models import get_embedding, get_llm from agent.profiles import ProfileName, create_profile_graphs from agent.profiles.base import InputState, OutputState +from util.embedding_environment import EmbeddingEnvironment from util.logging import logging LANGGRAPH_DB_URI = f"postgresql://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}@postgres:5432/{os.getenv('POSTGRES_LANGGRAPH_DB')}?sslmode=disable" @@ -24,13 +25,66 @@ logging.warning("POSTGRES_LANGGRAPH_DB undefined; falling back to MemorySaver.") +DEFAULT_EMBEDDING_MODEL = "text-embedding-3-large" + + +def resolve_embedding_model() -> str: + """Pick the embedding model, defaulting to whatever built the installed bundles. + + A query is embedded with this model and compared against vectors already in + Chroma, so it has to match the model that produced them. The bundle path + records that -- `openai/text-embedding-3-large/reactome/Release95` -- which + makes it the source of truth rather than a constant kept in sync by hand. + + The default used to be a literal "bge-m3". That is the right model for the + Plant Reactome deployment, which serves it from a self-hosted + OpenAI-compatible endpoint via OPENAI_BASE_URL, but it is wrong for every + bundle published for Reactome: against api.openai.com it is a 404 on the + first query. Hardcoding either one breaks the other deployment, so neither + is hardcoded. + + AgentGraph builds a single embedding shared by every profile, so all + installed bundles must agree on the model. Disagreement is a real + misconfiguration and is reported rather than silently resolved. + """ + configured = os.getenv("EMBEDDING_MODEL") + + # Bundle paths are "///"; the provider is + # supplied separately to get_embedding, so only the model is wanted here. + installed = { + bundle.parent.parent.name for bundle in EmbeddingEnvironment.get_dict().values() + } + + if not installed: + return configured or DEFAULT_EMBEDDING_MODEL + + if len(installed) > 1: + logging.error( + f"Installed bundles were built with different embedding models " + f"({', '.join(sorted(installed))}), but one embedding is shared by " + "every profile. Retrieval will be meaningless for whichever does not " + "match. Install bundles built with the same model." + ) + return configured or DEFAULT_EMBEDDING_MODEL + + bundle_model = installed.pop() + if configured and configured != bundle_model: + logging.error( + f"EMBEDDING_MODEL is {configured!r} but the installed bundle was built " + f"with {bundle_model!r}. Queries would be embedded with a different " + "model than the stored vectors, making retrieval meaningless. Unset " + "EMBEDDING_MODEL, or install a matching bundle." + ) + return configured or bundle_model + + class AgentGraph: def __init__( self, profiles: list[ProfileName], ) -> None: # Get base models - embedding_model = os.getenv("EMBEDDING_MODEL", "bge-m3") + embedding_model = resolve_embedding_model() llm_model = os.getenv("LLM_MODEL", "gpt-4o-mini") llm_base_url = os.getenv("LLM_BASE_URL", None) llm: BaseChatModel = get_llm( diff --git a/tests/agent/test_embedding_model_resolution.py b/tests/agent/test_embedding_model_resolution.py new file mode 100644 index 0000000..f8682c8 --- /dev/null +++ b/tests/agent/test_embedding_model_resolution.py @@ -0,0 +1,78 @@ +"""The query embedding model must match the one that built the bundle. + +Queries are embedded and compared against vectors already in Chroma. If the two +models differ the comparison is meaningless, and if the model does not exist at +the configured endpoint the first query fails outright. + +The default used to be the literal "bge-m3". That is correct for the Plant +Reactome deployment, which serves it from a self-hosted OpenAI-compatible +endpoint via OPENAI_BASE_URL, and wrong for every bundle published for Reactome, +where it 404s against api.openai.com. Hardcoding either breaks the other. +""" + +from pathlib import Path + +import pytest + +import util.embedding_environment as ee +from agent.graph import DEFAULT_EMBEDDING_MODEL, resolve_embedding_model + +REACTOME = "openai/text-embedding-3-large/reactome/Release95" +PLANT = "openai/bge-m3/plantreactome/Release68" + + +@pytest.fixture +def archive(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.setattr(ee, "EM_ARCHIVE", tmp_path) + monkeypatch.setattr(ee, "EM_CURRENT", tmp_path / "current") + monkeypatch.delenv("EMBEDDING_MODEL", raising=False) + return tmp_path + + +def test_reactome_bundle_resolves_to_its_own_model(archive: Path) -> None: + (archive / "current").write_text(REACTOME) + assert resolve_embedding_model() == "text-embedding-3-large" + + +def test_plantreactome_bundle_resolves_to_bge_m3(archive: Path) -> None: + """Plant Reactome legitimately uses bge-m3; it must not be overridden.""" + (archive / "current").write_text(PLANT) + assert resolve_embedding_model() == "bge-m3" + + +def test_never_silently_defaults_to_a_model_no_bundle_used(archive: Path) -> None: + """The regression: "bge-m3" was the default regardless of what was installed.""" + (archive / "current").write_text(REACTOME) + assert resolve_embedding_model() != "bge-m3" + + +def test_env_override_is_respected( + archive: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + (archive / "current").write_text(REACTOME) + monkeypatch.setenv("EMBEDDING_MODEL", "text-embedding-3-small") + assert resolve_embedding_model() == "text-embedding-3-small" + + +def test_override_that_disagrees_with_the_bundle_is_reported( + archive: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + (archive / "current").write_text(REACTOME) + monkeypatch.setenv("EMBEDDING_MODEL", "text-embedding-3-small") + with caplog.at_level("ERROR"): + resolve_embedding_model() + assert "text-embedding-3-large" in caplog.text + + +def test_bundles_built_with_different_models_are_reported( + archive: Path, caplog: pytest.LogCaptureFixture +) -> None: + """One embedding is shared by every profile, so the bundles must agree.""" + (archive / "current").write_text(f"{REACTOME}:{PLANT}") + with caplog.at_level("ERROR"): + resolve_embedding_model() + assert "different embedding models" in caplog.text + + +def test_falls_back_when_nothing_is_installed(archive: Path) -> None: + assert resolve_embedding_model() == DEFAULT_EMBEDDING_MODEL