fix(kb): retry-with-backoff on embed_content, fixing intra-file 429 bursts

knowledge-base/scripts/_test_clients/fault_injection.py

@@ -1,4 +1,4 @@
-"""Fault-injecting Gemini client for testing mmrag._retry_generate_content.
+"""Fault-injecting Gemini client for testing mmrag._retry_with_backoff.
 
 Wired via two env vars consumed by mmrag.get_genai_client:
 
@@ -56,34 +56,61 @@ def __init__(self, text):
         self.usage_metadata = None
 
 
+class _StubEmbedding:
+    def __init__(self, values):
+        self.values = values
+
+
+class _StubEmbedResponse:
+    def __init__(self, values):
+        self.embeddings = [_StubEmbedding(values)]
+
+
+def _next_scripted(script, index, kind):
+    if index >= len(script):
+        raise RuntimeError(
+            f"fault_injection: {kind} script exhausted at attempt {index + 1} "
+            f"(scripted {len(script)} responses)"
+        )
+    return script[index]
+
+
 class _StubModels:
-    def __init__(self, script):
+    def __init__(self, script, embed_script=None):
         self._script = list(script)
         self._index = 0
+        # embed_content is scripted SEPARATELY from generate_content — added
+        # 2026-09-03 (task_1788420454462_38838015, _retry_embed_content) so
+        # a test can drive one call type without needing to also script the
+        # other. Kept as None (raises like before) when no test needs it.
+        self._embed_script = list(embed_script) if embed_script is not None else None
+        self._embed_index = 0
 
     def generate_content(self, model=None, contents=None, **kwargs):
-        if self._index >= len(self._script):
-            raise RuntimeError(
-                f"fault_injection: script exhausted at attempt {self._index + 1} "
-                f"(scripted {len(self._script)} responses)"
-            )
-        code, message = self._script[self._index]
+        code, message = _next_scripted(self._script, self._index, "generate_content")
         self._index += 1
         if code == 200:
             return _StubResponse(message or "[stub] fault-injection success")
         status = _STATUS_FOR_CODE.get(code, "UNKNOWN")
         raise _InjectedAPIError(code, status, message or f"injected {code} {status}")
 
-    def embed_content(self, *a, **kw):
-        raise RuntimeError(
-            "fault_injection: embed_content is not scripted. Tests should target "
-            "_retry_generate_content directly, not the full ingest_pdf pipeline."
-        )
+    def embed_content(self, model=None, contents=None, config=None, **kwargs):
+        if self._embed_script is None:
+            raise RuntimeError(
+                "fault_injection: embed_content is not scripted for this client. "
+                "Construct FaultInjectionClient(script, embed_script=...) to test it."
+            )
+        code, message = _next_scripted(self._embed_script, self._embed_index, "embed_content")
+        self._embed_index += 1
+        if code == 200:
+            return _StubEmbedResponse([0.1, 0.2, 0.3])
+        status = _STATUS_FOR_CODE.get(code, "UNKNOWN")
+        raise _InjectedAPIError(code, status, message or f"injected {code} {status}")
 
 
 class FaultInjectionClient:
-    def __init__(self, script):
-        self.models = _StubModels(script)
+    def __init__(self, script, embed_script=None):
+        self.models = _StubModels(script, embed_script=embed_script)
 
 
 def _parse_script(spec):

knowledge-base/scripts/_test_clients/test_retry.py

@@ -1,15 +1,26 @@
-"""Behavioral tests for mmrag._retry_generate_content.
+"""Behavioral tests for mmrag._retry_with_backoff.
 
 Run from knowledge-base/scripts:
 
     python -m _test_clients.test_retry
 
-Exits 0 on all-pass, 1 on any failure. Three scenarios:
+Exits 0 on all-pass, 1 on any failure.
 
-  1. transient_then_success: 503 → 200 → returns response, no raise
-  2. all_exhausted: 503 → 503 → 503 → raises last APIError
-  3. fail_fast_nontransient: 403 (with '503' in body) → raises immediately;
-     proves the predicate is structural (.code / .status), not textual.
+_retry_with_backoff is the single retry implementation shared by both
+generate_content and embed_content call sites (consolidated 2026-09-03,
+task_1788420454462_38838015, from two near-identical ~25-line copies —
+embed_content originally had no retry at all while generate_content did;
+adding a second copy for embed_content was the wrong depth for that fix).
+Each scenario below is run against BOTH a generate_content-shaped call and
+an embed_content-shaped call, via the shared `_run` helper, so the suite
+proves the one implementation behaves identically regardless of which real
+call site wraps it — that is the property the consolidation is supposed to
+guarantee, and a single-call-shape test would not exercise it.
+
+  1. transient_then_success: 429/503 -> 200 -> returns response, no raise
+  2. all_exhausted: 3x transient -> raises last APIError
+  3. fail_fast_nontransient: 403 -> raises immediately; proves the predicate
+     is structural (.code / .status), not textual
 
 backoffs is passed as (0, 0, 0) so tests run in milliseconds.
 """
@@ -37,90 +48,101 @@ def _check(label, cond, detail=""):
         FAILURES.append(label)
 
 
-def test_transient_then_success():
-    print("\n[test 1/3] transient_then_success: 503 -> 200")
-    client = fault_injection.FaultInjectionClient(
-        fault_injection._parse_script("503:gemini busy,200:hello world")
-    )
-    response = mmrag._retry_generate_content(
-        client, model="x", contents=["x"], backoffs=(0, 0, 0)
-    )
-    _check("returns response after one transient", response is not None)
-    _check(
-        "response.text matches scripted message",
-        getattr(response, "text", None) == "hello world",
-        detail=f"got {getattr(response, 'text', None)!r}",
-    )
-    _check(
-        "consumed exactly 2 attempts",
-        client.models._index == 2,
-        detail=f"got {client.models._index}",
-    )
-
-
-def test_all_exhausted():
-    print("\n[test 2/3] all_exhausted: 503 -> 503 -> 503 -> re-raise")
-    client = fault_injection.FaultInjectionClient(
-        fault_injection._parse_script("503,503,503")
-    )
+# Each entry: (call-shape name, script -> client, fn(client) -> the callable to
+# retry, attempts-consumed accessor). Covers both real call sites so the shared
+# retry implementation is proven against both shapes, not just one.
+CALL_SHAPES = {
+    "generate_content": dict(
+        make_client=lambda script: fault_injection.FaultInjectionClient(script),
+        make_fn=lambda client: (lambda: client.models.generate_content(model="x", contents=["x"])),
+        attempts=lambda client: client.models._index,
+        response_ok=lambda r: getattr(r, "text", None) == "hello world",
+    ),
+    "embed_content": dict(
+        make_client=lambda script: fault_injection.FaultInjectionClient([], embed_script=script),
+        make_fn=lambda client: (lambda: client.models.embed_content(model="x", contents="x", config=None)),
+        attempts=lambda client: client.models._embed_index,
+        response_ok=lambda r: getattr(r.embeddings[0], "values", None) == [0.1, 0.2, 0.3],
+    ),
+}
+
+
+def test_transient_then_success(shape_name, shape):
+    print(f"\n[{shape_name}] transient_then_success: 429 -> 200")
+    client = shape["make_client"](fault_injection._parse_script("429:quota exhausted,200:hello world"))
+    response = mmrag._retry_with_backoff(shape["make_fn"](client), label=shape_name, backoffs=(0, 0, 0))
+    _check(f"{shape_name}: returns response after one transient", response is not None)
+    _check(f"{shape_name}: response matches scripted success", shape["response_ok"](response))
+    _check(f"{shape_name}: consumed exactly 2 attempts", shape["attempts"](client) == 2,
+           detail=f"got {shape['attempts'](client)}")
+
+
+def test_all_exhausted(shape_name, shape):
+    print(f"\n[{shape_name}] all_exhausted: 429 -> 429 -> 429 -> re-raise")
+    client = shape["make_client"](fault_injection._parse_script("429,429,429"))
     raised = None
     try:
-        mmrag._retry_generate_content(
-            client, model="x", contents=["x"], backoffs=(0, 0, 0)
-        )
+        mmrag._retry_with_backoff(shape["make_fn"](client), label=shape_name, backoffs=(0, 0, 0))
     except Exception as e:
         raised = e
-    _check("raises after all attempts exhausted", raised is not None)
+    _check(f"{shape_name}: raises after all attempts exhausted", raised is not None)
     if raised is not None:
-        _check("raised.code is 503", getattr(raised, "code", None) == 503)
-        _check(
-            "raised.status is UNAVAILABLE",
-            getattr(raised, "status", None) == "UNAVAILABLE",
-        )
-    _check(
-        "consumed exactly 3 attempts",
-        client.models._index == 3,
-        detail=f"got {client.models._index}",
-    )
+        _check(f"{shape_name}: raised.code is 429", getattr(raised, "code", None) == 429)
+        _check(f"{shape_name}: raised.status is RESOURCE_EXHAUSTED",
+               getattr(raised, "status", None) == "RESOURCE_EXHAUSTED")
+    _check(f"{shape_name}: consumed exactly 3 attempts", shape["attempts"](client) == 3,
+           detail=f"got {shape['attempts'](client)}")
 
 
-def test_fail_fast_nontransient():
-    print("\n[test 3/3] fail_fast_nontransient: 403 (with '503' in body) -> raises immediately")
-    client = fault_injection.FaultInjectionClient(
-        fault_injection._parse_script(
-            "403:Permission denied for resource ID 503-pseudo,200:should not reach"
-        )
-    )
+def test_fail_fast_nontransient(shape_name, shape):
+    print(f"\n[{shape_name}] fail_fast_nontransient: 403 -> raises immediately")
+    client = shape["make_client"](fault_injection._parse_script("403:Permission denied,200:should not reach"))
     raised = None
     try:
-        mmrag._retry_generate_content(
-            client, model="x", contents=["x"], backoffs=(0, 0, 0)
-        )
+        mmrag._retry_with_backoff(shape["make_fn"](client), label=shape_name, backoffs=(0, 0, 0))
     except Exception as e:
         raised = e
-    _check("raises immediately on non-transient", raised is not None)
+    _check(f"{shape_name}: raises immediately on non-transient", raised is not None)
     if raised is not None:
-        _check("raised.code is 403", getattr(raised, "code", None) == 403)
-        _check(
-            "raised.status is PERMISSION_DENIED",
-            getattr(raised, "status", None) == "PERMISSION_DENIED",
-        )
-    _check(
-        "did NOT consume the second scripted attempt (predicate is structural, not textual)",
-        client.models._index == 1,
-        detail=f"got {client.models._index}",
+        _check(f"{shape_name}: raised.code is 403", getattr(raised, "code", None) == 403)
+        _check(f"{shape_name}: raised.status is PERMISSION_DENIED",
+               getattr(raised, "status", None) == "PERMISSION_DENIED")
+    _check(f"{shape_name}: did NOT consume the second scripted attempt (predicate is structural)",
+           shape["attempts"](client) == 1, detail=f"got {shape['attempts'](client)}")
+
+
+def test_shapes_are_independently_scriptable():
+    """Guards the fault-injection double itself: scripting one call shape must
+    not leak into or block the other on the same client."""
+    print("\n[cross-shape] generate_content and embed_content script independently on one client")
+    client = fault_injection.FaultInjectionClient(
+        fault_injection._parse_script("200:generate ok"),
+        embed_script=fault_injection._parse_script("200"),
+    )
+    gen = mmrag._retry_with_backoff(
+        lambda: client.models.generate_content(model="x", contents=["x"]),
+        label="generate_content", backoffs=(0, 0, 0),
+    )
+    _check("generate_content returns its own scripted response", getattr(gen, "text", None) == "generate ok")
+    emb = mmrag._retry_with_backoff(
+        lambda: client.models.embed_content(model="x", contents="y", config=None),
+        label="embed_content", backoffs=(0, 0, 0),
     )
+    _check("embed_content returns its own separately-scripted response",
+           getattr(emb.embeddings[0], "values", None) == [0.1, 0.2, 0.3])
 
 
 if __name__ == "__main__":
-    test_transient_then_success()
-    test_all_exhausted()
-    test_fail_fast_nontransient()
+    for shape_name, shape in CALL_SHAPES.items():
+        test_transient_then_success(shape_name, shape)
+        test_all_exhausted(shape_name, shape)
+        test_fail_fast_nontransient(shape_name, shape)
+    test_shapes_are_independently_scriptable()
     print()
     if FAILURES:
         print(f"FAILED: {len(FAILURES)} assertion(s)")
         for f in FAILURES:
             print(f"  - {f}")
         sys.exit(1)
-    print(f"ALL PASS (3 scenarios)")
+    print(f"ALL PASS ({3 * len(CALL_SHAPES) + 1} scenarios)")
     sys.exit(0)

knowledge-base/scripts/mmrag.py

@@ -264,44 +264,53 @@ def get_genai_client(api_key):
     return genai.Client(api_key=api_key)
 
 
-def _retry_generate_content(client, *, model, contents, backoffs=(5, 15, 45)):
-    """Call client.models.generate_content with bounded retries on transient APIErrors.
+def _retry_with_backoff(fn, *, label, backoffs=(5, 15, 45)):
+    """Call the zero-arg `fn` with bounded retries on transient APIErrors.
 
     Retries on HTTP code in TRANSIENT_HTTP_CODES or status name in
     TRANSIENT_STATUS_NAMES; re-raises immediately on any other APIError (auth,
     malformed request, etc.); re-raises last_err after all attempts exhausted.
+    `label` is used only in log lines, to say which call is retrying.
 
     backoffs is a tuple of sleep seconds between attempts. len(backoffs) is the
     attempt count. Tests pass (0, 0, 0) to skip sleeps.
+
+    Single implementation shared by generate_content and embed_content (was two
+    near-identical ~25-line copies until 2026-09-03, task_1788420454462_38838015
+    — embed_content had no retry at all, generate_content already did; adding a
+    second copy for embed_content was caught in review as the wrong depth for
+    the fix and collapsed into this shared version instead. A third API surface
+    needing retry needs only a `fn` callable, not a third copy of this loop).
     """
     from google.genai import errors as _genai_errors
     last_err = None
     for attempt, backoff in enumerate(backoffs, start=1):
         try:
-            return client.models.generate_content(model=model, contents=contents)
+            return fn()
         except _genai_errors.APIError as e:
             last_err = e
             is_transient = (e.code in TRANSIENT_HTTP_CODES) or (e.status in TRANSIENT_STATUS_NAMES)
             if not is_transient:
                 raise
             if attempt < len(backoffs):
-                print(f"    Transient error (HTTP {e.code} {e.status or ''}); retrying in {backoff}s (attempt {attempt}/{len(backoffs)})")
+                print(f"    Transient error (HTTP {e.code} {e.status or ''}) on {label}; retrying in {backoff}s (attempt {attempt}/{len(backoffs)})")
                 time.sleep(backoff)
             else:
-                print(f"    Exhausted retries on transient error: HTTP {e.code} {e.status or ''}")
+                print(f"    Exhausted retries on transient {label} error: HTTP {e.code} {e.status or ''}")
     raise last_err if last_err else RuntimeError("retry loop completed without response or error")
 
 
 def embed_content(client, config, content, task_type="RETRIEVAL_DOCUMENT"):
     """Embed content using Gemini Embedding 2. Content can be text string or list of Parts."""
     from google.genai import types
-    result = client.models.embed_content(
-        model=config.get("embedding_model", "gemini-embedding-2-preview"),
-        contents=content,
-        config=types.EmbedContentConfig(
-            output_dimensionality=config.get("embedding_dimensions", DEFAULT_EMBEDDING_DIMENSIONS),
-            task_type=task_type,
-        ),
+    model = config.get("embedding_model", "gemini-embedding-2-preview")
+    embed_config = types.EmbedContentConfig(
+        output_dimensionality=config.get("embedding_dimensions", DEFAULT_EMBEDDING_DIMENSIONS),
+        task_type=task_type,
+    )
+    result = _retry_with_backoff(
+        lambda: client.models.embed_content(model=model, contents=content, config=embed_config),
+        label="embed_content",
     )
     if _tracker:
         _tracker.track_embedding(content)
@@ -848,7 +857,7 @@ def ingest_pdf(client, config, collection, file_path):
     print(f"  Analyzing PDF: {file_path.name}...")
 
     # Gemini Flash returns 503 UNAVAILABLE during high-demand windows. Without
-    # retries, a single 503 kills the ingest. _retry_generate_content wraps the
+    # retries, a single 503 kills the ingest. _retry_with_backoff wraps the
     # call with bounded retries on transient SDK conditions (HTTP 429/500/503,
     # status UNAVAILABLE/RESOURCE_EXHAUSTED) and fails fast on everything else.
     extraction_prompt = (
@@ -860,13 +869,14 @@ def ingest_pdf(client, config, collection, file_path):
         "Separate each page's content with '=== PAGE N ===' markers.\n"
         "Be thorough - this will be used for search and retrieval."
     )
-    response = _retry_generate_content(
-        client,
-        model=config.get("gemini_model", "gemini-2.5-flash"),
-        contents=[
-            types.Part.from_bytes(data=data, mime_type="application/pdf"),
-            extraction_prompt,
-        ],
+    gen_model = config.get("gemini_model", "gemini-2.5-flash")
+    gen_contents = [
+        types.Part.from_bytes(data=data, mime_type="application/pdf"),
+        extraction_prompt,
+    ]
+    response = _retry_with_backoff(
+        lambda: client.models.generate_content(model=gen_model, contents=gen_contents),
+        label="generate_content",
     )
     if _tracker:
         _tracker.track_generation(response)