Harden knowledge graph build (isolation, status reset, embed fallback)

app.py

@@ -437,6 +437,9 @@ async def graph_build(
                 embed=do_embed,
             ),
         )
+        # Set the terminal state INSIDE the try so the finally below sees it.
+        job.knowledge_graph = graph
+        job.kg_status = "ready"
     except RuntimeError as exc:
         # A Gemini/key/quota/network failure — actionable, single-line message.
         job.kg_status = "error"
@@ -448,15 +451,24 @@ async def graph_build(
         raise HTTPException(
             status_code=502, detail="Knowledge-graph build failed unexpectedly."
         )
+    finally:
+        # A client disconnect raises asyncio.CancelledError — a BaseException
+        # that bypasses the except clauses above (they catch only Exception).
+        # Without this, kg_status would stay stuck on "building" and 409-lock
+        # every future rebuild of this document until it's evicted. Reset it.
+        if job.kg_status == "building":
+            job.kg_status = "error"
+            job.kg_error = "Build was interrupted before it finished; try again."
 
-    job.knowledge_graph = graph
-    job.kg_status = "ready"
     return {
         "job_id": job_id,
         "status": "ready",
         "nodes": graph.num_nodes,
         "edges": graph.num_edges,
         "has_vectors": graph.has_vectors,
+        "pages_built": graph.pages_built,
+        "pages_failed": graph.pages_failed,
+        "embed_error": graph.embed_error,
     }
 
 

pipeline/kg.py

@@ -37,6 +37,7 @@ Public API
     retrieval and returns ranked, explainable results.
 """
 
+import hashlib
 import json
 import re
 import time
@@ -63,6 +64,7 @@ from pipeline.online_ocr import (
 # ---------------------------------------------------------------------------
 _GENERATE_URL = _API_BASE + "/{model}:generateContent"
 _BATCH_EMBED_URL = _API_BASE + "/{model}:batchEmbedContents"
+_EMBED_URL = _API_BASE + "/{model}:embedContent"  # single-item fallback
 
 # 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.
@@ -147,6 +149,19 @@ _TRIPLE_SCHEMA = {
 # ---------------------------------------------------------------------------
 # Low-level HTTP (stdlib) — mirrors online_ocr's request/error handling.
 # ---------------------------------------------------------------------------
+class GeminiHTTPError(RuntimeError):
+    """A friendly, KEY-FREE Gemini HTTP error that also carries the status code.
+
+    Subclasses RuntimeError so every existing ``except RuntimeError`` still
+    catches it; the ``.code`` lets callers branch (e.g. fall back from
+    batchEmbedContents to embedContent on a 404 method-not-found).
+    """
+
+    def __init__(self, message, code=None):
+        super().__init__(message)
+        self.code = code
+
+
 def _post_json(url: str, body: dict, api_key: str, *, timeout: float) -> dict:
     """POST a JSON body to Gemini and return the parsed JSON response.
 
@@ -177,7 +192,7 @@ def _post_json(url: str, body: dict, api_key: str, *, timeout: float) -> dict:
             if err.code in _RETRY_STATUSES and attempt < _RETRY_ATTEMPTS - 1:
                 time.sleep(_RETRY_BASE_DELAY * (2 ** attempt))
                 continue
-            raise _friendly_http_error(err) from None
+            raise GeminiHTTPError(str(_friendly_http_error(err)), code=err.code) from None
         except urllib.error.URLError as err:
             raise RuntimeError(
                 "Could not reach the Gemini API ({}). Check your internet "
@@ -188,7 +203,10 @@ def _post_json(url: str, body: dict, api_key: str, *, timeout: float) -> dict:
                 "The Gemini request timed out after {}s.".format(timeout)
             ) from None
     else:  # pragma: no cover - loop always breaks or raises
-        raise _friendly_http_error(last_http_err)
+        raise GeminiHTTPError(
+            str(_friendly_http_error(last_http_err)),
+            code=getattr(last_http_err, "code", None),
+        )
 
     try:
         payload = json.loads(raw.decode("utf-8", "replace"))
@@ -213,11 +231,12 @@ 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
+# Process-level cache of the discovered embedding model, keyed by a HASH of the
+# API key (never the key itself). Different keys/projects can have access to
+# different models — on the multi-user demo Space the first visitor's key must
+# NOT pin the model for everyone. Only a SUCCESSFUL discovery is cached, so a
+# transient failure that fell back to the default is retried next time.
+_EMBED_MODEL_RESOLVED: dict = {}
 
 
 def _list_embedding_models(api_key, *, timeout=30) -> list:
@@ -273,11 +292,12 @@ def _pick_embedding_model(api_key, requested=None, *, timeout=30) -> str:
     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
+    key_hash = hashlib.sha256(_clean_key(api_key).encode("utf-8")).hexdigest()
+    cached = _EMBED_MODEL_RESOLVED.get(key_hash)
+    if cached:
+        return cached
     try:
         available = set(_list_embedding_models(api_key, timeout=timeout))
     except RuntimeError:
@@ -291,7 +311,7 @@ def _pick_embedding_model(api_key, requested=None, *, timeout=30) -> str:
             sorted(available)[0],
         )
     if chosen:
-        _EMBED_MODEL_RESOLVED = chosen  # cache only a real discovery
+        _EMBED_MODEL_RESOLVED[key_hash] = chosen  # cache only a real discovery
         return chosen
     return DEFAULT_EMBED_MODEL
 
@@ -338,11 +358,15 @@ def _parse_triple_payload(payload: dict) -> dict:
     text = _response_text(payload).strip()
     if not text:
         return {"entities": [], "triples": []}
-    # Strip an accidental ```json ... ``` fence if the model added one.
+    # Strip a ```json ... ``` fence if the model wrapped its JSON in one — but
+    # only the fence markers, so backticks legitimately inside the content
+    # (e.g. a value with a code span) are preserved.
     if text.startswith("```"):
-        text = text.strip("`")
+        text = text[3:]
         if text[:4].lower() == "json":
             text = text[4:]
+        if text.endswith("```"):
+            text = text[:-3]
         text = text.strip()
     try:
         obj = json.loads(text)
@@ -402,6 +426,24 @@ def _parse_embed_payload(payload: dict, expected: int) -> list:
     return out
 
 
+def _parse_single_embed_payload(payload: dict) -> list:
+    """Parse a single-item embedContent response: ``{"embedding": {"values": [...]}}``."""
+    emb = payload.get("embedding") if isinstance(payload, dict) else None
+    values = emb.get("values") if isinstance(emb, dict) else None
+    if not isinstance(values, list) or not values:
+        raise RuntimeError("Gemini returned an empty embedding vector.")
+    return [float(v) for v in values]
+
+
+def _embed_item(wire_model: str, text: str, task_type: str) -> dict:
+    """One embedding request item (same shape for batch list and single body)."""
+    return {
+        "model": wire_model,
+        "content": {"parts": [{"text": text}]},
+        "taskType": task_type,
+    }
+
+
 # ---------------------------------------------------------------------------
 # Public extraction / embedding calls.
 # ---------------------------------------------------------------------------
@@ -467,23 +509,29 @@ def embed_texts(
     api_key = _clean_key(api_key)
     model_id = _resolve_model(model, DEFAULT_EMBED_MODEL)
     wire_model = "models/" + model_id
-    url = _BATCH_EMBED_URL.format(model=model_id)
+    batch_url = _BATCH_EMBED_URL.format(model=model_id)
+    single_url = _EMBED_URL.format(model=model_id)
 
     vectors: list = []
+    use_batch = True
     for start in range(0, len(items), _EMBED_BATCH):
         chunk = items[start:start + _EMBED_BATCH]
-        body = {
-            "requests": [
-                {
-                    "model": wire_model,
-                    "content": {"parts": [{"text": t}]},
-                    "taskType": task_type,
-                }
-                for t in chunk
-            ]
-        }
-        payload = _post_json(url, body, api_key, timeout=timeout)
-        vectors.extend(_parse_embed_payload(payload, len(chunk)))
+        if use_batch:
+            body = {"requests": [_embed_item(wire_model, t, task_type) for t in chunk]}
+            try:
+                payload = _post_json(batch_url, body, api_key, timeout=timeout)
+                vectors.extend(_parse_embed_payload(payload, len(chunk)))
+                continue
+            except GeminiHTTPError as exc:
+                # 404 = this model/version doesn't expose batchEmbedContents.
+                # Degrade to per-item embedContent instead of failing the build.
+                if exc.code != 404:
+                    raise
+                use_batch = False
+        for t in chunk:
+            payload = _post_json(single_url, _embed_item(wire_model, t, task_type),
+                                 api_key, timeout=timeout)
+            vectors.append(_parse_single_embed_payload(payload))
 
     return np.asarray(vectors, dtype=np.float32)
 
@@ -521,6 +569,10 @@ class KnowledgeGraph:
     _index: dict = field(default_factory=dict, repr=False)  # norm_key -> node_id
     vectors: Optional[np.ndarray] = field(default=None, repr=False)  # (n, dim) L2-normalized
     embed_model: Optional[str] = None
+    # Build telemetry (surfaced so partial builds / degraded search aren't silent).
+    pages_built: int = 0      # pages whose triple extraction succeeded
+    pages_failed: int = 0     # pages skipped due to a per-page Gemini failure
+    embed_error: Optional[str] = None  # set if embedding failed -> lexical-only
 
     # --- construction helpers ------------------------------------------------
     def _add_node(self, name: str, etype: str = "OTHER", page: Optional[int] = None) -> int:
@@ -604,34 +656,33 @@ class KnowledgeGraph:
         return self.vectors is not None and self.vectors.shape[0] == len(self.nodes)
 
     # --- retrieval -----------------------------------------------------------
-    def _semantic_scores(self, query: str, query_vec: Optional[np.ndarray]) -> np.ndarray:
+    def _semantic_scores(self, query: str, query_vec: Optional[np.ndarray]):
         """Per-node relevance to the query in [0,1].
 
-        Uses cosine similarity when vectors + a query vector are available;
-        otherwise falls back to lexical token overlap so the feature still works
-        offline / when embeddings are unavailable.
+        Returns ``(scores, used_vectors)``. Uses cosine similarity when vectors +
+        a usable query vector are available; otherwise falls back to lexical
+        token overlap (``used_vectors=False``) so the caller can report the mode
+        accurately even when it silently fell back (e.g. a dim mismatch).
         """
         n = len(self.nodes)
         if n == 0:
-            return np.zeros((0,), dtype=np.float32)
-        if self.has_vectors and query_vec is not None and query_vec.size:
+            return np.zeros((0,), dtype=np.float32), False
+        if self.has_vectors and query_vec is not None and getattr(query_vec, "size", 0):
             q = np.asarray(query_vec, dtype=np.float32).reshape(-1)
             qn = np.linalg.norm(q)
-            if qn > 0:
-                q = q / qn
-                if q.shape[0] == self.vectors.shape[1]:
-                    sims = self.vectors @ q  # cosine; both are L2-normalized
-                    return ((sims + 1.0) / 2.0).astype(np.float32)  # map [-1,1]->[0,1]
+            if qn > 0 and q.shape[0] == self.vectors.shape[1]:
+                sims = self.vectors @ (q / qn)  # cosine; both are L2-normalized
+                return ((sims + 1.0) / 2.0).astype(np.float32), True  # [-1,1]->[0,1]
         # Lexical fallback.
         q_tokens = _tokens(query)
         scores = np.zeros((n,), dtype=np.float32)
         if not q_tokens:
-            return scores
+            return scores, False
         for i, node in enumerate(self.nodes):
             nt = _tokens(node["name"])
             if nt:
                 scores[i] = len(q_tokens & nt) / float(len(q_tokens | nt))
-        return scores
+        return scores, False
 
     def _multi_source_paths(self, seeds: list, hops: int) -> dict:
         """Multi-source BFS from ``seeds`` up to ``hops``.
@@ -722,8 +773,8 @@ class KnowledgeGraph:
             except RuntimeError:
                 query_vec = None
 
-        sem = self._semantic_scores(query, query_vec)
-        out["mode"] = "semantic" if (self.has_vectors and query_vec is not None) else "lexical"
+        sem, used_vectors = self._semantic_scores(query, query_vec)
+        out["mode"] = "semantic" if used_vectors else "lexical"
         if sem.size == 0 or float(sem.max()) <= 0:
             return out
 
@@ -876,26 +927,49 @@ def build_graph(
         usable = usable[:max_pages]
 
     graph = KnowledgeGraph()
+    last_error: Optional[Exception] = None
     for page_no, text in usable:
-        data = extract_triples(text, api_key=api_key, model=triple_model, timeout=timeout)
+        # Isolate each page: a single rate-limited / safety-blocked page (common
+        # on the free tier) must NOT discard every page already processed (and
+        # its API spend). Skip it and keep going.
+        try:
+            data = extract_triples(text, api_key=api_key, model=triple_model, timeout=timeout)
+        except RuntimeError as exc:
+            graph.pages_failed += 1
+            last_error = exc
+            continue
+        graph.pages_built += 1
         # Register typed entities first so types are known, then triples.
         for e in data.get("entities", []):
             graph._add_node(e["name"], e.get("type", "OTHER"), page=page_no)
         for t in data.get("triples", []):
             graph._add_triple(t["subject"], t["predicate"], t["object"], page=page_no)
 
+    # If EVERY page failed, this is a systemic failure (bad key / quota / network),
+    # not "a document with no facts" — surface it instead of returning an empty graph.
+    if usable and graph.pages_failed == len(usable):
+        raise last_error or RuntimeError(
+            "Knowledge-graph extraction failed for every page."
+        )
+
     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=resolved,
-            task_type="RETRIEVAL_DOCUMENT",
-            timeout=timeout,
-        )
-        graph.set_vectors(vectors, model=resolved)
+        # Embeddings are the SEMANTIC half: if they fail (quota / retired model),
+        # keep the graph for LEXICAL search rather than throwing away the (costly)
+        # triple extraction. The degradation is surfaced via has_vectors/embed_error.
+        try:
+            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=resolved,
+                task_type="RETRIEVAL_DOCUMENT",
+                timeout=timeout,
+            )
+            graph.set_vectors(vectors, model=resolved)
+        except RuntimeError as exc:
+            graph.embed_error = str(exc)
 
     return graph

static/app.js

@@ -1183,10 +1183,17 @@
         return;
       }
       setKgStatus(job, data.nodes + " entities · " + data.edges + " facts", "good");
-      el.kgBuildNote.textContent =
-        data.has_vectors
-          ? "Graph ready — ask a question below (semantic + graph search)."
-          : "Graph ready — ask a question below (graph search).";
+      let note = data.has_vectors
+        ? "Graph ready — ask a question below (semantic + graph search)."
+        : "Graph ready — ask a question below (keyword graph search).";
+      if (data.pages_failed) {
+        note += " " + data.pages_failed + " page(s) were skipped (rate limit or block).";
+      }
+      if (data.embed_error) {
+        note += " Semantic search unavailable (" + shortErr(data.embed_error) +
+          ") — using keyword search.";
+      }
+      el.kgBuildNote.textContent = note;
       if (el.kgBuildBtn) el.kgBuildBtn.hidden = true;
       el.kgQueryWrap.hidden = false;
       showToast("Knowledge graph: " + data.nodes + " entities, " + data.edges + " facts");