Source code for langchain_aws.chat_models.bedrock_converse

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.pydantic_v1 import BaseModel, Field, root_validator
from langchain_core.runnables import Runnable, RunnableMap, RunnablePassthrough
from langchain_core.tools import BaseTool
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 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 langchain_core.pydantic_v1 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 langchain_core.pydantic_v1 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 """ 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. """ class Config: """Configuration for this pydantic object.""" extra = "forbid" allow_population_by_field_name = True @root_validator(pre=True) def set_disable_streaming(cls, values: Dict) -> Dict: values["provider"] = ( values.get("provider") or (values.get("model_id", values["model"])).split(".")[0] ) # As of 08/05/24 only Anthropic models support streamed tool calling if "disable_streaming" not in values: values["disable_streaming"] = ( False if "anthropic" in values["provider"] else "tool_calling" ) return values @root_validator(pre=False, skip_on_failure=True) def validate_environment(cls, values: Dict) -> Dict: """Validate that AWS credentials to and python package exists in environment.""" if values["client"] is not None: return values try: if values["credentials_profile_name"] is not None: session = boto3.Session(profile_name=values["credentials_profile_name"]) else: session = boto3.Session() except ValueError as e: raise ValueError(f"Error raised by bedrock service: {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: {e}" ) from e values["region_name"] = ( values.get("region_name") or os.getenv("AWS_DEFAULT_REGION") or session.region_name ) client_params = {} if values["region_name"]: client_params["region_name"] = values["region_name"] if values["endpoint_url"]: client_params["endpoint_url"] = values["endpoint_url"] if values["config"]: client_params["config"] = values["config"] try: values["client"] = session.client("bedrock-runtime", **client_params) except ValueError as e: raise ValueError(f"Error raised by bedrock service: {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: {e}" ) from e # 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 values["supports_tool_choice_values"] is None: if "claude-3" in values["model_id"]: values["supports_tool_choice_values"] = ("auto", "any", "tool") elif "mistral-large" in values["model_id"]: values["supports_tool_choice_values"] = ("auto", "any") else: values["supports_tool_choice_values"] = () return values 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 bind_tools( self, tools: Sequence[Union[Dict[str, Any], TypeBaseModel, Callable, BaseTool]], *, tool_choice: Optional[Union[dict, str, Literal["auto", "any"]]] = None, **kwargs: Any, ) -> Runnable[LanguageModelInput, BaseMessage]: if tool_choice: tool_choice = _format_tool_choice(tool_choice) tool_choice_type = list(tool_choice.keys())[0] if tool_choice_type not in list(self.supports_tool_choice_values or []): if self.supports_tool_choice_values: supported = ( f"Model {self.model_id} does not currently support tool_choice " f"of type {tool_choice_type}. The following tool_choice types " f"are supported: {self.supports_tool_choice_values}." ) else: supported = ( f"Model {self.model_id} does not currently support tool_choice." ) raise ValueError( f"{supported} Please see " f"https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html " # noqa: E501 f"for the latest documentation on models that support tool choice." ) kwargs["tool_choice"] = _format_tool_choice(tool_choice) return self.bind(tools=_format_tools(tools), **kwargs)
[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"
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() return bedrock_messages, bedrock_system 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=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_function(tool) 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"))}, }