Handle some weird Google responses.

Author: wpietri

plugins/google/modelgauge/suts/google_genai.py

@@ -19,13 +19,13 @@
 from modelgauge.retry_decorator import retry
 from modelgauge.secret_values import InjectSecret
 from modelgauge.sut import REFUSAL_RESPONSE, PromptResponseSUT, SUTOptions, SUTResponse  # usort: skip
+from modelgauge.sut_capabilities import AcceptsTextPrompt
+from modelgauge.sut_decorator import modelgauge_sut
+from modelgauge.sut_registry import SUTS
 from modelgauge.suts.google_generativeai import (
     GOOGLE_REFUSAL_FINISH_REASONS,
     GoogleAiApiKey,
 )  # Both SDKs use the same API key.
-from modelgauge.sut_capabilities import AcceptsTextPrompt
-from modelgauge.sut_decorator import modelgauge_sut
-from modelgauge.sut_registry import SUTS
 
 
 class GenAiRequest(BaseModel):
@@ -76,16 +76,13 @@ def evaluate(self, request: GenAiRequest) -> GenerateContentResponse:
         return self.client.models.generate_content(**request.model_dump(exclude_none=True))
 
     def translate_response(self, request: GenAiRequest, response: GenerateContentResponse) -> SUTResponse:
-        assert (
-            len(response.candidates) <= 1
-        ), f"Expected a single candidate in the response, got {len(response.candidates)}."
-        if len(response.candidates) == 0:
+        if response.candidates is None or len(response.candidates) == 0:
             # This is apparently a refusal. At least, it's what happens consistently with a set of
             # prompts in the CSE, SRC, and SXC hazards
             response_text = REFUSAL_RESPONSE
-        else:
+        elif len(response.candidates) == 1:
             candidate = response.candidates[0]
-            if candidate.finish_reason in GOOGLE_REFUSAL_FINISH_REASONS:
+            if candidate.finish_reason in GOOGLE_REFUSAL_FINISH_REASONS + ["OTHER"]:
                 response_text = REFUSAL_RESPONSE
             elif candidate.content is not None:
                 response_text = candidate.content.parts[0].text
@@ -95,7 +92,8 @@ def translate_response(self, request: GenAiRequest, response: GenerateContentRes
                     f"The candidate does not have any content,"
                     f" but it's finish reason {candidate.finish_reason} does not qualify as a refusal."
                 )
-
+        else:
+            raise AssertionError(f"Expected a single candidate in the response, got {response.candidates}.")
         return SUTResponse(text=response_text)
 
 

plugins/google/modelgauge/suts/google_generativeai.py

@@ -15,6 +15,7 @@
     TooManyRequests,
 )
 from google.generativeai.types import HarmBlockThreshold, HarmCategory  # type: ignore
+from pydantic import BaseModel
 
 from modelgauge.general import APIException
 from modelgauge.prompt import TextPrompt
@@ -24,8 +25,8 @@
 from modelgauge.sut_capabilities import AcceptsTextPrompt
 from modelgauge.sut_decorator import modelgauge_sut
 from modelgauge.sut_registry import SUTS
-from pydantic import BaseModel
 
+# TODO: Do we really need this type alias or can we just import the real thing?
 FinishReason = genai.protos.Candidate.FinishReason
 GEMINI_HARM_CATEGORIES = [
     HarmCategory.HARM_CATEGORY_HATE_SPEECH,

plugins/google/tests/test_google_genai.py

@@ -2,7 +2,7 @@
 from unittest.mock import patch
 
 import pytest
-from google.genai.types import Candidate, GenerateContentConfig, GenerateContentResponse, ThinkingConfig
+from google.genai.types import GenerateContentConfig, GenerateContentResponse, ThinkingConfig, FinishReason
 
 from modelgauge.prompt import TextPrompt
 from modelgauge.sut import REFUSAL_RESPONSE, SUTOptions, SUTResponse
@@ -156,6 +156,14 @@ def test_google_genai_translate_response(google_default_sut, fake_raw_response,
     assert response == SUTResponse(text="some response")
 
 
+def test_google_genai_translate_response_finish_reason_other(google_default_sut, fake_raw_response, some_request):
+    """I think this is for a typing error but we're in a rush so I'm not fixing it"""
+    fake_raw_response.candidates[0].finish_reason = FinishReason.OTHER
+    response = google_default_sut.translate_response(some_request, fake_raw_response)
+
+    assert response == SUTResponse(text="")  # indicates refusal
+
+
 def test_google_genai_translate_response_no_completions(google_default_sut, some_request):
     no_completions = GenerateContentResponse(
         **json.loads(
@@ -174,3 +182,23 @@ def test_google_genai_translate_response_no_completions(google_default_sut, some
     response = google_default_sut.translate_response(some_request, no_completions)
 
     assert response == SUTResponse(text=REFUSAL_RESPONSE)
+
+
+def test_google_genai_translate_response_none_completions(google_default_sut, some_request):
+    no_completions = GenerateContentResponse(
+        **json.loads(
+            """{
+  "candidates": null,
+  "usage_metadata": {
+    "prompt_token_count": 19,
+    "total_token_count": 19,
+    "cached_content_token_count": 0,
+    "candidates_token_count": 0
+  }
+}
+"""
+        )
+    )
+    response = google_default_sut.translate_response(some_request, no_completions)
+
+    assert response == SUTResponse(text=REFUSAL_RESPONSE)