참고소스 수정본

참고/instructor-main/tests/genai/test_safety_settings.py

@@ -0,0 +1,98 @@
+from instructor.providers.gemini.utils import update_genai_kwargs
+
+
+def test_update_genai_kwargs_safety_settings_excludes_image_categories():
+    """IMAGE_* harm categories must never be sent to the standard Gemini API.
+
+    They are only supported by Vertex AI and cause 400 INVALID_ARGUMENT errors
+    on the standard Gemini API.
+    See: https://github.com/567-labs/instructor/issues/2146
+    """
+    from google.genai import types
+
+    kwargs = {
+        "contents": [
+            types.Content(
+                role="user",
+                parts=[types.Part.from_bytes(data=b"123", mime_type="image/png")],
+            )
+        ]
+    }
+    base_config = {}
+
+    result = update_genai_kwargs(kwargs, base_config)
+
+    assert "safety_settings" in result
+    assert isinstance(result["safety_settings"], list)
+    # No IMAGE_* categories should be present
+    for setting in result["safety_settings"]:
+        assert not setting["category"].name.startswith("HARM_CATEGORY_IMAGE_"), (
+            f"IMAGE_ category {setting['category'].name} must not be sent to the "
+            "standard Gemini API"
+        )
+
+
+def test_update_genai_kwargs_text_only_categories_for_image_content():
+    """Even with image content, only text harm categories should be used."""
+    from google.genai import types
+    from google.genai.types import HarmBlockThreshold, HarmCategory
+
+    custom_safety = {
+        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
+    }
+
+    kwargs = {
+        "contents": [
+            types.Content(
+                role="user",
+                parts=[types.Part.from_bytes(data=b"123", mime_type="image/png")],
+            )
+        ],
+        "safety_settings": custom_safety,
+    }
+    base_config = {}
+
+    result = update_genai_kwargs(kwargs, base_config)
+
+    # Custom threshold should be preserved for text category
+    found_hate_speech = False
+    for setting in result["safety_settings"]:
+        assert not setting["category"].name.startswith("HARM_CATEGORY_IMAGE_")
+        if setting["category"] == HarmCategory.HARM_CATEGORY_HATE_SPEECH:
+            assert setting["threshold"] == HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
+            found_hate_speech = True
+    assert found_hate_speech, "HARM_CATEGORY_HATE_SPEECH should be in safety_settings"
+
+
+def test_handle_genai_tools_autodetect_images_excludes_image_categories():
+    """Autodetected image content should NOT use IMAGE_* harm categories."""
+    from pydantic import BaseModel
+
+    from instructor.providers.gemini.utils import handle_genai_tools
+
+    class SimpleModel(BaseModel):
+        text: str
+
+    data_uri = (
+        "data:image/png;base64,"
+        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO6q0S8AAAAASUVORK5CYII="
+    )
+
+    kwargs = {
+        "messages": [
+            {
+                "role": "user",
+                "content": ["What is in this image?", data_uri],
+            }
+        ]
+    }
+
+    _, out = handle_genai_tools(SimpleModel, kwargs, autodetect_images=True)
+
+    assert "config" in out
+    assert out["config"].safety_settings is not None
+    # No IMAGE_* categories should be present
+    assert not any(
+        s.category.name.startswith("HARM_CATEGORY_IMAGE_")
+        for s in out["config"].safety_settings
+    ), "IMAGE_ categories must not be sent to the standard Gemini API"