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

agno/models/cohere/chat.py

@@ -1,41 +1,37 @@
-import json
 from dataclasses import dataclass
 from os import getenv
-from typing import Any, Dict, Iterator, List, Optional, Tuple
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, Type, Union
 
-from agno.models.base import Model, StreamData
+import httpx
+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.tools.function import FunctionCall
-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.http import get_default_async_client, get_default_sync_client
+from agno.utils.log import log_debug, log_error, log_warning
+from agno.utils.models.cohere import format_messages
 
 try:
-    from cohere import Client as CohereClient
-    from cohere.types.api_meta import ApiMeta
-    from cohere.types.api_meta_tokens import ApiMetaTokens
-    from cohere.types.non_streamed_chat_response import NonStreamedChatResponse
-    from cohere.types.streamed_chat_response import (
-        StreamedChatResponse,
-        StreamEndStreamedChatResponse,
-        StreamStartStreamedChatResponse,
-        TextGenerationStreamedChatResponse,
-        ToolCallsChunkStreamedChatResponse,
-        ToolCallsGenerationStreamedChatResponse,
-    )
-    from cohere.types.tool import Tool as CohereTool
-    from cohere.types.tool_call import ToolCall
-    from cohere.types.tool_parameter_definitions_value import (
-        ToolParameterDefinitionsValue,
-    )
-    from cohere.types.tool_result import ToolResult
-except (ModuleNotFoundError, ImportError):
+    from cohere import AsyncClientV2 as CohereAsyncClient
+    from cohere import ClientV2 as CohereClient
+    from cohere.v2.types.v2chat_response import V2ChatResponse
+    from cohere.v2.types.v2chat_stream_response import V2ChatStreamResponse
+except ImportError:
     raise ImportError("`cohere` not installed. Please install using `pip install cohere`")
 
 
 @dataclass
 class Cohere(Model):
+    """
+    A class representing the Cohere model.
+
+    For more information, see: https://docs.cohere.com/docs/chat-api
+    """
+
     id: str = "command-r-plus"
     name: str = "cohere"
     provider: str = "Cohere"
@@ -45,565 +41,383 @@ class Cohere(Model):
     max_tokens: Optional[int] = None
     top_k: Optional[int] = None
     top_p: Optional[float] = None
+    seed: Optional[int] = None
     frequency_penalty: Optional[float] = None
     presence_penalty: Optional[float] = None
+    logprobs: Optional[bool] = None
     request_params: Optional[Dict[str, Any]] = None
+    strict_tools: bool = False
     # Add chat history to the cohere messages instead of using the conversation_id
     add_chat_history: bool = False
     # -*- Client parameters
     api_key: Optional[str] = None
     client_params: Optional[Dict[str, Any]] = None
+    http_client: Optional[Union[httpx.Client, httpx.AsyncClient]] = None
     # -*- Provide the Cohere client manually
-    cohere_client: Optional[CohereClient] = None
+    client: Optional[CohereClient] = None
+    async_client: Optional[CohereAsyncClient] = None
 
     def get_client(self) -> CohereClient:
-        if self.cohere_client:
-            return self.cohere_client
+        if self.client:
+            return self.client
 
         _client_params: Dict[str, Any] = {}
 
         self.api_key = self.api_key or getenv("CO_API_KEY")
         if not self.api_key:
-            logger.error("CO_API_KEY not set. Please set the CO_API_KEY environment variable.")
+            log_error("CO_API_KEY not set. Please set the CO_API_KEY environment variable.")
+
+        _client_params["api_key"] = self.api_key
+
+        if self.http_client:
+            if isinstance(self.http_client, httpx.Client):
+                _client_params["httpx_client"] = self.http_client
+            else:
+                log_warning("http_client is not an instance of httpx.Client. Using default global httpx.Client.")
+                # Use global sync client when user http_client is invalid
+                _client_params["httpx_client"] = get_default_sync_client()
+        else:
+            # Use global sync client when no custom http_client is provided
+            _client_params["httpx_client"] = get_default_sync_client()
+
+        self.client = CohereClient(**_client_params)
+        return self.client  # type: ignore
 
-        if self.api_key:
-            _client_params["api_key"] = self.api_key
-        return CohereClient(**_client_params)
+    def get_async_client(self) -> CohereAsyncClient:
+        if self.async_client:
+            return self.async_client
 
-    @property
-    def request_kwargs(self) -> Dict[str, Any]:
+        _client_params: Dict[str, Any] = {}
+
+        self.api_key = self.api_key or getenv("CO_API_KEY")
+
+        if not self.api_key:
+            raise ModelProviderError(
+                message="CO_API_KEY not set. Please set the CO_API_KEY environment variable.",
+                model_name=self.name,
+                model_id=self.id,
+            )
+
+        _client_params["api_key"] = self.api_key
+
+        if self.http_client:
+            if isinstance(self.http_client, httpx.AsyncClient):
+                _client_params["httpx_client"] = self.http_client
+            else:
+                log_warning(
+                    "http_client is not an instance of httpx.AsyncClient. Using default global httpx.AsyncClient."
+                )
+                # Use global async client when user http_client is invalid
+                _client_params["httpx_client"] = get_default_async_client()
+        else:
+            # Use global async client when no custom http_client is provided
+            _client_params["httpx_client"] = get_default_async_client()
+        self.async_client = CohereAsyncClient(**_client_params)
+        return self.async_client  # type: ignore
+
+    def get_request_params(
+        self,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+    ) -> Dict[str, Any]:
         _request_params: Dict[str, Any] = {}
-        if self.session_id is not None and not self.add_chat_history:
-            _request_params["conversation_id"] = self.session_id
         if self.temperature:
             _request_params["temperature"] = self.temperature
         if self.max_tokens:
             _request_params["max_tokens"] = self.max_tokens
         if self.top_k:
-            _request_params["top_k"] = self.top_k
+            _request_params["k"] = self.top_k
         if self.top_p:
-            _request_params["top_p"] = self.top_p
+            _request_params["p"] = self.top_p
+        if self.seed:
+            _request_params["seed"] = self.seed
+        if self.logprobs:
+            _request_params["logprobs"] = self.logprobs
         if self.frequency_penalty:
             _request_params["frequency_penalty"] = self.frequency_penalty
         if self.presence_penalty:
             _request_params["presence_penalty"] = self.presence_penalty
-        if self.request_params:
-            _request_params.update(self.request_params)
-        return _request_params
 
-    def _get_tools(self) -> Optional[List[CohereTool]]:
-        """
-        Get the tools in the format supported by the Cohere API.
-
-        Returns:
-            Optional[List[CohereTool]]: The list of tools.
-        """
-        if not self._functions:
-            return None
-
-        # Returns the tools in the format supported by the Cohere API
-        return [
-            CohereTool(
-                name=f_name,
-                description=function.description or "",
-                parameter_definitions={
-                    param_name: ToolParameterDefinitionsValue(
-                        type=param_info["type"] if isinstance(param_info["type"], str) else param_info["type"][0],
-                        required="null" not in param_info["type"],
-                    )
-                    for param_name, param_info in function.parameters.get("properties", {}).items()
-                },
-            )
-            for f_name, function in self._functions.items()
-        ]
-
-    def _prepare_for_invoke(
-        self, messages: List[Message], tool_results: Optional[List[ToolResult]] = None
-    ) -> Tuple[str, Dict[str, Any]]:
-        api_kwargs: Dict[str, Any] = self.request_kwargs
-        chat_message: str = ""
-
-        if self.add_chat_history:
-            logger.debug("Providing chat_history to cohere")
-            chat_history: List = []
-            for m in messages:
-                if m.role == "system" and "preamble" not in api_kwargs:
-                    api_kwargs["preamble"] = m.content
-                elif m.role == "user":
-                    # Update the chat_message to the new user message
-                    chat_message = m.get_content_string()
-                    chat_history.append({"role": "USER", "message": chat_message})
+        if response_format:
+            _request_params["response_format"] = response_format
+
+        # Format tools
+        if tools is not None and len(tools) > 0:
+            formatted_tools = []
+            for tool in tools:
+                if tool.get("type") == "function" and "function" in tool:
+                    # Extract only the fields that Cohere supports
+                    filtered_tool = {
+                        "type": "function",
+                        "function": {
+                            "name": tool["function"]["name"],
+                            "description": tool["function"]["description"],
+                            "parameters": tool["function"]["parameters"],
+                        },
+                    }
+                    formatted_tools.append(filtered_tool)
                 else:
-                    chat_history.append({"role": "CHATBOT", "message": m.get_content_string() or ""})
-            if chat_history[-1].get("role") == "USER":
-                chat_history.pop()
-            api_kwargs["chat_history"] = chat_history
-        else:
-            # Set first system message as preamble
-            for m in messages:
-                if m.role == "system" and "preamble" not in api_kwargs:
-                    api_kwargs["preamble"] = m.get_content_string()
-                    break
-            # Set last user message as chat_message
-            for m in reversed(messages):
-                if m.role == "user":
-                    chat_message = m.get_content_string()
-                    break
-
-        if self.tools:
-            api_kwargs["tools"] = self._get_tools()
+                    # For non-function tools, pass them through as-is
+                    formatted_tools.append(tool)
+
+            _request_params["tools"] = formatted_tools
+            # Fix optional parameters where the "type" is [type, null]
+            for tool in _request_params["tools"]:  # type: ignore
+                if "parameters" in tool["function"] and "properties" in tool["function"]["parameters"]:  # type: ignore
+                    params = tool["function"]["parameters"]
+                    # Cohere requires at least one required parameter, so if unset, set it to all parameters
+                    if len(params.get("required", [])) == 0:
+                        params["required"] = list(params["properties"].keys())
+            _request_params["strict_tools"] = self.strict_tools
 
-        if tool_results:
-            api_kwargs["tool_results"] = tool_results
+        if self.request_params:
+            _request_params.update(self.request_params)
 
-        return chat_message, api_kwargs
+        if _request_params:
+            log_debug(f"Calling {self.provider} with request parameters: {_request_params}", log_level=2)
+        return _request_params
 
     def invoke(
-        self, messages: List[Message], tool_results: Optional[List[ToolResult]] = None
-    ) -> NonStreamedChatResponse:
+        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 a non-streamed chat response from the Cohere API.
+        """
+        request_kwargs = self.get_request_params(response_format=response_format, tools=tools)
 
-        Args:
-            messages (List[Message]): The list of messages.
-            tool_results (Optional[List[ToolResult]]): The list of tool results.
+        try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
 
-        Returns:
-            NonStreamedChatResponse: The non-streamed chat response.
-        """
-        chat_message, api_kwargs = self._prepare_for_invoke(messages, tool_results)
+            assistant_message.metrics.start_timer()
+            provider_response = self.get_client().chat(
+                model=self.id,
+                messages=format_messages(messages, compress_tool_results),  # type: ignore
+                **request_kwargs,
+            )  # type: ignore
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+
+            return model_response
 
-        return self.get_client().chat(model=self.id, message=chat_message, **api_kwargs)
+        except Exception as e:
+            log_error(f"Unexpected error calling Cohere 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], tool_results: Optional[List[ToolResult]] = None
-    ) -> Iterator[StreamedChatResponse]:
+        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]:
         """
         Invoke a streamed chat response from the Cohere API.
+        """
+        request_kwargs = self.get_request_params(response_format=response_format, tools=tools)
 
-        Args:
-            messages (List[Message]): The list of messages.
-            tool_results (Optional[List[ToolResult]]): The list of tool results.
+        try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
 
-        Returns:
-            Iterator[StreamedChatResponse]: An iterator of streamed chat responses.
-        """
-        chat_message, api_kwargs = self._prepare_for_invoke(messages, tool_results)
+            tool_use: Dict[str, Any] = {}
 
-        return self.get_client().chat_stream(model=self.id, message=chat_message, **api_kwargs)
+            assistant_message.metrics.start_timer()
 
-    def _prepare_function_calls(self, agent_message: Message) -> Tuple[List[FunctionCall], List[Message]]:
-        """
-        Prepares function calls based on tool calls in the agent message.
+            for response in self.get_client().chat_stream(
+                model=self.id,
+                messages=format_messages(messages, compress_tool_results),  # type: ignore
+                **request_kwargs,
+            ):
+                model_response, tool_use = self._parse_provider_response_delta(response, tool_use=tool_use)
+                yield model_response
 
-        This method processes tool calls, matches them with available functions,
-        and prepares them for execution. It also handles errors if functions
-        are not found or if there are issues with the function calls.
+            assistant_message.metrics.stop_timer()
 
-        Args:
-            agent_message (Message): The message containing tool calls to process.
+        except Exception as e:
+            log_error(f"Unexpected error calling Cohere API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
 
-        Returns:
-            Tuple[List[FunctionCall], List[Message]]: A tuple containing a list of
-            prepared function calls and a list of error messages.
+    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:
         """
-        function_calls_to_run: List[FunctionCall] = []
-        error_messages: List[Message] = []
-
-        # Check if tool_calls is None or empty
-        if not agent_message.tool_calls:
-            return function_calls_to_run, error_messages
+        Asynchronously invoke a non-streamed chat response from the Cohere API.
+        """
+        request_kwargs = self.get_request_params(response_format=response_format, tools=tools)
 
-        # Process each tool call in the agent message
-        for tool_call in agent_message.tool_calls:
-            # Attempt to get a function call for the tool call
-            _function_call = get_function_call_for_tool_call(tool_call, self._functions)
+        try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
 
-            # Handle cases where function call cannot be created
-            if _function_call is None:
-                error_messages.append(Message(role="user", content="Could not find function to call."))
-                continue
+            assistant_message.metrics.start_timer()
+            provider_response = await self.get_async_client().chat(
+                model=self.id,
+                messages=format_messages(messages, compress_tool_results),  # type: ignore
+                **request_kwargs,
+            )
+            assistant_message.metrics.stop_timer()
 
-            # Handle cases where function call has an error
-            if _function_call.error is not None:
-                error_messages.append(Message(role="user", content=_function_call.error))
-                continue
+            model_response = self._parse_provider_response(provider_response, response_format=response_format)
 
-            # Add valid function calls to the list
-            function_calls_to_run.append(_function_call)
+            return model_response
 
-        return function_calls_to_run, error_messages
+        except Exception as e:
+            log_error(f"Unexpected error calling Cohere API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
 
-    def _handle_tool_calls(
+    async def ainvoke_stream(
         self,
-        assistant_message: Message,
         messages: List[Message],
-        response_tool_calls: Optional[List[ToolCall]],
-        model_response: ModelResponse,
-    ) -> Optional[Any]:
+        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]:
         """
-        Handle tool calls in the assistant message.
-
-        Args:
-            assistant_message (Message): The assistant message.
-            messages (List[Message]): The list of messages.
-            response_tool_calls (List[Any]): The list of response tool calls.
-            model_response (ModelResponse): The model response.
-
-        Returns:
-            Optional[Any]: The tool results.
+        Asynchronously invoke a streamed chat response from the Cohere API.
         """
+        request_kwargs = self.get_request_params(response_format=response_format, tools=tools)
 
-        model_response.content = ""
-        tool_role: str = "tool"
-        function_calls_to_run: List[FunctionCall] = []
-        function_call_results: List[Message] = []
-        if assistant_message.tool_calls is None:
-            return None
-
-        if model_response.tool_calls is None:
-            model_response.tool_calls = []
-
-        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)
-
-        model_response.content = assistant_message.get_content_string() + "\n\n"
+        try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
 
-        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"
+            tool_use: Dict[str, Any] = {}
 
-        function_calls_to_run, error_messages = self._prepare_function_calls(assistant_message)
+            assistant_message.metrics.start_timer()
 
-        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
+            async for response in self.get_async_client().chat_stream(
+                model=self.id,
+                messages=format_messages(messages, compress_tool_results),  # type: ignore
+                **request_kwargs,
             ):
-                model_response.tool_calls.extend(function_call_response.tool_calls)
+                model_response, tool_use = self._parse_provider_response_delta(response, tool_use=tool_use)
+                yield model_response
 
-        if len(function_call_results) > 0:
-            messages.extend(function_call_results)
-
-        # Prepare tool results for the next API call
-        if response_tool_calls:
-            tool_results = [
-                ToolResult(
-                    call=tool_call,
-                    outputs=[tool_call.parameters, {"result": fn_result.content}],
-                )
-                for tool_call, fn_result in zip(response_tool_calls, function_call_results)
-            ]
-        else:
-            tool_results = None
+            assistant_message.metrics.stop_timer()
 
-        return tool_results
+        except Exception as e:
+            log_error(f"Unexpected error calling Cohere API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
 
-    def _create_assistant_message(self, response: NonStreamedChatResponse) -> Message:
+    def _parse_provider_response(self, response: V2ChatResponse, **kwargs) -> ModelResponse:
         """
-        Create an assistant message from the response.
+        Parse the model provider response.
 
         Args:
-            response (NonStreamedChatResponse): The response from the Cohere API.
-
-        Returns:
-            Message: The assistant message.
-        """
-        response_content = response.text
-        return Message(role="assistant", content=response_content)
-
-    def response(self, messages: List[Message], tool_results: Optional[List[ToolResult]] = None) -> ModelResponse:
-        """
-        Send a chat completion request to the Cohere API.
-
-        Args:
-            messages (List[Message]): A list of message objects representing the conversation.
-
-        Returns:
-            ModelResponse: The model response from the API.
+            response (V2ChatResponse): The response from the Cohere API.
         """
-        logger.debug("---------- Cohere Response Start ----------")
-        self._log_messages(messages)
         model_response = ModelResponse()
 
-        # Timer for response
-        response_timer = Timer()
-        response_timer.start()
-        logger.debug(f"Tool Results: {tool_results}")
-        response: NonStreamedChatResponse = self.invoke(messages=messages, tool_results=tool_results)
-        response_timer.stop()
-        logger.debug(f"Time to generate response: {response_timer.elapsed:.4f}s")
-
-        assistant_message = self._create_assistant_message(response)
-
-        # Process tool calls if present
-        response_tool_calls = response.tool_calls
-        if response_tool_calls:
-            tool_calls = [
-                {
-                    "type": "function",
-                    "function": {
-                        "name": tools.name,
-                        "arguments": json.dumps(tools.parameters),
-                    },
-                }
-                for tools in response_tool_calls
-            ]
-            assistant_message.tool_calls = tool_calls
-
-        # Handle tool calls if present and tool running is enabled
-        if assistant_message.tool_calls:
-            tool_results = self._handle_tool_calls(
-                assistant_message=assistant_message,
-                messages=messages,
-                response_tool_calls=response_tool_calls,
-                model_response=model_response,
-            )
+        model_response.role = response.message.role
 
-            # Make a recursive call with tool results if available
-            if tool_results:
-                # Cohere doesn't allow tool calls in the same message as the user's message, so we add a new user message with empty content
-                messages.append(Message(role="user", content=""))
+        response_message = response.message
+        if response_message.content is not None:
+            full_content = ""
+            for item in response_message.content:
+                if hasattr(item, "text") and item.text is not None:  # type: ignore
+                    full_content += item.text  # type: ignore
+            model_response.content = full_content
 
-            response_after_tool_calls = self.response(messages=messages, tool_results=tool_results)
-            if response_after_tool_calls.content:
-                if model_response.content is None:
-                    model_response.content = ""
-                model_response.content += response_after_tool_calls.content
-            return model_response
+        if response_message.tool_calls is not None:
+            model_response.tool_calls = [t.model_dump() for t in response_message.tool_calls]
 
-        # If no tool calls, return the agent message content
-        if assistant_message.content:
-            model_response.content = assistant_message.get_content_string()
+        if response.usage is not None:
+            model_response.response_usage = self._get_metrics(response.usage)
 
-        logger.debug("---------- Cohere Response End ----------")
         return model_response
 
-    def _update_stream_metrics(self, stream_data: StreamData, assistant_message: Message):
+    def _parse_provider_response_delta(
+        self, response: V2ChatStreamResponse, tool_use: Dict[str, Any]
+    ) -> Tuple[ModelResponse, Dict[str, Any]]:  # type: ignore
         """
-        Update the metrics for the streaming response.
+        Parse the streaming response from the model provider into ModelResponse objects.
 
         Args:
-            stream_data (StreamData): The streaming data
-            assistant_message (Message): The assistant message.
+            response: Raw response chunk from the model provider
+
+        Returns:
+            ModelResponse: Parsed response delta
         """
-        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}"
-            )
+        model_response = ModelResponse()
 
-        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
+        # 1. Add content
+        if (
+            response.type == "content-delta"
+            and response.delta is not None
+            and response.delta.message is not None
+            and response.delta.message.content is not None
+            and response.delta.message.content.text is not None
+        ):
+            model_response.content = response.delta.message.content.text
 
-        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
+        # 2. Add tool calls information
 
-        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
+        # 2.1 Add starting tool call
+        elif response.type == "tool-call-start" and response.delta is not None:
+            if response.delta.message is not None and response.delta.message.tool_calls is not None:
+                tool_use = response.delta.message.tool_calls.model_dump()
 
-    def response_stream(
-        self, messages: List[Message], tool_results: Optional[List[ToolResult]] = None
-    ) -> Iterator[ModelResponse]:
-        logger.debug("---------- Cohere Response Start ----------")
-        # -*- Log messages for debugging
-        self._log_messages(messages)
-
-        stream_data: StreamData = StreamData()
-        stream_data.response_timer.start()
-
-        stream_data.response_content = ""
-        tool_calls: List[Dict[str, Any]] = []
-        stream_data.response_tool_calls = []
-        last_delta: Optional[NonStreamedChatResponse] = None
-
-        for response in self.invoke_stream(messages=messages, tool_results=tool_results):
-            if isinstance(response, StreamStartStreamedChatResponse):
-                pass
-
-            if isinstance(response, TextGenerationStreamedChatResponse):
-                if response.text is not None:
-                    stream_data.response_content += response.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=response.text)
-
-            if isinstance(response, ToolCallsChunkStreamedChatResponse):
-                if response.tool_call_delta is None:
-                    yield ModelResponse(content=response.text)
-
-            # Detect if response is a tool call
-            if isinstance(response, ToolCallsGenerationStreamedChatResponse):
-                for tc in response.tool_calls:
-                    stream_data.response_tool_calls.append(tc)
-                    tool_calls.append(
-                        {
-                            "type": "function",
-                            "function": {
-                                "name": tc.name,
-                                "arguments": json.dumps(tc.parameters),
-                            },
-                        }
-                    )
-
-            if isinstance(response, StreamEndStreamedChatResponse):
-                last_delta = response.response
-
-        yield ModelResponse(content="\n\n")
-
-        stream_data.response_timer.stop()
-        logger.debug(f"Time to generate response: {stream_data.response_timer.elapsed:.4f}s")
-
-        # -*- Create assistant message
-        assistant_message = Message(role="assistant", content=stream_data.response_content)
-        # -*- Add tool calls to assistant message
-        if len(stream_data.response_tool_calls) > 0:
-            assistant_message.tool_calls = tool_calls
-
-        # -*- Update usage metrics
-        # Add response time to metrics
-        assistant_message.metrics["time"] = stream_data.response_timer.elapsed
-        if "response_times" not in self.metrics:
-            self.metrics["response_times"] = []
-        self.metrics["response_times"].append(stream_data.response_timer.elapsed)
-
-        # Add token usage to metrics
-        meta: Optional[ApiMeta] = last_delta.meta if last_delta else None
-        tokens: Optional[ApiMetaTokens] = meta.tokens if meta else None
-
-        if tokens:
-            input_tokens = tokens.input_tokens
-            output_tokens = tokens.output_tokens
-
-            if input_tokens is not None:
-                assistant_message.metrics["input_tokens"] = input_tokens
-                self.metrics["input_tokens"] = self.metrics.get("input_tokens", 0) + input_tokens
-
-            if output_tokens is not None:
-                assistant_message.metrics["output_tokens"] = output_tokens
-                self.metrics["output_tokens"] = self.metrics.get("output_tokens", 0) + output_tokens
-
-            if input_tokens is not None and output_tokens is not None:
-                assistant_message.metrics["total_tokens"] = input_tokens + output_tokens
-                self.metrics["total_tokens"] = self.metrics.get("total_tokens", 0) + input_tokens + output_tokens
-
-        # -*- Add assistant message to messages
-        self._update_stream_metrics(stream_data=stream_data, assistant_message=assistant_message)
-        messages.append(assistant_message)
-        assistant_message.log()
-        logger.debug(f"Assistant Message: {assistant_message}")
-
-        # -*- Parse and run function call
-        if assistant_message.tool_calls is not None and len(assistant_message.tool_calls) > 0:
-            tool_role: str = "tool"
-            function_calls_to_run: List[FunctionCall] = []
-            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_role,
-                            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_role,
-                            tool_call_id=_tool_call_id,
-                            content=_function_call.error,
-                        )
-                    )
-                    continue
-                function_calls_to_run.append(_function_call)
-
-            if self.show_tool_calls:
-                if len(function_calls_to_run) == 1:
-                    yield ModelResponse(content=f"- Running: {function_calls_to_run[0].get_call_str()}\n\n")
-                elif len(function_calls_to_run) > 1:
-                    yield ModelResponse(content="Running:")
-                    for _f in function_calls_to_run:
-                        yield ModelResponse(content=f"\n - {_f.get_call_str()}")
-                    yield ModelResponse(content="\n\n")
-
-            for intermediate_model_response in self.run_function_calls(
-                function_calls=function_calls_to_run, function_call_results=function_call_results, tool_role=tool_role
+        # 2.2 Add tool call delta
+        elif response.type == "tool-call-delta" and response.delta is not None:
+            if (
+                response.delta.message is not None
+                and response.delta.message.tool_calls is not None
+                and response.delta.message.tool_calls.function is not None
+                and response.delta.message.tool_calls.function.arguments is not None
             ):
-                yield intermediate_model_response
-
-            if len(function_call_results) > 0:
-                messages.extend(function_call_results)
+                tool_use["function"]["arguments"] += response.delta.message.tool_calls.function.arguments
+
+        # 2.3 Add ending tool call
+        elif response.type == "tool-call-end":
+            model_response.tool_calls = [tool_use]
+
+        # 3. Add metrics
+        elif (
+            response.type == "message-end"
+            and response.delta is not None
+            and response.delta.usage is not None
+            and response.delta.usage.tokens is not None
+        ):
+            model_response.response_usage = self._get_metrics(response.delta.usage)  # type: ignore
 
-            # Making sure the length of tool calls and function call results are the same to avoid unexpected behavior
-            if stream_data.response_tool_calls is not None:
-                # Constructs a list named tool_results, where each element is a dictionary that contains details of tool calls and their outputs.
-                # It pairs each tool call in response_tool_calls with its corresponding result in function_call_results.
-                tool_results = [
-                    ToolResult(call=tool_call, outputs=[tool_call.parameters, {"result": fn_result.content}])
-                    for tool_call, fn_result in zip(stream_data.response_tool_calls, function_call_results)
-                ]
-                messages.append(Message(role="user", content=""))
+        return model_response, tool_use
 
-            # -*- Yield new response using results of tool calls
-            yield from self.response_stream(messages=messages, tool_results=tool_results)
-        logger.debug("---------- Cohere Response End ----------")
+    def _get_metrics(self, response_usage) -> Metrics:
+        """
+        Parse the given Cohere usage into an Agno Metrics object.
 
-    async def ainvoke(self, *args, **kwargs) -> Any:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
+        Args:
+            response_usage: Usage data from Cohere
 
-    async def ainvoke_stream(self, *args, **kwargs) -> Any:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
+        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.tokens.input_tokens or 0
+        metrics.output_tokens = response_usage.tokens.output_tokens 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