Adding bedrock chat completion for anthropic models

python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py

@@ -1,14 +1,20 @@
 from ._anthropic_client import AnthropicChatCompletionClient, BaseAnthropicChatCompletionClient
 from .config import (
+    AnthropicBedrockClientConfiguration,
+    AnthropicBedrockClientConfigurationConfigModel,
     AnthropicClientConfiguration,
     AnthropicClientConfigurationConfigModel,
+    BedrockInfo,
     CreateArgumentsConfigModel,
 )
 
 __all__ = [
     "AnthropicChatCompletionClient",
     "BaseAnthropicChatCompletionClient",
     "AnthropicClientConfiguration",
+    "AnthropicBedrockClientConfiguration",
     "AnthropicClientConfigurationConfigModel",
+    "AnthropicBedrockClientConfigurationConfigModel",
     "CreateArgumentsConfigModel",
+    "BedrockInfo",
 ]

python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py

@@ -25,7 +25,7 @@
 )
 
 import tiktoken
-from anthropic import AsyncAnthropic, AsyncStream
+from anthropic import AnthropicBedrock, AsyncAnthropic, AsyncStream
 from anthropic.types import (
     Base64ImageSourceParam,
     ContentBlock,
@@ -67,7 +67,13 @@
 from typing_extensions import Self, Unpack
 
 from . import _model_info
-from .config import AnthropicClientConfiguration, AnthropicClientConfigurationConfigModel
+from .config import (
+    AnthropicBedrockClientConfiguration,
+    AnthropicBedrockClientConfigurationConfigModel,
+    AnthropicClientConfiguration,
+    AnthropicClientConfigurationConfigModel,
+    BedrockInfo,
+)
 
 logger = logging.getLogger(EVENT_LOGGER_NAME)
 trace_logger = logging.getLogger(TRACE_LOGGER_NAME)
@@ -410,7 +416,7 @@
 class BaseAnthropicChatCompletionClient(ChatCompletionClient):
     def __init__(
         self,
-        client: AsyncAnthropic,
+        client: Any,
         *,
         create_args: Dict[str, Any],
         model_info: Optional[ModelInfo] = None,
@@ -1109,3 +1115,138 @@
             copied_config["api_key"] = config.api_key.get_secret_value()
 
         return cls(**copied_config)
+
+
+class AnthropicBedrockChatCompletionClient(
+    BaseAnthropicChatCompletionClient, Component[AnthropicBedrockClientConfigurationConfigModel]
+):
+    """
+    Chat completion client for Anthropic's Claude models.
+
+    Args:
+        model (str): The Claude model to use (e.g., "claude-3-sonnet-20240229", "claude-3-opus-20240229")
+        api_key (str, optional): Anthropic API key. Required if not in environment variables.
+        base_url (str, optional): Override the default API endpoint.
+        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.
+        bedrock_info (BedrockInfo, optional): The capabilities of the model in bedrock. Required if using a model from AWS bedrock.
+
+    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 AnthropicBedrockChatCompletionClient
+        from autogen_core.models import UserMessage
+
+
+        async def main():
+            config = {
+                "model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
+                "temperature": 0.1,
+                "model_info": {
+                    "vision": True,
+                    "function_calling": True,
+                    "json_output": True,
+                    "family": ModelFamily.CLAUDE_3_5_SONNET,
+                },
+                "bedrock_info": {
+                    "aws_access_key": "<aws_access_key>",
+                    "aws_secret_key": "<aws_secret_key>",
+                    "aws_session_token": "<aws_session_token>",
+                    "aws_region": "<aws_region>",
+                },
+            }
+            anthropic_client = AnthropicBedrockChatCompletionClient(**config)
+
+            result = await anthropic_client.create([UserMessage(content="What is the capital of France?", source="user")])  # type: ignore
+            print(result)
+
+
+        if __name__ == "__main__":
+            asyncio.run(main())
+    """
+
+    component_type = "model"
+    component_config_schema = AnthropicBedrockClientConfigurationConfigModel
+    component_provider_override = "autogen_ext.models.anthropic.AnthropicChatCompletionClient"
+
+    def __init__(self, **kwargs: Unpack[AnthropicBedrockClientConfiguration]):
+        if "model" not in kwargs:
+            raise ValueError("model is required for AnthropicChatCompletionClient")
+
+        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"]
+
+        bedrock_info: Optional[BedrockInfo] = None
+        if "bedrock_info" in kwargs:
+            bedrock_info = kwargs["bedrock_info"]
+
+        if bedrock_info is None:
+            raise ValueError("bedrock_info is required for AnthropicBedrockChatCompletionClient")
+
+        # Handle bedrock_info as secretestr
+        aws_region = bedrock_info["aws_region"]
+        aws_access_key = bedrock_info["aws_access_key"].get_secret_value()
+        aws_secret_key = bedrock_info["aws_secret_key"].get_secret_value()
+        aws_session_token = bedrock_info["aws_session_token"].get_secret_value()
+
+        client = AnthropicBedrock(
+            aws_access_key=aws_access_key,
+            aws_secret_key=aws_secret_key,
+            aws_session_token=aws_session_token,
+            aws_region=aws_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)
+        self._client = _anthropic_client_from_config(state["_raw_config"])
+
+    def _to_config(self) -> AnthropicBedrockClientConfigurationConfigModel:
+        copied_config = self._raw_config.copy()
+        return AnthropicBedrockClientConfigurationConfigModel(**copied_config)
+
+    @classmethod
+    def _from_config(cls, config: AnthropicBedrockClientConfigurationConfigModel) -> Self:
+        copied_config = config.model_copy().model_dump(exclude_none=True)
+
+        # Handle api_key as SecretStr
+        if "api_key" in copied_config and isinstance(config.api_key, SecretStr):
+            copied_config["api_key"] = config.api_key.get_secret_value()
+
+        # Handle bedrock_info as SecretStr
+        if "bedrock_info" in copied_config and isinstance(config.bedrock_info, dict):
+            copied_config["bedrock_info"] = {
+                "aws_access_key": config.bedrock_info["aws_access_key"].get_secret_value(),
+                "aws_secret_key": config.bedrock_info["aws_secret_key"].get_secret_value(),
+                "aws_session_token": config.bedrock_info["aws_session_token"].get_secret_value(),
+                "aws_region": config.bedrock_info["aws_region"],
+            }
+
+        return cls(**copied_config)

python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py

@@ -2,7 +2,7 @@
 
 from autogen_core.models import ModelCapabilities, ModelInfo  # type: ignore
 from pydantic import BaseModel, SecretStr
-from typing_extensions import TypedDict
+from typing_extensions import Required, TypedDict
 
 
 class ResponseFormat(TypedDict):
@@ -20,6 +20,22 @@ class CreateArguments(TypedDict, total=False):
     metadata: Optional[Dict[str, str]]
 
 
+class BedrockInfo(TypedDict):
+    """BedrockInfo is a dictionary that contains information about a bedrock's properties.
+    It is expected to be used in the bedrock_info property of a model client.
+
+    """
+
+    aws_access_key: Required[SecretStr]
+    """Access key for the aws account to gain bedrock model access"""
+    aws_secret_key: Required[SecretStr]
+    """Access secret key for the aws account to gain bedrock model access"""
+    aws_session_token: Required[SecretStr]
+    """aws session token for the aws account to gain bedrock model access"""
+    aws_region: Required[str]
+    """aws region for the aws account to gain bedrock model access"""
+
+
 class BaseAnthropicClientConfiguration(CreateArguments, total=False):
     api_key: str
     base_url: Optional[str]
@@ -36,6 +52,10 @@ class AnthropicClientConfiguration(BaseAnthropicClientConfiguration, total=False
     tool_choice: Optional[Union[Literal["auto", "any", "none"], Dict[str, Any]]]
 
 
+class AnthropicBedrockClientConfiguration(AnthropicClientConfiguration, total=False):
+    bedrock_info: BedrockInfo
+
+
 # Pydantic equivalents of the above TypedDicts
 class CreateArgumentsConfigModel(BaseModel):
     model: str
@@ -61,3 +81,7 @@ class BaseAnthropicClientConfigurationConfigModel(CreateArgumentsConfigModel):
 class AnthropicClientConfigurationConfigModel(BaseAnthropicClientConfigurationConfigModel):
     tools: List[Dict[str, Any]] | None = None
     tool_choice: Union[Literal["auto", "any", "none"], Dict[str, Any]] | None = None
+
+
+class AnthropicBedrockClientConfigurationConfigModel(AnthropicClientConfigurationConfigModel):
+    bedrock_info: BedrockInfo | None = None