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

agno/models/huggingface/huggingface.py

@@ -5,12 +5,15 @@ from os import getenv
 from typing import Any, Dict, Iterator, List, Optional, Type, Union
 
 import httpx
+from huggingface_hub import ChatCompletionInputStreamOptions
 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
 
 try:
@@ -228,19 +231,29 @@ class HuggingFace(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[ChatCompletionOutput]:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Send a chat completion request to the HuggingFace Hub.
         """
         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],
                 **self.get_request_params(tools=tools, tool_choice=tool_choice),
             )
+            assistant_message.metrics.stop_timer()
+
+            return self._parse_provider_response(provider_response, response_format=response_format)
+
         except InferenceTimeoutError as e:
             log_error(f"Error invoking HuggingFace model: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -251,20 +264,29 @@ class HuggingFace(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[ChatCompletionOutput]:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Sends an asynchronous chat completion request to the HuggingFace Hub Inference.
         """
         try:
-            async with self.get_async_client() as client:
-                return await client.chat.completions.create(
-                    model=self.id,
-                    messages=[self._format_message(m) for m in messages],
-                    **self.get_request_params(tools=tools, tool_choice=tool_choice),
-                )
+            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.completions.create(
+                model=self.id,
+                messages=[self._format_message(m) for m in messages],
+                **self.get_request_params(tools=tools, tool_choice=tool_choice),
+            )
+            assistant_message.metrics.stop_timer()
+
+            return self._parse_provider_response(provider_response, response_format=response_format)
+
         except InferenceTimeoutError as e:
             log_error(f"Error invoking HuggingFace model: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -275,21 +297,34 @@ class HuggingFace(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[ChatCompletionStreamOutput]:
+        run_response: Optional[RunOutput] = None,
+    ) -> Iterator[ModelResponse]:
         """
         Send a streaming chat completion request to the HuggingFace API.
         """
         try:
-            yield from 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()
+
+            stream = self.get_client().chat.completions.create(
                 model=self.id,
                 messages=[self._format_message(m) for m in messages],
                 stream=True,
-                stream_options={"include_usage": True},
+                stream_options=ChatCompletionInputStreamOptions(include_usage=True),  # type: ignore
                 **self.get_request_params(tools=tools, tool_choice=tool_choice),
-            )  # type: ignore
+            )
+
+            for chunk in stream:
+                yield self._parse_provider_response_delta(chunk)
+
+            assistant_message.metrics.stop_timer()
+
         except InferenceTimeoutError as e:
             log_error(f"Error invoking HuggingFace model: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -300,24 +335,33 @@ class HuggingFace(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[Any]:
         """
         Sends an asynchronous streaming chat completion request to the HuggingFace API.
         """
         try:
-            async with self.get_async_client() as client:
-                stream = await client.chat.completions.create(
-                    model=self.id,
-                    messages=[self._format_message(m) for m in messages],
-                    stream=True,
-                    stream_options={"include_usage": True},
-                    **self.get_request_params(tools=tools, tool_choice=tool_choice),
-                )
-                async for chunk in stream:
-                    yield chunk
+            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.completions.create(
+                model=self.id,
+                messages=[self._format_message(m) for m in messages],
+                stream=True,
+                stream_options=ChatCompletionInputStreamOptions(include_usage=True),  # type: ignore
+                **self.get_request_params(tools=tools, tool_choice=tool_choice),
+            )
+
+            async for chunk in provider_response:
+                yield self._parse_provider_response_delta(chunk)
+
+            assistant_message.metrics.stop_timer()
+
         except InferenceTimeoutError as e:
             log_error(f"Error invoking HuggingFace model: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -366,7 +410,7 @@ class HuggingFace(Model):
                     tool_call_entry["type"] = _tool_call_type
         return tool_calls
 
-    def parse_provider_response(
+    def _parse_provider_response(
         self,
         response: ChatCompletionOutput,
         response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
@@ -401,11 +445,11 @@ class HuggingFace(Model):
             log_warning(f"Error retrieving structured outputs: {e}")
 
         if response.usage is not None:
-            model_response.response_usage = response.usage
+            model_response.response_usage = self._get_metrics(response)
 
         return model_response
 
-    def parse_provider_response_delta(self, response_delta: ChatCompletionStreamOutput) -> ModelResponse:
+    def _parse_provider_response_delta(self, response_delta: ChatCompletionStreamOutput) -> ModelResponse:
         """
         Parse the provider response delta into a ModelResponse.
         """
@@ -420,6 +464,27 @@ class HuggingFace(Model):
             if response_delta_message.tool_calls is not None and len(response_delta_message.tool_calls) > 0:
                 model_response.tool_calls = [response_delta_message.tool_calls]  # type: ignore
         if response_delta.usage is not None:
-            model_response.response_usage = response_delta.usage
+            model_response.response_usage = self._get_metrics(response_delta)
 
         return model_response
+
+    def _get_metrics(self, response: Union[ChatCompletionOutput, ChatCompletionStreamOutput]) -> Metrics:
+        """
+        Parse the given HuggingFace-specific usage into an Agno Metrics object.
+
+        Args:
+            response: The HuggingFace response to parse.
+
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
+
+        if not response.usage:
+            return 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/ibm/watsonx.py

@@ -1,13 +1,15 @@
 from dataclasses import dataclass
 from os import getenv
-from typing import Any, AsyncGenerator, 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, log_warning
 from agno.utils.models.watsonx import format_images_for_message
 
@@ -154,14 +156,19 @@ class WatsonX(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,
-    ) -> Any:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Send a chat completion request to the WatsonX API.
         """
         try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
             client = self.get_client()
 
             formatted_messages = [self._format_message(m) for m in messages]
@@ -169,9 +176,13 @@ class WatsonX(Model):
                 response_format=response_format, tools=tools, tool_choice=tool_choice
             )
 
-            # Call chat method
+            assistant_message.metrics.start_timer()
             response = client.chat(messages=formatted_messages, **request_params)
-            return response
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(response, response_format=response_format)
+
+            return model_response
 
         except Exception as e:
             log_error(f"Error calling WatsonX API: {str(e)}")
@@ -180,14 +191,19 @@ class WatsonX(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,
     ) -> Any:
         """
         Sends an asynchronous chat completion request to the WatsonX API.
         """
         try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
             client = self.get_client()
             formatted_messages = [self._format_message(m) for m in messages]
 
@@ -195,7 +211,13 @@ class WatsonX(Model):
                 response_format=response_format, tools=tools, tool_choice=tool_choice
             )
 
-            return await client.achat(messages=formatted_messages, **request_params)
+            assistant_message.metrics.start_timer()
+            provider_response = await client.achat(messages=formatted_messages, **request_params)
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+
+            return model_response
 
         except Exception as e:
             log_error(f"Error calling WatsonX API: {str(e)}")
@@ -204,10 +226,12 @@ class WatsonX(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]:
         """
         Send a streaming chat completion request to the WatsonX API.
         """
@@ -219,7 +243,15 @@ class WatsonX(Model):
                 response_format=response_format, tools=tools, tool_choice=tool_choice
             )
 
-            yield from client.chat_stream(messages=formatted_messages, **request_params)
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            assistant_message.metrics.start_timer()
+
+            for chunk in client.chat_stream(messages=formatted_messages, **request_params):
+                yield self._parse_provider_response_delta(chunk)
+
+            assistant_message.metrics.stop_timer()
 
         except Exception as e:
             log_error(f"Error calling WatsonX API: {str(e)}")
@@ -228,14 +260,19 @@ class WatsonX(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,
-    ) -> AsyncGenerator[Any, None]:
+        run_response: Optional[RunOutput] = None,
+    ) -> AsyncIterator[ModelResponse]:
         """
         Sends an asynchronous streaming chat completion request to the WatsonX API.
         """
         try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
             client = self.get_client()
             formatted_messages = [self._format_message(m) for m in messages]
 
@@ -244,9 +281,13 @@ class WatsonX(Model):
                 response_format=response_format, tools=tools, tool_choice=tool_choice
             )
 
+            assistant_message.metrics.start_timer()
+
             async_stream = await client.achat_stream(messages=formatted_messages, **request_params)
             async for chunk in async_stream:
-                yield chunk
+                yield self._parse_provider_response_delta(chunk)
+
+            assistant_message.metrics.stop_timer()
 
         except Exception as e:
             log_error(f"Error in async streaming from WatsonX API: {str(e)}")
@@ -293,7 +334,7 @@ class WatsonX(Model):
                     tool_call_entry["type"] = _tool_call_type
         return tool_calls
 
-    def parse_provider_response(
+    def _parse_provider_response(
         self,
         response: Dict[str, Any],
         response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
@@ -334,11 +375,11 @@ class WatsonX(Model):
                 log_warning(f"Error processing tool calls: {e}")
 
         if response.get("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: Dict[str, Any]) -> ModelResponse:
+    def _parse_provider_response_delta(self, response_delta: Dict[str, Any]) -> ModelResponse:
         """
         Parse the OpenAI streaming response into a ModelResponse.
         """
@@ -358,6 +399,24 @@ class WatsonX(Model):
 
         # Add usage metrics if present
         if response_delta.get("usage") is not None:
-            model_response.response_usage = response_delta["usage"]
+            model_response.response_usage = self._get_metrics(response_delta["usage"])
 
         return model_response
+
+    def _get_metrics(self, response_usage: Dict[str, Any]) -> Metrics:
+        """
+        Parse the given WatsonX usage into an Agno Metrics object.
+
+        Args:
+            response_usage: Usage data from WatsonX
+
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
+
+        metrics.input_tokens = response_usage.get("prompt_tokens") or 0
+        metrics.output_tokens = response_usage.get("completion_tokens") or 0
+        metrics.total_tokens = metrics.input_tokens + metrics.output_tokens
+
+        return metrics

agno/models/litellm/chat.py

@@ -1,13 +1,15 @@
 import json
 from dataclasses import dataclass
 from os import getenv
-from typing import Any, AsyncIterator, Dict, Iterator, List, Mapping, Optional, Type, Union
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Type, Union
 
 from pydantic import BaseModel
 
 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 _format_file_for_message, audio_to_message, images_to_message
 
@@ -162,65 +164,112 @@ class LiteLLM(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:
         """Sends a chat completion request to the LiteLLM API."""
         completion_kwargs = self.get_request_params(tools=tools)
         completion_kwargs["messages"] = self._format_messages(messages)
-        return self.get_client().completion(**completion_kwargs)
+
+        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().completion(**completion_kwargs)
+
+        assistant_message.metrics.stop_timer()
+
+        model_response = self._parse_provider_response(provider_response, response_format=response_format)
+        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 completion request to the LiteLLM API."""
         completion_kwargs = self.get_request_params(tools=tools)
         completion_kwargs["messages"] = self._format_messages(messages)
         completion_kwargs["stream"] = True
         completion_kwargs["stream_options"] = {"include_usage": True}
-        return self.get_client().completion(**completion_kwargs)
+
+        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().completion(**completion_kwargs):
+            yield self._parse_provider_response_delta(chunk)
+
+        assistant_message.metrics.stop_timer()
 
     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 completion request to the LiteLLM API."""
         completion_kwargs = self.get_request_params(tools=tools)
         completion_kwargs["messages"] = self._format_messages(messages)
-        return await self.get_client().acompletion(**completion_kwargs)
+
+        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().acompletion(**completion_kwargs)
+
+        assistant_message.metrics.stop_timer()
+
+        model_response = self._parse_provider_response(provider_response, response_format=response_format)
+        return model_response
 
     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,
-    ) -> AsyncIterator[Any]:
+        run_response: Optional[RunOutput] = None,
+    ) -> AsyncIterator[ModelResponse]:
         """Sends an asynchronous streaming chat request to the LiteLLM API."""
         completion_kwargs = self.get_request_params(tools=tools)
         completion_kwargs["messages"] = self._format_messages(messages)
         completion_kwargs["stream"] = True
         completion_kwargs["stream_options"] = {"include_usage": True}
 
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+
         try:
             # litellm.acompletion returns a coroutine that resolves to an async iterator
             # We need to await it first to get the actual async iterator
             async_stream = await self.get_client().acompletion(**completion_kwargs)
             async for chunk in async_stream:
-                yield chunk
+                yield self._parse_provider_response_delta(chunk)
+
+            assistant_message.metrics.stop_timer()
+
         except Exception as e:
             log_error(f"Error in streaming response: {e}")
             raise
 
-    def parse_provider_response(self, response: Any, **kwargs) -> ModelResponse:
+    def _parse_provider_response(self, response: Any, **kwargs) -> ModelResponse:
         """Parse the provider response."""
         model_response = ModelResponse()
 
@@ -241,11 +290,11 @@ class LiteLLM(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: Any) -> ModelResponse:
+    def _parse_provider_response_delta(self, response_delta: Any) -> ModelResponse:
         """Parse the provider response delta for streaming responses."""
         model_response = ModelResponse()
 
@@ -284,7 +333,7 @@ class LiteLLM(Model):
 
         # Add usage metrics if present in streaming response
         if hasattr(response_delta, "usage") and response_delta.usage is not None:
-            model_response.response_usage = response_delta.usage
+            model_response.response_usage = self._get_metrics(response_delta.usage)
 
         return model_response
 
@@ -377,3 +426,26 @@ class LiteLLM(Model):
             result.append(tc_copy)
 
         return result
+
+    def _get_metrics(self, response_usage: Any) -> Metrics:
+        """
+        Parse the given LiteLLM usage into an Agno Metrics object.
+
+        Args:
+            response_usage: Usage data from LiteLLM
+
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
+
+        if isinstance(response_usage, dict):
+            metrics.input_tokens = response_usage.get("prompt_tokens") or 0
+            metrics.output_tokens = response_usage.get("completion_tokens") or 0
+        else:
+            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