Fix knowledge graph embedding model discovery

pipeline/kg.py

@@ -67,9 +67,20 @@ _BATCH_EMBED_URL = _API_BASE + "/{model}:batchEmbedContents"
 # Triple extraction must default to a FREE-TIER model (Flash). Pro/preview
 # models return HTTP 429 limit:0 on the free tier — see online_ocr notes.
 DEFAULT_TRIPLE_MODEL = "gemini-2.5-flash"
-# Embeddings: text-embedding-004 is the stable, free, 768-dim model. The wire
-# format wants the "models/" prefix in the per-request "model" field.
-DEFAULT_EMBED_MODEL = "text-embedding-004"
+
+# Embeddings: the wire format wants the "models/" prefix in the per-request
+# "model" field. The embedding model id has CHANGED over time — Google retired
+# `text-embedding-004` on v1beta (it now 404s) in favour of `gemini-embedding-001`.
+# So we DON'T trust a single hard-coded id: build_graph DISCOVERS the model the
+# key can actually use via ModelService.ListModels (`_pick_embedding_model`),
+# trying this preference order and falling back to the default only if discovery
+# fails. Order newest→oldest so a key that still has an old model also works.
+DEFAULT_EMBED_MODEL = "gemini-embedding-001"
+EMBED_MODEL_PREFERENCE = (
+    "gemini-embedding-001",   # current GA model (default ~3072-dim; we L2-normalize)
+    "text-embedding-004",     # legacy 768-dim (retired on v1beta for many keys)
+    "embedding-001",          # older fallback
+)
 
 # Gemini caps batchEmbedContents at 100 requests per call; chunk to stay under.
 _EMBED_BATCH = 100
@@ -202,6 +213,89 @@ def _resolve_model(model, default: str) -> str:
     return name or default
 
 
+# Process-level cache of the discovered embedding model. Only a SUCCESSFUL
+# discovery is cached (so a transient failure that fell back to the default is
+# retried next time). Not keyed by the API key — the available models are a
+# property of the deployment, not a secret, and we never store the key itself.
+_EMBED_MODEL_RESOLVED = None
+
+
+def _list_embedding_models(api_key, *, timeout=30) -> list:
+    """Return bare ids of models the key can use for embedding.
+
+    Calls ``GET /v1beta/models`` and keeps those whose
+    ``supportedGenerationMethods`` include ``batchEmbedContents`` (what we call)
+    or ``embedContent``. Used to pick a model that actually exists, instead of
+    hard-coding an id that Google may retire (the cause of the text-embedding-004
+    404). Mirrors ``online_ocr.list_models`` but filters for embedding support.
+    """
+    api_key = _clean_key(api_key)
+    request = urllib.request.Request(
+        _API_BASE, method="GET", headers={_API_KEY_HEADER: api_key}
+    )
+    try:
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            raw = response.read()
+    except urllib.error.HTTPError as err:
+        raise _friendly_http_error(err) from None
+    except urllib.error.URLError as err:
+        raise RuntimeError(
+            "Could not reach the Gemini API ({}). Check your internet "
+            "connection.".format(err.reason)
+        ) from None
+    except TimeoutError:
+        raise RuntimeError(
+            "The Gemini request timed out after {}s.".format(timeout)
+        ) from None
+    try:
+        payload = json.loads(raw.decode("utf-8", "replace"))
+    except (ValueError, AttributeError):
+        return []
+    out = []
+    for entry in (payload.get("models") or []):
+        if not isinstance(entry, dict):
+            continue
+        methods = entry.get("supportedGenerationMethods") or []
+        if "batchEmbedContents" in methods or "embedContent" in methods:
+            name = str(entry.get("name") or "")
+            if name.startswith("models/"):
+                name = name[len("models/"):]
+            if name:
+                out.append(name)
+    return out
+
+
+def _pick_embedding_model(api_key, requested=None, *, timeout=30) -> str:
+    """Choose an embedding model id the key can actually use.
+
+    An explicit ``requested`` model wins. Otherwise discover the available
+    embedding models and pick by :data:`EMBED_MODEL_PREFERENCE`; on a transient
+    discovery failure fall back to :data:`DEFAULT_EMBED_MODEL` (without caching,
+    so it's retried). A successful discovery is cached for the process.
+    """
+    global _EMBED_MODEL_RESOLVED
+    if requested:
+        return _resolve_model(requested, DEFAULT_EMBED_MODEL)
+    if _EMBED_MODEL_RESOLVED:
+        return _EMBED_MODEL_RESOLVED
+    try:
+        available = set(_list_embedding_models(api_key, timeout=timeout))
+    except RuntimeError:
+        available = set()
+    chosen = next((m for m in EMBED_MODEL_PREFERENCE if m in available), None)
+    if not chosen and available:
+        # An unfamiliar but embedding-capable model — prefer one that looks like
+        # an embedding model, else just take the first deterministically.
+        chosen = next(
+            (m for m in sorted(available) if "embedding" in m.lower()),
+            sorted(available)[0],
+        )
+    if chosen:
+        _EMBED_MODEL_RESOLVED = chosen  # cache only a real discovery
+        return chosen
+    return DEFAULT_EMBED_MODEL
+
+
 # ---------------------------------------------------------------------------
 # Response parsers (pure functions — unit-tested with canned dicts, no network).
 # ---------------------------------------------------------------------------
@@ -356,9 +450,11 @@ def embed_texts(
     """Embed a list of strings into an ``(n, dim)`` float32 matrix via Gemini.
 
     ``task_type`` should be ``RETRIEVAL_DOCUMENT`` for the graph's nodes and
-    ``RETRIEVAL_QUERY`` for the search query — text-embedding-004 uses it to
+    ``RETRIEVAL_QUERY`` for the search query — the embedding model uses it to
     tune the vectors for retrieval, a real accuracy win. Batches at 100 inputs
-    per request (Gemini's cap) and concatenates in order.
+    per request (Gemini's cap) and concatenates in order. Pass an explicit
+    ``model`` (typically resolved by :func:`_pick_embedding_model`); the default
+    is only a last resort since Google retires embedding-model ids over time.
     """
     if not is_configured(api_key):
         raise RuntimeError(
@@ -789,14 +885,17 @@ def build_graph(
             graph._add_triple(t["subject"], t["predicate"], t["object"], page=page_no)
 
     if embed and graph.num_nodes:
+        # Resolve the embedding model the key can actually use (Google retires
+        # ids over time), so search uses the SAME model the nodes were built with.
+        resolved = _pick_embedding_model(api_key, embed_model, timeout=min(timeout, 30))
         node_texts = [graph._node_embed_text(i) for i in range(graph.num_nodes)]
         vectors = embed_texts(
             node_texts,
             api_key=api_key,
-            model=embed_model,
+            model=resolved,
             task_type="RETRIEVAL_DOCUMENT",
             timeout=timeout,
         )
-        graph.set_vectors(vectors, model=_resolve_model(embed_model, DEFAULT_EMBED_MODEL))
+        graph.set_vectors(vectors, model=resolved)
 
     return graph