From 535dfc737c976fa73f0294e96e15d92bf4a166e4 Mon Sep 17 00:00:00 2001 From: Rudy Celekli <47457359+rudycelekli@users.noreply.github.com> Date: Mon, 5 Oct 2026 04:25:44 -0400 Subject: [PATCH 1/2] fix(providers): reject native Gemini and Bedrock judge cutoffs Signed-off-by: Rudy Celekli <47457359+rudycelekli@users.noreply.github.com> --- ifixai/providers/bedrock.py | 9 ++- ifixai/providers/gemini.py | 7 ++ ifixai/tests/test_native_judge_cutoff.py | 95 ++++++++++++++++++++++++ 3 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 ifixai/tests/test_native_judge_cutoff.py diff --git a/ifixai/providers/bedrock.py b/ifixai/providers/bedrock.py index 6c897246..6a640a2d 100644 --- a/ifixai/providers/bedrock.py +++ b/ifixai/providers/bedrock.py @@ -12,6 +12,7 @@ ProviderRateLimitError, ProviderResponseError, ProviderTimeoutError, + raise_if_truncated, ) from ifixai.providers.schemas import ConversePayload @@ -84,6 +85,7 @@ async def send_message( system_prompts, converse_messages, inference_config, + config.reject_truncated, ), timeout=float(config.timeout), ) @@ -160,6 +162,7 @@ def _invoke_converse( system_prompts: list[dict], messages: list[dict], inference_config: dict, + reject_truncated: bool = False, ) -> str: converse_kwargs: dict = { "modelId": model_id, @@ -174,6 +177,11 @@ def _invoke_converse( output = response.get("output", {}) message = output.get("message", {}) content_blocks = message.get("content", []) + text_parts = [block["text"] for block in content_blocks if "text" in block] + if reject_truncated: + raise_if_truncated( + "bedrock", "", response.get("stopReason", ""), "\n".join(text_parts) + ) if not content_blocks: raise ProviderEmptyContentError( @@ -182,7 +190,6 @@ def _invoke_converse( details="Empty content in Bedrock converse response", ) - text_parts = [block["text"] for block in content_blocks if "text" in block] if not text_parts: raise ProviderResponseError( provider="bedrock", diff --git a/ifixai/providers/gemini.py b/ifixai/providers/gemini.py index 034f69d6..e70cdf7a 100644 --- a/ifixai/providers/gemini.py +++ b/ifixai/providers/gemini.py @@ -12,6 +12,7 @@ ProviderRateLimitError, ProviderResponseError, ProviderTimeoutError, + raise_if_truncated, ) from ifixai.providers.schemas import GeminiMessages @@ -87,6 +88,12 @@ async def send_message( for part in candidate.content.parts if hasattr(part, "text") and part.text ] + if config.reject_truncated: + raise_if_truncated( + "gemini", endpoint, + getattr(candidate.finish_reason, "name", ""), + "\n".join(text_parts), + ) if not text_parts: raise ProviderEmptyContentError( provider="gemini", diff --git a/ifixai/tests/test_native_judge_cutoff.py b/ifixai/tests/test_native_judge_cutoff.py new file mode 100644 index 00000000..0dc11928 --- /dev/null +++ b/ifixai/tests/test_native_judge_cutoff.py @@ -0,0 +1,95 @@ +"""Judge cutoff handling for Gemini SDK candidates and native Bedrock HTTP.""" + +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest + +from ifixai.core.types import ChatMessage, ProviderConfig +from ifixai.providers.base import ProviderTruncatedError + +PARTIAL = '{"passed": true' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reject,cutoff,text", [(True, True, PARTIAL), (False, True, PARTIAL), (True, False, PARTIAL), (True, True, "")]) +async def test_gemini_sdk_finish_reason_controls_judge_cutoff(monkeypatch, reject, cutoff, text): + genai = pytest.importorskip("google.generativeai") + glm = pytest.importorskip("google.ai.generativelanguage") + from google.generativeai import client + + from ifixai.providers.gemini import GeminiProvider + + class OwnedTransport: + async def generate_content(self, request, **kwargs): + assert request.contents[0].parts[0].text == "judge this" + return glm.GenerateContentResponse(candidates=[glm.Candidate( + content=glm.Content(parts=[glm.Part(text=text)]), + finish_reason=glm.Candidate.FinishReason.MAX_TOKENS if cutoff else glm.Candidate.FinishReason.STOP, + )]) + + # Use the installed GenerativeModel and its native response conversion. + # Only the outbound transport is synthetic; no Google service is called. + monkeypatch.setattr(client, "get_default_generative_async_client", OwnedTransport) + config = ProviderConfig(provider="gemini", api_key="synthetic-local-key", reject_truncated=reject, max_retries=0) + call = GeminiProvider().send_message([ChatMessage(content="judge this")], config) + if reject and cutoff: + with pytest.raises(ProviderTruncatedError): + await call + else: + assert await call == text + assert genai is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reject,cutoff,text", [(True, True, PARTIAL), (False, True, PARTIAL), (True, False, PARTIAL), (True, True, "")]) +async def test_bedrock_native_http_stop_reason_controls_judge_cutoff(monkeypatch, tmp_path, reject, cutoff, text): + pytest.importorskip("boto3") + from ifixai.providers.bedrock import BedrockProvider + + # Explicit owned credentials avoid SDK discovery of any user credentials. + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "synthetic-local-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "synthetic-local-secret") + monkeypatch.setenv("AWS_SESSION_TOKEN", "") + monkeypatch.delenv("AWS_PROFILE", raising=False) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "absent-credentials")) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "absent-config")) + seen = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + seen.append((self.path, json.loads(self.rfile.read(int(self.headers["Content-Length"]))))) + payload = json.dumps({ + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": "max_tokens" if cutoff else "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + "metrics": {"latencyMs": 1}, + }).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + config = ProviderConfig(provider="bedrock", endpoint=f"http://127.0.0.1:{server.server_port}", model="owned-model", reject_truncated=reject, max_retries=0) + call = BedrockProvider().send_message([ChatMessage(content="judge this")], config) + if reject and cutoff: + with pytest.raises(ProviderTruncatedError): + await call + else: + assert await call == text + assert len(seen) == 1 + assert seen[0][0] == "/model/owned-model/converse" + assert seen[0][1]["messages"][0]["content"] == [{"text": "judge this"}] + finally: + server.shutdown() + server.server_close() + thread.join() From d03b9876fd061a6eb83e7498dcf5d65b2219dcdb Mon Sep 17 00:00:00 2001 From: Rudy Celekli <47457359+rudycelekli@users.noreply.github.com> Date: Mon, 5 Oct 2026 04:50:01 -0400 Subject: [PATCH 2/2] test(gemini): keep cutoff transports local across client lifecycles Signed-off-by: Rudy Celekli <47457359+rudycelekli@users.noreply.github.com> --- ifixai/tests/test_native_judge_cutoff.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/ifixai/tests/test_native_judge_cutoff.py b/ifixai/tests/test_native_judge_cutoff.py index 0dc11928..d7f52ed5 100644 --- a/ifixai/tests/test_native_judge_cutoff.py +++ b/ifixai/tests/test_native_judge_cutoff.py @@ -19,9 +19,15 @@ async def test_gemini_sdk_finish_reason_controls_judge_cutoff(monkeypatch, rejec glm = pytest.importorskip("google.ai.generativelanguage") from google.generativeai import client - from ifixai.providers.gemini import GeminiProvider + from ifixai.providers import gemini class OwnedTransport: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + async def generate_content(self, request, **kwargs): assert request.contents[0].parts[0].text == "judge this" return glm.GenerateContentResponse(candidates=[glm.Candidate( @@ -32,8 +38,9 @@ async def generate_content(self, request, **kwargs): # Use the installed GenerativeModel and its native response conversion. # Only the outbound transport is synthetic; no Google service is called. monkeypatch.setattr(client, "get_default_generative_async_client", OwnedTransport) + monkeypatch.setattr(gemini, "GenerativeServiceAsyncClient", lambda **kwargs: OwnedTransport(), raising=False) config = ProviderConfig(provider="gemini", api_key="synthetic-local-key", reject_truncated=reject, max_retries=0) - call = GeminiProvider().send_message([ChatMessage(content="judge this")], config) + call = gemini.GeminiProvider().send_message([ChatMessage(content="judge this")], config) if reject and cutoff: with pytest.raises(ProviderTruncatedError): await call