From e305021d12f3de9e38a6723e00ca7c3f5be96231 Mon Sep 17 00:00:00 2001 From: Srikar Mannepalli Date: Thu, 25 Sep 2025 12:01:38 +0530 Subject: [PATCH 1/3] anthropic vertex client --- .../autogen_ext/models/anthropic/__init__.py | 8 + .../models/anthropic/_anthropic_client.py | 137 ++++++++++- .../models/anthropic/config/__init__.py | 27 +++ .../models/test_anthropic_model_client.py | 222 ++++++++++++++++++ 4 files changed, 393 insertions(+), 1 deletion(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py index f31e7b1c0b72..d1c29a9c82f4 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py @@ -1,6 +1,7 @@ from ._anthropic_client import ( AnthropicBedrockChatCompletionClient, AnthropicChatCompletionClient, + AnthropicVertexChatCompletionClient, BaseAnthropicChatCompletionClient, ) from .config import ( @@ -8,18 +9,25 @@ AnthropicBedrockClientConfigurationConfigModel, AnthropicClientConfiguration, AnthropicClientConfigurationConfigModel, + AnthropicVertexClientConfiguration, + AnthropicVertexClientConfigurationConfigModel, BedrockInfo, CreateArgumentsConfigModel, + VertexInfo, ) __all__ = [ "AnthropicChatCompletionClient", "AnthropicBedrockChatCompletionClient", + "AnthropicVertexChatCompletionClient", "BaseAnthropicChatCompletionClient", "AnthropicClientConfiguration", "AnthropicBedrockClientConfiguration", + "AnthropicVertexClientConfiguration", "AnthropicClientConfigurationConfigModel", "AnthropicBedrockClientConfigurationConfigModel", + "AnthropicVertexClientConfigurationConfigModel", "CreateArgumentsConfigModel", "BedrockInfo", + "VertexInfo", ] diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py index 3bc989878b59..0062af6ff877 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py @@ -23,7 +23,7 @@ ) import tiktoken -from anthropic import AsyncAnthropic, AsyncAnthropicBedrock, AsyncStream +from anthropic import AsyncAnthropic, AsyncAnthropicBedrock, AsyncAnthropicVertex, AsyncStream from anthropic.types import ( Base64ImageSourceParam, ContentBlock, @@ -71,7 +71,10 @@ AnthropicBedrockClientConfigurationConfigModel, AnthropicClientConfiguration, AnthropicClientConfigurationConfigModel, + AnthropicVertexClientConfiguration, + AnthropicVertexClientConfigurationConfigModel, BedrockInfo, + VertexInfo, ) logger = logging.getLogger(EVENT_LOGGER_NAME) @@ -1345,3 +1348,135 @@ def _from_config(cls, config: AnthropicBedrockClientConfigurationConfigModel) -> } return cls(**copied_config) + + + +class AnthropicVertexChatCompletionClient( + BaseAnthropicChatCompletionClient, Component[AnthropicVertexClientConfigurationConfigModel] +): + """ + Chat completion client for Anthropic's Claude models via Google Vertex AI. + + Args: + model (str): The Claude model to use (e.g., "claude-3-sonnet-20240229", "claude-3-opus-20240229") + vertex_info (VertexInfo): Configuration for Vertex AI including project_id and region + max_tokens (int, optional): Maximum tokens in the response. Default is 4096. + temperature (float, optional): Controls randomness. Lower is more deterministic. Default is 1.0. + top_p (float, optional): Controls diversity via nucleus sampling. Default is 1.0. + top_k (int, optional): Controls diversity via top-k sampling. Default is -1 (disabled). + model_info (ModelInfo, optional): The capabilities of the model. Required if using a custom model. + + To use this client, you must install the Anthropic extension: + + .. code-block:: bash + + pip install "autogen-ext[anthropic]" + + Example: + + .. code-block:: python + + import asyncio + from autogen_ext.models.anthropic import AnthropicVertexChatCompletionClient, VertexInfo + from autogen_core.models import UserMessage + + async def main(): + vertex_client = AnthropicVertexChatCompletionClient( + model="claude-3-sonnet-20240229", + vertex_info=VertexInfo( + project_id="your-gcp-project-id", + region="us-east5" + ) + ) + + result = await vertex_client.create([UserMessage(content="What is the capital of France?", source="user")]) # type: ignore + print(result) + + if __name__ == "__main__": + asyncio.run(main()) + + To load the client from a configuration: + + .. code-block:: python + + from autogen_core.models import ChatCompletionClient + + config = { + "provider": "AnthropicVertexChatCompletionClient", + "config": { + "model": "claude-3-sonnet-20240229", + "vertex_info": { + "project_id": "your-gcp-project-id", + "region": "us-east5" + } + }, + } + + client = ChatCompletionClient.load_component(config) + """ + + component_type = "model" + component_config_schema = AnthropicVertexClientConfigurationConfigModel + component_provider_override = "autogen_ext.models.anthropic.AnthropicVertexChatCompletionClient" + + def __init__(self, **kwargs: Unpack[AnthropicVertexClientConfiguration]): + if "model" not in kwargs: + raise ValueError("model is required for AnthropicVertexChatCompletionClient") + + self._raw_config: Dict[str, Any] = dict(kwargs).copy() + copied_args = dict(kwargs).copy() + + model_info: Optional[ModelInfo] = None + if "model_info" in kwargs: + model_info = kwargs["model_info"] + del copied_args["model_info"] + + vertex_info: Optional[VertexInfo] = None + if "vertex_info" in kwargs: + vertex_info = kwargs["vertex_info"] + + if vertex_info is None: + raise ValueError("vertex_info is required for AnthropicVertexChatCompletionClient") + + # Handle vertex_info + project_id = vertex_info["project_id"] + region = vertex_info["region"] + + client = AsyncAnthropicVertex( + project_id=project_id, + region=region, + ) + create_args = _create_args_from_config(copied_args) + + super().__init__( + client=client, + create_args=create_args, + model_info=model_info, + ) + + def __getstate__(self) -> Dict[str, Any]: + state = self.__dict__.copy() + state["_client"] = None + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + self.__dict__.update(state) + # Recreate the client from raw config + vertex_info = state["_raw_config"]["vertex_info"] + self._client = AsyncAnthropicVertex( + project_id=vertex_info["project_id"], + region=vertex_info["region"], + ) + + def _to_config(self) -> AnthropicVertexClientConfigurationConfigModel: + copied_config = self._raw_config.copy() + return AnthropicVertexClientConfigurationConfigModel(**copied_config) + + @classmethod + def _from_config(cls, config: AnthropicVertexClientConfigurationConfigModel) -> Self: + copied_config = config.model_copy().model_dump(exclude_none=True) + + # Handle vertex_info properly - no secret values to extract like bedrock + # vertex_info contains project_id and region which are not secret + + return cls(**copied_config) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py index 22dfc067f318..ddc12eb7e1c6 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py @@ -36,6 +36,18 @@ class BedrockInfo(TypedDict): """aws region for the aws account to gain bedrock model access""" +class VertexInfo(TypedDict): + """VertexInfo is a dictionary that contains information about Vertex AI configuration. + It is expected to be used in the vertex_info property of a model client. + + """ + + project_id: Required[str] + """GCP project ID for Vertex AI access""" + region: Required[str] + """GCP region for Vertex AI access""" + + class BaseAnthropicClientConfiguration(CreateArguments, total=False): api_key: str base_url: Optional[str] @@ -56,6 +68,10 @@ class AnthropicBedrockClientConfiguration(AnthropicClientConfiguration, total=Fa bedrock_info: BedrockInfo +class AnthropicVertexClientConfiguration(AnthropicClientConfiguration, total=False): + vertex_info: VertexInfo + + # Pydantic equivalents of the above TypedDicts class CreateArgumentsConfigModel(BaseModel): model: str @@ -95,3 +111,14 @@ class BedrockInfoConfigModel(TypedDict): class AnthropicBedrockClientConfigurationConfigModel(AnthropicClientConfigurationConfigModel): bedrock_info: BedrockInfoConfigModel | None = None + + +class VertexInfoConfigModel(TypedDict): + project_id: Required[str] + """GCP project ID for Vertex AI access""" + region: Required[str] + """GCP region for Vertex AI access""" + + +class AnthropicVertexClientConfigurationConfigModel(AnthropicClientConfigurationConfigModel): + vertex_info: VertexInfoConfigModel | None = None diff --git a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py index 13bff87cad2b..bedb2ca49e23 100644 --- a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py @@ -21,8 +21,10 @@ from autogen_ext.models.anthropic import ( AnthropicBedrockChatCompletionClient, AnthropicChatCompletionClient, + AnthropicVertexChatCompletionClient, BaseAnthropicChatCompletionClient, BedrockInfo, + VertexInfo, ) @@ -1107,3 +1109,223 @@ async def test_streaming_tool_usage_with_arguments(provider: str) -> None: assert isinstance(content, FunctionCall) assert content.name == "add_numbers" assert json.loads(content.arguments) is not None + + +# ===== Vertex Client Tests ===== + + +def test_vertex_client_initialization_success(): + """Test successful initialization of AnthropicVertexChatCompletionClient.""" + vertex_info = VertexInfo( + project_id="test-project-123", + region="us-east5" + ) + + client = AnthropicVertexChatCompletionClient( + model="claude-3-sonnet-20240229", + vertex_info=vertex_info + ) + + assert client._create_args["model"] == "claude-3-sonnet-20240229" + assert client._raw_config["vertex_info"]["project_id"] == "test-project-123" + assert client._raw_config["vertex_info"]["region"] == "us-east5" + + +def test_vertex_client_initialization_missing_model(): + """Test that missing model raises ValueError.""" + vertex_info = VertexInfo( + project_id="test-project-123", + region="us-east5" + ) + + with pytest.raises(ValueError, match="model is required for AnthropicVertexChatCompletionClient"): + AnthropicVertexChatCompletionClient(vertex_info=vertex_info) + + +def test_vertex_client_initialization_missing_vertex_info(): + """Test that missing vertex_info raises ValueError.""" + with pytest.raises(ValueError, match="vertex_info is required for AnthropicVertexChatCompletionClient"): + AnthropicVertexChatCompletionClient(model="claude-3-sonnet-20240229") + + +def test_vertex_client_with_model_info(): + """Test initialization with custom model_info.""" + vertex_info = VertexInfo( + project_id="test-project-123", + region="us-east5" + ) + + custom_model_info = ModelInfo( + vision=True, + function_calling=True, + json_output=True, + family="test-family", + structured_output=False + ) + + client = AnthropicVertexChatCompletionClient( + model="custom-claude-model", + vertex_info=vertex_info, + model_info=custom_model_info + ) + + assert client.model_info == custom_model_info + + +@pytest.mark.asyncio +async def test_vertex_client_mock_create(): + """Test basic create functionality with mocked client.""" + from anthropic.types import TextBlock + + # Create mock client and response + mock_client = AsyncMock() + mock_message = MagicMock() + mock_message.content = [TextBlock(type="text", text="Hello! I'm Claude.")] + mock_message.usage.input_tokens = 10 + mock_message.usage.output_tokens = 5 + mock_message.stop_reason = "end_turn" + + mock_client.messages.create.return_value = mock_message + + # Create real client but patch the underlying Anthropic client + vertex_info = VertexInfo(project_id="test-project", region="us-east5") + client = AnthropicVertexChatCompletionClient( + model="claude-3-sonnet-20240229", + vertex_info=vertex_info + ) + + messages: List[LLMMessage] = [ + UserMessage(content="Hello", source="user"), + ] + + with patch.object(client, "_client", mock_client): + result = await client.create(messages=messages) + + # Verify the call was made + mock_client.messages.create.assert_called_once() + + # Verify the result + assert result.content == "Hello! I'm Claude." + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + + +@pytest.mark.asyncio +async def test_vertex_client_mock_tool_usage(): + """Test tool usage with mocked client.""" + from anthropic.types import ToolUseBlock + + # Create mock client and response + mock_client = AsyncMock() + mock_message = MagicMock() + mock_message.content = [ToolUseBlock(type="tool_use", name="add_numbers", input={"a": 1, "b": 2}, id="call_123")] + mock_message.usage.input_tokens = 15 + mock_message.usage.output_tokens = 8 + mock_message.stop_reason = "tool_use" + + mock_client.messages.create.return_value = mock_message + + # Create real client + vertex_info = VertexInfo(project_id="test-project", region="us-east5") + client = AnthropicVertexChatCompletionClient( + model="claude-3-sonnet-20240229", + vertex_info=vertex_info + ) + + # Define tool + add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") + + messages: List[LLMMessage] = [ + UserMessage(content="What is 1 + 2?", source="user"), + ] + + with patch.object(client, "_client", mock_client): + result = await client.create( + messages=messages, + tools=[add_tool], + tool_choice="auto" + ) + + # Verify the call + mock_client.messages.create.assert_called_once() + call_args = mock_client.messages.create.call_args + assert "tools" in call_args.kwargs + assert "tool_choice" in call_args.kwargs + + # Verify tool call result + assert isinstance(result.content, list) + assert len(result.content) == 1 + tool_call = result.content[0] + assert isinstance(tool_call, FunctionCall) + assert tool_call.name == "add_numbers" + + +def test_vertex_client_serialization(): + """Test that client can be serialized and deserialized.""" + import pickle + + vertex_info = VertexInfo( + project_id="test-project-456", + region="us-central1" + ) + + client = AnthropicVertexChatCompletionClient( + model="claude-3-haiku-20240307", + vertex_info=vertex_info, + temperature=0.7, + max_tokens=1000 + ) + + # Serialize and deserialize + serialized = pickle.dumps(client) + deserialized = pickle.loads(serialized) + + # Verify the deserialized client has same config + assert deserialized._create_args["model"] == "claude-3-haiku-20240307" + assert deserialized._create_args["temperature"] == 0.7 + assert deserialized._create_args["max_tokens"] == 1000 + assert deserialized._raw_config["vertex_info"]["project_id"] == "test-project-456" + assert deserialized._raw_config["vertex_info"]["region"] == "us-central1" + + +def test_vertex_client_config_conversion(): + """Test _to_config and _from_config methods.""" + from autogen_ext.models.anthropic import AnthropicVertexClientConfigurationConfigModel + + vertex_info = VertexInfo( + project_id="config-test-project", + region="europe-west1" + ) + + client = AnthropicVertexChatCompletionClient( + model="claude-3-opus-20240229", + vertex_info=vertex_info, + temperature=0.5 + ) + + # Convert to config + config = client._to_config() + assert isinstance(config, AnthropicVertexClientConfigurationConfigModel) + + # Create new client from config + new_client = AnthropicVertexChatCompletionClient._from_config(config) + + # Verify equivalence + assert new_client._create_args["model"] == "claude-3-opus-20240229" + assert new_client._create_args["temperature"] == 0.5 + assert new_client._raw_config["vertex_info"]["project_id"] == "config-test-project" + assert new_client._raw_config["vertex_info"]["region"] == "europe-west1" + + +def test_vertex_client_component_attributes(): + """Test that component attributes are set correctly.""" + vertex_info = VertexInfo(project_id="test", region="us-east5") + client = AnthropicVertexChatCompletionClient( + model="claude-3-sonnet-20240229", + vertex_info=vertex_info + ) + + assert client.component_type == "model" + assert hasattr(client, "component_config_schema") + assert hasattr(client, "component_provider_override") + assert client.component_provider_override == "autogen_ext.models.anthropic.AnthropicVertexChatCompletionClient" From 8c73091cb48e1ed28e580b7d831590871ca49f70 Mon Sep 17 00:00:00 2001 From: Srikar Mannepalli Date: Fri, 3 Oct 2025 13:37:42 +0530 Subject: [PATCH 2/3] poe check fixes --- .../models/anthropic/_anthropic_client.py | 13 +- .../models/test_anthropic_model_client.py | 146 ++++++------------ 2 files changed, 54 insertions(+), 105 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py index 74553a4e5ce2..f9eb41f94b41 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py @@ -1419,7 +1419,6 @@ def _from_config(cls, config: AnthropicBedrockClientConfigurationConfigModel) -> return cls(**copied_config) - class AnthropicVertexChatCompletionClient( BaseAnthropicChatCompletionClient, Component[AnthropicVertexClientConfigurationConfigModel] ): @@ -1449,18 +1448,17 @@ class AnthropicVertexChatCompletionClient( from autogen_ext.models.anthropic import AnthropicVertexChatCompletionClient, VertexInfo from autogen_core.models import UserMessage + async def main(): vertex_client = AnthropicVertexChatCompletionClient( model="claude-3-sonnet-20240229", - vertex_info=VertexInfo( - project_id="your-gcp-project-id", - region="us-east5" - ) + vertex_info=VertexInfo(project_id="your-gcp-project-id", region="us-east5"), ) result = await vertex_client.create([UserMessage(content="What is the capital of France?", source="user")]) # type: ignore print(result) + if __name__ == "__main__": asyncio.run(main()) @@ -1474,10 +1472,7 @@ async def main(): "provider": "AnthropicVertexChatCompletionClient", "config": { "model": "claude-3-sonnet-20240229", - "vertex_info": { - "project_id": "your-gcp-project-id", - "region": "us-east5" - } + "vertex_info": {"project_id": "your-gcp-project-id", "region": "us-east5"}, }, } diff --git a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py index 959b38eff295..4015ae9d0e31 100644 --- a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py @@ -1114,66 +1114,48 @@ async def test_streaming_tool_usage_with_arguments(provider: str) -> None: # ===== Vertex Client Tests ===== -def test_vertex_client_initialization_success(): +def test_vertex_client_initialization_success() -> None: """Test successful initialization of AnthropicVertexChatCompletionClient.""" - vertex_info = VertexInfo( - project_id="test-project-123", - region="us-east5" - ) + vertex_info = VertexInfo(project_id="test-project-123", region="us-east5") - client = AnthropicVertexChatCompletionClient( - model="claude-3-sonnet-20240229", - vertex_info=vertex_info - ) + client = AnthropicVertexChatCompletionClient(model="claude-3-sonnet-20240229", vertex_info=vertex_info) - assert client._create_args["model"] == "claude-3-sonnet-20240229" - assert client._raw_config["vertex_info"]["project_id"] == "test-project-123" - assert client._raw_config["vertex_info"]["region"] == "us-east5" + assert client._create_args["model"] == "claude-3-sonnet-20240229" # pyright: ignore[reportPrivateUsage] + assert client._raw_config["vertex_info"]["project_id"] == "test-project-123" # pyright: ignore[reportPrivateUsage] + assert client._raw_config["vertex_info"]["region"] == "us-east5" # pyright: ignore[reportPrivateUsage] -def test_vertex_client_initialization_missing_model(): +def test_vertex_client_initialization_missing_model() -> None: """Test that missing model raises ValueError.""" - vertex_info = VertexInfo( - project_id="test-project-123", - region="us-east5" - ) + vertex_info = VertexInfo(project_id="test-project-123", region="us-east5") with pytest.raises(ValueError, match="model is required for AnthropicVertexChatCompletionClient"): AnthropicVertexChatCompletionClient(vertex_info=vertex_info) -def test_vertex_client_initialization_missing_vertex_info(): +def test_vertex_client_initialization_missing_vertex_info() -> None: """Test that missing vertex_info raises ValueError.""" with pytest.raises(ValueError, match="vertex_info is required for AnthropicVertexChatCompletionClient"): AnthropicVertexChatCompletionClient(model="claude-3-sonnet-20240229") -def test_vertex_client_with_model_info(): +def test_vertex_client_with_model_info() -> None: """Test initialization with custom model_info.""" - vertex_info = VertexInfo( - project_id="test-project-123", - region="us-east5" - ) + vertex_info = VertexInfo(project_id="test-project-123", region="us-east5") custom_model_info = ModelInfo( - vision=True, - function_calling=True, - json_output=True, - family="test-family", - structured_output=False + vision=True, function_calling=True, json_output=True, family="test-family", structured_output=False ) client = AnthropicVertexChatCompletionClient( - model="custom-claude-model", - vertex_info=vertex_info, - model_info=custom_model_info + model="custom-claude-model", vertex_info=vertex_info, model_info=custom_model_info ) assert client.model_info == custom_model_info @pytest.mark.asyncio -async def test_vertex_client_mock_create(): +async def test_vertex_client_mock_create() -> None: """Test basic create functionality with mocked client.""" from anthropic.types import TextBlock @@ -1189,10 +1171,7 @@ async def test_vertex_client_mock_create(): # Create real client but patch the underlying Anthropic client vertex_info = VertexInfo(project_id="test-project", region="us-east5") - client = AnthropicVertexChatCompletionClient( - model="claude-3-sonnet-20240229", - vertex_info=vertex_info - ) + client = AnthropicVertexChatCompletionClient(model="claude-3-sonnet-20240229", vertex_info=vertex_info) messages: List[LLMMessage] = [ UserMessage(content="Hello", source="user"), @@ -1211,7 +1190,7 @@ async def test_vertex_client_mock_create(): @pytest.mark.asyncio -async def test_vertex_client_mock_tool_usage(): +async def test_vertex_client_mock_tool_usage() -> None: """Test tool usage with mocked client.""" from anthropic.types import ToolUseBlock @@ -1229,7 +1208,20 @@ async def test_vertex_client_mock_tool_usage(): vertex_info = VertexInfo(project_id="test-project", region="us-east5") client = AnthropicVertexChatCompletionClient( model="claude-3-sonnet-20240229", - vertex_info=vertex_info + vertex_info=vertex_info, + ) + + # Patch the client and make a call + messages = [UserMessage(content="Calculate 1 + 2", source="user")] + with patch.object(client, "_client", mock_client): + result = await client.create(messages=messages) + + # Verify tool use was returned + assert len(result.content) == 1 + assert isinstance(result.content[0], FunctionCall) + assert result.content[0].arguments == '{"a": 1, "b": 2}' # pyright: ignore[reportUnknownMemberType] + + def test_mock_thinking_config_validation() -> None: """Test thinking configuration handling logic.""" client = AnthropicChatCompletionClient( @@ -1359,44 +1351,31 @@ async def test_anthropic_thinking_mode_with_tools() -> None: add_tool = FunctionTool(_add_numbers, description="Add two numbers together", name="add_numbers") messages: List[LLMMessage] = [ - UserMessage(content="What is 1 + 2?", source="user"), + UserMessage(content="What is 1 + 2? Think step by step.", source="user"), ] - with patch.object(client, "_client", mock_client): - result = await client.create( - messages=messages, - tools=[add_tool], - tool_choice="auto" - ) + thinking_config = {"thinking": {"type": "enabled", "budget_tokens": 2000}} - # Verify the call - mock_client.messages.create.assert_called_once() - call_args = mock_client.messages.create.call_args - assert "tools" in call_args.kwargs - assert "tool_choice" in call_args.kwargs + result = await client.create( + messages=messages, tools=[add_tool], tool_choice="auto", extra_create_args=thinking_config + ) # Verify tool call result assert isinstance(result.content, list) - assert len(result.content) == 1 + assert len(result.content) >= 1 tool_call = result.content[0] assert isinstance(tool_call, FunctionCall) assert tool_call.name == "add_numbers" -def test_vertex_client_serialization(): +def test_vertex_client_serialization() -> None: """Test that client can be serialized and deserialized.""" import pickle - vertex_info = VertexInfo( - project_id="test-project-456", - region="us-central1" - ) + vertex_info = VertexInfo(project_id="test-project-456", region="us-central1") client = AnthropicVertexChatCompletionClient( - model="claude-3-haiku-20240307", - vertex_info=vertex_info, - temperature=0.7, - max_tokens=1000 + model="claude-3-haiku-20240307", vertex_info=vertex_info, temperature=0.7, max_tokens=1000 ) # Serialize and deserialize @@ -1411,61 +1390,36 @@ def test_vertex_client_serialization(): assert deserialized._raw_config["vertex_info"]["region"] == "us-central1" -def test_vertex_client_config_conversion(): +def test_vertex_client_config_conversion() -> None: """Test _to_config and _from_config methods.""" from autogen_ext.models.anthropic import AnthropicVertexClientConfigurationConfigModel - vertex_info = VertexInfo( - project_id="config-test-project", - region="europe-west1" - ) + vertex_info = VertexInfo(project_id="config-test-project", region="europe-west1") client = AnthropicVertexChatCompletionClient( - model="claude-3-opus-20240229", - vertex_info=vertex_info, - temperature=0.5 + model="claude-3-opus-20240229", vertex_info=vertex_info, temperature=0.5 ) # Convert to config - config = client._to_config() + config = client._to_config() # pyright: ignore[reportPrivateUsage] assert isinstance(config, AnthropicVertexClientConfigurationConfigModel) # Create new client from config - new_client = AnthropicVertexChatCompletionClient._from_config(config) + new_client = AnthropicVertexChatCompletionClient._from_config(config) # pyright: ignore[reportPrivateUsage] # Verify equivalence - assert new_client._create_args["model"] == "claude-3-opus-20240229" - assert new_client._create_args["temperature"] == 0.5 - assert new_client._raw_config["vertex_info"]["project_id"] == "config-test-project" - assert new_client._raw_config["vertex_info"]["region"] == "europe-west1" + assert new_client._create_args["model"] == "claude-3-opus-20240229" # pyright: ignore[reportPrivateUsage] + assert new_client._create_args["temperature"] == 0.5 # pyright: ignore[reportPrivateUsage] + assert new_client._raw_config["vertex_info"]["project_id"] == "config-test-project" # pyright: ignore[reportPrivateUsage] + assert new_client._raw_config["vertex_info"]["region"] == "europe-west1" # pyright: ignore[reportPrivateUsage] -def test_vertex_client_component_attributes(): +def test_vertex_client_component_attributes() -> None: """Test that component attributes are set correctly.""" vertex_info = VertexInfo(project_id="test", region="us-east5") - client = AnthropicVertexChatCompletionClient( - model="claude-3-sonnet-20240229", - vertex_info=vertex_info - ) + client = AnthropicVertexChatCompletionClient(model="claude-3-sonnet-20240229", vertex_info=vertex_info) assert client.component_type == "model" assert hasattr(client, "component_config_schema") assert hasattr(client, "component_provider_override") assert client.component_provider_override == "autogen_ext.models.anthropic.AnthropicVertexChatCompletionClient" - messages = [ - UserMessage(content="I need to add 25 and 17. Use the add tool after thinking about it.", source="test") - ] - - thinking_config = {"thinking": {"type": "enabled", "budget_tokens": 2000}} - - result = await client.create(messages, tools=[add_tool], extra_create_args=thinking_config) - - # Should get tool calls - assert isinstance(result.content, list) - assert len(result.content) >= 1 - assert isinstance(result.content[0], FunctionCall) - assert result.content[0].name == "add_numbers" - - # Should have thinking content even with tool calls - assert result.thought is not None - assert len(result.thought) > 10 From 461c79b8c2e212b79199cff7bc1688a1a1d95f0d Mon Sep 17 00:00:00 2001 From: Srikar Mannepalli Date: Fri, 3 Oct 2025 13:40:22 +0530 Subject: [PATCH 3/3] format --- .../autogen-ext/tests/models/test_anthropic_model_client.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py index 4015ae9d0e31..66ba71d50877 100644 --- a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py @@ -1111,9 +1111,6 @@ async def test_streaming_tool_usage_with_arguments(provider: str) -> None: assert json.loads(content.arguments) is not None -# ===== Vertex Client Tests ===== - - def test_vertex_client_initialization_success() -> None: """Test successful initialization of AnthropicVertexChatCompletionClient.""" vertex_info = VertexInfo(project_id="test-project-123", region="us-east5")