diff --git a/knowledge-base/scripts/_test_clients/fault_injection.py b/knowledge-base/scripts/_test_clients/fault_injection.py index 5320c38e8a..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: @@ -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.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/mmrag.py b/knowledge-base/scripts/mmrag.py index 9eea792658..f19e6d77ac 100755 --- a/knowledge-base/scripts/mmrag.py +++ b/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)