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

agno/models/google/gemini.py

@@ -1,9 +1,10 @@
 import json
 import time
+from collections.abc import AsyncIterator
 from dataclasses import dataclass
 from os import getenv
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import Any, Dict, Iterator, List, Optional, Type, Union
 from uuid import uuid4
 
 from pydantic import BaseModel
@@ -11,8 +12,10 @@ from pydantic import BaseModel
 from agno.exceptions import ModelProviderError
 from agno.media import Audio, File, ImageArtifact, Video
 from agno.models.base import Model
-from agno.models.message import Citations, Message, MessageMetrics, UrlCitation
+from agno.models.message import Citations, Message, UrlCitation
+from agno.models.metrics import Metrics
 from agno.models.response import ModelResponse
+from agno.run.agent import RunOutput
 from agno.utils.gemini import convert_schema, format_function_definitions, format_image_for_message
 from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.utils.models.schema_utils import get_response_schema_for_provider
@@ -262,21 +265,33 @@ class Gemini(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,
-    ):
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Invokes the model with a list of messages and returns the response.
         """
         formatted_messages, system_message = self._format_messages(messages)
         request_kwargs = self.get_request_params(system_message, response_format=response_format, tools=tools)
         try:
-            return self.get_client().models.generate_content(
+            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().models.generate_content(
                 model=self.id,
                 contents=formatted_messages,
                 **request_kwargs,
             )
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+
+            return model_response
+
         except (ClientError, ServerError) as e:
             log_error(f"Error from Gemini API: {e}")
             error_message = str(e.response) if hasattr(e, "response") else str(e)
@@ -293,10 +308,12 @@ class Gemini(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,
-    ):
+        run_response: Optional[RunOutput] = None,
+    ) -> Iterator[ModelResponse]:
         """
         Invokes the model with a list of messages and returns the response as a stream.
         """
@@ -304,11 +321,19 @@ class Gemini(Model):
 
         request_kwargs = self.get_request_params(system_message, response_format=response_format, tools=tools)
         try:
-            yield from self.get_client().models.generate_content_stream(
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+            for response in self.get_client().models.generate_content_stream(
                 model=self.id,
                 contents=formatted_messages,
                 **request_kwargs,
-            )
+            ):
+                yield self._parse_provider_response_delta(response)
+
+            assistant_message.metrics.stop_timer()
+
         except (ClientError, ServerError) as e:
             log_error(f"Error from Gemini API: {e}")
             raise ModelProviderError(
@@ -324,10 +349,12 @@ class Gemini(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,
-    ):
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Invokes the model with a list of messages and returns the response.
         """
@@ -336,11 +363,21 @@ class Gemini(Model):
         request_kwargs = self.get_request_params(system_message, response_format=response_format, tools=tools)
 
         try:
-            return await self.get_client().aio.models.generate_content(
+            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_client().aio.models.generate_content(
                 model=self.id,
                 contents=formatted_messages,
                 **request_kwargs,
             )
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+
+            return model_response
+
         except (ClientError, ServerError) as e:
             log_error(f"Error from Gemini API: {e}")
             raise ModelProviderError(
@@ -356,10 +393,12 @@ class Gemini(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,
-    ):
+        run_response: Optional[RunOutput] = None,
+    ) -> AsyncIterator[ModelResponse]:
         """
         Invokes the model with a list of messages and returns the response as a stream.
         """
@@ -368,13 +407,21 @@ class Gemini(Model):
         request_kwargs = self.get_request_params(system_message, response_format=response_format, tools=tools)
 
         try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+
             async_stream = await self.get_client().aio.models.generate_content_stream(
                 model=self.id,
                 contents=formatted_messages,
                 **request_kwargs,
             )
             async for chunk in async_stream:
-                yield chunk
+                yield self._parse_provider_response_delta(chunk)
+
+            assistant_message.metrics.stop_timer()
+
         except (ClientError, ServerError) as e:
             log_error(f"Error from Gemini API: {e}")
             raise ModelProviderError(
@@ -696,7 +743,7 @@ class Gemini(Model):
         """
         combined_content: List = []
         combined_function_result: List = []
-        message_metrics = MessageMetrics()
+        message_metrics = Metrics()
         if len(function_call_results) > 0:
             for result in function_call_results:
                 combined_content.append(result.content)
@@ -710,7 +757,7 @@ class Gemini(Model):
                 )
             )
 
-    def parse_provider_response(self, response: GenerateContentResponse, **kwargs) -> ModelResponse:
+    def _parse_provider_response(self, response: GenerateContentResponse, **kwargs) -> ModelResponse:
         """
         Parse the OpenAI response into a ModelResponse.
 
@@ -765,8 +812,12 @@ class Gemini(Model):
                                 model_response.content += content_str
 
                 if hasattr(part, "inline_data") and part.inline_data is not None:
-                    model_response.image = ImageArtifact(
-                        id=str(uuid4()), content=part.inline_data.data, mime_type=part.inline_data.mime_type
+                    if model_response.images is None:
+                        model_response.images = []
+                    model_response.images.append(
+                        ImageArtifact(
+                            id=str(uuid4()), content=part.inline_data.data, mime_type=part.inline_data.mime_type
+                        )
                     )
 
                 # Extract function call if present
@@ -836,18 +887,7 @@ class Gemini(Model):
 
         # Extract usage metadata if present
         if hasattr(response, "usage_metadata") and response.usage_metadata is not None:
-            usage: GenerateContentResponseUsageMetadata = response.usage_metadata
-
-            output_tokens = usage.candidates_token_count or 0
-            if hasattr(usage, "thoughts_token_count") and usage.thoughts_token_count is not None:
-                output_tokens += usage.thoughts_token_count or 0
-
-            model_response.response_usage = {
-                "input_tokens": usage.prompt_token_count or 0,
-                "output_tokens": output_tokens,
-                "total_tokens": usage.total_token_count or 0,
-                "cached_tokens": usage.cached_content_token_count or 0,
-            }
+            model_response.response_usage = self._get_metrics(response.usage_metadata)
 
         # If we have no content but have a role, add a default empty content
         if model_response.role and model_response.content is None and not model_response.tool_calls:
@@ -855,7 +895,7 @@ class Gemini(Model):
 
         return model_response
 
-    def parse_provider_response_delta(self, response_delta: GenerateContentResponse) -> ModelResponse:
+    def _parse_provider_response_delta(self, response_delta: GenerateContentResponse) -> ModelResponse:
         model_response = ModelResponse()
 
         if response_delta.candidates and len(response_delta.candidates) > 0:
@@ -886,8 +926,12 @@ class Gemini(Model):
                                 model_response.content += text_content
 
                     if hasattr(part, "inline_data") and part.inline_data is not None:
-                        model_response.image = ImageArtifact(
-                            id=str(uuid4()), content=part.inline_data.data, mime_type=part.inline_data.mime_type
+                        if model_response.images is None:
+                            model_response.images = []
+                        model_response.images.append(
+                            ImageArtifact(
+                                id=str(uuid4()), content=part.inline_data.data, mime_type=part.inline_data.mime_type
+                            )
                         )
 
                     # Extract function call if present
@@ -932,18 +976,7 @@ class Gemini(Model):
 
             # Extract usage metadata if present
             if hasattr(response_delta, "usage_metadata") and response_delta.usage_metadata is not None:
-                usage: GenerateContentResponseUsageMetadata = response_delta.usage_metadata
-
-                output_tokens = usage.candidates_token_count or 0
-                if hasattr(usage, "thoughts_token_count") and usage.thoughts_token_count is not None:
-                    output_tokens += usage.thoughts_token_count or 0
-
-                model_response.response_usage = {
-                    "input_tokens": usage.prompt_token_count or 0,
-                    "output_tokens": output_tokens,
-                    "total_tokens": usage.total_token_count or 0,
-                    "cached_tokens": usage.cached_content_token_count or 0,
-                }
+                model_response.response_usage = self._get_metrics(response_delta.usage_metadata)
 
         return model_response
 
@@ -984,3 +1017,28 @@ class Gemini(Model):
         setattr(new_instance, "client", None)
 
         return new_instance
+
+    def _get_metrics(self, response_usage: GenerateContentResponseUsageMetadata) -> Metrics:
+        """
+        Parse the given Google Gemini usage into an Agno Metrics object.
+
+        Args:
+            response_usage: Usage data from Google Gemini
+
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
+
+        metrics.input_tokens = response_usage.prompt_token_count or 0
+        metrics.output_tokens = response_usage.candidates_token_count or 0
+        metrics.total_tokens = metrics.input_tokens + metrics.output_tokens
+        if response_usage.thoughts_token_count is not None:
+            metrics.output_tokens += response_usage.thoughts_token_count or 0
+
+        metrics.cache_read_tokens = response_usage.cached_content_token_count or 0
+
+        if response_usage.traffic_type is not None:
+            metrics.provider_metrics = {"traffic_type": response_usage.traffic_type}
+
+        return metrics

agno/models/groq/groq.py

@@ -1,3 +1,4 @@
+from collections.abc import AsyncIterator
 from dataclasses import dataclass
 from os import getenv
 from typing import Any, Dict, Iterator, List, Optional, Type, Union
@@ -8,7 +9,9 @@ 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, log_warning
 from agno.utils.openai import images_to_message
 
@@ -18,6 +21,7 @@ try:
     from groq import Groq as GroqClient
     from groq.types.chat import ChatCompletion
     from groq.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDelta, ChoiceDeltaToolCall
+    from groq.types.completion_usage import CompletionUsage
 except ImportError:
     raise ImportError("`groq` not installed. Please install using `pip install groq`")
 
@@ -250,19 +254,31 @@ class Groq(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,
-    ) -> ChatCompletion:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Send a chat completion request to the Groq API.
         """
         try:
-            return self.get_client().chat.completions.create(
+            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.completions.create(
                 model=self.id,
                 messages=[self.format_message(m) for m in messages],  # type: ignore
                 **self.get_request_params(response_format=response_format, tools=tools, tool_choice=tool_choice),
             )
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+
+            return model_response
+
         except (APIResponseValidationError, APIStatusError) as e:
             log_error(f"Error calling Groq API: {str(e)}")
             raise ModelProviderError(
@@ -278,19 +294,31 @@ class Groq(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,
-    ) -> ChatCompletion:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Sends an asynchronous chat completion request to the Groq API.
         """
         try:
-            return await self.get_async_client().chat.completions.create(
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+            response = await self.get_async_client().chat.completions.create(
                 model=self.id,
                 messages=[self.format_message(m) for m in messages],  # type: ignore
                 **self.get_request_params(response_format=response_format, tools=tools, tool_choice=tool_choice),
             )
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(response, response_format=response_format)
+
+            return model_response
+
         except (APIResponseValidationError, APIStatusError) as e:
             log_error(f"Error calling Groq API: {str(e)}")
             raise ModelProviderError(
@@ -306,20 +334,31 @@ class Groq(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[ChatCompletionChunk]:
+        run_response: Optional[RunOutput] = None,
+    ) -> Iterator[ModelResponse]:
         """
         Send a streaming chat completion request to the Groq API.
         """
         try:
-            return self.get_client().chat.completions.create(
+            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.completions.create(
                 model=self.id,
                 messages=[self.format_message(m) for m in messages],  # type: ignore
                 stream=True,
                 **self.get_request_params(response_format=response_format, tools=tools, tool_choice=tool_choice),
-            )
+            ):
+                yield self._parse_provider_response_delta(chunk)  # type: ignore
+
+            assistant_message.metrics.stop_timer()
+
         except (APIResponseValidationError, APIStatusError) as e:
             log_error(f"Error calling Groq API: {str(e)}")
             raise ModelProviderError(
@@ -335,23 +374,33 @@ class Groq(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]:
         """
         Sends an asynchronous streaming chat completion request to the Groq API.
         """
 
         try:
-            stream = await self.get_async_client().chat.completions.create(
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+
+            async_stream = await self.get_async_client().chat.completions.create(
                 model=self.id,
                 messages=[self.format_message(m) for m in messages],  # type: ignore
                 stream=True,
                 **self.get_request_params(response_format=response_format, tools=tools, tool_choice=tool_choice),
             )
-            async for chunk in stream:  # type: ignore
-                yield chunk
+            async for chunk in async_stream:  # type: ignore
+                yield self._parse_provider_response_delta(chunk)  # type: ignore
+
+            assistant_message.metrics.stop_timer()
+
         except (APIResponseValidationError, APIStatusError) as e:
             log_error(f"Error calling Groq API: {str(e)}")
             raise ModelProviderError(
@@ -405,7 +454,7 @@ class Groq(Model):
                     tool_call_entry["type"] = _tool_call_type
         return tool_calls
 
-    def parse_provider_response(self, response: ChatCompletion, **kwargs) -> ModelResponse:
+    def _parse_provider_response(self, response: ChatCompletion, **kwargs) -> ModelResponse:
         """
         Parse the Groq response into a ModelResponse.
 
@@ -437,20 +486,11 @@ class Groq(Model):
 
         # Add usage metrics if present
         if response.usage is not None:
-            model_response.response_usage = {
-                "input_tokens": response.usage.prompt_tokens,
-                "output_tokens": response.usage.completion_tokens,
-                "total_tokens": response.usage.total_tokens,
-                "additional_metrics": {
-                    "completion_time": response.usage.completion_time,
-                    "prompt_time": response.usage.prompt_time,
-                    "queue_time": response.usage.queue_time,
-                    "total_time": response.usage.total_time,
-                },
-            }
+            model_response.response_usage = self._get_metrics(response.usage)
+
         return model_response
 
-    def parse_provider_response_delta(self, response: ChatCompletionChunk) -> ModelResponse:
+    def _parse_provider_response_delta(self, response: ChatCompletionChunk) -> ModelResponse:
         """
         Parse the Groq streaming response into ModelResponse objects.
 
@@ -476,16 +516,38 @@ class Groq(Model):
 
         # Add usage metrics if present
         if response.x_groq is not None and response.x_groq.usage is not None:
-            model_response.response_usage = {
-                "input_tokens": response.x_groq.usage.prompt_tokens,
-                "output_tokens": response.x_groq.usage.completion_tokens,
-                "total_tokens": response.x_groq.usage.total_tokens,
-                "additional_metrics": {
-                    "completion_time": response.x_groq.usage.completion_time,
-                    "prompt_time": response.x_groq.usage.prompt_time,
-                    "queue_time": response.x_groq.usage.queue_time,
-                    "total_time": response.x_groq.usage.total_time,
-                },
-            }
+            model_response.response_usage = self._get_metrics(response.x_groq.usage)
 
         return model_response
+
+    def _get_metrics(self, response_usage: CompletionUsage) -> Metrics:
+        """
+        Parse the given Groq usage into an Agno Metrics object.
+
+        Args:
+            response_usage: Usage data from Groq
+
+        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
+
+        # Additional time metrics offered by Groq
+        if completion_time := response_usage.completion_time:
+            metrics.provider_metrics = metrics.provider_metrics or {}
+            metrics.provider_metrics["completion_time"] = completion_time
+        if prompt_time := response_usage.prompt_time:
+            metrics.provider_metrics = metrics.provider_metrics or {}
+            metrics.provider_metrics["prompt_time"] = prompt_time
+        if queue_time := response_usage.queue_time:
+            metrics.provider_metrics = metrics.provider_metrics or {}
+            metrics.provider_metrics["queue_time"] = queue_time
+        if total_time := response_usage.total_time:
+            metrics.provider_metrics = metrics.provider_metrics or {}
+            metrics.provider_metrics["total_time"] = total_time
+
+        return metrics