test: Add test for content_to_message_param

Author: Happy-ryan

tests/unittests/models/test_anthropic_llm.py

@@ -20,6 +20,7 @@
 from google.adk import version as adk_version
 from google.adk.models import anthropic_llm
 from google.adk.models.anthropic_llm import Claude
+from google.adk.models.anthropic_llm import content_to_message_param
 from google.adk.models.anthropic_llm import function_declaration_to_tool_param
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
@@ -32,69 +33,69 @@
 
 @pytest.fixture
 def generate_content_response():
-  return anthropic_types.Message(
-      id="msg_vrtx_testid",
-      content=[
-          anthropic_types.TextBlock(
-              citations=None, text="Hi! How can I help you today?", type="text"
-          )
-      ],
-      model="claude-3-5-sonnet-v2-20241022",
-      role="assistant",
-      stop_reason="end_turn",
-      stop_sequence=None,
-      type="message",
-      usage=anthropic_types.Usage(
-          cache_creation_input_tokens=0,
-          cache_read_input_tokens=0,
-          input_tokens=13,
-          output_tokens=12,
-          server_tool_use=None,
-          service_tier=None,
-      ),
-  )
+    return anthropic_types.Message(
+        id="msg_vrtx_testid",
+        content=[
+            anthropic_types.TextBlock(
+                citations=None, text="Hi! How can I help you today?", type="text"
+            )
+        ],
+        model="claude-3-5-sonnet-v2-20241022",
+        role="assistant",
+        stop_reason="end_turn",
+        stop_sequence=None,
+        type="message",
+        usage=anthropic_types.Usage(
+            cache_creation_input_tokens=0,
+            cache_read_input_tokens=0,
+            input_tokens=13,
+            output_tokens=12,
+            server_tool_use=None,
+            service_tier=None,
+        ),
+    )
 
 
 @pytest.fixture
 def generate_llm_response():
-  return LlmResponse.create(
-      types.GenerateContentResponse(
-          candidates=[
-              types.Candidate(
-                  content=Content(
-                      role="model",
-                      parts=[Part.from_text(text="Hello, how can I help you?")],
-                  ),
-                  finish_reason=types.FinishReason.STOP,
-              )
-          ]
-      )
-  )
+    return LlmResponse.create(
+        types.GenerateContentResponse(
+            candidates=[
+                types.Candidate(
+                    content=Content(
+                        role="model",
+                        parts=[Part.from_text(text="Hello, how can I help you?")],
+                    ),
+                    finish_reason=types.FinishReason.STOP,
+                )
+            ]
+        )
+    )
 
 
 @pytest.fixture
 def claude_llm():
-  return Claude(model="claude-3-5-sonnet-v2@20241022")
+    return Claude(model="claude-3-5-sonnet-v2@20241022")
 
 
 @pytest.fixture
 def llm_request():
-  return LlmRequest(
-      model="claude-3-5-sonnet-v2@20241022",
-      contents=[Content(role="user", parts=[Part.from_text(text="Hello")])],
-      config=types.GenerateContentConfig(
-          temperature=0.1,
-          response_modalities=[types.Modality.TEXT],
-          system_instruction="You are a helpful assistant",
-      ),
-  )
+    return LlmRequest(
+        model="claude-3-5-sonnet-v2@20241022",
+        contents=[Content(role="user", parts=[Part.from_text(text="Hello")])],
+        config=types.GenerateContentConfig(
+            temperature=0.1,
+            response_modalities=[types.Modality.TEXT],
+            system_instruction="You are a helpful assistant",
+        ),
+    )
 
 
 def test_supported_models():
-  models = Claude.supported_models()
-  assert len(models) == 2
-  assert models[0] == r"claude-3-.*"
-  assert models[1] == r"claude-.*-4.*"
+    models = Claude.supported_models()
+    assert len(models) == 2
+    assert models[0] == r"claude-3-.*"
+    assert models[1] == r"claude-.*-4.*"
 
 
 function_declaration_test_cases = [
@@ -133,9 +134,7 @@ def test_supported_models():
                 "properties": {
                     "location": {
                         "type": "string",
-                        "description": (
-                            "City and state, e.g., San Francisco, CA"
-                        ),
+                        "description": ("City and state, e.g., San Francisco, CA"),
                     }
                 },
             },
@@ -284,65 +283,142 @@ def test_supported_models():
 async def test_function_declaration_to_tool_param(
     _, function_declaration, expected_tool_param
 ):
-  """Test function_declaration_to_tool_param."""
-  assert (
-      function_declaration_to_tool_param(function_declaration)
-      == expected_tool_param
-  )
+    """Test function_declaration_to_tool_param."""
+    assert (
+        function_declaration_to_tool_param(function_declaration) == expected_tool_param
+    )
 
 
 @pytest.mark.asyncio
 async def test_generate_content_async(
     claude_llm, llm_request, generate_content_response, generate_llm_response
 ):
-  with mock.patch.object(claude_llm, "_anthropic_client") as mock_client:
-    with mock.patch.object(
-        anthropic_llm,
-        "message_to_generate_content_response",
-        return_value=generate_llm_response,
-    ):
-      # Create a mock coroutine that returns the generate_content_response.
-      async def mock_coro():
-        return generate_content_response
+    with mock.patch.object(claude_llm, "_anthropic_client") as mock_client:
+        with mock.patch.object(
+            anthropic_llm,
+            "message_to_generate_content_response",
+            return_value=generate_llm_response,
+        ):
+            # Create a mock coroutine that returns the generate_content_response.
+            async def mock_coro():
+                return generate_content_response
 
-      # Assign the coroutine to the mocked method
-      mock_client.messages.create.return_value = mock_coro()
+            # Assign the coroutine to the mocked method
+            mock_client.messages.create.return_value = mock_coro()
 
-      responses = [
-          resp
-          async for resp in claude_llm.generate_content_async(
-              llm_request, stream=False
-          )
-      ]
-      assert len(responses) == 1
-      assert isinstance(responses[0], LlmResponse)
-      assert responses[0].content.parts[0].text == "Hello, how can I help you?"
+            responses = [
+                resp
+                async for resp in claude_llm.generate_content_async(
+                    llm_request, stream=False
+                )
+            ]
+            assert len(responses) == 1
+            assert isinstance(responses[0], LlmResponse)
+            assert responses[0].content.parts[0].text == "Hello, how can I help you?"
 
 
 @pytest.mark.asyncio
 async def test_generate_content_async_with_max_tokens(
     llm_request, generate_content_response, generate_llm_response
 ):
-  claude_llm = Claude(model="claude-3-5-sonnet-v2@20241022", max_tokens=4096)
-  with mock.patch.object(claude_llm, "_anthropic_client") as mock_client:
-    with mock.patch.object(
-        anthropic_llm,
-        "message_to_generate_content_response",
-        return_value=generate_llm_response,
-    ):
-      # Create a mock coroutine that returns the generate_content_response.
-      async def mock_coro():
-        return generate_content_response
+    claude_llm = Claude(model="claude-3-5-sonnet-v2@20241022", max_tokens=4096)
+    with mock.patch.object(claude_llm, "_anthropic_client") as mock_client:
+        with mock.patch.object(
+            anthropic_llm,
+            "message_to_generate_content_response",
+            return_value=generate_llm_response,
+        ):
+            # Create a mock coroutine that returns the generate_content_response.
+            async def mock_coro():
+                return generate_content_response
+
+            # Assign the coroutine to the mocked method
+            mock_client.messages.create.return_value = mock_coro()
+
+            _ = [
+                resp
+                async for resp in claude_llm.generate_content_async(
+                    llm_request, stream=False
+                )
+            ]
+            mock_client.messages.create.assert_called_once()
+            _, kwargs = mock_client.messages.create.call_args
+            assert kwargs["max_tokens"] == 4096
+
+
+content_to_message_param_test_cases = [
+    (
+        "user_role_with_text_and_image",
+        Content(
+            role="user",
+            parts=[
+                Part.from_text(text="What's in this image?"),
+                Part(
+                    inline_data=types.Blob(
+                        mime_type="image/jpeg", data=b"fake_image_data"
+                    )
+                ),
+            ],
+        ),
+        "user",
+        2,  # Expected content length
+        False,  # Should not log warning
+    ),
+    (
+        "model_role_with_text_and_image",
+        Content(
+            role="model",
+            parts=[
+                Part.from_text(text="I see a cat."),
+                Part(
+                    inline_data=types.Blob(
+                        mime_type="image/png", data=b"fake_image_data"
+                    )
+                ),
+            ],
+        ),
+        "assistant",
+        1,  # Image filtered out, only text remains
+        True,  # Should log warning
+    ),
+    (
+        "assistant_role_with_text_and_image",
+        Content(
+            role="assistant",
+            parts=[
+                Part.from_text(text="Here's what I found."),
+                Part(
+                    inline_data=types.Blob(
+                        mime_type="image/webp", data=b"fake_image_data"
+                    )
+                ),
+            ],
+        ),
+        "assistant",
+        1,  # Image filtered out, only text remains
+        True,  # Should log warning
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "_, content, expected_role, expected_content_length, should_log_warning",
+    content_to_message_param_test_cases,
+    ids=[case[0] for case in content_to_message_param_test_cases],
+)
+def test_content_to_message_param_with_images(
+    _, content, expected_role, expected_content_length, should_log_warning
+):
+    """Test content_to_message_param handles images correctly based on role."""
+    with mock.patch("google.adk.models.anthropic_llm.logger") as mock_logger:
+        result = content_to_message_param(content)
 
-      # Assign the coroutine to the mocked method
-      mock_client.messages.create.return_value = mock_coro()
+        assert result["role"] == expected_role
+        assert len(result["content"]) == expected_content_length
 
-      _ = [
-          resp
-          async for resp in claude_llm.generate_content_async(
-              llm_request, stream=False
-          )
-      ]
-      mock_client.messages.create.assert_called_once()
-      _, kwargs = mock_client.messages.create.call_args
-      assert kwargs["max_tokens"] == 4096
+        if should_log_warning:
+            mock_logger.warning.assert_called_once_with(
+                "Image data is not supported in Claude for model turns."
+            )
+        else:
+            mock_logger.warning.assert_not_called()