From 17af3c40fd661260f83a71364706faa5a5c2d93a Mon Sep 17 00:00:00 2001 From: Aaron Sachs <898627+asachs01@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:27:32 +0000 Subject: [PATCH 1/3] fix(kb): add retry-with-backoff to embed_content, mirroring generate_content MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit embed_content() had zero retry logic while generate_content already retries transient 429/500/503 errors via _retry_generate_content. A large text file (e.g. an agent's MEMORY.md) embeds every chunk back-to-back with no pacing, so once the file's own chunk volume trips the shared embedding quota, every remaining chunk in that run also 429s with nothing pausing for recovery — the ingest dies partway through on every attempt, independent of any cross-agent contention (confirmed via isolated reproduction and adoption's MEMORY.md failing 5 consecutive staggered-schedule ingests). Adds _retry_embed_content mirroring the existing _retry_generate_content pattern exactly, wires embed_content() through it. Extends the fault_injection test harness to script embed_content independently of generate_content, with 4 new regression tests. --- .../scripts/_test_clients/fault_injection.py | 55 +++++-- .../scripts/_test_clients/test_retry_embed.py | 151 ++++++++++++++++++ knowledge-base/scripts/mmrag.py | 42 ++++- 3 files changed, 232 insertions(+), 16 deletions(-) create mode 100644 knowledge-base/scripts/_test_clients/test_retry_embed.py diff --git a/knowledge-base/scripts/_test_clients/fault_injection.py b/knowledge-base/scripts/_test_clients/fault_injection.py index 5320c38e8a..3c7d5bcd3e 100644 --- a/knowledge-base/scripts/_test_clients/fault_injection.py +++ b/knowledge-base/scripts/_test_clients/fault_injection.py @@ -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): diff --git a/knowledge-base/scripts/_test_clients/test_retry_embed.py b/knowledge-base/scripts/_test_clients/test_retry_embed.py new file mode 100644 index 0000000000..fbbfa61519 --- /dev/null +++ b/knowledge-base/scripts/_test_clients/test_retry_embed.py @@ -0,0 +1,151 @@ +"""Behavioral tests for mmrag._retry_embed_content. + +Run from knowledge-base/scripts: + + python -m _test_clients.test_retry_embed + +Exits 0 on all-pass, 1 on any failure. Mirrors test_retry.py's three scenarios +against embed_content instead of generate_content — added 2026-09-03 +(task_1788420454462_38838015) alongside _retry_embed_content itself, since +embed_content previously had NO retry at all and a large file's own chunk +volume could self-trip a 429 with zero other-agent contention (adoption's +MEMORY.md: 5 consecutive failed ingests across 3 already-staggered cycles). + + 1. transient_then_success: 429 -> 200 -> returns embedding, no raise + 2. all_exhausted: 429 -> 429 -> 429 -> raises last APIError + 3. fail_fast_nontransient: 403 -> raises immediately, predicate is structural + +backoffs is passed as (0, 0, 0) so tests run in milliseconds. +""" + +import os +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +PARENT = os.path.dirname(HERE) +if PARENT not in sys.path: + sys.path.insert(0, PARENT) + +import mmrag +from _test_clients import fault_injection + + +FAILURES = [] + + +def _check(label, cond, detail=""): + if cond: + print(f" PASS {label}") + else: + print(f" FAIL {label}: {detail}") + FAILURES.append(label) + + +def test_transient_then_success(): + print("\n[test 1/3] transient_then_success: 429 -> 200") + client = fault_injection.FaultInjectionClient( + [], embed_script=fault_injection._parse_script("429:quota exhausted,200:ok") + ) + response = mmrag._retry_embed_content( + client, model="x", contents="chunk text", embed_config=None, backoffs=(0, 0, 0) + ) + _check("returns response after one transient", response is not None) + _check( + "response.embeddings[0].values is the stubbed vector", + getattr(response.embeddings[0], "values", None) == [0.1, 0.2, 0.3], + detail=f"got {getattr(response.embeddings[0], 'values', None)!r}", + ) + _check( + "consumed exactly 2 attempts", + client.models._embed_index == 2, + detail=f"got {client.models._embed_index}", + ) + + +def test_all_exhausted(): + print("\n[test 2/3] all_exhausted: 429 -> 429 -> 429 -> re-raise") + client = fault_injection.FaultInjectionClient( + [], embed_script=fault_injection._parse_script("429,429,429") + ) + raised = None + try: + mmrag._retry_embed_content( + client, model="x", contents="chunk text", embed_config=None, backoffs=(0, 0, 0) + ) + except Exception as e: + raised = e + _check("raises after all attempts exhausted", raised is not None) + if raised is not None: + _check("raised.code is 429", getattr(raised, "code", None) == 429) + _check( + "raised.status is RESOURCE_EXHAUSTED", + getattr(raised, "status", None) == "RESOURCE_EXHAUSTED", + ) + _check( + "consumed exactly 3 attempts", + client.models._embed_index == 3, + detail=f"got {client.models._embed_index}", + ) + + +def test_fail_fast_nontransient(): + print("\n[test 3/3] fail_fast_nontransient: 403 -> raises immediately") + client = fault_injection.FaultInjectionClient( + [], embed_script=fault_injection._parse_script( + "403:Permission denied,200:should not reach" + ) + ) + raised = None + try: + mmrag._retry_embed_content( + client, model="x", contents="chunk text", embed_config=None, backoffs=(0, 0, 0) + ) + except Exception as e: + raised = e + _check("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._embed_index == 1, + detail=f"got {client.models._embed_index}", + ) + + +def test_generate_content_still_unscripted_by_default(): + """Guards against a regression where embed_script leaks into generate_content's + script or vice versa — the two must stay independently scriptable.""" + print("\n[test 4/4] generate_content path is untouched by embed_script wiring") + client = fault_injection.FaultInjectionClient( + fault_injection._parse_script("200:generate ok"), + embed_script=fault_injection._parse_script("200"), + ) + gen_response = mmrag._retry_generate_content( + client, model="x", contents=["x"], backoffs=(0, 0, 0) + ) + _check("generate_content still returns its own scripted response", + getattr(gen_response, "text", None) == "generate ok") + embed_response = mmrag._retry_embed_content( + client, model="x", contents="y", embed_config=None, backoffs=(0, 0, 0) + ) + _check("embed_content returns its own separately-scripted response", + getattr(embed_response.embeddings[0], "values", None) == [0.1, 0.2, 0.3]) + + +if __name__ == "__main__": + test_transient_then_success() + test_all_exhausted() + test_fail_fast_nontransient() + test_generate_content_still_unscripted_by_default() + print() + if FAILURES: + print(f"FAILED: {len(FAILURES)} assertion(s)") + for f in FAILURES: + print(f" - {f}") + sys.exit(1) + print(f"ALL PASS (4 scenarios)") + sys.exit(0) diff --git a/knowledge-base/scripts/mmrag.py b/knowledge-base/scripts/mmrag.py index 9eea792658..d668799735 100755 --- a/knowledge-base/scripts/mmrag.py +++ b/knowledge-base/scripts/mmrag.py @@ -292,13 +292,51 @@ def _retry_generate_content(client, *, model, contents, backoffs=(5, 15, 45)): raise last_err if last_err else RuntimeError("retry loop completed without response or error") +def _retry_embed_content(client, *, model, contents, embed_config, backoffs=(5, 15, 45)): + """Call client.models.embed_content with bounded retries on transient APIErrors. + + Mirrors _retry_generate_content above (same TRANSIENT_HTTP_CODES/TRANSIENT_STATUS_NAMES + classification, same default backoff shape) — added 2026-09-03 (task_1788420454462_38838015) + because embed_content had NO retry at all, unlike generate_content. This matters specifically + for a large text file: chunk_text() on a file the size of MEMORY.md produces 200+ chunks + embedded back-to-back in a tight loop with zero pacing, so a 429 tripped by the file's OWN + chunk volume (confirmed 2026-09-03: a 2nd chunk-sized embed_content call 429'd with ZERO other + agents contending) previously killed the entire rest of that ingest run — every subsequent + chunk in the same loop also 429'd, since nothing paused for the quota to recover. adoption's + MEMORY.md failed 5 consecutive ingests across 3 already-staggered heartbeat cycles, proving + stagger alone does not fix an intra-file burst. This does not fix that structural gap either + (it still bails after backoffs is exhausted) — it turns "the rest of this file's chunks are all + doomed once one 429s" into "this one chunk pauses and retries", which is a materially different + and much better failure mode for a multi-hundred-chunk file. Verify empirically against + adoption's next MEMORY.md ingest attempts; this backoff shape is inherited from + _retry_generate_content, not independently tuned for embedding's burst characteristics. + """ + from google.genai import errors as _genai_errors + last_err = None + for attempt, backoff in enumerate(backoffs, start=1): + try: + return client.models.embed_content(model=model, contents=contents, config=embed_config) + 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 ''}) on embed_content; retrying in {backoff}s (attempt {attempt}/{len(backoffs)})") + time.sleep(backoff) + else: + print(f" Exhausted retries on transient embed_content 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( + result = _retry_embed_content( + client, model=config.get("embedding_model", "gemini-embedding-2-preview"), contents=content, - config=types.EmbedContentConfig( + embed_config=types.EmbedContentConfig( output_dimensionality=config.get("embedding_dimensions", DEFAULT_EMBEDDING_DIMENSIONS), task_type=task_type, ), From 9478ab6051fdf547c30b6d5a1b3fece5cbd8a200 Mon Sep 17 00:00:00 2001 From: Aaron Sachs <898627+asachs01@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:32:19 +0000 Subject: [PATCH 2/3] refactor(kb): consolidate embed/generate retry into one _retry_with_backoff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /simplify review of the prior commit found _retry_embed_content was a near-verbatim copy of _retry_generate_content (same loop, same transient-code classification, same backoff shape — only the wrapped call differed). Collapse both into a single _retry_with_backoff(fn, label=...) taking a zero-arg callable, so a third retry-needing call site only needs a lambda, not a third copy of the loop. Also merges test_retry_embed.py into test_retry.py: since there is now one implementation, testing it via two near-duplicate suites (one per call shape) re-introduced the same duplication at the test layer. test_retry.py now runs each scenario against both a generate_content-shaped and an embed_content- shaped call through a shared table, plus a cross-shape independence check. Reviewed and deliberately NOT addressed here (would require changes well outside this diff): proactive inter-chunk pacing in the ingest loops themselves (retry-after-failure treats the symptom; the tight per-chunk loop is what generates the burst) — flagged by /simplify's altitude pass as the deeper fix, tracked for follow-up rather than expanding this diff's scope. Also not changed: the backoffs=(5,15,45) default is borrowed from generate_content's tuning, not independently measured for embedding's burst window — real-world validation is pending against a live large-file ingest. --- .../scripts/_test_clients/test_retry.py | 162 ++++++++++-------- .../scripts/_test_clients/test_retry_embed.py | 151 ---------------- knowledge-base/scripts/mmrag.py | 84 +++------ 3 files changed, 120 insertions(+), 277 deletions(-) delete mode 100644 knowledge-base/scripts/_test_clients/test_retry_embed.py diff --git a/knowledge-base/scripts/_test_clients/test_retry.py b/knowledge-base/scripts/_test_clients/test_retry.py index 022b12cf75..3824678f4d 100644 --- a/knowledge-base/scripts/_test_clients/test_retry.py +++ b/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) diff --git a/knowledge-base/scripts/_test_clients/test_retry_embed.py b/knowledge-base/scripts/_test_clients/test_retry_embed.py deleted file mode 100644 index fbbfa61519..0000000000 --- a/knowledge-base/scripts/_test_clients/test_retry_embed.py +++ /dev/null @@ -1,151 +0,0 @@ -"""Behavioral tests for mmrag._retry_embed_content. - -Run from knowledge-base/scripts: - - python -m _test_clients.test_retry_embed - -Exits 0 on all-pass, 1 on any failure. Mirrors test_retry.py's three scenarios -against embed_content instead of generate_content — added 2026-09-03 -(task_1788420454462_38838015) alongside _retry_embed_content itself, since -embed_content previously had NO retry at all and a large file's own chunk -volume could self-trip a 429 with zero other-agent contention (adoption's -MEMORY.md: 5 consecutive failed ingests across 3 already-staggered cycles). - - 1. transient_then_success: 429 -> 200 -> returns embedding, no raise - 2. all_exhausted: 429 -> 429 -> 429 -> raises last APIError - 3. fail_fast_nontransient: 403 -> raises immediately, predicate is structural - -backoffs is passed as (0, 0, 0) so tests run in milliseconds. -""" - -import os -import sys - -HERE = os.path.dirname(os.path.abspath(__file__)) -PARENT = os.path.dirname(HERE) -if PARENT not in sys.path: - sys.path.insert(0, PARENT) - -import mmrag -from _test_clients import fault_injection - - -FAILURES = [] - - -def _check(label, cond, detail=""): - if cond: - print(f" PASS {label}") - else: - print(f" FAIL {label}: {detail}") - FAILURES.append(label) - - -def test_transient_then_success(): - print("\n[test 1/3] transient_then_success: 429 -> 200") - client = fault_injection.FaultInjectionClient( - [], embed_script=fault_injection._parse_script("429:quota exhausted,200:ok") - ) - response = mmrag._retry_embed_content( - client, model="x", contents="chunk text", embed_config=None, backoffs=(0, 0, 0) - ) - _check("returns response after one transient", response is not None) - _check( - "response.embeddings[0].values is the stubbed vector", - getattr(response.embeddings[0], "values", None) == [0.1, 0.2, 0.3], - detail=f"got {getattr(response.embeddings[0], 'values', None)!r}", - ) - _check( - "consumed exactly 2 attempts", - client.models._embed_index == 2, - detail=f"got {client.models._embed_index}", - ) - - -def test_all_exhausted(): - print("\n[test 2/3] all_exhausted: 429 -> 429 -> 429 -> re-raise") - client = fault_injection.FaultInjectionClient( - [], embed_script=fault_injection._parse_script("429,429,429") - ) - raised = None - try: - mmrag._retry_embed_content( - client, model="x", contents="chunk text", embed_config=None, backoffs=(0, 0, 0) - ) - except Exception as e: - raised = e - _check("raises after all attempts exhausted", raised is not None) - if raised is not None: - _check("raised.code is 429", getattr(raised, "code", None) == 429) - _check( - "raised.status is RESOURCE_EXHAUSTED", - getattr(raised, "status", None) == "RESOURCE_EXHAUSTED", - ) - _check( - "consumed exactly 3 attempts", - client.models._embed_index == 3, - detail=f"got {client.models._embed_index}", - ) - - -def test_fail_fast_nontransient(): - print("\n[test 3/3] fail_fast_nontransient: 403 -> raises immediately") - client = fault_injection.FaultInjectionClient( - [], embed_script=fault_injection._parse_script( - "403:Permission denied,200:should not reach" - ) - ) - raised = None - try: - mmrag._retry_embed_content( - client, model="x", contents="chunk text", embed_config=None, backoffs=(0, 0, 0) - ) - except Exception as e: - raised = e - _check("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._embed_index == 1, - detail=f"got {client.models._embed_index}", - ) - - -def test_generate_content_still_unscripted_by_default(): - """Guards against a regression where embed_script leaks into generate_content's - script or vice versa — the two must stay independently scriptable.""" - print("\n[test 4/4] generate_content path is untouched by embed_script wiring") - client = fault_injection.FaultInjectionClient( - fault_injection._parse_script("200:generate ok"), - embed_script=fault_injection._parse_script("200"), - ) - gen_response = mmrag._retry_generate_content( - client, model="x", contents=["x"], backoffs=(0, 0, 0) - ) - _check("generate_content still returns its own scripted response", - getattr(gen_response, "text", None) == "generate ok") - embed_response = mmrag._retry_embed_content( - client, model="x", contents="y", embed_config=None, backoffs=(0, 0, 0) - ) - _check("embed_content returns its own separately-scripted response", - getattr(embed_response.embeddings[0], "values", None) == [0.1, 0.2, 0.3]) - - -if __name__ == "__main__": - test_transient_then_success() - test_all_exhausted() - test_fail_fast_nontransient() - test_generate_content_still_unscripted_by_default() - print() - if FAILURES: - print(f"FAILED: {len(FAILURES)} assertion(s)") - for f in FAILURES: - print(f" - {f}") - sys.exit(1) - print(f"ALL PASS (4 scenarios)") - sys.exit(0) diff --git a/knowledge-base/scripts/mmrag.py b/knowledge-base/scripts/mmrag.py index d668799735..fa9e4ffccc 100755 --- a/knowledge-base/scripts/mmrag.py +++ b/knowledge-base/scripts/mmrag.py @@ -264,82 +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. - """ - 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) - 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)})") - time.sleep(backoff) - else: - print(f" Exhausted retries on transient error: HTTP {e.code} {e.status or ''}") - raise last_err if last_err else RuntimeError("retry loop completed without response or error") - -def _retry_embed_content(client, *, model, contents, embed_config, backoffs=(5, 15, 45)): - """Call client.models.embed_content with bounded retries on transient APIErrors. - - Mirrors _retry_generate_content above (same TRANSIENT_HTTP_CODES/TRANSIENT_STATUS_NAMES - classification, same default backoff shape) — added 2026-09-03 (task_1788420454462_38838015) - because embed_content had NO retry at all, unlike generate_content. This matters specifically - for a large text file: chunk_text() on a file the size of MEMORY.md produces 200+ chunks - embedded back-to-back in a tight loop with zero pacing, so a 429 tripped by the file's OWN - chunk volume (confirmed 2026-09-03: a 2nd chunk-sized embed_content call 429'd with ZERO other - agents contending) previously killed the entire rest of that ingest run — every subsequent - chunk in the same loop also 429'd, since nothing paused for the quota to recover. adoption's - MEMORY.md failed 5 consecutive ingests across 3 already-staggered heartbeat cycles, proving - stagger alone does not fix an intra-file burst. This does not fix that structural gap either - (it still bails after backoffs is exhausted) — it turns "the rest of this file's chunks are all - doomed once one 429s" into "this one chunk pauses and retries", which is a materially different - and much better failure mode for a multi-hundred-chunk file. Verify empirically against - adoption's next MEMORY.md ingest attempts; this backoff shape is inherited from - _retry_generate_content, not independently tuned for embedding's burst characteristics. + 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.embed_content(model=model, contents=contents, config=embed_config) + 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 ''}) on embed_content; 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 embed_content 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 = _retry_embed_content( - client, - model=config.get("embedding_model", "gemini-embedding-2-preview"), - contents=content, - embed_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) @@ -898,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) From b4fc01c4efa5ab25fad5254df02698baf4f1f1b1 Mon Sep 17 00:00:00 2001 From: Aaron Sachs <898627+asachs01@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:36:16 +0000 Subject: [PATCH 3/3] docs(kb): fix two stale _retry_generate_content references after rename Caught by analyst's independent review of #170 (ran the test suite in an isolated worktree rather than trusting the PR description): mmrag.py:860's comment and fault_injection.py's module docstring still named the pre- consolidation function. Cosmetic only, no behavior change. --- knowledge-base/scripts/_test_clients/fault_injection.py | 2 +- knowledge-base/scripts/mmrag.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/knowledge-base/scripts/_test_clients/fault_injection.py b/knowledge-base/scripts/_test_clients/fault_injection.py index 3c7d5bcd3e..a83fdaa1c1 100644 --- a/knowledge-base/scripts/_test_clients/fault_injection.py +++ b/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: diff --git a/knowledge-base/scripts/mmrag.py b/knowledge-base/scripts/mmrag.py index fa9e4ffccc..f19e6d77ac 100755 --- a/knowledge-base/scripts/mmrag.py +++ b/knowledge-base/scripts/mmrag.py @@ -857,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 = (