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

agno/models/mistral/mistral.py

@@ -1,25 +1,41 @@
 from dataclasses import dataclass
 from os import getenv
-from typing import Any, Dict, Iterator, List, Optional, Union
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Type, Union
 
-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.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.log import log_debug, log_error
+from agno.utils.models.mistral import format_messages
 
 try:
+    from mistralai import CompletionEvent
     from mistralai import Mistral as MistralClient
-    from mistralai.models import AssistantMessage, SystemMessage, ToolMessage, UserMessage
-    from mistralai.models.chatcompletionresponse import ChatCompletionResponse
+    from mistralai.extra import response_format_from_pydantic_model
+    from mistralai.extra.struct_chat import ParsedChatCompletionResponse
+    from mistralai.models import (
+        AssistantMessage,
+        HTTPValidationError,
+        SDKError,
+        SystemMessage,
+        ToolMessage,
+        UserMessage,
+    )
+    from mistralai.models.chatcompletionresponse import (
+        ChatCompletionResponse,
+    )
     from mistralai.models.deltamessage import DeltaMessage
     from mistralai.types.basemodel import Unset
-except (ModuleNotFoundError, ImportError):
-    raise ImportError("`mistralai` not installed. Please install using `pip install mistralai`")
 
-MistralMessage = Union[UserMessage, AssistantMessage, SystemMessage, ToolMessage]
+    MistralMessage = Union[UserMessage, AssistantMessage, SystemMessage, ToolMessage]
+
+except ImportError:
+    raise ImportError("`mistralai` not installed. Please install using `pip install mistralai`")
 
 
 @dataclass
@@ -27,31 +43,15 @@ class MistralChat(Model):
     """
     MistralChat is a model that uses the Mistral API to generate responses to messages.
 
-    Args:
-        id (str): The ID of the model.
-        name (str): The name of the model.
-        provider (str): The provider of the model.
-        temperature (Optional[float]): The temperature of the model.
-        max_tokens (Optional[int]): The maximum number of tokens to generate.
-        top_p (Optional[float]): The top p of the model.
-        random_seed (Optional[int]): The random seed of the model.
-        safe_mode (bool): The safe mode of the model.
-        safe_prompt (bool): The safe prompt of the model.
-        response_format (Optional[Union[Dict[str, Any], ChatCompletionResponse]]): The response format of the model.
-        request_params (Optional[Dict[str, Any]]): The request parameters of the model.
-        api_key (Optional[str]): The API key of the model.
-        endpoint (Optional[str]): The endpoint of the model.
-        max_retries (Optional[int]): The maximum number of retries of the model.
-        timeout (Optional[int]): The timeout of the model.
-        client_params (Optional[Dict[str, Any]]): The client parameters of the model.
-        mistral_client (Optional[Mistral]): The Mistral client of the model.
-
+    For more information, see the Mistral API documentation: https://docs.mistral.ai/capabilities/completion/
     """
 
     id: str = "mistral-large-latest"
     name: str = "MistralChat"
     provider: str = "Mistral"
 
+    supports_native_structured_outputs: bool = True
+
     # -*- Request parameters
     temperature: Optional[float] = None
     max_tokens: Optional[int] = None
@@ -59,7 +59,6 @@ class MistralChat(Model):
     random_seed: Optional[int] = None
     safe_mode: bool = False
     safe_prompt: bool = False
-    response_format: Optional[Union[Dict[str, Any], ChatCompletionResponse]] = None
     request_params: Optional[Dict[str, Any]] = None
     # -*- Client parameters
     api_key: Optional[str] = None
@@ -70,36 +69,51 @@ class MistralChat(Model):
     # -*- Provide the Mistral Client manually
     mistral_client: Optional[MistralClient] = None
 
-    @property
-    def client(self) -> MistralClient:
+    def get_client(self) -> MistralClient:
         """
         Get the Mistral client.
 
         Returns:
-            MistralChat: The Mistral client.
+            MistralClient: The Mistral client instance.
         """
         if self.mistral_client:
             return self.mistral_client
 
+        _client_params = self._get_client_params()
+        self.mistral_client = MistralClient(**_client_params)
+        return self.mistral_client
+
+    def _get_client_params(self) -> Dict[str, Any]:
+        """
+        Get the client parameters for initializing Mistral clients.
+
+        Returns:
+            Dict[str, Any]: The client parameters.
+        """
+        client_params: Dict[str, Any] = {}
+
         self.api_key = self.api_key or getenv("MISTRAL_API_KEY")
         if not self.api_key:
-            logger.error("MISTRAL_API_KEY not set. Please set the MISTRAL_API_KEY environment variable.")
-
-        _client_params: Dict[str, Any] = {}
-        if self.api_key:
-            _client_params["api_key"] = self.api_key
-        if self.endpoint:
-            _client_params["endpoint"] = self.endpoint
-        if self.max_retries:
-            _client_params["max_retries"] = self.max_retries
-        if self.timeout:
-            _client_params["timeout"] = self.timeout
-        if self.client_params:
-            _client_params.update(self.client_params)
-        return MistralClient(**_client_params)
-
-    @property
-    def api_kwargs(self) -> Dict[str, Any]:
+            log_error("MISTRAL_API_KEY not set. Please set the MISTRAL_API_KEY environment variable.")
+
+        client_params.update(
+            {
+                "api_key": self.api_key,
+                "endpoint": self.endpoint,
+                "max_retries": self.max_retries,
+                "timeout": self.timeout,
+            }
+        )
+
+        if self.client_params is not None:
+            client_params.update(self.client_params)
+
+        # Remove None values
+        return {k: v for k, v in client_params.items() if v is not None}
+
+    def get_request_params(
+        self, tools: Optional[List[Dict[str, Any]]] = None, tool_choice: Optional[Union[str, Dict[str, Any]]] = None
+    ) -> Dict[str, Any]:
         """
         Get the API kwargs for the Mistral model.
 
@@ -119,14 +133,17 @@ class MistralChat(Model):
             _request_params["safe_mode"] = self.safe_mode
         if self.safe_prompt:
             _request_params["safe_prompt"] = self.safe_prompt
-        if self.tools:
-            _request_params["tools"] = self.tools
-            if self.tool_choice is None:
+        if tools:
+            _request_params["tools"] = tools
+            if tool_choice is None:
                 _request_params["tool_choice"] = "auto"
             else:
-                _request_params["tool_choice"] = self.tool_choice
+                _request_params["tool_choice"] = tool_choice
         if self.request_params:
             _request_params.update(self.request_params)
+
+        if _request_params:
+            log_debug(f"Calling {self.provider} with request parameters: {_request_params}", log_level=2)
         return _request_params
 
     def to_dict(self) -> Dict[str, Any]:
@@ -144,393 +161,276 @@ class MistralChat(Model):
                 "random_seed": self.random_seed,
                 "safe_mode": self.safe_mode,
                 "safe_prompt": self.safe_prompt,
-                "response_format": self.response_format,
             }
         )
         cleaned_dict = {k: v for k, v in _dict.items() if v is not None}
         return cleaned_dict
 
-    def _prepare_messages(self, messages: List[Message]) -> List[MistralMessage]:
-        mistral_messages: List[MistralMessage] = []
-        for m in messages:
-            mistral_message: MistralMessage
-            if m.role == "user":
-                mistral_message = UserMessage(role=m.role, content=m.content)
-            elif m.role == "assistant":
-                if m.tool_calls is not None:
-                    mistral_message = AssistantMessage(role=m.role, content=m.content, tool_calls=m.tool_calls)
-                else:
-                    mistral_message = AssistantMessage(role=m.role, content=m.content)
-            elif m.role == "system":
-                mistral_message = SystemMessage(role=m.role, content=m.content)
-            elif m.role == "tool":
-                mistral_message = ToolMessage(name=m.name, content=m.content, tool_call_id=m.tool_call_id)
-            else:
-                raise ValueError(f"Unknown role: {m.role}")
-            mistral_messages.append(mistral_message)
-        return mistral_messages
-
-    def invoke(self, messages: List[Message]) -> ChatCompletionResponse:
+    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:
         """
         Send a chat completion request to the Mistral model.
+        """
+        mistral_messages = format_messages(messages, compress_tool_results)
+        try:
+            response: Union[ChatCompletionResponse, ParsedChatCompletionResponse]
+            if (
+                response_format is not None
+                and isinstance(response_format, type)
+                and issubclass(response_format, BaseModel)
+            ):
+                if run_response and run_response.metrics:
+                    run_response.metrics.set_time_to_first_token()
 
-        Args:
-            messages (List[Message]): The messages to send to the model.
+                assistant_message.metrics.start_timer()
 
-        Returns:
-            ChatCompletionResponse: The response from the model.
-        """
-        mistral_messages = self._prepare_messages(messages)
-        logger.debug(f"Mistral messages: {mistral_messages}")
-        response = self.client.chat.complete(
-            messages=mistral_messages,
-            model=self.id,
-            **self.api_kwargs,
-        )
-        if response is None:
-            raise ValueError("Chat completion returned None")
-        return response
+                response = self.get_client().chat.complete(
+                    model=self.id,
+                    messages=mistral_messages,
+                    response_format=response_format_from_pydantic_model(response_format),
+                    **self.get_request_params(tools=tools, tool_choice=tool_choice),
+                )
+            else:
+                if run_response and run_response.metrics:
+                    run_response.metrics.set_time_to_first_token()
 
-    def invoke_stream(self, messages: List[Message]) -> Iterator[Any]:
-        """
-        Stream the response from the Mistral model.
+                assistant_message.metrics.start_timer()
+                response = self.get_client().chat.complete(
+                    model=self.id,
+                    messages=mistral_messages,
+                    **self.get_request_params(tools=tools, tool_choice=tool_choice),
+                )
 
-        Args:
-            messages (List[Message]): The messages to send to the model.
+            assistant_message.metrics.stop_timer()
 
-        Returns:
-            Iterator[Any]: The streamed response.
+            model_response = self._parse_provider_response(response, response_format=response_format)
+
+            return model_response
+
+        except HTTPValidationError as e:
+            log_error(f"HTTPValidationError from Mistral: {e}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
+        except SDKError as e:
+            log_error(f"SDKError from Mistral: {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]:
         """
-        mistral_messages = self._prepare_messages(messages)
-        logger.debug(f"Mistral messages sending to stream endpoint: {mistral_messages}")
-        response = self.client.chat.stream(
-            messages=mistral_messages,
-            model=self.id,
-            **self.api_kwargs,
-        )
-        if response is None:
-            raise ValueError("Chat stream returned None")
-        # Since response is a generator, use 'yield from' to yield its items
-        yield from response
-
-    def _handle_tool_calls(
-        self, assistant_message: Message, messages: List[Message], model_response: ModelResponse
-    ) -> Optional[ModelResponse]:
+        Stream the response from the Mistral model.
         """
-        Handle tool calls in the assistant message.
+        mistral_messages = format_messages(messages, compress_tool_results)
 
-        Args:
-            assistant_message (Message): The assistant message.
-            messages (List[Message]): The messages to send to the model.
-            model_response (ModelResponse): The model response.
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
 
-        Returns:
-            Optional[ModelResponse]: The model response after handling tool calls.
+        assistant_message.metrics.start_timer()
+
+        try:
+            for chunk in self.get_client().chat.stream(
+                model=self.id,
+                messages=mistral_messages,
+                **self.get_request_params(tools=tools, tool_choice=tool_choice),
+            ):
+                yield self._parse_provider_response_delta(chunk)
+
+            assistant_message.metrics.stop_timer()
+
+        except HTTPValidationError as e:
+            log_error(f"HTTPValidationError from Mistral: {e}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
+        except SDKError as e:
+            log_error(f"SDKError from Mistral: {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:
         """
-        if assistant_message.tool_calls is not None and len(assistant_message.tool_calls) > 0:
-            if model_response.tool_calls is None:
-                model_response.tool_calls = []
-            model_response.content = ""
-            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["type"] = "function"
-                _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, tool_call_error=True, content=_function_call.error
-                        )
-                    )
-                    continue
-                function_calls_to_run.append(_function_call)
-
-            if self.show_tool_calls:
-                if len(function_calls_to_run) == 1:
-                    model_response.content += f"\n - Running: {function_calls_to_run[0].get_call_str()}\n\n"
-                elif len(function_calls_to_run) > 1:
-                    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
+        Send an asynchronous chat completion request to the Mistral API.
+        """
+        mistral_messages = format_messages(messages, compress_tool_results)
+        try:
+            response: Union[ChatCompletionResponse, ParsedChatCompletionResponse]
+            if (
+                response_format is not None
+                and isinstance(response_format, type)
+                and issubclass(response_format, BaseModel)
             ):
-                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 run_response and run_response.metrics:
+                    run_response.metrics.set_time_to_first_token()
+                assistant_message.metrics.start_timer()
+                response = await self.get_client().chat.complete_async(
+                    model=self.id,
+                    messages=mistral_messages,
+                    response_format=response_format_from_pydantic_model(response_format),
+                    **self.get_request_params(tools=tools, tool_choice=tool_choice),
+                )
+            else:
+                if run_response and run_response.metrics:
+                    run_response.metrics.set_time_to_first_token()
+                assistant_message.metrics.start_timer()
+                response = await self.get_client().chat.complete_async(
+                    model=self.id,
+                    messages=mistral_messages,
+                    **self.get_request_params(tools=tools, tool_choice=tool_choice),
+                )
 
-            if len(function_call_results) > 0:
-                messages.extend(function_call_results)
+            assistant_message.metrics.stop_timer()
 
-            return model_response
-        return None
+            model_response = self._parse_provider_response(response, response_format=response_format)
 
-    def _create_assistant_message(self, response: ChatCompletionResponse) -> Message:
+            return model_response
+        except HTTPValidationError as e:
+            log_error(f"HTTPValidationError from Mistral: {e}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
+        except SDKError as e:
+            log_error(f"SDKError from Mistral: {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]:
         """
-        Create an assistant message from the response.
-
-        Args:
-            response (ChatCompletionResponse): The response from the model.
-
-        Returns:
-            Message: The assistant message.
+        Stream an asynchronous response from the Mistral API.
         """
-        if response.choices is None or len(response.choices) == 0:
-            raise ValueError("The response does not contain any choices.")
+        mistral_messages = format_messages(messages, compress_tool_results)
+        try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
 
-        response_message: AssistantMessage = response.choices[0].message
+            assistant_message.metrics.start_timer()
 
-        # Create assistant message
-        assistant_message = Message(
-            role=response_message.role or "assistant",
-            content=response_message.content,
-        )
+            async for chunk in await self.get_client().chat.stream_async(
+                model=self.id,
+                messages=mistral_messages,
+                **self.get_request_params(tools=tools, tool_choice=tool_choice),
+            ):
+                yield self._parse_provider_response_delta(chunk)
 
-        if isinstance(response_message.tool_calls, list) and len(response_message.tool_calls) > 0:
-            assistant_message.tool_calls = [t.model_dump() for t in response_message.tool_calls]
+            assistant_message.metrics.stop_timer()
 
-        return assistant_message
+        except HTTPValidationError as e:
+            log_error(f"HTTPValidationError from Mistral: {e}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
+        except SDKError as e:
+            log_error(f"SDKError from Mistral: {e}")
+            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
 
-    def _update_usage_metrics(
-        self, assistant_message: Message, response: ChatCompletionResponse, response_timer: Timer
-    ) -> None:
+    def _parse_provider_response(self, response: ChatCompletionResponse, **kwargs) -> ModelResponse:
         """
-        Update the usage metrics for the response.
+        Parse the response from the Mistral model.
 
         Args:
-            assistant_message (Message): The assistant message.
             response (ChatCompletionResponse): The response from the model.
-            response_timer (Timer): The timer for the response.
         """
-        # -*- Update usage metrics
-        # 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
-        self.metrics.update(response.usage.model_dump())
-
-    def response(self, messages: List[Message]) -> ModelResponse:
-        """
-        Send a chat completion request to the Mistral model.
-
-        Args:
-            messages (List[Message]): The messages to send to the model.
-
-        Returns:
-            ModelResponse: The response from the model.
-        """
-        logger.debug("---------- Mistral Response Start ----------")
-        # -*- Log messages for debugging
-        self._log_messages(messages)
         model_response = ModelResponse()
+        if response.choices is not None and len(response.choices) > 0:
+            response_message: AssistantMessage = response.choices[0].message
 
-        response_timer = Timer()
-        response_timer.start()
-        response: ChatCompletionResponse = self.invoke(messages=messages)
-        response_timer.stop()
-        logger.debug(f"Time to generate response: {response_timer.elapsed:.4f}s")
-
-        # -*- Ensure response.choices is not None
-        if response.choices is None or len(response.choices) == 0:
-            raise ValueError("Chat completion response has no choices")
-
-        # -*- Create assistant message
-        assistant_message = self._create_assistant_message(response)
-
-        # -*- Update usage metrics
-        self._update_usage_metrics(assistant_message, response, response_timer)
-
-        # -*- Add assistant message to messages
-        messages.append(assistant_message)
-        assistant_message.log()
+            # -*- Set content
+            model_response.content = response_message.content  # type: ignore
 
-        # -*- Parse and run tool calls
-        logger.debug(f"Functions: {self._functions}")
+            # -*- Set role
+            model_response.role = response_message.role
 
-        # -*- Handle tool calls
-        if self._handle_tool_calls(assistant_message, messages, model_response):
-            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
+            # -*- Set tool calls
+            if isinstance(response_message.tool_calls, list) and len(response_message.tool_calls) > 0:
+                model_response.tool_calls = []
+                for tool_call in response_message.tool_calls:
+                    model_response.tool_calls.append(
+                        {
+                            "id": tool_call.id,
+                            "type": "function",
+                            "function": tool_call.function.model_dump(),
+                        }
+                    )
 
-        # -*- Add content to model response
-        if assistant_message.content is not None:
-            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("---------- Mistral Response End ----------")
         return model_response
 
-    def _update_stream_metrics(self, stream_data: StreamData, assistant_message: Message):
+    def _parse_provider_response_delta(self, response_delta: CompletionEvent) -> ModelResponse:
         """
-        Update the metrics for the streaming response.
-
-        Args:
-            stream_data (StreamData): The streaming data
-            assistant_message (Message): The assistant message.
+        Parse the response delta from the Mistral model.
         """
-        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)
-
-        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
-
-        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
+        model_response = ModelResponse()
+
+        delta_message: DeltaMessage = response_delta.data.choices[0].delta
+        if delta_message.role is not None and not isinstance(delta_message.role, Unset):
+            model_response.role = delta_message.role  # type: ignore
+
+        if (
+            delta_message.content is not None
+            and not isinstance(delta_message.content, Unset)
+            and isinstance(delta_message.content, str)
+        ):
+            model_response.content = delta_message.content
+
+        if delta_message.tool_calls is not None:
+            model_response.tool_calls = []
+            for tool_call in delta_message.tool_calls:
+                model_response.tool_calls.append(
+                    {
+                        "id": tool_call.id,  # type: ignore
+                        "type": "function",
+                        "function": {
+                            "name": tool_call.function.name,  # type: ignore
+                            "arguments": tool_call.function.arguments,  # type: ignore
+                        },
+                    }
+                )
+
+        if response_delta.data.usage is not None:
+            model_response.response_usage = self._get_metrics(response_delta.data.usage)
 
-        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
+        return model_response
 
-    def response_stream(self, messages: List[Message]) -> Iterator[ModelResponse]:
+    def _get_metrics(self, response_usage: Any) -> Metrics:
         """
-        Stream the response from the Mistral model.
+        Parse the given Mistral usage into an Agno Metrics object.
 
         Args:
-            messages (List[Message]): The messages to send to the model.
+            response_usage: Usage data from Mistral
 
         Returns:
-            Iterator[ModelResponse]: The streamed response.
+            Metrics: Parsed metrics data
         """
-        logger.debug("---------- Mistral Response Start ----------")
-        # -*- Log messages for debugging
-        self._log_messages(messages)
-
-        stream_data: StreamData = StreamData()
-        stream_data.response_timer.start()
-
-        assistant_message_role = None
-        for response in self.invoke_stream(messages=messages):
-            # -*- Parse response
-            response_delta: DeltaMessage = response.data.choices[0].delta
-            if assistant_message_role is None and response_delta.role is not None:
-                assistant_message_role = response_delta.role
-
-            response_content: Optional[str] = None
-            if (
-                response_delta.content is not None
-                and not isinstance(response_delta.content, Unset)
-                and isinstance(response_delta.content, str)
-            ):
-                response_content = response_delta.content
-            response_tool_calls = response_delta.tool_calls
-
-            # -*- Return content if present, otherwise get tool call
-            if response_content is not None:
-                stream_data.response_content += response_content
-                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_content)
-
-            # -*- Parse tool calls
-            if response_tool_calls is not None:
-                if stream_data.response_tool_calls is None:
-                    stream_data.response_tool_calls = []
-                stream_data.response_tool_calls.extend(response_tool_calls)
-
-        stream_data.response_timer.stop()
-        completion_tokens = stream_data.completion_tokens
-        if completion_tokens > 0:
-            logger.debug(f"Time per output token: {stream_data.response_timer.elapsed / completion_tokens:.4f}s")
-            logger.debug(f"Throughput: {completion_tokens / stream_data.response_timer.elapsed:.4f} tokens/s")
-
-        # -*- Create assistant message
-        assistant_message = Message(role=(assistant_message_role or "assistant"))
-        if stream_data.response_content != "":
-            assistant_message.content = stream_data.response_content
-
-        # -*- Add tool calls to assistant message
-        if stream_data.response_tool_calls is not None:
-            assistant_message.tool_calls = [t.model_dump() for t in stream_data.response_tool_calls]
-
-        # -*- Update usage metrics
-        self._update_stream_metrics(stream_data, assistant_message)
-        messages.append(assistant_message)
-        assistant_message.log()
-
-        # -*- Parse and run tool calls
-        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")
-                tool_call["type"] = "function"
-                _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, tool_call_error=True, 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"\n - Running: {function_calls_to_run[0].get_call_str()}\n\n")
-                elif len(function_calls_to_run) > 1:
-                    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 intermediate_model_response in self.run_function_calls(
-                function_calls=function_calls_to_run, function_call_results=function_call_results, tool_role=tool_role
-            ):
-                yield intermediate_model_response
-
-            if len(function_call_results) > 0:
-                messages.extend(function_call_results)
-
-            yield from self.response_stream(messages=messages)
-        logger.debug("---------- Mistral Response End ----------")
-
-    async def ainvoke(self, *args, **kwargs) -> Any:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
-
-    async def ainvoke_stream(self, *args, **kwargs) -> Any:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
+        metrics = Metrics()
 
-    async def aresponse(self, messages: List[Message]) -> ModelResponse:
-        raise NotImplementedError(f"Async not supported on {self.name}.")
+        metrics.input_tokens = response_usage.prompt_tokens or 0
+        metrics.output_tokens = response_usage.completion_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