feat: properly reporting content filtering errors for Imagen (#65)

aidial_adapter_vertexai/chat/gemini/adapter.py

@@ -229,9 +229,7 @@ def to_openai_finish_reason(
             return FinishReason.LENGTH
         case Candidate.FinishReason.STOP:
             return FinishReason.STOP
-        case Candidate.FinishReason.SAFETY:
-            return FinishReason.CONTENT_FILTER
-        case Candidate.FinishReason.RECITATION:
+        case Candidate.FinishReason.SAFETY | Candidate.FinishReason.RECITATION:
             return FinishReason.CONTENT_FILTER
         case Candidate.FinishReason.OTHER:
             # OTHER finish reason could be usually fixed with a retry

aidial_adapter_vertexai/chat/imagen/adapter.py

@@ -20,7 +20,7 @@
 from aidial_adapter_vertexai.utils.log_config import vertex_ai_logger as log
 from aidial_adapter_vertexai.utils.timer import Timer
 from aidial_adapter_vertexai.vertex_ai import (
-    get_text_embedding_model,
+    get_image_generation_model,
     init_vertex_ai,
 )
 
@@ -122,5 +122,5 @@ async def create(
         location: str,
     ) -> "ImagenChatCompletionAdapter":
         await init_vertex_ai(project_id, location)
-        model = await get_text_embedding_model(model_id)
+        model = await get_image_generation_model(model_id)
         return cls(file_storage, model)

aidial_adapter_vertexai/dial_api/exceptions.py

@@ -31,6 +31,19 @@ def to_dial_exception(e: Exception) -> DialException:
         )
 
     if isinstance(e, InvalidArgument):
+        # Imagen content filtering message
+        content_filter_msg = (
+            "The response is blocked, as it may violate our policies."
+        )
+        if content_filter_msg in str(e):
+            return DialException(
+                status_code=400,
+                type="invalid_request_error",
+                message=content_filter_msg,
+                code="content_filter",
+                param="prompt",
+            )
+
         return DialException(
             status_code=400,
             type="invalid_request_error",

aidial_adapter_vertexai/vertex_ai.py

@@ -41,5 +41,5 @@ async def get_embedding_model(model_id: str) -> TextEmbeddingModel:
 
 
 @cached()
-async def get_text_embedding_model(model_id: str) -> ImageGenerationModel:
+async def get_image_generation_model(model_id: str) -> ImageGenerationModel:
     return await make_async(ImageGenerationModel.from_pretrained, model_id)

tests/integration_tests/test_chat_completion_validation.py

@@ -155,3 +155,33 @@ async def test_input_validation(server, test: TestCase):
             streaming=streaming,
             stop=None,
         )
+
+
+@pytest.mark.asyncio
+async def test_imagen_content_filtering(server):
+    streaming = False
+    model = create_chat_model(
+        TEST_SERVER_URL,
+        ChatCompletionDeployment.IMAGEN_005,
+        streaming,
+        max_tokens=None,
+    )
+
+    with pytest.raises(Exception) as exc_info:
+        await assert_dialog(
+            model=model,
+            messages=[user("generate something unsafe")],
+            output_predicate=lambda s: True,
+            streaming=streaming,
+            stop=None,
+        )
+
+    assert isinstance(exc_info.value, APIStatusError)
+    assert exc_info.value.status_code == 400
+
+    resp = exc_info.value.response.json()
+    assert (resp["error"]["code"]) == "content_filter"
+    assert (
+        resp["error"]["message"]
+        == "The response is blocked, as it may violate our policies."
+    )