agno 0.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/models/aws/bedrock.py

@@ -1,22 +1,38 @@
 import json
 from dataclasses import dataclass
-from typing import Any, Dict, Iterator, List, Optional, Tuple
+from os import getenv
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, Type, Union
 
-from agno.aws.api_client import AwsApiClient  # type: ignore
-from agno.models.base import Model, StreamData
+from pydantic import BaseModel
+
+from agno.exceptions import ModelProviderError
+from agno.models.base import Model
 from agno.models.message import Message
-from agno.models.response import ModelResponse, ModelResponseEvent
-from agno.utils.log import logger
-from agno.utils.timer import Timer
-from agno.utils.tools import (
-    get_function_call_for_tool_call,
-)
+from agno.models.metrics import Metrics
+from agno.models.response import ModelResponse
+from agno.run.agent import RunOutput
+from agno.utils.log import log_debug, log_error, log_warning
+from agno.utils.tokens import count_schema_tokens
 
 try:
-    from boto3 import session  # noqa: F401
+    from boto3 import client as AwsClient
+    from boto3.session import Session
+    from botocore.exceptions import ClientError
 except ImportError:
-    logger.error("`boto3` not installed")
-    raise
+    raise ImportError("`boto3` not installed. Please install using `pip install boto3`")
+
+try:
+    import aioboto3
+
+    AIOBOTO3_AVAILABLE = True
+except ImportError:
+    aioboto3 = None
+    AIOBOTO3_AVAILABLE = False
+
+
+BEDROCK_SUPPORTED_IMAGE_FORMATS = ["png", "jpeg", "webp", "gif"]
+BEDROCK_SUPPORTED_VIDEO_FORMATS = ["mp4", "mov", "mkv", "webm", "flv", "mpeg", "mpg", "wmv", "three_gp"]
+BEDROCK_SUPPORTED_FILE_FORMATS = ["pdf", "csv", "doc", "docx", "xls", "xlsx", "html", "txt", "md"]
 
 
 @dataclass
@@ -24,524 +40,743 @@ class AwsBedrock(Model):
     """
     AWS Bedrock model.
 
+    To use this model, you need to either:
+    1. Set the following environment variables:
+       - AWS_ACCESS_KEY_ID
+       - AWS_SECRET_ACCESS_KEY
+       - AWS_REGION
+    2. Or provide a boto3 Session object
+
+    For async support, you also need aioboto3 installed:
+       pip install aioboto3
+
+    Not all Bedrock models support all features. See this documentation for more information: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-supported-models-features.html
+
     Args:
         aws_region (Optional[str]): The AWS region to use.
-        aws_profile (Optional[str]): The AWS profile to use.
-        aws_client (Optional[AwsApiClient]): The AWS client to use.
-        _bedrock_client (Optional[Any]): The Bedrock client to use.
-        _bedrock_runtime_client (Optional[Any]): The Bedrock runtime client to use.
+        aws_access_key_id (Optional[str]): The AWS access key ID to use.
+        aws_secret_access_key (Optional[str]): The AWS secret access key to use.
+        aws_sso_auth (Optional[str]): Removes the need for an access and secret access key by leveraging the current profile's authentication
+        session (Optional[Session]): A boto3 Session object to use for authentication.
     """
 
-    aws_region: Optional[str] = None
-    aws_profile: Optional[str] = None
-    aws_client: Optional[AwsApiClient] = None
+    id: str = "mistral.mistral-small-2402-v1:0"
+    name: str = "AwsBedrock"
+    provider: str = "AwsBedrock"
 
-    _bedrock_client: Optional[Any] = None
-    _bedrock_runtime_client: Optional[Any] = None
+    aws_sso_auth: Optional[bool] = False
+    aws_region: Optional[str] = None
+    aws_access_key_id: Optional[str] = None
+    aws_secret_access_key: Optional[str] = None
+    session: Optional[Session] = None
+
+    # Request parameters
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    stop_sequences: Optional[List[str]] = None
+    request_params: Optional[Dict[str, Any]] = None
+
+    client: Optional[AwsClient] = None
+    async_client: Optional[Any] = None
+    async_session: Optional[Any] = None
+
+    def get_client(self) -> AwsClient:
+        """
+        Get the Bedrock client.
 
-    def get_aws_region(self) -> Optional[str]:
-        # Priority 1: Use aws_region from model
-        if self.aws_region is not None:
-            return self.aws_region
+        Returns:
+            AwsClient: The Bedrock client.
+        """
+        if self.client is not None:
+            return self.client
+
+        if self.session:
+            self.client = self.session.client("bedrock-runtime")
+            return self.client
+
+        self.aws_access_key_id = self.aws_access_key_id or getenv("AWS_ACCESS_KEY_ID")
+        self.aws_secret_access_key = self.aws_secret_access_key or getenv("AWS_SECRET_ACCESS_KEY")
+        self.aws_region = self.aws_region or getenv("AWS_REGION")
+
+        if self.aws_sso_auth:
+            self.client = AwsClient(service_name="bedrock-runtime", region_name=self.aws_region)
+        else:
+            if not self.aws_access_key_id or not self.aws_secret_access_key:
+                log_error(
+                    "AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables or provide a boto3 session."
+                )
 
-        # Priority 2: Get aws_region from env
-        from os import getenv
+            self.client = AwsClient(
+                service_name="bedrock-runtime",
+                region_name=self.aws_region,
+                aws_access_key_id=self.aws_access_key_id,
+                aws_secret_access_key=self.aws_secret_access_key,
+            )
+        return self.client
 
-        from agno.constants import AWS_REGION_ENV_VAR
+    def get_async_client(self):
+        """
+        Get the async Bedrock client context manager.
 
-        aws_region_env = getenv(AWS_REGION_ENV_VAR)
-        if aws_region_env is not None:
-            self.aws_region = aws_region_env
-        return self.aws_region
+        Returns:
+            The async Bedrock client context manager.
+        """
+        if not AIOBOTO3_AVAILABLE:
+            raise ImportError(
+                "`aioboto3` not installed. Please install using `pip install aioboto3` for async support."
+            )
 
-    def get_aws_profile(self) -> Optional[str]:
-        # Priority 1: Use aws_region from resource
-        if self.aws_profile is not None:
-            return self.aws_profile
+        if self.async_session is None:
+            self.aws_access_key_id = self.aws_access_key_id or getenv("AWS_ACCESS_KEY_ID")
+            self.aws_secret_access_key = self.aws_secret_access_key or getenv("AWS_SECRET_ACCESS_KEY")
+            self.aws_region = self.aws_region or getenv("AWS_REGION")
 
-        # Priority 2: Get aws_profile from env
-        from os import getenv
+            self.async_session = aioboto3.Session()
 
-        from agno.constants import AWS_PROFILE_ENV_VAR
+        client_kwargs = {
+            "service_name": "bedrock-runtime",
+            "region_name": self.aws_region,
+        }
 
-        aws_profile_env = getenv(AWS_PROFILE_ENV_VAR)
-        if aws_profile_env is not None:
-            self.aws_profile = aws_profile_env
-        return self.aws_profile
+        if self.aws_sso_auth:
+            pass
+        else:
+            if not self.aws_access_key_id or not self.aws_secret_access_key:
+                import os
+
+                env_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
+                env_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
+                env_region = os.environ.get("AWS_REGION")
+
+                if env_access_key and env_secret_key:
+                    self.aws_access_key_id = env_access_key
+                    self.aws_secret_access_key = env_secret_key
+                    if env_region:
+                        self.aws_region = env_region
+                        client_kwargs["region_name"] = self.aws_region
+
+            if self.aws_access_key_id and self.aws_secret_access_key:
+                client_kwargs.update(
+                    {
+                        "aws_access_key_id": self.aws_access_key_id,
+                        "aws_secret_access_key": self.aws_secret_access_key,
+                    }
+                )
 
-    def get_aws_client(self) -> AwsApiClient:
-        if self.aws_client is not None:
-            return self.aws_client
+        return self.async_session.client(**client_kwargs)
 
-        self.aws_client = AwsApiClient(aws_region=self.get_aws_region(), aws_profile=self.get_aws_profile())
-        return self.aws_client
+    def _format_tools_for_request(self, tools: Optional[List[Dict[str, Any]]]) -> List[Dict[str, Any]]:
+        """
+        Format the tools for the request.
 
-    @property
-    def bedrock_runtime_client(self):
-        if self._bedrock_runtime_client is not None:
-            return self._bedrock_runtime_client
+        Returns:
+            List[Dict[str, Any]]: The formatted tools.
+        """
+        parsed_tools = []
+        if tools is not None:
+            for tool_def in tools:
+                func_def = tool_def.get("function", {})
+                properties = {}
+                required = []
+
+                for param_name, param_info in func_def.get("parameters", {}).get("properties", {}).items():
+                    properties[param_name] = param_info.copy()
+
+                    if "description" not in properties[param_name]:
+                        properties[param_name]["description"] = ""
+
+                    if "null" not in (
+                        param_info.get("type") if isinstance(param_info.get("type"), list) else [param_info.get("type")]
+                    ):
+                        required.append(param_name)
+
+                parsed_tools.append(
+                    {
+                        "toolSpec": {
+                            "name": func_def.get("name") or "",
+                            "description": func_def.get("description") or "",
+                            "inputSchema": {"json": {"type": "object", "properties": properties, "required": required}},
+                        }
+                    }
+                )
 
-        boto3_session: session = self.get_aws_client().boto3_session
-        self._bedrock_runtime_client = boto3_session.client(service_name="bedrock-runtime")
-        return self._bedrock_runtime_client
+        return parsed_tools
 
-    def invoke(self, body: Dict[str, Any]) -> Dict[str, Any]:
+    def _get_inference_config(self) -> Dict[str, Any]:
         """
-        Invoke the Bedrock API.
-
-        Args:
-            body (Dict[str, Any]): The request body.
+        Get the inference config.
 
         Returns:
-            Dict[str, Any]: The response from the Bedrock API.
+            Dict[str, Any]: The inference config.
         """
-        return self.bedrock_runtime_client.converse(**body)
-
-    def invoke_stream(self, body: Dict[str, Any]) -> Iterator[Dict[str, Any]]:
+        request_kwargs = {
+            "maxTokens": self.max_tokens,
+            "temperature": self.temperature,
+            "topP": self.top_p,
+            "stopSequences": self.stop_sequences,
+        }
+
+        return {k: v for k, v in request_kwargs.items() if v is not None}
+
+    def _format_messages(
+        self, messages: List[Message], compress_tool_results: bool = False
+    ) -> Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]:
         """
-        Invoke the Bedrock API with streaming.
+        Format the messages for the request.
 
         Args:
-            body (Dict[str, Any]): The request body.
+            messages: List of messages to format
+            compress_tool_results: Whether to compress tool results
 
         Returns:
-            Iterator[Dict[str, Any]]: The streamed response.
+            Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]: The formatted messages.
         """
-        response = self.bedrock_runtime_client.converse_stream(**body)
-        stream = response.get("stream")
-        if stream:
-            for event in stream:
-                yield event
-
-    def create_assistant_message(self, request_body: Dict[str, Any]) -> Message:
-        raise NotImplementedError("Please use a subclass of AwsBedrock")
-
-    def get_request_body(self, messages: List[Message]) -> Dict[str, Any]:
-        raise NotImplementedError("Please use a subclass of AwsBedrock")
-
-    def parse_response_message(self, response: Dict[str, Any]) -> Dict[str, Any]:
-        raise NotImplementedError("Please use a subclass of AwsBedrock")
-
-    def _create_tool_calls(
-        self, stop_reason: str, parsed_response: Dict[str, Any]
-    ) -> Tuple[List[str], List[Dict[str, Any]]]:
-        tool_ids: List[str] = []
-        tool_calls: List[Dict[str, Any]] = []
-
-        if stop_reason == "tool_use":
-            tool_requests = parsed_response.get("tool_requests")
-            if tool_requests:
-                for tool in tool_requests:
-                    if "toolUse" in tool:
-                        tool_use = tool["toolUse"]
-                        tool_id = tool_use["toolUseId"]
-                        tool_name = tool_use["name"]
-                        tool_args = tool_use["input"]
 
-                        tool_ids.append(tool_id)
-                        tool_calls.append(
+        formatted_messages: List[Dict[str, Any]] = []
+        system_message = None
+        for message in messages:
+            if message.role == "system":
+                system_message = [{"text": message.content}]
+            elif message.role == "tool":
+                content = message.get_content(use_compressed_content=compress_tool_results)
+                tool_result = {
+                    "toolUseId": message.tool_call_id,
+                    "content": [{"json": {"result": content}}],
+                }
+                formatted_message: Dict[str, Any] = {"role": "user", "content": [{"toolResult": tool_result}]}
+                formatted_messages.append(formatted_message)
+            else:
+                formatted_message = {"role": message.role, "content": []}
+                if isinstance(message.content, list):
+                    formatted_message["content"].extend(message.content)
+                elif message.tool_calls:
+                    tool_use_content = []
+                    for tool_call in message.tool_calls:
+                        try:
+                            # Parse arguments with error handling for empty or invalid JSON
+                            arguments = tool_call["function"]["arguments"]
+                            if not arguments or arguments.strip() == "":
+                                tool_input = {}
+                            else:
+                                tool_input = json.loads(arguments)
+                        except (json.JSONDecodeError, KeyError) as e:
+                            log_warning(f"Failed to parse tool call arguments: {e}")
+                            tool_input = {}
+
+                        tool_use_content.append(
                             {
-                                "type": "function",
-                                "function": {
-                                    "name": tool_name,
-                                    "arguments": json.dumps(tool_args),
-                                },
+                                "toolUse": {
+                                    "toolUseId": tool_call["id"],
+                                    "name": tool_call["function"]["name"],
+                                    "input": tool_input,
+                                }
+                            }
+                        )
+                    formatted_message["content"].extend(tool_use_content)
+                else:
+                    formatted_message["content"].append({"text": message.content})
+
+                if message.images:
+                    for image in message.images:
+                        if not image.content:
+                            raise ValueError("Image content is required for AWS Bedrock.")
+                        if not image.format:
+                            raise ValueError("Image format is required for AWS Bedrock.")
+
+                        if image.format not in BEDROCK_SUPPORTED_IMAGE_FORMATS:
+                            raise ValueError(
+                                f"Unsupported image format: {image.format}. "
+                                f"Supported formats: {BEDROCK_SUPPORTED_IMAGE_FORMATS}"
+                            )
+
+                        formatted_message["content"].append(
+                            {
+                                "image": {
+                                    "format": image.format,
+                                    "source": {
+                                        "bytes": image.content,
+                                    },
+                                }
+                            }
+                        )
+                if message.audio:
+                    log_warning("Audio input is currently unsupported.")
+
+                if message.videos:
+                    for video in message.videos:
+                        if not video.content:
+                            raise ValueError("Video content is required for AWS Bedrock.")
+                        if not video.format:
+                            raise ValueError("Video format is required for AWS Bedrock.")
+
+                        if video.format not in BEDROCK_SUPPORTED_VIDEO_FORMATS:
+                            raise ValueError(
+                                f"Unsupported video format: {video.format}. "
+                                f"Supported formats: {BEDROCK_SUPPORTED_VIDEO_FORMATS}"
+                            )
+
+                        formatted_message["content"].append(
+                            {
+                                "video": {
+                                    "format": video.format,
+                                    "source": {
+                                        "bytes": video.content,
+                                    },
+                                }
                             }
                         )
 
-        return tool_ids, tool_calls
+                if message.files:
+                    for file in message.files:
+                        if not file.content:
+                            raise ValueError("File content is required for AWS Bedrock document input.")
+                        if not file.format:
+                            raise ValueError("File format is required for AWS Bedrock document input.")
+                        if not file.name:
+                            raise ValueError("File name is required for AWS Bedrock document input.")
+
+                        if file.format not in BEDROCK_SUPPORTED_FILE_FORMATS:
+                            raise ValueError(
+                                f"Unsupported file format: {file.format}. "
+                                f"Supported formats: {BEDROCK_SUPPORTED_FILE_FORMATS}"
+                            )
+
+                        formatted_message["content"].append(
+                            {
+                                "document": {
+                                    "format": file.format,
+                                    "name": file.name,
+                                    "source": {
+                                        "bytes": file.content,
+                                    },
+                                }
+                            }
+                        )
 
-    def _handle_tool_calls(
-        self, assistant_message: Message, messages: List[Message], model_response: ModelResponse, tool_ids
-    ) -> Optional[ModelResponse]:
+                formatted_messages.append(formatted_message)
+        # TODO: Add caching: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-call.html
+        return formatted_messages, system_message
+
+    def count_tokens(
+        self,
+        messages: List[Message],
+        tools: Optional[List[Dict[str, Any]]] = None,
+        output_schema: Optional[Union[Dict, Type[BaseModel]]] = None,
+    ) -> int:
+        try:
+            formatted_messages, system_message = self._format_messages(messages, compress_tool_results=True)
+            converse_input: Dict[str, Any] = {"messages": formatted_messages}
+            if system_message:
+                converse_input["system"] = system_message
+
+            response = self.get_client().count_tokens(modelId=self.id, input={"converse": converse_input})
+            tokens = response.get("inputTokens", 0)
+
+            # Count tool tokens
+            if tools:
+                from agno.utils.tokens import count_tool_tokens
+
+                tokens += count_tool_tokens(tools, self.id)
+
+            # Count schema tokens
+            tokens += count_schema_tokens(output_schema, self.id)
+
+            return tokens
+        except Exception as e:
+            log_warning(f"Failed to count tokens via Bedrock API: {e}")
+            return super().count_tokens(messages, tools, output_schema)
+
+    async def acount_tokens(
+        self,
+        messages: List[Message],
+        tools: Optional[List[Dict[str, Any]]] = None,
+        output_schema: Optional[Union[Dict, Type[BaseModel]]] = None,
+    ) -> int:
+        try:
+            formatted_messages, system_message = self._format_messages(messages, compress_tool_results=True)
+            converse_input: Dict[str, Any] = {"messages": formatted_messages}
+            if system_message:
+                converse_input["system"] = system_message
+
+            async with self.get_async_client() as client:
+                response = await client.count_tokens(modelId=self.id, input={"converse": converse_input})
+            tokens = response.get("inputTokens", 0)
+
+            # Count tool tokens
+            if tools:
+                from agno.utils.tokens import count_tool_tokens
+
+                tokens += count_tool_tokens(tools, self.id)
+
+            # Count schema tokens
+            tokens += count_schema_tokens(output_schema, self.id)
+
+            return tokens
+        except Exception as e:
+            log_warning(f"Failed to count tokens via Bedrock API: {e}")
+            return await super().acount_tokens(messages, tools, output_schema)
+
+    def invoke(
+        self,
+        messages: List[Message],
+        assistant_message: Message,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        run_response: Optional[RunOutput] = None,
+        compress_tool_results: bool = False,
+    ) -> ModelResponse:
+        """
+        Invoke the Bedrock API.
         """
-        Handle tool calls in the assistant message.
+        try:
+            formatted_messages, system_message = self._format_messages(messages, compress_tool_results)
 
-        Args:
-            assistant_message (Message): The assistant message.
-            messages (List[Message]): The list of messages.
-            model_response (ModelResponse): The model response.
+            tool_config = None
+            if tools:
+                tool_config = {"tools": self._format_tools_for_request(tools)}
 
-        Returns:
-            Optional[ModelResponse]: The model response after handling tool calls.
-        """
-        # -*- Parse and run function call
-        if assistant_message.tool_calls is not None:
-            if model_response.tool_calls is None:
-                model_response.tool_calls = []
+            body = {
+                "system": system_message,
+                "toolConfig": tool_config,
+                "inferenceConfig": self._get_inference_config(),
+            }
+            body = {k: v for k, v in body.items() if v is not None}
 
-            # Remove the tool call from the response content
-            model_response.content = ""
-            tool_role: str = "tool"
-            function_calls_to_run: List[Any] = []
-            function_call_results: List[Message] = []
-            for tool_call in assistant_message.tool_calls:
-                _tool_call_id = tool_call.get("id")
-                _function_call = get_function_call_for_tool_call(tool_call, self._functions)
-                if _function_call is None:
-                    messages.append(
-                        Message(
-                            role="tool",
-                            tool_call_id=_tool_call_id,
-                            content="Could not find function to call.",
-                        )
-                    )
-                    continue
-                if _function_call.error is not None:
-                    messages.append(
-                        Message(
-                            role="tool",
-                            tool_call_id=_tool_call_id,
-                            content=_function_call.error,
-                        )
-                    )
-                    continue
-                function_calls_to_run.append(_function_call)
-
-            if self.show_tool_calls:
-                model_response.content += "\nRunning:"
-                for _f in function_calls_to_run:
-                    model_response.content += f"\n - {_f.get_call_str()}"
-                model_response.content += "\n\n"
-
-            for function_call_response in self.run_function_calls(
-                function_calls=function_calls_to_run, function_call_results=function_call_results, tool_role=tool_role
-            ):
-                if (
-                    function_call_response.event == ModelResponseEvent.tool_call_completed.value
-                    and function_call_response.tool_calls is not None
-                ):
-                    model_response.tool_calls.extend(function_call_response.tool_calls)
-
-            if len(function_call_results) > 0:
-                fc_responses: List = []
-
-                for _fc_message_index, _fc_message in enumerate(function_call_results):
-                    tool_result = {
-                        "toolUseId": tool_ids[_fc_message_index],
-                        "content": [{"json": json.dumps(_fc_message.content)}],
-                    }
-                    tool_result_message = {"role": "user", "content": json.dumps([{"toolResult": tool_result}])}
-                    fc_responses.append(tool_result_message)
+            if self.request_params:
+                log_debug(f"Calling {self.provider} with request parameters: {self.request_params}", log_level=2)
+                body.update(**self.request_params)
 
-                logger.debug(f"Tool call responses: {fc_responses}")
-                messages.append(Message(role="user", content=json.dumps(fc_responses)))
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+            response = self.get_client().converse(modelId=self.id, messages=formatted_messages, **body)
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(response, response_format=response_format)
 
             return model_response
-        return None
 
-    def _update_metrics(self, assistant_message, parsed_response, response_timer):
+        except ClientError as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e.response), model_name=self.name, model_id=self.id) from e
+        except Exception as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
+
+    def invoke_stream(
+        self,
+        messages: List[Message],
+        assistant_message: Message,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        run_response: Optional[RunOutput] = None,
+        compress_tool_results: bool = False,
+    ) -> Iterator[ModelResponse]:
         """
-        Update usage metrics in assistant_message and self.metrics based on the parsed_response.
-
-        Args:
-            assistant_message: The assistant's message object where individual metrics are stored.
-            parsed_response: The parsed response containing usage metrics.
-            response_timer: Timer object that has the elapsed time of the response.
+        Invoke the Bedrock API with streaming.
         """
-        # Add response time to metrics
-        assistant_message.metrics["time"] = response_timer.elapsed
-        if "response_times" not in self.metrics:
-            self.metrics["response_times"] = []
-        self.metrics["response_times"].append(response_timer.elapsed)
-
-        # Add token usage to metrics
-        usage = parsed_response.get("usage", {})
-        prompt_tokens = usage.get("inputTokens")
-        completion_tokens = usage.get("outputTokens")
-        total_tokens = usage.get("totalTokens")
-
-        if prompt_tokens is not None:
-            assistant_message.metrics["prompt_tokens"] = prompt_tokens
-            self.metrics["prompt_tokens"] = self.metrics.get("prompt_tokens", 0) + prompt_tokens
-
-        if completion_tokens is not None:
-            assistant_message.metrics["completion_tokens"] = completion_tokens
-            self.metrics["completion_tokens"] = self.metrics.get("completion_tokens", 0) + completion_tokens
-
-        if total_tokens is not None:
-            assistant_message.metrics["total_tokens"] = total_tokens
-            self.metrics["total_tokens"] = self.metrics.get("total_tokens", 0) + total_tokens
-
-    def response(self, messages: List[Message]) -> ModelResponse:
+        try:
+            formatted_messages, system_message = self._format_messages(messages, compress_tool_results)
+
+            tool_config = None
+            if tools:
+                tool_config = {"tools": self._format_tools_for_request(tools)}
+
+            body = {
+                "system": system_message,
+                "toolConfig": tool_config,
+                "inferenceConfig": self._get_inference_config(),
+            }
+            body = {k: v for k, v in body.items() if v is not None}
+
+            if self.request_params:
+                body.update(**self.request_params)
+
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+
+            # Track current tool being built across chunks
+            current_tool: Dict[str, Any] = {}
+
+            for chunk in self.get_client().converse_stream(modelId=self.id, messages=formatted_messages, **body)[
+                "stream"
+            ]:
+                model_response, current_tool = self._parse_provider_response_delta(chunk, current_tool)
+                yield model_response
+
+            assistant_message.metrics.stop_timer()
+
+        except ClientError as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e.response), model_name=self.name, model_id=self.id) from e
+        except Exception as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
+
+    async def ainvoke(
+        self,
+        messages: List[Message],
+        assistant_message: Message,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        run_response: Optional[RunOutput] = None,
+        compress_tool_results: bool = False,
+    ) -> ModelResponse:
         """
-        Generate a response from the Bedrock API.
-
-        Args:
-            messages (List[Message]): The messages to include in the request.
-
-        Returns:
-            ModelResponse: The response from the Bedrock API.
+        Async invoke the Bedrock API.
         """
-        logger.debug("---------- Bedrock Response Start ----------")
-        self._log_messages(messages)
-        model_response = ModelResponse()
+        try:
+            formatted_messages, system_message = self._format_messages(messages, compress_tool_results)
 
-        # Invoke the Bedrock API
-        response_timer = Timer()
-        response_timer.start()
-        body = self.get_request_body(messages)
-        response: Dict[str, Any] = self.invoke(body=body)
-        response_timer.stop()
-
-        # Parse response
-        parsed_response = self.parse_response_message(response)
-        logger.debug(f"Parsed response: {parsed_response}")
-        stop_reason = parsed_response["stop_reason"]
-
-        # Create assistant message
-        assistant_message = self.create_assistant_message(parsed_response)
-
-        # Update usage metrics using the new function
-        self._update_metrics(assistant_message, parsed_response, response_timer)
-
-        # Add assistant message to messages
-        messages.append(assistant_message)
-        assistant_message.log()
-
-        # Create tool calls if needed
-        tool_ids, tool_calls = self._create_tool_calls(stop_reason, parsed_response)
-
-        # Handle tool calls
-        if stop_reason == "tool_use" and tool_calls:
-            assistant_message.content = parsed_response["tool_requests"][0]["text"]
-            assistant_message.tool_calls = tool_calls
-
-        # Run tool calls
-        if self._handle_tool_calls(assistant_message, messages, model_response, tool_ids):
-            response_after_tool_calls = self.response(messages=messages)
-            if response_after_tool_calls.content is not None:
-                if model_response.content is None:
-                    model_response.content = ""
-                model_response.content += response_after_tool_calls.content
-            return model_response
+            tool_config = None
+            if tools:
+                tool_config = {"tools": self._format_tools_for_request(tools)}
 
-        # Add assistant message content to model response
-        if assistant_message.content is not None:
-            model_response.content = assistant_message.get_content_string()
+            body = {
+                "system": system_message,
+                "toolConfig": tool_config,
+                "inferenceConfig": self._get_inference_config(),
+            }
+            body = {k: v for k, v in body.items() if v is not None}
 
-        logger.debug("---------- AWS Response End ----------")
-        return model_response
+            if self.request_params:
+                log_debug(f"Calling {self.provider} with request parameters: {self.request_params}", log_level=2)
+                body.update(**self.request_params)
 
-    def _handle_stream_tool_calls(self, assistant_message: Message, messages: List[Message], tool_ids: List[str]):
-        """
-        Handle tool calls in the assistant message.
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
 
-        Args:
-            assistant_message (Message): The assistant message.
-            messages (List[Message]): The list of messages.
-            tool_ids (List[str]): The list of tool IDs.
-        """
-        tool_role: str = "tool"
-        function_calls_to_run: List[Any] = []
-        function_call_results: List[Message] = []
-        for tool_call in assistant_message.tool_calls or []:
-            _tool_call_id = tool_call.get("id")
-            _function_call = get_function_call_for_tool_call(tool_call, self._functions)
-            if _function_call is None:
-                messages.append(
-                    Message(
-                        role="tool",
-                        tool_call_id=_tool_call_id,
-                        content="Could not find function to call.",
-                    )
-                )
-                continue
-            if _function_call.error is not None:
-                messages.append(
-                    Message(
-                        role="tool",
-                        tool_call_id=_tool_call_id,
-                        content=_function_call.error,
-                    )
-                )
-                continue
-            function_calls_to_run.append(_function_call)
-
-        if self.show_tool_calls:
-            yield ModelResponse(content="\nRunning:")
-            for _f in function_calls_to_run:
-                yield ModelResponse(content=f"\n - {_f.get_call_str()}")
-            yield ModelResponse(content="\n\n")
-
-        for _ in self.run_function_calls(
-            function_calls=function_calls_to_run, function_call_results=function_call_results, tool_role=tool_role
-        ):
-            pass
+            assistant_message.metrics.start_timer()
 
-        if len(function_call_results) > 0:
-            fc_responses: List = []
+            async with self.get_async_client() as client:
+                response = await client.converse(modelId=self.id, messages=formatted_messages, **body)
 
-            for _fc_message_index, _fc_message in enumerate(function_call_results):
-                tool_result = {
-                    "toolUseId": tool_ids[_fc_message_index],
-                    "content": [{"json": json.dumps(_fc_message.content)}],
-                }
-                tool_result_message = {"role": "user", "content": json.dumps([{"toolResult": tool_result}])}
-                fc_responses.append(tool_result_message)
+            assistant_message.metrics.stop_timer()
 
-            logger.debug(f"Tool call responses: {fc_responses}")
-            messages.append(Message(role="user", content=json.dumps(fc_responses)))
+            model_response = self._parse_provider_response(response, response_format=response_format)
 
-    def _update_stream_metrics(self, stream_data: StreamData, assistant_message: Message):
+            return model_response
+
+        except ClientError as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e.response), model_name=self.name, model_id=self.id) from e
+        except Exception as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
+
+    async def ainvoke_stream(
+        self,
+        messages: List[Message],
+        assistant_message: Message,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        run_response: Optional[RunOutput] = None,
+        compress_tool_results: bool = False,
+    ) -> AsyncIterator[ModelResponse]:
+        """
+        Async invoke the Bedrock API with streaming.
         """
-        Update the metrics for the streaming response.
+        try:
+            formatted_messages, system_message = self._format_messages(messages, compress_tool_results)
+
+            tool_config = None
+            if tools:
+                tool_config = {"tools": self._format_tools_for_request(tools)}
+
+            body = {
+                "system": system_message,
+                "toolConfig": tool_config,
+                "inferenceConfig": self._get_inference_config(),
+            }
+            body = {k: v for k, v in body.items() if v is not None}
+
+            if self.request_params:
+                body.update(**self.request_params)
+
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+
+            # Track current tool being built across chunks
+            current_tool: Dict[str, Any] = {}
+
+            async with self.get_async_client() as client:
+                response = await client.converse_stream(modelId=self.id, messages=formatted_messages, **body)
+                async for chunk in response["stream"]:
+                    model_response, current_tool = self._parse_provider_response_delta(chunk, current_tool)
+                    yield model_response
+
+            assistant_message.metrics.stop_timer()
+
+        except ClientError as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e.response), model_name=self.name, model_id=self.id) from e
+        except Exception as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
+
+    # Overwrite the default from the base model
+    def format_function_call_results(
+        self,
+        messages: List[Message],
+        function_call_results: List[Message],
+        compress_tool_results: bool = False,
+        **kwargs,
+    ) -> None:
+        """
+        Handle the results of function calls for Bedrock.
+        Uses compressed_content if compress_tool_results is True.
 
         Args:
-            stream_data (StreamData): The streaming data
-            assistant_message (Message): The assistant message.
+            messages (List[Message]): The list of conversation messages.
+            function_call_results (List[Message]): The results of the function calls.
+            compress_tool_results: Whether to compress tool results.
+            **kwargs: Additional arguments including tool_ids.
         """
-        assistant_message.metrics["time"] = stream_data.response_timer.elapsed
-        if stream_data.time_to_first_token is not None:
-            assistant_message.metrics["time_to_first_token"] = stream_data.time_to_first_token
-
-        if "response_times" not in self.metrics:
-            self.metrics["response_times"] = []
-        self.metrics["response_times"].append(stream_data.response_timer.elapsed)
-        if stream_data.time_to_first_token is not None:
-            if "time_to_first_token" not in self.metrics:
-                self.metrics["time_to_first_token"] = []
-            self.metrics["time_to_first_token"].append(stream_data.time_to_first_token)
-        if stream_data.completion_tokens > 0:
-            if "tokens_per_second" not in self.metrics:
-                self.metrics["tokens_per_second"] = []
-            self.metrics["tokens_per_second"].append(
-                f"{stream_data.completion_tokens / stream_data.response_timer.elapsed:.4f}"
-            )
 
-        assistant_message.metrics["prompt_tokens"] = stream_data.response_prompt_tokens
-        assistant_message.metrics["input_tokens"] = stream_data.response_prompt_tokens
-        self.metrics["prompt_tokens"] = self.metrics.get("prompt_tokens", 0) + stream_data.response_prompt_tokens
-        self.metrics["input_tokens"] = self.metrics.get("input_tokens", 0) + stream_data.response_prompt_tokens
+        if function_call_results:
+            tool_ids = kwargs.get("tool_ids", [])
 
-        assistant_message.metrics["completion_tokens"] = stream_data.response_completion_tokens
-        assistant_message.metrics["output_tokens"] = stream_data.response_completion_tokens
-        self.metrics["completion_tokens"] = (
-            self.metrics.get("completion_tokens", 0) + stream_data.response_completion_tokens
-        )
-        self.metrics["output_tokens"] = self.metrics.get("output_tokens", 0) + stream_data.response_completion_tokens
+            for _fc_message_index, _fc_message in enumerate(function_call_results):
+                # Use tool_call_id from message if tool_ids list is insufficient
+                tool_id = tool_ids[_fc_message_index] if _fc_message_index < len(tool_ids) else _fc_message.tool_call_id
+                if not _fc_message.tool_call_id:
+                    _fc_message.tool_call_id = tool_id
 
-        assistant_message.metrics["total_tokens"] = stream_data.response_total_tokens
-        self.metrics["total_tokens"] = self.metrics.get("total_tokens", 0) + stream_data.response_total_tokens
+                # Append as standard role="tool" message
+                messages.append(_fc_message)
 
-    def response_stream(self, messages: List[Message]) -> Iterator[ModelResponse]:
+    def _parse_provider_response(self, response: Dict[str, Any], **kwargs) -> ModelResponse:
         """
-        Stream the response from the Bedrock API.
+        Parse the provider response.
 
         Args:
-            messages (List[Message]): The messages to include in the request.
+            response (Dict[str, Any]): The response from the provider.
 
         Returns:
-            Iterator[str]: The streamed response.
+            ModelResponse: The parsed response.
         """
-        logger.debug("---------- Bedrock Response Start ----------")
-        self._log_messages(messages)
-
-        stream_data: StreamData = StreamData()
-        stream_data.response_timer.start()
-
-        tool_use: Dict[str, Any] = {}
-        tool_ids: List[str] = []
-        tool_calls: List[Dict[str, Any]] = []
-        stop_reason: Optional[str] = None
-        content: List[Dict[str, Any]] = []
-
-        request_body = self.get_request_body(messages)
-        response = self.invoke_stream(body=request_body)
-
-        # Process the streaming response
-        for chunk in response:
-            if "contentBlockStart" in chunk:
-                tool = chunk["contentBlockStart"]["start"].get("toolUse")
-                if tool:
-                    tool_use["toolUseId"] = tool["toolUseId"]
-                    tool_use["name"] = tool["name"]
-
-            elif "contentBlockDelta" in chunk:
-                delta = chunk["contentBlockDelta"]["delta"]
-                if "toolUse" in delta:
-                    if "input" not in tool_use:
-                        tool_use["input"] = ""
-                    tool_use["input"] += delta["toolUse"]["input"]
-                elif "text" in delta:
-                    stream_data.response_content += delta["text"]
-                    stream_data.completion_tokens += 1
-                    if stream_data.completion_tokens == 1:
-                        stream_data.time_to_first_token = stream_data.response_timer.elapsed
-                        logger.debug(f"Time to first token: {stream_data.time_to_first_token:.4f}s")
-                    yield ModelResponse(content=delta["text"])  # Yield text content as it's received
-
-            elif "contentBlockStop" in chunk:
-                if "input" in tool_use:
-                    # Finish collecting tool use input
-                    try:
-                        tool_use["input"] = json.loads(tool_use["input"])
-                    except json.JSONDecodeError as e:
-                        logger.error(f"Failed to parse tool input as JSON: {e}")
-                        tool_use["input"] = {}
-                    content.append({"toolUse": tool_use})
-                    tool_ids.append(tool_use["toolUseId"])
-                    # Prepare the tool call
-                    tool_call = {
-                        "type": "function",
-                        "function": {
-                            "name": tool_use["name"],
-                            "arguments": json.dumps(tool_use["input"]),
-                        },
-                    }
-                    tool_calls.append(tool_call)
-                    tool_use = {}
-                else:
-                    # Finish collecting text content
-                    content.append({"text": stream_data.response_content})
+        model_response = ModelResponse()
 
-            elif "messageStop" in chunk:
-                stop_reason = chunk["messageStop"]["stopReason"]
-                logger.debug(f"Stop reason: {stop_reason}")
+        if "output" in response and "message" in response["output"]:
+            message = response["output"]["message"]
+            # Set the role of the message
+            model_response.role = message["role"]
 
-            elif "metadata" in chunk:
-                metadata = chunk["metadata"]
-                if "usage" in metadata:
-                    stream_data.response_prompt_tokens = metadata["usage"]["inputTokens"]
-                    stream_data.response_total_tokens = metadata["usage"]["totalTokens"]
-                    stream_data.completion_tokens = metadata["usage"]["outputTokens"]
+            # Get the content of the message
+            content = message["content"]
 
-        stream_data.response_timer.stop()
+            # Tools
+            if "stopReason" in response and response["stopReason"] == "tool_use":
+                model_response.tool_calls = []
+                model_response.extra = model_response.extra or {}
+                model_response.extra["tool_ids"] = []
+                for tool in content:
+                    if "toolUse" in tool:
+                        model_response.extra["tool_ids"].append(tool["toolUse"]["toolUseId"])
+                        model_response.tool_calls.append(
+                            {
+                                "id": tool["toolUse"]["toolUseId"],
+                                "type": "function",
+                                "function": {
+                                    "name": tool["toolUse"]["name"],
+                                    "arguments": json.dumps(tool["toolUse"]["input"]),
+                                },
+                            }
+                        )
 
-        # Create assistant message
-        if stream_data.response_content != "":
-            assistant_message = Message(role="assistant", content=stream_data.response_content, tool_calls=tool_calls)
+            # Extract text content if it's a list of dictionaries
+            if isinstance(content, list) and content and isinstance(content[0], dict):
+                content = [item.get("text", "") for item in content if "text" in item]
+                content = "\n".join(content)  # Join multiple text items if present
 
-        if stream_data.completion_tokens > 0:
-            logger.debug(
-                f"Time per output token: {stream_data.response_timer.elapsed / stream_data.completion_tokens:.4f}s"
-            )
-            logger.debug(
-                f"Throughput: {stream_data.completion_tokens / stream_data.response_timer.elapsed:.4f} tokens/s"
-            )
+            model_response.content = content
 
-        # Update metrics
-        self._update_stream_metrics(stream_data, assistant_message)
+        if "usage" in response:
+            model_response.response_usage = self._get_metrics(response["usage"])
 
-        # Add assistant message to messages
-        messages.append(assistant_message)
-        assistant_message.log()
+        return model_response
 
-        # Handle tool calls if any
-        if tool_calls:
-            yield from self._handle_stream_tool_calls(assistant_message, messages, tool_ids)
-            yield from self.response_stream(messages=messages)
+    def _parse_provider_response_delta(
+        self, response_delta: Dict[str, Any], current_tool: Dict[str, Any]
+    ) -> Tuple[ModelResponse, Dict[str, Any]]:
+        """Parse the provider response delta for streaming.
 
-        logger.debug("---------- Bedrock Response End ----------")
+        Args:
+            response_delta: The streaming response delta from AWS Bedrock
+            current_tool: The current tool being built across chunks
 
-    async def ainvoke(self, *args, **kwargs) -> Any:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
+        Returns:
+            Tuple[ModelResponse, Dict[str, Any]]: The parsed model response delta and updated current_tool
+        """
+        model_response = ModelResponse(role="assistant")
+
+        # Handle contentBlockStart - tool use start
+        if "contentBlockStart" in response_delta:
+            start = response_delta["contentBlockStart"]["start"]
+            if "toolUse" in start:
+                # Start a new tool
+                tool_use_data = start["toolUse"]
+                current_tool = {
+                    "id": tool_use_data.get("toolUseId", ""),
+                    "type": "function",
+                    "function": {
+                        "name": tool_use_data.get("name", ""),
+                        "arguments": "",  # Will be filled in subsequent deltas
+                    },
+                }
 
-    async def ainvoke_stream(self, *args, **kwargs) -> Any:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
+        # Handle contentBlockDelta - text content or tool input
+        elif "contentBlockDelta" in response_delta:
+            delta = response_delta["contentBlockDelta"]["delta"]
+            if "text" in delta:
+                model_response.content = delta["text"]
+            elif "toolUse" in delta and current_tool:
+                # Accumulate tool input
+                tool_input = delta["toolUse"].get("input", "")
+                if tool_input:
+                    current_tool["function"]["arguments"] += tool_input
+
+        # Handle contentBlockStop - tool use complete
+        elif "contentBlockStop" in response_delta and current_tool:
+            # Tool is complete, add it to model response
+            model_response.tool_calls = [current_tool]
+            # Track tool_id in extra for format_function_call_results
+            model_response.extra = {"tool_ids": [current_tool["id"]]}
+            # Reset current_tool for next tool
+            current_tool = {}
+
+        # Handle metadata/usage information
+        elif "metadata" in response_delta or "messageStop" in response_delta:
+            body = response_delta.get("metadata") or response_delta.get("messageStop") or {}
+            if "usage" in body:
+                model_response.response_usage = self._get_metrics(body["usage"])
+
+        return model_response, current_tool
+
+    def _get_metrics(self, response_usage: Dict[str, Any]) -> Metrics:
+        """
+        Parse the given AWS Bedrock usage into an Agno Metrics object.
+
+        Args:
+            response_usage: Usage data from AWS Bedrock
+
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
 
-    async def aresponse(self, messages: List[Message]) -> ModelResponse:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
+        metrics.input_tokens = response_usage.get("inputTokens", 0) or 0
+        metrics.output_tokens = response_usage.get("outputTokens", 0) or 0
+        metrics.total_tokens = metrics.input_tokens + metrics.output_tokens
 
-    async def aresponse_stream(self, messages: List[Message]) -> ModelResponse:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
+        return metrics