Fix: the embedding model default breaks one deployment or the other

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=

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=

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 "<provider>/<model>/<database>/<version>"; 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(

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