Skip to content

Commit

Permalink
Fix Anthropic Bedrock support (#3210)
Browse files Browse the repository at this point in the history
* Added _configure_openai_config_for_bedrock to include aws variables in openai_config, necessary for setting AnthropicBedrock as client.

* Removed aws_session_token from required_keys

* Removed check for aws_session_token

* Removed all checks for aws_session_token

* Ran pre-commit

---------

Co-authored-by: Chi Wang <wang.chi@microsoft.com>
Co-authored-by: HRUSHIKESH DOKALA <96101829+Hk669@users.noreply.github.com>
  • Loading branch information
3 people authored Jul 25, 2024
1 parent e08fd68 commit 2b29274
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 8 deletions.
8 changes: 1 addition & 7 deletions autogen/oai/anthropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,17 +99,11 @@ def __init__(self, **kwargs: Any):
if not self._aws_secret_key:
self._aws_secret_key = os.getenv("AWS_SECRET_KEY")

if not self._aws_session_token:
self._aws_session_token = os.getenv("AWS_SESSION_TOKEN")

if not self._aws_region:
self._aws_region = os.getenv("AWS_REGION")

if self._api_key is None and (
self._aws_access_key is None
or self._aws_secret_key is None
or self._aws_session_token is None
or self._aws_region is None
self._aws_access_key is None or self._aws_secret_key is None or self._aws_region is None
):
raise ValueError("API key or AWS credentials are required to use the Anthropic API.")

Expand Down
12 changes: 11 additions & 1 deletion autogen/oai/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -454,12 +454,20 @@ def _configure_azure_openai(self, config: Dict[str, Any], openai_config: Dict[st
azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
)

def _configure_openai_config_for_bedrock(self, config: Dict[str, Any], openai_config: Dict[str, Any]) -> None:
"""Update openai_config with AWS credentials from config."""
required_keys = ["aws_access_key", "aws_secret_key", "aws_region"]

for key in required_keys:
if key in config:
openai_config[key] = config[key]

def _register_default_client(self, config: Dict[str, Any], openai_config: Dict[str, Any]) -> None:
"""Create a client with the given config to override openai_config,
after removing extra kwargs.
For Azure models/deployment names there's a convenience modification of model removing dots in
the it's value (Azure deploment names can't have dots). I.e. if you have Azure deployment name
the it's value (Azure deployment names can't have dots). I.e. if you have Azure deployment name
"gpt-35-turbo" and define model "gpt-3.5-turbo" in the config the function will remove the dot
from the name and create a client that connects to "gpt-35-turbo" Azure deployment.
"""
Expand All @@ -485,6 +493,8 @@ def _register_default_client(self, config: Dict[str, Any], openai_config: Dict[s
client = GeminiClient(**openai_config)
self._clients.append(client)
elif api_type is not None and api_type.startswith("anthropic"):
if "api_key" not in config:
self._configure_openai_config_for_bedrock(config, openai_config)
if anthropic_import_exception:
raise ImportError("Please install `anthropic` to use Anthropic API.")
client = AnthropicClient(**openai_config)
Expand Down

0 comments on commit 2b29274

Please sign in to comment.