add bedrock support

This commit is contained in:
LeonOstrez
2024-09-19 18:44:43 +02:00
parent 3e0917a476
commit a2f1768b7c
2 changed files with 35 additions and 2 deletions

View File

@@ -465,13 +465,40 @@ class ConfigLoader:
loader = ConfigLoader()
def adapt_for_bedrock(config: Config) -> Config:
"""
Adapt the configuration for use with Bedrock.
:param config: Configuration to adapt.
:return: Adapted configuration.
"""
if "anthropic" not in config.llm:
return config
if config.llm["anthropic"].base_url is None or "bedrock/anthropic" not in config.llm["anthropic"].base_url:
return config
replacement_map = {
"claude-3-5-sonnet-20240620": "anthropic.claude-3-5-sonnet-20240620-v1:0",
"claude-3-sonnet-20240229": "anthropic.claude-3-sonnet-20240229-v1:0",
"claude-3-haiku-20240307": "anthropic.claude-3-haiku-20240307-v1:0",
"claude-3-opus-20240229": "anthropic.claude-3-opus-20240229-v1:0",
}
for agent in config.agent:
if config.agent[agent].model in replacement_map:
config.agent[agent].model = replacement_map[config.agent[agent].model]
return config
def get_config() -> Config:
"""
Return current configuration.
:return: Current configuration object.
"""
return loader.config
return adapt_for_bedrock(loader.config)
__all__ = ["loader", "get_config"]

View File

@@ -74,8 +74,14 @@ class AnthropicClient(BaseLLMClient):
"temperature": self.config.temperature if temperature is None else temperature,
}
if "bedrock/anthropic" in self.config.base_url:
completion_kwargs["extra_headers"] = {"anthropic-version": "bedrock-2023-05-31"}
if "sonnet" in self.config.model:
completion_kwargs["extra_headers"] = {"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"}
if "extra_headers" in completion_kwargs:
completion_kwargs["extra_headers"]["anthropic-beta"] = "max-tokens-3-5-sonnet-2024-07-15"
else:
completion_kwargs["extra_headers"] = {"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"}
completion_kwargs["max_tokens"] = MAX_TOKENS_SONNET
if json_mode: