import base64
import json
import os
import re
from operator import itemgetter
from typing import (
Any,
Callable,
Dict,
Iterator,
List,
Literal,
Optional,
Sequence,
Tuple,
Type,
TypeVar,
Union,
cast,
)
import boto3
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import BaseChatModel, LanguageModelInput
from langchain_core.language_models.chat_models import LangSmithParams
from langchain_core.messages import (
AIMessage,
BaseMessage,
BaseMessageChunk,
HumanMessage,
HumanMessageChunk,
SystemMessage,
ToolCall,
ToolMessage,
merge_message_runs,
)
from langchain_core.messages.ai import AIMessageChunk, UsageMetadata
from langchain_core.messages.tool import tool_call as create_tool_call
from langchain_core.messages.tool import tool_call_chunk
from langchain_core.output_parsers import JsonOutputKeyToolsParser, PydanticToolsParser
from langchain_core.output_parsers.base import OutputParserLike
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_core.runnables import Runnable, RunnableMap, RunnablePassthrough
from langchain_core.tools import BaseTool
from langchain_core.utils import secret_from_env
from langchain_core.utils.function_calling import (
convert_to_openai_function,
convert_to_openai_tool,
)
from langchain_core.utils.pydantic import TypeBaseModel, is_basemodel_subclass
from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
from typing_extensions import Self
from langchain_aws.function_calling import ToolsOutputParser
_BM = TypeVar("_BM", bound=BaseModel)
_DictOrPydanticClass = Union[Dict[str, Any], Type[_BM], Type]
[docs]
class ChatBedrockConverse(BaseChatModel):
"""Bedrock chat model integration built on the Bedrock converse API.
This implementation will eventually replace the existing ChatBedrock implementation
once the Bedrock converse API has feature parity with older Bedrock API.
Specifically the converse API does not yet support custom Bedrock models.
Setup:
To use Amazon Bedrock make sure you've gone through all the steps described
here: https://docs.aws.amazon.com/bedrock/latest/userguide/setting-up.html
Once that's completed, install the LangChain integration:
.. code-block:: bash
pip install -U langchain-aws
Key init args — completion params:
model: str
Name of BedrockConverse model to use.
temperature: float
Sampling temperature.
max_tokens: Optional[int]
Max number of tokens to generate.
Key init args — client params:
region_name: Optional[str]
AWS region to use, e.g. 'us-west-2'.
base_url: Optional[str]
Bedrock endpoint to use. Needed if you don't want to default to us-east-
1 endpoint.
credentials_profile_name: Optional[str]
The name of the profile in the ~/.aws/credentials or ~/.aws/config files.
See full list of supported init args and their descriptions in the params section.
Instantiate:
.. code-block:: python
from langchain_aws import ChatBedrockConverse
llm = ChatBedrockConverse(
model="anthropic.claude-3-sonnet-20240229-v1:0",
temperature=0,
max_tokens=None,
# other params...
)
Invoke:
.. code-block:: python
messages = [
("system", "You are a helpful translator. Translate the user sentence to French."),
("human", "I love programming."),
]
llm.invoke(messages)
.. code-block:: python
AIMessage(content=[{'type': 'text', 'text': "J'aime la programmation."}], response_metadata={'ResponseMetadata': {'RequestId': '9ef1e313-a4c1-4f79-b631-171f658d3c0e', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Sat, 15 Jun 2024 01:19:24 GMT', 'content-type': 'application/json', 'content-length': '205', 'connection': 'keep-alive', 'x-amzn-requestid': '9ef1e313-a4c1-4f79-b631-171f658d3c0e'}, 'RetryAttempts': 0}, 'stopReason': 'end_turn', 'metrics': {'latencyMs': 609}}, id='run-754e152b-2b41-4784-9538-d40d71a5c3bc-0', usage_metadata={'input_tokens': 25, 'output_tokens': 11, 'total_tokens': 36})
Stream:
.. code-block:: python
for chunk in llm.stream(messages):
print(chunk)
.. code-block:: python
AIMessageChunk(content=[], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'type': 'text', 'text': 'J', 'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'text': "'", 'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'text': 'a', 'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'text': 'ime', 'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'text': ' la', 'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'text': ' programm', 'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'text': 'ation', 'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'text': '.', 'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[{'index': 0}], id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[], response_metadata={'stopReason': 'end_turn'}, id='run-da3c2606-4792-440a-ac66-72e0d1f6d117')
AIMessageChunk(content=[], response_metadata={'metrics': {'latencyMs': 581}}, id='run-da3c2606-4792-440a-ac66-72e0d1f6d117', usage_metadata={'input_tokens': 25, 'output_tokens': 11, 'total_tokens': 36})
.. code-block:: python
stream = llm.stream(messages)
full = next(stream)
for chunk in stream:
full += chunk
full
.. code-block:: python
AIMessageChunk(content=[{'type': 'text', 'text': "J'aime la programmation.", 'index': 0}], response_metadata={'stopReason': 'end_turn', 'metrics': {'latencyMs': 554}}, id='run-56a5a5e0-de86-412b-9835-624652dc3539', usage_metadata={'input_tokens': 25, 'output_tokens': 11, 'total_tokens': 36})
Tool calling:
.. code-block:: python
from pydantic import BaseModel, Field
class GetWeather(BaseModel):
'''Get the current weather in a given location'''
location: str = Field(..., description="The city and state, e.g. San Francisco, CA")
class GetPopulation(BaseModel):
'''Get the current population in a given location'''
location: str = Field(..., description="The city and state, e.g. San Francisco, CA")
llm_with_tools = llm.bind_tools([GetWeather, GetPopulation])
ai_msg = llm_with_tools.invoke("Which city is hotter today and which is bigger: LA or NY?")
ai_msg.tool_calls
.. code-block:: python
[{'name': 'GetWeather',
'args': {'location': 'Los Angeles, CA'},
'id': 'tooluse_Mspi2igUTQygp-xbX6XGVw'},
{'name': 'GetWeather',
'args': {'location': 'New York, NY'},
'id': 'tooluse_tOPHiDhvR2m0xF5_5tyqWg'},
{'name': 'GetPopulation',
'args': {'location': 'Los Angeles, CA'},
'id': 'tooluse__gcY_klbSC-GqB-bF_pxNg'},
{'name': 'GetPopulation',
'args': {'location': 'New York, NY'},
'id': 'tooluse_-1HSoGX0TQCSaIg7cdFy8Q'}]
See ``ChatBedrockConverse.bind_tools()`` method for more.
Structured output:
.. code-block:: python
from typing import Optional
from pydantic import BaseModel, Field
class Joke(BaseModel):
'''Joke to tell user.'''
setup: str = Field(description="The setup of the joke")
punchline: str = Field(description="The punchline to the joke")
rating: Optional[int] = Field(description="How funny the joke is, from 1 to 10")
structured_llm = llm.with_structured_output(Joke)
structured_llm.invoke("Tell me a joke about cats")
.. code-block:: python
Joke(setup='What do you call a cat that gets all dressed up?', punchline='A purrfessional!', rating=7)
See ``ChatBedrockConverse.with_structured_output()`` for more.
Image input:
.. code-block:: python
import base64
import httpx
from langchain_core.messages import HumanMessage
image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
image_data = base64.b64encode(httpx.get(image_url).content).decode("utf-8")
message = HumanMessage(
content=[
{"type": "text", "text": "describe the weather in this image"},
{
"type": "image",
"source": {"type": "base64", "media_type": "image/jpeg", "data": image_data},
},
],
)
ai_msg = llm.invoke([message])
ai_msg.content
.. code-block:: python
[{'type': 'text',
'text': 'The image depicts a sunny day with a partly cloudy sky. The sky is a brilliant blue color with scattered white clouds drifting across. The lighting and cloud patterns suggest pleasant, mild weather conditions. The scene shows an open grassy field or meadow, indicating warm temperatures conducive for vegetation growth. Overall, the weather portrayed in this scenic outdoor image appears to be sunny with some clouds, likely representing a nice, comfortable day.'}]
Token usage:
.. code-block:: python
ai_msg = llm.invoke(messages)
ai_msg.usage_metadata
.. code-block:: python
{'input_tokens': 25, 'output_tokens': 11, 'total_tokens': 36}
Response metadata
.. code-block:: python
ai_msg = llm.invoke(messages)
ai_msg.response_metadata
.. code-block:: python
{'ResponseMetadata': {'RequestId': '776a2a26-5946-45ae-859e-82dc5f12017c',
'HTTPStatusCode': 200,
'HTTPHeaders': {'date': 'Mon, 17 Jun 2024 01:37:05 GMT',
'content-type': 'application/json',
'content-length': '206',
'connection': 'keep-alive',
'x-amzn-requestid': '776a2a26-5946-45ae-859e-82dc5f12017c'},
'RetryAttempts': 0},
'stopReason': 'end_turn',
'metrics': {'latencyMs': 1290}}
""" # noqa: E501
client: Any = Field(default=None, exclude=True) #: :meta private:
model_id: str = Field(alias="model")
"""Id of the model to call.
e.g., ``"anthropic.claude-3-sonnet-20240229-v1:0"``. This is equivalent to the
modelID property in the list-foundation-models api. For custom and provisioned
models, an ARN value is expected. See
https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html#model-ids-arns
for a list of all supported built-in models.
"""
max_tokens: Optional[int] = None
"""Max tokens to generate."""
stop_sequences: Optional[List[str]] = Field(default=None, alias="stop")
"""Stop generation if any of these substrings occurs."""
temperature: Optional[float] = None
"""Sampling temperature. Must be 0 to 1."""
top_p: Optional[float] = None
"""The percentage of most-likely candidates that are considered for the next token.
Must be 0 to 1.
For example, if you choose a value of 0.8 for topP, the model selects from
the top 80% of the probability distribution of tokens that could be next in the
sequence."""
region_name: Optional[str] = None
"""The aws region, e.g., `us-west-2`.
Falls back to AWS_DEFAULT_REGION env variable or region specified in ~/.aws/config
in case it is not provided here.
"""
credentials_profile_name: Optional[str] = Field(default=None, exclude=True)
"""The name of the profile in the ~/.aws/credentials or ~/.aws/config files.
Profile should either have access keys or role information specified.
If not specified, the default credential profile or, if on an EC2 instance,
credentials from IMDS will be used.
See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
"""
aws_access_key_id: Optional[SecretStr] = Field(
default_factory=secret_from_env("AWS_ACCESS_KEY_ID", default=None)
)
"""AWS access key id.
If provided, aws_secret_access_key must also be provided.
If not specified, the default credential profile or, if on an EC2 instance,
credentials from IMDS will be used.
See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
If not provided, will be read from 'AWS_ACCESS_KEY_ID' environment variable.
"""
aws_secret_access_key: Optional[SecretStr] = Field(
default_factory=secret_from_env("AWS_SECRET_ACCESS_KEY", default=None)
)
"""AWS secret_access_key.
If provided, aws_access_key_id must also be provided.
If not specified, the default credential profile or, if on an EC2 instance,
credentials from IMDS will be used.
See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
If not provided, will be read from 'AWS_SECRET_ACCESS_KEY' environment variable.
"""
aws_session_token: Optional[SecretStr] = Field(
default_factory=secret_from_env("AWS_SESSION_TOKEN", default=None)
)
"""AWS session token.
If provided, aws_access_key_id and aws_secret_access_key must
also be provided. Not required unless using temporary credentials.
See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
If not provided, will be read from 'AWS_SESSION_TOKEN' environment variable.
"""
provider: str = ""
"""The model provider, e.g., amazon, cohere, ai21, etc.
When not supplied, provider is extracted from the first part of the model_id, e.g.
'amazon' in 'amazon.titan-text-express-v1'. This value should be provided for model
ids that do not have the provider in them, like custom and provisioned models that
have an ARN associated with them.
"""
endpoint_url: Optional[str] = Field(default=None, alias="base_url")
"""Needed if you don't want to default to us-east-1 endpoint"""
config: Any = None
"""An optional botocore.config.Config instance to pass to the client."""
guardrail_config: Optional[Dict[str, Any]] = Field(default=None, alias="guardrails")
"""Configuration information for a guardrail that you want to use in the request."""
additional_model_request_fields: Optional[Dict[str, Any]] = None
"""Additional inference parameters that the model supports.
Parameters beyond the base set of inference parameters that Converse supports in the
inferenceConfig field.
"""
additional_model_response_field_paths: Optional[List[str]] = None
"""Additional model parameters field paths to return in the response.
Converse returns the requested fields as a JSON Pointer object in the
additionalModelResponseFields field. The following is example JSON for
additionalModelResponseFieldPaths.
"""
supports_tool_choice_values: Optional[
Sequence[Literal["auto", "any", "tool"]]
] = None
"""Which types of tool_choice values the model supports.
Inferred if not specified. Inferred as ('auto', 'any', 'tool') if a 'claude-3'
model is used, ('auto', 'any') if a 'mistral-large' model is used, empty otherwise.
"""
model_config = ConfigDict(
extra="forbid",
populate_by_name=True,
)
@model_validator(mode="before")
@classmethod
def set_disable_streaming(cls, values: Dict) -> Any:
model_id = values.get("model_id", values.get("model"))
model_parts = model_id.split(".")
values["provider"] = values.get("provider") or (
model_parts[-2] if len(model_parts) > 1 else model_parts[0]
)
# As of 09/15/24 Anthropic and Cohere models support streamed tool calling
if "disable_streaming" not in values:
values["disable_streaming"] = (
False
if values["provider"] in ["anthropic", "cohere"]
else "tool_calling"
)
return values
@model_validator(mode="after")
def validate_environment(self) -> Self:
"""Validate that AWS credentials to and python package exists in environment."""
# As of 08/05/24 only claude-3 and mistral-large models support tool choice:
# https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
if self.supports_tool_choice_values is None:
if "claude-3" in self.model_id:
self.supports_tool_choice_values = ("auto", "any", "tool")
elif "mistral-large" in self.model_id:
self.supports_tool_choice_values = ("auto", "any")
else:
self.supports_tool_choice_values = ()
# Skip creating new client if passed in constructor
if self.client is None:
creds = {
"aws_access_key_id": self.aws_access_key_id,
"aws_secret_access_key": self.aws_secret_access_key,
"aws_session_token": self.aws_session_token,
}
if creds["aws_access_key_id"] and creds["aws_secret_access_key"]:
session_params = {
k: v.get_secret_value() for k, v in creds.items() if v
}
elif any(creds.values()):
raise ValueError(
f"If any of aws_access_key_id, aws_secret_access_key, or "
f"aws_session_token are specified then both aws_access_key_id and "
f"aws_secret_access_key must be specified. Only received "
f"{(k for k, v in creds.items() if v)}."
)
elif self.credentials_profile_name is not None:
session_params = {"profile_name": self.credentials_profile_name}
else:
# use default credentials
session_params = {}
try:
session = boto3.Session(**session_params)
self.region_name = (
self.region_name
or os.getenv("AWS_DEFAULT_REGION")
or session.region_name
)
client_params = {
"endpoint_url": self.endpoint_url,
"config": self.config,
"region_name": self.region_name,
}
client_params = {k: v for k, v in client_params.items() if v}
self.client = session.client("bedrock-runtime", **client_params)
except ValueError as e:
raise ValueError(f"Error raised by bedrock service:\n\n{e}") from e
except Exception as e:
raise ValueError(
"Could not load credentials to authenticate with AWS client. "
"Please check that credentials in the specified "
f"profile name are valid. Bedrock error:\n\n{e}"
) from e
return self
def _generate(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> ChatResult:
"""Top Level call"""
bedrock_messages, system = _messages_to_bedrock(messages)
params = self._converse_params(
stop=stop, **_snake_to_camel_keys(kwargs, excluded_keys={"inputSchema"})
)
response = self.client.converse(
messages=bedrock_messages, system=system, **params
)
response_message = _parse_response(response)
return ChatResult(generations=[ChatGeneration(message=response_message)])
def _stream(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> Iterator[ChatGenerationChunk]:
bedrock_messages, system = _messages_to_bedrock(messages)
params = self._converse_params(
stop=stop, **_snake_to_camel_keys(kwargs, excluded_keys={"inputSchema"})
)
response = self.client.converse_stream(
messages=bedrock_messages, system=system, **params
)
for event in response["stream"]:
if message_chunk := _parse_stream_event(event):
yield ChatGenerationChunk(message=message_chunk)
# TODO: Add async support once there are async bedrock.converse methods.
[docs]
def with_structured_output(
self,
schema: _DictOrPydanticClass,
*,
include_raw: bool = False,
**kwargs: Any,
) -> Runnable[LanguageModelInput, Union[Dict, BaseModel]]:
supports_tool_choice_values = self.supports_tool_choice_values or ()
if "tool" in supports_tool_choice_values:
tool_choice = convert_to_openai_function(schema)["name"]
elif "any" in supports_tool_choice_values:
tool_choice = "any"
else:
tool_choice = None
llm = self.bind_tools([schema], tool_choice=tool_choice)
if isinstance(schema, type) and is_basemodel_subclass(schema):
if self.disable_streaming:
output_parser: OutputParserLike = ToolsOutputParser(
first_tool_only=True, pydantic_schemas=[schema]
)
else:
output_parser = PydanticToolsParser(
tools=[schema],
first_tool_only=True,
)
else:
tool_name = convert_to_openai_tool(schema)["function"]["name"]
if self.disable_streaming:
output_parser = ToolsOutputParser(first_tool_only=True, args_only=True)
else:
output_parser = JsonOutputKeyToolsParser(
key_name=tool_name, first_tool_only=True
)
if include_raw:
parser_assign = RunnablePassthrough.assign(
parsed=itemgetter("raw") | output_parser, parsing_error=lambda _: None
)
parser_none = RunnablePassthrough.assign(parsed=lambda _: None)
parser_with_fallback = parser_assign.with_fallbacks(
[parser_none], exception_key="parsing_error"
)
return RunnableMap(raw=llm) | parser_with_fallback
else:
return llm | output_parser
def _converse_params(
self,
*,
stop: Optional[List[str]] = None,
stopSequences: Optional[List[str]] = None,
maxTokens: Optional[List[str]] = None,
temperature: Optional[float] = None,
topP: Optional[float] = None,
tools: Optional[List] = None,
toolChoice: Optional[dict] = None,
modelId: Optional[str] = None,
inferenceConfig: Optional[dict] = None,
toolConfig: Optional[dict] = None,
additionalModelRequestFields: Optional[dict] = None,
additionalModelResponseFieldPaths: Optional[List[str]] = None,
guardrailConfig: Optional[dict] = None,
) -> Dict[str, Any]:
if not inferenceConfig:
inferenceConfig = {
"maxTokens": maxTokens or self.max_tokens,
"temperature": temperature or self.temperature,
"topP": self.top_p or topP,
"stopSequences": stop or stopSequences or self.stop_sequences,
}
if not toolConfig and tools:
toolChoice = _format_tool_choice(toolChoice) if toolChoice else None
toolConfig = {"tools": _format_tools(tools), "toolChoice": toolChoice}
return _drop_none(
{
"modelId": modelId or self.model_id,
"inferenceConfig": inferenceConfig,
"toolConfig": toolConfig,
"additionalModelRequestFields": additionalModelRequestFields
or self.additional_model_request_fields,
"additionalModelResponseFieldPaths": additionalModelResponseFieldPaths
or self.additional_model_response_field_paths,
"guardrailConfig": guardrailConfig or self.guardrail_config,
}
)
def _get_ls_params(
self, stop: Optional[List[str]] = None, **kwargs: Any
) -> LangSmithParams:
"""Get standard params for tracing."""
params = self._get_invocation_params(stop=stop, **kwargs)
ls_params = LangSmithParams(
ls_provider="amazon_bedrock",
ls_model_name=self.model_id,
ls_model_type="chat",
ls_temperature=params.get("temperature", self.temperature),
)
if ls_max_tokens := params.get("max_tokens", self.max_tokens):
ls_params["ls_max_tokens"] = ls_max_tokens
if ls_stop := stop or params.get("stop", None):
ls_params["ls_stop"] = ls_stop
return ls_params
@property
def _llm_type(self) -> str:
"""Return type of chat model."""
return "amazon_bedrock_converse_chat"
@classmethod
def is_lc_serializable(cls) -> bool:
return True
@classmethod
def get_lc_namespace(cls) -> list[str]:
return ["langchain_aws", "chat_models"]
@property
def lc_secrets(self) -> Dict[str, str]:
return {
"aws_access_key_id": "AWS_ACCESS_KEY_ID",
"aws_secret_access_key": "AWS_SECRET_ACCESS_KEY",
"aws_session_token": "AWS_SESSION_TOKEN",
}
def _messages_to_bedrock(
messages: List[BaseMessage],
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
"""Handle Bedrock converse and Anthropic style content blocks"""
bedrock_messages: List[Dict[str, Any]] = []
bedrock_system: List[Dict[str, Any]] = []
# Merge system, human, ai message runs because Anthropic expects (at most) 1
# system message then alternating human/ai messages.
messages = merge_message_runs(messages)
for msg in messages:
content = _anthropic_to_bedrock(msg.content)
if isinstance(msg, HumanMessage):
# If there's a human, tool, human message sequence, the
# tool message will be merged with the first human message, so the second
# human message will now be preceded by a human message and should also
# be merged with it.
if bedrock_messages and bedrock_messages[-1]["role"] == "user":
bedrock_messages[-1]["content"].extend(content)
else:
bedrock_messages.append({"role": "user", "content": content})
elif isinstance(msg, AIMessage):
content = _upsert_tool_calls_to_bedrock_content(content, msg.tool_calls)
bedrock_messages.append({"role": "assistant", "content": content})
elif isinstance(msg, SystemMessage):
bedrock_system.extend(content)
elif isinstance(msg, ToolMessage):
if bedrock_messages and bedrock_messages[-1]["role"] == "user":
curr = bedrock_messages.pop()
else:
curr = {"role": "user", "content": []}
curr["content"].append(
{
"toolResult": {
"content": content,
"toolUseId": msg.tool_call_id,
"status": msg.status,
}
}
)
bedrock_messages.append(curr)
else:
raise ValueError(f"Unsupported message type {type(msg)}")
return bedrock_messages, bedrock_system
def _extract_response_metadata(response: Dict[str, Any]) -> Dict[str, Any]:
response_metadata = response
# response_metadata only supports string, list or dict
if "metrics" in response and "latencyMs" in response["metrics"]:
response_metadata["metrics"]["latencyMs"] = [response["metrics"]["latencyMs"]]
return response_metadata
def _parse_response(response: Dict[str, Any]) -> AIMessage:
anthropic_content = _bedrock_to_anthropic(
response.pop("output")["message"]["content"]
)
tool_calls = _extract_tool_calls(anthropic_content)
usage = UsageMetadata(_camel_to_snake_keys(response.pop("usage"))) # type: ignore[misc]
return AIMessage(
content=_str_if_single_text_block(anthropic_content), # type: ignore[arg-type]
usage_metadata=usage,
response_metadata=_extract_response_metadata(response),
tool_calls=tool_calls,
)
def _parse_stream_event(event: Dict[str, Any]) -> Optional[BaseMessageChunk]:
if "messageStart" in event:
# TODO: needed?
return (
AIMessageChunk(content=[])
if event["messageStart"]["role"] == "assistant"
else HumanMessageChunk(content=[])
)
elif "contentBlockStart" in event:
block = {
**_bedrock_to_anthropic([event["contentBlockStart"]["start"]])[0],
"index": event["contentBlockStart"]["contentBlockIndex"],
}
tool_call_chunks = []
if block["type"] == "tool_use":
tool_call_chunks.append(
tool_call_chunk(
name=block.get("name"),
id=block.get("id"),
args=block.get("input"),
index=event["contentBlockStart"]["contentBlockIndex"],
)
)
return AIMessageChunk(content=[block], tool_call_chunks=tool_call_chunks)
elif "contentBlockDelta" in event:
block = {
**_bedrock_to_anthropic([event["contentBlockDelta"]["delta"]])[0],
"index": event["contentBlockDelta"]["contentBlockIndex"],
}
tool_call_chunks = []
if block["type"] == "tool_use":
tool_call_chunks.append(
tool_call_chunk(
name=block.get("name"),
id=block.get("id"),
args=block.get("input"),
index=event["contentBlockDelta"]["contentBlockIndex"],
)
)
return AIMessageChunk(content=[block], tool_call_chunks=tool_call_chunks)
elif "contentBlockStop" in event:
# TODO: needed?
return AIMessageChunk(
content=[{"index": event["contentBlockStop"]["contentBlockIndex"]}]
)
elif "messageStop" in event:
# TODO: snake case response metadata?
return AIMessageChunk(content=[], response_metadata=event["messageStop"])
elif "metadata" in event:
usage = UsageMetadata(_camel_to_snake_keys(event["metadata"].pop("usage"))) # type: ignore[misc]
return AIMessageChunk(
content=[], response_metadata=event["metadata"], usage_metadata=usage
)
elif "Exception" in list(event.keys())[0]:
name, info = list(event.items())[0]
raise ValueError(
f"Received AWS exception {name}:\n\n{json.dumps(info, indent=2)}"
)
else:
raise ValueError(f"Received unsupported stream event:\n\n{event}")
def _anthropic_to_bedrock(
content: Union[str, List[Union[str, Dict[str, Any]]]],
) -> List[Dict[str, Any]]:
if isinstance(content, str):
content = [{"text": content}]
bedrock_content: List[Dict[str, Any]] = []
for block in _snake_to_camel_keys(content):
if isinstance(block, str):
bedrock_content.append({"text": block})
# Assume block is already in bedrock format.
elif "type" not in block:
bedrock_content.append(block)
elif block["type"] == "text":
bedrock_content.append({"text": block["text"]})
elif block["type"] == "image":
# Assume block is already in bedrock format.
if "image" in block:
bedrock_content.append({"image": block["image"]})
else:
bedrock_content.append(
{
"image": {
"format": block["source"]["mediaType"].split("/")[1],
"source": {
"bytes": _b64str_to_bytes(block["source"]["data"])
},
}
}
)
elif block["type"] == "image_url":
# Support OpenAI image format as well.
bedrock_content.append(
{"image": _format_openai_image_url(block["imageUrl"]["url"])}
)
elif block["type"] == "document":
# Assume block in bedrock document format
bedrock_content.append({"document": block["document"]})
elif block["type"] == "tool_use":
bedrock_content.append(
{
"toolUse": {
"toolUseId": block["id"],
"input": block["input"],
"name": block["name"],
}
}
)
elif block["type"] == "tool_result":
bedrock_content.append(
{
"toolResult": {
"toolUseId": block["toolUseId"],
"content": _anthropic_to_bedrock(block["content"]),
"status": "error" if block.get("isError") else "success",
}
}
)
# Only needed for tool_result content blocks.
elif block["type"] == "json":
bedrock_content.append({"json": block["json"]})
elif block["type"] == "guard_content":
bedrock_content.append({"guardContent": {"text": {"text": block["text"]}}})
else:
raise ValueError(f"Unsupported content block type:\n{block}")
# drop empty text blocks
return [block for block in bedrock_content if block.get("text", True)]
def _bedrock_to_anthropic(content: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
anthropic_content = []
for block in _camel_to_snake_keys(content):
if "text" in block:
anthropic_content.append({"type": "text", "text": block["text"]})
elif "tool_use" in block:
block["tool_use"]["id"] = block["tool_use"].pop("tool_use_id", None)
anthropic_content.append({"type": "tool_use", **block["tool_use"]})
elif "image" in block:
anthropic_content.append(
{
"type": "image",
"source": {
"media_type": f"image/{block['image']['format']}",
"type": "base64",
"data": _bytes_to_b64_str(block["image"]["source"]["bytes"]),
},
}
)
elif "tool_result" in block:
anthropic_content.append(
{
"type": "tool_result",
"tool_use_id": block["tool_result"]["tool_use_id"],
"is_error": block["tool_result"].get("status") == "error",
"content": _bedrock_to_anthropic(block["tool_result"]["content"]),
}
)
# Only occurs in content blocks of a tool_result:
elif "json" in block:
anthropic_content.append({"type": "json", **block})
elif "guard_content" in block:
anthropic_content.append(
{
"type": "guard_content",
"guard_content": {
"type": "text",
"text": block["guard_content"]["text"]["text"],
},
}
)
else:
raise ValueError(
"Unexpected content block type in content. Expected to have one of "
"'text', 'tool_use', 'image', or 'tool_result' keys. Received:\n\n"
f"{block}"
)
return anthropic_content
def _format_tools(
tools: Sequence[Union[Dict[str, Any], TypeBaseModel, Callable, BaseTool],],
) -> List[Dict[Literal["toolSpec"], Dict[str, Union[Dict[str, Any], str]]]]:
formatted_tools: List = []
for tool in tools:
if isinstance(tool, dict) and "toolSpec" in tool:
formatted_tools.append(tool)
else:
spec = convert_to_openai_tool(tool)["function"]
spec["inputSchema"] = {"json": spec.pop("parameters")}
formatted_tools.append({"toolSpec": spec})
return formatted_tools
def _format_tool_choice(
tool_choice: Union[Dict[str, Dict], Literal["auto", "any"], str],
) -> Dict[str, Dict[str, str]]:
if isinstance(tool_choice, dict):
return tool_choice
elif tool_choice in ("auto", "any"):
return {tool_choice: {}}
else:
return {"tool": {"name": tool_choice}}
def _extract_tool_calls(anthropic_content: List[dict]) -> List[ToolCall]:
tool_calls = []
for block in anthropic_content:
if block["type"] == "tool_use":
tool_calls.append(
create_tool_call(
name=block["name"], args=block["input"], id=block["id"]
)
)
return tool_calls
def _snake_to_camel(text: str) -> str:
split = text.split("_")
return "".join(split[:1] + [s.title() for s in split[1:]])
def _camel_to_snake(text: str) -> str:
pattern = re.compile(r"(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])")
return pattern.sub("_", text).lower()
_T = TypeVar("_T")
def _camel_to_snake_keys(obj: _T) -> _T:
if isinstance(obj, list):
return cast(_T, [_camel_to_snake_keys(e) for e in obj])
elif isinstance(obj, dict):
return cast(
_T, {_camel_to_snake(k): _camel_to_snake_keys(v) for k, v in obj.items()}
)
else:
return obj
def _snake_to_camel_keys(obj: _T, excluded_keys: set = set()) -> _T:
if isinstance(obj, list):
return cast(
_T, [_snake_to_camel_keys(e, excluded_keys=excluded_keys) for e in obj]
)
elif isinstance(obj, dict):
_dict = {}
for k, v in obj.items():
if k in excluded_keys:
_dict[k] = v
else:
_dict[_snake_to_camel(k)] = _snake_to_camel_keys(
v, excluded_keys=excluded_keys
)
return cast(_T, _dict)
else:
return obj
def _drop_none(obj: Any) -> Any:
if isinstance(obj, dict):
new = {k: _drop_none(v) for k, v in obj.items() if _drop_none(v) is not None}
return new
else:
return obj
def _b64str_to_bytes(base64_str: str) -> bytes:
return base64.b64decode(base64_str.encode("utf-8"))
def _bytes_to_b64_str(bytes_: bytes) -> str:
return base64.b64encode(bytes_).decode("utf-8")
def _str_if_single_text_block(
anthropic_content: List[Dict[str, Any]],
) -> Union[str, List[Dict[str, Any]]]:
if len(anthropic_content) == 1 and anthropic_content[0]["type"] == "text":
return anthropic_content[0]["text"]
return anthropic_content
def _upsert_tool_calls_to_bedrock_content(
content: List[Dict[str, Any]], tool_calls: List[ToolCall]
) -> List[Dict[str, Any]]:
existing_tc_blocks = [block for block in content if "toolUse" in block]
for tool_call in tool_calls:
if tool_call["id"] in [
block["toolUse"]["toolUseId"] for block in existing_tc_blocks
]:
tc_block = next(
block
for block in existing_tc_blocks
if block["toolUse"]["toolUseId"] == tool_call["id"]
)
tc_block["toolUse"]["input"] = tool_call["args"]
tc_block["toolUse"]["name"] = tool_call["name"]
else:
content.append(
{
"toolUse": {
"toolUseId": tool_call["id"],
"input": tool_call["args"],
"name": tool_call["name"],
}
}
)
return content
def _format_openai_image_url(image_url: str) -> Dict:
"""
Formats an image of format data:image/jpeg;base64,{b64_string}
to a dict for bedrock api.
And throws an error if url is not a b64 image.
"""
regex = r"^data:image/(?P<media_type>.+);base64,(?P<data>.+)$"
match = re.match(regex, image_url)
if match is None:
raise ValueError(
"Bedrock does not currently support OpenAI-format image URLs, only "
"base64-encoded images. Example: data:image/png;base64,'/9j/4AAQSk'..."
)
return {
"format": match.group("media_type"),
"source": {"bytes": _b64str_to_bytes(match.group("data"))},
}
