fix: Refactor Anthropic integration to support both direct API and Vertex AI

This change introduces an `AnthropicLlm` base class for direct Anthropic API calls using `AsyncAnthropic`. The existing `Claude` class now inherits from `AnthropicLlm` and is specialized to use `AsyncAnthropicVertex` for models hosted on Vertex AI. The `messages.create` call is now properly awaited

Merge: https://github.com/google/adk-python/pull/2904

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 838851026
This commit is contained in:
Eitan Yarmush
2025-12-01 11:04:51 -08:00
committed by Copybara-Service
co-authored by George Weale
parent 7edd7ea9b7
commit 8e82838f1e
2 changed files with 57 additions and 8 deletions
@@ -19,6 +19,7 @@ from unittest import mock
from anthropic import types as anthropic_types
from google.adk import version as adk_version
from google.adk.models import anthropic_llm
from google.adk.models.anthropic_llm import AnthropicLlm
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
@@ -359,6 +360,37 @@ async def test_generate_content_async(
assert responses[0].content.parts[0].text == "Hello, how can I help you?"
@pytest.mark.asyncio
async def test_anthropic_llm_generate_content_async(
llm_request, generate_content_response, generate_llm_response
):
anthropic_llm_instance = AnthropicLlm(model="claude-sonnet-4-20250514")
with mock.patch.object(
anthropic_llm_instance, "_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()
responses = [
resp
async for resp in anthropic_llm_instance.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