agno 1.8.0__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

agno/models/mistral/mistral.py

@@ -1,13 +1,15 @@
 from dataclasses import dataclass
 from os import getenv
-from typing import Any, Dict, Iterator, List, Optional, Type, Union
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Type, Union
 
 from pydantic import BaseModel
 
 from agno.exceptions import ModelProviderError
 from agno.models.base import Model
 from agno.models.message import Message
+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
 
@@ -24,7 +26,9 @@ try:
         ToolMessage,
         UserMessage,
     )
-    from mistralai.models.chatcompletionresponse import ChatCompletionResponse
+    from mistralai.models.chatcompletionresponse import (
+        ChatCompletionResponse,
+    )
     from mistralai.models.deltamessage import DeltaMessage
     from mistralai.types.basemodel import Unset
 
@@ -165,10 +169,12 @@ class MistralChat(Model):
     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,
-    ) -> Union[ChatCompletionResponse, ParsedChatCompletionResponse]:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Send a chat completion request to the Mistral model.
         """
@@ -180,6 +186,11 @@ class MistralChat(Model):
                 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()
+
+                assistant_message.metrics.start_timer()
+
                 response = self.get_client().chat.complete(
                     model=self.id,
                     messages=mistral_messages,
@@ -187,12 +198,21 @@ class MistralChat(Model):
                     **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 = self.get_client().chat.complete(
                     model=self.id,
                     messages=mistral_messages,
                     **self.get_request_params(tools=tools, tool_choice=tool_choice),
                 )
-            return response
+
+            assistant_message.metrics.stop_timer()
+
+            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}")
@@ -204,21 +224,32 @@ class MistralChat(Model):
     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,
-    ) -> Iterator[Any]:
+        run_response: Optional[RunOutput] = None,
+    ) -> Iterator[ModelResponse]:
         """
         Stream the response from the Mistral model.
         """
         mistral_messages = format_messages(messages)
+
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+
         try:
-            stream = self.get_client().chat.stream(
+            for chunk in self.get_client().chat.stream(
                 model=self.id,
                 messages=mistral_messages,
                 **self.get_request_params(tools=tools, tool_choice=tool_choice),
-            )
-            return stream
+            ):
+                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
@@ -229,10 +260,12 @@ class MistralChat(Model):
     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,
-    ) -> Union[ChatCompletionResponse, ParsedChatCompletionResponse]:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Send an asynchronous chat completion request to the Mistral API.
         """
@@ -244,6 +277,9 @@ class MistralChat(Model):
                 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()
+                assistant_message.metrics.start_timer()
                 response = await self.get_client().chat.complete_async(
                     model=self.id,
                     messages=mistral_messages,
@@ -251,12 +287,20 @@ class MistralChat(Model):
                     **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),
                 )
-            return response
+
+            assistant_message.metrics.stop_timer()
+
+            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
@@ -267,24 +311,31 @@ class MistralChat(Model):
     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,
-    ) -> Any:
+        run_response: Optional[RunOutput] = None,
+    ) -> AsyncIterator[ModelResponse]:
         """
         Stream an asynchronous response from the Mistral API.
         """
         mistral_messages = format_messages(messages)
         try:
-            stream = await self.get_client().chat.stream_async(
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+
+            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),
-            )
-            if stream is None:
-                raise ValueError("Chat stream returned None")
-            async for chunk in stream:
-                yield chunk
+            ):
+                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
@@ -292,7 +343,7 @@ class MistralChat(Model):
             log_error(f"SDKError from Mistral: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
 
-    def parse_provider_response(self, response: ChatCompletionResponse, **kwargs) -> ModelResponse:
+    def _parse_provider_response(self, response: ChatCompletionResponse, **kwargs) -> ModelResponse:
         """
         Parse the response from the Mistral model.
 
@@ -322,11 +373,11 @@ class MistralChat(Model):
                     )
 
         if response.usage is not None:
-            model_response.response_usage = response.usage
+            model_response.response_usage = self._get_metrics(response.usage)
 
         return model_response
 
-    def parse_provider_response_delta(self, response_delta: CompletionEvent) -> Optional[ModelResponse]:
+    def _parse_provider_response_delta(self, response_delta: CompletionEvent) -> ModelResponse:
         """
         Parse the response delta from the Mistral model.
         """
@@ -358,6 +409,24 @@ class MistralChat(Model):
                 )
 
         if response_delta.data.usage is not None:
-            model_response.response_usage = response_delta.data.usage
+            model_response.response_usage = self._get_metrics(response_delta.data.usage)
 
         return model_response
+
+    def _get_metrics(self, response_usage: Any) -> Metrics:
+        """
+        Parse the given Mistral usage into an Agno Metrics object.
+
+        Args:
+            response_usage: Usage data from Mistral
+
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
+
+        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
+
+        return metrics

agno/models/ollama/__init__.py

@@ -1,7 +1,5 @@
 from agno.models.ollama.chat import Ollama
-from agno.models.ollama.tools import OllamaTools
 
 __all__ = [
     "Ollama",
-    "OllamaTools",
 ]

agno/models/ollama/chat.py

@@ -1,11 +1,13 @@
 import json
 from dataclasses import dataclass
-from typing import Any, Dict, Iterator, List, Mapping, Optional, Type, Union
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Type, Union
 
 from pydantic import BaseModel
 
+from agno.agent import RunOutput
 from agno.models.base import Model
 from agno.models.message import Message
+from agno.models.metrics import Metrics
 from agno.models.response import ModelResponse
 from agno.utils.log import log_debug, log_warning
 
@@ -182,86 +184,124 @@ class Ollama(Model):
     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,
-    ) -> Mapping[str, Any]:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Send a chat request to the Ollama API.
         """
         request_kwargs = self._prepare_request_kwargs_for_invoke(response_format=response_format, tools=tools)
 
-        return self.get_client().chat(
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+
+        provider_response = self.get_client().chat(
             model=self.id.strip(),
             messages=[self._format_message(m) for m in messages],  # type: ignore
             **request_kwargs,
         )  # type: ignore
 
+        assistant_message.metrics.stop_timer()
+
+        model_response = self._parse_provider_response(provider_response)  # type: ignore
+        return model_response
+
     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,
-    ) -> Mapping[str, Any]:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Sends an asynchronous chat request to the Ollama API.
         """
         request_kwargs = self._prepare_request_kwargs_for_invoke(response_format=response_format, tools=tools)
 
-        return await self.get_async_client().chat(
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+
+        provider_response = await self.get_async_client().chat(
             model=self.id.strip(),
             messages=[self._format_message(m) for m in messages],  # type: ignore
             **request_kwargs,
         )  # type: ignore
 
+        assistant_message.metrics.stop_timer()
+
+        model_response = self._parse_provider_response(provider_response)  # type: ignore
+        return model_response
+
     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,
-    ) -> Iterator[Mapping[str, Any]]:
+        run_response: Optional[RunOutput] = None,
+    ) -> Iterator[ModelResponse]:
         """
         Sends a streaming chat request to the Ollama API.
         """
-        yield from self.get_client().chat(
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+
+        for chunk in self.get_client().chat(
             model=self.id,
             messages=[self._format_message(m) for m in messages],  # type: ignore
             stream=True,
             **self.get_request_params(tools=tools),
-        )  # type: ignore
+        ):
+            yield self._parse_provider_response_delta(chunk)
+
+        assistant_message.metrics.stop_timer()
 
     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,
-    ) -> Any:
+        run_response: Optional[RunOutput] = None,
+    ) -> AsyncIterator[ModelResponse]:
         """
         Sends an asynchronous streaming chat completion request to the Ollama API.
         """
-        async_stream = await self.get_async_client().chat(
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+
+        async for chunk in await self.get_async_client().chat(
             model=self.id.strip(),
             messages=[self._format_message(m) for m in messages],  # type: ignore
             stream=True,
             **self.get_request_params(tools=tools),
-        )
-        async for chunk in async_stream:  # type: ignore
-            yield chunk
+        ):
+            yield self._parse_provider_response_delta(chunk)
 
-    def parse_provider_response(
-        self,
-        response: ChatResponse,
-        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
-    ) -> ModelResponse:
+        assistant_message.metrics.stop_timer()
+
+    def _parse_provider_response(self, response: dict) -> ModelResponse:
         """
         Parse the provider response.
         """
         model_response = ModelResponse()
         # Get response message
-        response_message: OllamaMessage = response.get("message")
+        response_message: OllamaMessage = response.get("message")  # type: ignore
 
         if response_message.get("role") is not None:
             model_response.role = response_message.get("role")
@@ -272,7 +312,7 @@ class Ollama(Model):
         if response_message.get("tool_calls") is not None:
             if model_response.tool_calls is None:
                 model_response.tool_calls = []
-            for block in response_message.get("tool_calls"):
+            for block in response_message.get("tool_calls", []):
                 tool_call = block.get("function")
                 tool_name = tool_call.get("name")
                 tool_args = tool_call.get("arguments")
@@ -288,33 +328,23 @@ class Ollama(Model):
 
         # Get response usage
         if response.get("done"):
-            model_response.response_usage = {
-                "input_tokens": response.get("prompt_eval_count", 0),
-                "output_tokens": response.get("eval_count", 0),
-                "total_tokens": response.get("prompt_eval_count", 0) + response.get("eval_count", 0),
-                "additional_metrics": {
-                    "total_duration": response.get("total_duration", 0),
-                    "load_duration": response.get("load_duration", 0),
-                    "prompt_eval_duration": response.get("prompt_eval_duration", 0),
-                    "eval_duration": response.get("eval_duration", 0),
-                },
-            }
+            model_response.response_usage = self._get_metrics(response)
 
         return model_response
 
-    def parse_provider_response_delta(self, response_delta: ChatResponse) -> ModelResponse:
+    def _parse_provider_response_delta(self, response: ChatResponse) -> ModelResponse:
         """
         Parse the provider response delta.
 
         Args:
-            response_delta (ChatResponse): The response from the provider.
+            response (ChatResponse): The response from the provider.
 
         Returns:
             Iterator[ModelResponse]: An iterator of the model response.
         """
         model_response = ModelResponse()
 
-        response_message = response_delta.get("message")
+        response_message = response.get("message")
 
         if response_message is not None:
             content_delta = response_message.get("content")
@@ -333,17 +363,25 @@ class Ollama(Model):
                     }
                     model_response.tool_calls.append({"type": "function", "function": function_def})
 
-        if response_delta.get("done"):
-            model_response.response_usage = {
-                "input_tokens": response_delta.get("prompt_eval_count", 0),
-                "output_tokens": response_delta.get("eval_count", 0),
-                "total_tokens": response_delta.get("prompt_eval_count", 0) + response_delta.get("eval_count", 0),
-                "additional_metrics": {
-                    "total_duration": response_delta.get("total_duration", 0),
-                    "load_duration": response_delta.get("load_duration", 0),
-                    "prompt_eval_duration": response_delta.get("prompt_eval_duration", 0),
-                    "eval_duration": response_delta.get("eval_duration", 0),
-                },
-            }
+        if response.get("done"):
+            model_response.response_usage = self._get_metrics(response)
 
         return model_response
+
+    def _get_metrics(self, response: Union[dict, ChatResponse]) -> Metrics:
+        """
+        Parse the given Ollama usage into an Agno Metrics object.
+
+        Args:
+            response: The response from the provider.
+
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
+
+        metrics.input_tokens = response.get("prompt_eval_count", 0)
+        metrics.output_tokens = response.get("eval_count", 0)
+        metrics.total_tokens = metrics.input_tokens + metrics.output_tokens
+
+        return metrics