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

agno/models/base.py

@@ -5,7 +5,6 @@ from dataclasses import dataclass, field
 from types import AsyncGeneratorType, GeneratorType
 from typing import (
     Any,
-    AsyncGenerator,
     AsyncIterator,
     Dict,
     Iterator,
@@ -22,12 +21,13 @@ from uuid import uuid4
 from pydantic import BaseModel
 
 from agno.exceptions import AgentRunException
-from agno.media import AudioResponse, ImageArtifact
-from agno.models.message import Citations, Message, MessageMetrics
+from agno.media import Audio, Image, Video
+from agno.models.message import Citations, Message
+from agno.models.metrics import Metrics
 from agno.models.response import ModelResponse, ModelResponseEvent, ToolExecution
-from agno.run.response import RunResponseContentEvent, RunResponseEvent
-from agno.run.team import RunResponseContentEvent as TeamRunResponseContentEvent
-from agno.run.team import TeamRunResponseEvent
+from agno.run.agent import CustomEvent, RunContentEvent, RunOutput, RunOutputEvent
+from agno.run.team import RunContentEvent as TeamRunContentEvent
+from agno.run.team import TeamRunOutputEvent
 from agno.tools.function import Function, FunctionCall, FunctionExecutionResult, UserInputField
 from agno.utils.log import log_debug, log_error, log_warning
 from agno.utils.timer import Timer
@@ -38,13 +38,14 @@ from agno.utils.tools import get_function_call_for_tool_call, get_function_call_
 class MessageData:
     response_role: Optional[Literal["system", "user", "assistant", "tool"]] = None
     response_content: Any = ""
-    response_thinking: Any = ""
-    response_redacted_thinking: Any = ""
+    response_reasoning_content: Any = ""
+    response_redacted_reasoning_content: Any = ""
     response_citations: Optional[Citations] = None
     response_tool_calls: List[Dict[str, Any]] = field(default_factory=list)
 
-    response_audio: Optional[AudioResponse] = None
-    response_image: Optional[ImageArtifact] = None
+    response_audio: Optional[Audio] = None
+    response_image: Optional[Image] = None
+    response_video: Optional[Video] = None
 
     # Data from the provider that we might need on subsequent messages
     response_provider_data: Optional[Dict[str, Any]] = None
@@ -61,143 +62,17 @@ def _log_messages(messages: List[Message]) -> None:
         m.log(metrics=False)
 
 
-def _add_usage_metrics_to_assistant_message(assistant_message: Message, response_usage: Any) -> None:
-    """
-    Add usage metrics from the model provider to the assistant message.
-
-    Args:
-        assistant_message: Message to update with metrics
-        response_usage: Usage data from model provider
-    """
-
-    # Standard token metrics
-    if isinstance(response_usage, dict):
-        if "input_tokens" in response_usage and response_usage.get("input_tokens") is not None:
-            assistant_message.metrics.input_tokens = response_usage.get("input_tokens", 0)
-        if "output_tokens" in response_usage and response_usage.get("output_tokens") is not None:
-            assistant_message.metrics.output_tokens = response_usage.get("output_tokens", 0)
-        if "prompt_tokens" in response_usage and response_usage.get("prompt_tokens") is not None:
-            assistant_message.metrics.input_tokens = response_usage.get("prompt_tokens", 0)
-        if "completion_tokens" in response_usage and response_usage.get("completion_tokens") is not None:
-            assistant_message.metrics.output_tokens = response_usage.get("completion_tokens", 0)
-        if "cached_tokens" in response_usage and response_usage.get("cached_tokens") is not None:
-            assistant_message.metrics.cached_tokens = response_usage.get("cached_tokens", 0)
-        if "cache_write_tokens" in response_usage and response_usage.get("cache_write_tokens") is not None:
-            assistant_message.metrics.cache_write_tokens = response_usage.get("cache_write_tokens", 0)
-        if "total_tokens" in response_usage and response_usage.get("total_tokens") is not None:
-            assistant_message.metrics.total_tokens = response_usage.get("total_tokens", 0)
-        else:
-            assistant_message.metrics.total_tokens = (
-                assistant_message.metrics.input_tokens + assistant_message.metrics.output_tokens
-            )
-    else:
-        if hasattr(response_usage, "input_tokens") and response_usage.input_tokens:
-            assistant_message.metrics.input_tokens = response_usage.input_tokens
-        if hasattr(response_usage, "output_tokens") and response_usage.output_tokens:
-            assistant_message.metrics.output_tokens = response_usage.output_tokens
-        if hasattr(response_usage, "prompt_tokens") and response_usage.prompt_tokens is not None:
-            assistant_message.metrics.input_tokens = response_usage.prompt_tokens
-            assistant_message.metrics.prompt_tokens = response_usage.prompt_tokens
-        if hasattr(response_usage, "completion_tokens") and response_usage.completion_tokens is not None:
-            assistant_message.metrics.output_tokens = response_usage.completion_tokens
-            assistant_message.metrics.completion_tokens = response_usage.completion_tokens
-        if hasattr(response_usage, "total_tokens") and response_usage.total_tokens is not None:
-            assistant_message.metrics.total_tokens = response_usage.total_tokens
-        if hasattr(response_usage, "cached_tokens") and response_usage.cached_tokens is not None:
-            assistant_message.metrics.cached_tokens = response_usage.cached_tokens
-        if hasattr(response_usage, "cache_write_tokens") and response_usage.cache_write_tokens is not None:
-            assistant_message.metrics.cache_write_tokens = response_usage.cache_write_tokens
-
-    # If you didn't capture any total tokens
-    if not assistant_message.metrics.total_tokens:
-        if assistant_message.metrics.input_tokens is None:
-            assistant_message.metrics.input_tokens = 0
-        if assistant_message.metrics.output_tokens is None:
-            assistant_message.metrics.output_tokens = 0
-
-        assistant_message.metrics.total_tokens = (
-            assistant_message.metrics.input_tokens + assistant_message.metrics.output_tokens
-        )
-
-    # Additional metrics (e.g., from Groq, Ollama)
-    if isinstance(response_usage, dict) and "additional_metrics" in response_usage:
-        assistant_message.metrics.additional_metrics = response_usage["additional_metrics"]
-
-    # Token details (e.g., from OpenAI)
-    if hasattr(response_usage, "prompt_tokens_details"):
-        if isinstance(response_usage.prompt_tokens_details, dict):
-            assistant_message.metrics.prompt_tokens_details = response_usage.prompt_tokens_details
-            if (
-                "audio_tokens" in response_usage.prompt_tokens_details
-                and response_usage.prompt_tokens_details["audio_tokens"] is not None
-            ):
-                assistant_message.metrics.input_audio_tokens = response_usage.prompt_tokens_details["audio_tokens"]
-            if (
-                "cached_tokens" in response_usage.prompt_tokens_details
-                and response_usage.prompt_tokens_details["cached_tokens"] is not None
-            ):
-                assistant_message.metrics.cached_tokens = response_usage.prompt_tokens_details["cached_tokens"]
-        elif hasattr(response_usage.prompt_tokens_details, "model_dump"):
-            assistant_message.metrics.prompt_tokens_details = response_usage.prompt_tokens_details.model_dump(
-                exclude_none=True
-            )
-            if (
-                hasattr(response_usage.prompt_tokens_details, "audio_tokens")
-                and response_usage.prompt_tokens_details.audio_tokens is not None
-            ):
-                assistant_message.metrics.input_audio_tokens = response_usage.prompt_tokens_details.audio_tokens
-            if (
-                hasattr(response_usage.prompt_tokens_details, "cached_tokens")
-                and response_usage.prompt_tokens_details.cached_tokens is not None
-            ):
-                assistant_message.metrics.cached_tokens = response_usage.prompt_tokens_details.cached_tokens
-
-    if hasattr(response_usage, "completion_tokens_details"):
-        if isinstance(response_usage.completion_tokens_details, dict):
-            assistant_message.metrics.completion_tokens_details = response_usage.completion_tokens_details
-            if (
-                "audio_tokens" in response_usage.completion_tokens_details
-                and response_usage.completion_tokens_details["audio_tokens"] is not None
-            ):
-                assistant_message.metrics.output_audio_tokens = response_usage.completion_tokens_details["audio_tokens"]
-            if (
-                "reasoning_tokens" in response_usage.completion_tokens_details
-                and response_usage.completion_tokens_details["reasoning_tokens"] is not None
-            ):
-                assistant_message.metrics.reasoning_tokens = response_usage.completion_tokens_details[
-                    "reasoning_tokens"
-                ]
-        elif hasattr(response_usage.completion_tokens_details, "model_dump"):
-            assistant_message.metrics.completion_tokens_details = response_usage.completion_tokens_details.model_dump(
-                exclude_none=True
-            )
-            if (
-                hasattr(response_usage.completion_tokens_details, "audio_tokens")
-                and response_usage.completion_tokens_details.audio_tokens is not None
-            ):
-                assistant_message.metrics.output_audio_tokens = response_usage.completion_tokens_details.audio_tokens
-            if (
-                hasattr(response_usage.completion_tokens_details, "reasoning_tokens")
-                and response_usage.completion_tokens_details.reasoning_tokens is not None
-            ):
-                assistant_message.metrics.reasoning_tokens = response_usage.completion_tokens_details.reasoning_tokens
-
-    assistant_message.metrics.audio_tokens = (
-        assistant_message.metrics.input_audio_tokens + assistant_message.metrics.output_audio_tokens
-    )
-
-
-def _handle_agent_exception(a_exc: AgentRunException, additional_messages: Optional[List[Message]] = None) -> None:
+def _handle_agent_exception(a_exc: AgentRunException, additional_input: Optional[List[Message]] = None) -> None:
     """Handle AgentRunException and collect additional messages."""
-    if additional_messages is None:
-        additional_messages = []
+    if additional_input is None:
+        additional_input = []
     if a_exc.user_message is not None:
         msg = (
             Message(role="user", content=a_exc.user_message)
             if isinstance(a_exc.user_message, str)
             else a_exc.user_message
         )
-        additional_messages.append(msg)
+        additional_input.append(msg)
 
     if a_exc.agent_message is not None:
         msg = (
@@ -205,20 +80,20 @@ def _handle_agent_exception(a_exc: AgentRunException, additional_messages: Optio
             if isinstance(a_exc.agent_message, str)
             else a_exc.agent_message
         )
-        additional_messages.append(msg)
+        additional_input.append(msg)
 
     if a_exc.messages:
         for m in a_exc.messages:
             if isinstance(m, Message):
-                additional_messages.append(m)
+                additional_input.append(m)
             elif isinstance(m, dict):
                 try:
-                    additional_messages.append(Message(**m))
+                    additional_input.append(Message(**m))
                 except Exception as e:
                     log_warning(f"Failed to convert dict to Message: {e}")
 
     if a_exc.stop_execution:
-        for m in additional_messages:
+        for m in additional_input:
             m.stop_after_tool_call = True
 
 
@@ -270,23 +145,23 @@ class Model(ABC):
         return self.provider or self.name or self.__class__.__name__
 
     @abstractmethod
-    def invoke(self, *args, **kwargs) -> Any:
+    def invoke(self, *args, **kwargs) -> ModelResponse:
         pass
 
     @abstractmethod
-    async def ainvoke(self, *args, **kwargs) -> Any:
+    async def ainvoke(self, *args, **kwargs) -> ModelResponse:
         pass
 
     @abstractmethod
-    def invoke_stream(self, *args, **kwargs) -> Iterator[Any]:
+    def invoke_stream(self, *args, **kwargs) -> Iterator[ModelResponse]:
         pass
 
     @abstractmethod
-    async def ainvoke_stream(self, *args, **kwargs) -> AsyncGenerator[Any, None]:
+    def ainvoke_stream(self, *args, **kwargs) -> AsyncIterator[ModelResponse]:
         pass
 
     @abstractmethod
-    def parse_provider_response(self, response: Any, **kwargs) -> ModelResponse:
+    def _parse_provider_response(self, response: Any, **kwargs) -> ModelResponse:
         """
         Parse the raw response from the model provider into a ModelResponse.
 
@@ -299,7 +174,7 @@ class Model(ABC):
         pass
 
     @abstractmethod
-    def parse_provider_response_delta(self, response: Any) -> ModelResponse:
+    def _parse_provider_response_delta(self, response: Any) -> ModelResponse:
         """
         Parse the streaming response from the model provider into ModelResponse objects.
 
@@ -319,6 +194,7 @@ class Model(ABC):
         functions: Optional[Dict[str, Function]] = None,
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
         tool_call_limit: Optional[int] = None,
+        run_response: Optional[RunOutput] = None,
     ) -> ModelResponse:
         """
         Generate a response from the model.
@@ -342,6 +218,7 @@ class Model(ABC):
                 response_format=response_format,
                 tools=tools,
                 tool_choice=tool_choice or self._tool_choice,
+                run_response=run_response,
             )
 
             # Add assistant message to messages
@@ -369,6 +246,26 @@ class Model(ABC):
                     function_call_limit=tool_call_limit,
                 ):
                     if isinstance(function_call_response, ModelResponse):
+                        # The session state is updated by the function call
+                        if function_call_response.updated_session_state is not None:
+                            model_response.updated_session_state = function_call_response.updated_session_state
+
+                        # Media artifacts are generated by the function call
+                        if function_call_response.images is not None:
+                            if model_response.images is None:
+                                model_response.images = []
+                            model_response.images.extend(function_call_response.images)
+
+                        if function_call_response.audios is not None:
+                            if model_response.audios is None:
+                                model_response.audios = []
+                            model_response.audios.extend(function_call_response.audios)
+
+                        if function_call_response.videos is not None:
+                            if model_response.videos is None:
+                                model_response.videos = []
+                            model_response.videos.extend(function_call_response.videos)
+
                         if (
                             function_call_response.event
                             in [
@@ -395,6 +292,11 @@ class Model(ABC):
                 self.format_function_call_results(
                     messages=messages, function_call_results=function_call_results, **model_response.extra or {}
                 )
+
+                if any(msg.images or msg.videos or msg.audio for msg in function_call_results):
+                    # Handle function call media
+                    self._handle_function_call_media(messages=messages, function_call_results=function_call_results)
+
                 for function_call_result in function_call_results:
                     function_call_result.log(metrics=True)
 
@@ -480,6 +382,26 @@ class Model(ABC):
                     function_call_limit=tool_call_limit,
                 ):
                     if isinstance(function_call_response, ModelResponse):
+                        # The session state is updated by the function call
+                        if function_call_response.updated_session_state is not None:
+                            model_response.updated_session_state = function_call_response.updated_session_state
+
+                        # Media artifacts are generated by the function call
+                        if function_call_response.images is not None:
+                            if model_response.images is None:
+                                model_response.images = []
+                            model_response.images.extend(function_call_response.images)
+
+                        if function_call_response.audios is not None:
+                            if model_response.audios is None:
+                                model_response.audios = []
+                            model_response.audios.extend(function_call_response.audios)
+
+                        if function_call_response.videos is not None:
+                            if model_response.videos is None:
+                                model_response.videos = []
+                            model_response.videos.extend(function_call_response.videos)
+
                         if (
                             function_call_response.event
                             in [
@@ -505,6 +427,11 @@ class Model(ABC):
                 self.format_function_call_results(
                     messages=messages, function_call_results=function_call_results, **model_response.extra or {}
                 )
+
+                if any(msg.images or msg.videos or msg.audio for msg in function_call_results):
+                    # Handle function call media
+                    self._handle_function_call_media(messages=messages, function_call_results=function_call_results)
+
                 for function_call_result in function_call_results:
                     function_call_result.log(metrics=True)
 
@@ -541,6 +468,7 @@ class Model(ABC):
         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,
     ) -> None:
         """
         Process a single model response and return the assistant message and whether to continue.
@@ -549,21 +477,14 @@ class Model(ABC):
             Tuple[Message, bool]: (assistant_message, should_continue)
         """
         # Generate response
-        assistant_message.metrics.start_timer()
-        response = self.invoke(
+        provider_response = self.invoke(
+            assistant_message=assistant_message,
             messages=messages,
             response_format=response_format,
             tools=tools,
             tool_choice=tool_choice or self._tool_choice,
+            run_response=run_response,
         )
-        assistant_message.metrics.stop_timer()
-
-        # Parse provider response
-        provider_response: ModelResponse = self.parse_provider_response(response, response_format=response_format)
-
-        # Add parsed data to model response
-        if provider_response.parsed is not None:
-            model_response.parsed = provider_response.parsed
 
         # Populate the assistant message
         self._populate_assistant_message(assistant_message=assistant_message, provider_response=provider_response)
@@ -574,16 +495,19 @@ class Model(ABC):
                 model_response.content = assistant_message.get_content_string()
             else:
                 model_response.content += assistant_message.get_content_string()
-        if assistant_message.thinking is not None:
-            model_response.thinking = assistant_message.thinking
-        if assistant_message.redacted_thinking is not None:
-            model_response.redacted_thinking = assistant_message.redacted_thinking
+        if assistant_message.reasoning_content is not None:
+            model_response.reasoning_content = assistant_message.reasoning_content
+        if assistant_message.redacted_reasoning_content is not None:
+            model_response.redacted_reasoning_content = assistant_message.redacted_reasoning_content
         if assistant_message.citations is not None:
             model_response.citations = assistant_message.citations
         if assistant_message.audio_output is not None:
-            model_response.audio = assistant_message.audio_output
+            if isinstance(assistant_message.audio_output, Audio):
+                model_response.audio = assistant_message.audio_output
         if assistant_message.image_output is not None:
-            model_response.image = assistant_message.image_output
+            model_response.images = [assistant_message.image_output]
+        if assistant_message.video_output is not None:
+            model_response.videos = [assistant_message.video_output]
         if provider_response.extra is not None:
             if model_response.extra is None:
                 model_response.extra = {}
@@ -597,6 +521,7 @@ class Model(ABC):
         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,
     ) -> None:
         """
         Process a single async model response and return the assistant message and whether to continue.
@@ -605,21 +530,14 @@ class Model(ABC):
             Tuple[Message, bool]: (assistant_message, should_continue)
         """
         # Generate response
-        assistant_message.metrics.start_timer()
-        response = await self.ainvoke(
+        provider_response = await self.ainvoke(
             messages=messages,
             response_format=response_format,
             tools=tools,
             tool_choice=tool_choice or self._tool_choice,
+            assistant_message=assistant_message,
+            run_response=run_response,
         )
-        assistant_message.metrics.stop_timer()
-
-        # Parse provider response
-        provider_response: ModelResponse = self.parse_provider_response(response, response_format=response_format)
-
-        # Add parsed data to model response
-        if provider_response.parsed is not None:
-            model_response.parsed = provider_response.parsed
 
         # Populate the assistant message
         self._populate_assistant_message(assistant_message=assistant_message, provider_response=provider_response)
@@ -630,16 +548,19 @@ class Model(ABC):
                 model_response.content = assistant_message.get_content_string()
             else:
                 model_response.content += assistant_message.get_content_string()
-        if assistant_message.thinking is not None:
-            model_response.thinking = assistant_message.thinking
-        if assistant_message.redacted_thinking is not None:
-            model_response.redacted_thinking = assistant_message.redacted_thinking
+        if assistant_message.reasoning_content is not None:
+            model_response.reasoning_content = assistant_message.reasoning_content
+        if assistant_message.redacted_reasoning_content is not None:
+            model_response.redacted_reasoning_content = assistant_message.redacted_reasoning_content
         if assistant_message.citations is not None:
             model_response.citations = assistant_message.citations
         if assistant_message.audio_output is not None:
-            model_response.audio = assistant_message.audio_output
+            if isinstance(assistant_message.audio_output, Audio):
+                model_response.audio = assistant_message.audio_output
         if assistant_message.image_output is not None:
-            model_response.image = assistant_message.image_output
+            model_response.images = [assistant_message.image_output]
+        if assistant_message.video_output is not None:
+            model_response.videos = [assistant_message.video_output]
         if provider_response.extra is not None:
             if model_response.extra is None:
                 model_response.extra = {}
@@ -677,16 +598,22 @@ class Model(ABC):
             assistant_message.audio_output = provider_response.audio
 
         # Add image to assistant message
-        if provider_response.image is not None:
-            assistant_message.image_output = provider_response.image
+        if provider_response.images is not None:
+            if provider_response.images:
+                assistant_message.image_output = provider_response.images[-1]  # Taking last (most recent) image
+
+        # Add video to assistant message
+        if provider_response.videos is not None:
+            if provider_response.videos:
+                assistant_message.video_output = provider_response.videos[-1]  # Taking last (most recent) video
 
-        # Add thinking content to assistant message
-        if provider_response.thinking is not None:
-            assistant_message.thinking = provider_response.thinking
+        if provider_response.audios is not None:
+            if provider_response.audios:
+                assistant_message.audio_output = provider_response.audios[-1]  # Taking last (most recent) audio
 
         # Add redacted thinking content to assistant message
-        if provider_response.redacted_thinking is not None:
-            assistant_message.redacted_thinking = provider_response.redacted_thinking
+        if provider_response.redacted_reasoning_content is not None:
+            assistant_message.redacted_reasoning_content = provider_response.redacted_reasoning_content
 
         # Add reasoning content to assistant message
         if provider_response.reasoning_content is not None:
@@ -702,9 +629,7 @@ class Model(ABC):
 
         # Add usage metrics if provided
         if provider_response.response_usage is not None:
-            _add_usage_metrics_to_assistant_message(
-                assistant_message=assistant_message, response_usage=provider_response.response_usage
-            )
+            assistant_message.metrics += provider_response.response_usage
 
         return assistant_message
 
@@ -716,22 +641,28 @@ class Model(ABC):
         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]:
         """
         Process a streaming response from the model.
         """
-        assistant_message.metrics.start_timer()
+
         for response_delta in self.invoke_stream(
             messages=messages,
+            assistant_message=assistant_message,
             response_format=response_format,
             tools=tools,
             tool_choice=tool_choice or self._tool_choice,
+            run_response=run_response,
         ):
-            model_response_delta = self.parse_provider_response_delta(response_delta)
             yield from self._populate_stream_data_and_assistant_message(
-                stream_data=stream_data, assistant_message=assistant_message, model_response_delta=model_response_delta
+                stream_data=stream_data,
+                assistant_message=assistant_message,
+                model_response_delta=response_delta,
             )
-        assistant_message.metrics.stop_timer()
+
+        # Add final metrics to assistant message
+        self._populate_assistant_message(assistant_message=assistant_message, provider_response=response_delta)
 
     def response_stream(
         self,
@@ -742,7 +673,8 @@ class Model(ABC):
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
         tool_call_limit: Optional[int] = None,
         stream_model_response: bool = True,
-    ) -> Iterator[Union[ModelResponse, RunResponseEvent, TeamRunResponseEvent]]:
+        run_response: Optional[RunOutput] = None,
+    ) -> Iterator[Union[ModelResponse, RunOutputEvent, TeamRunOutputEvent]]:
         """
         Generate a streaming response from the model.
         """
@@ -766,15 +698,16 @@ class Model(ABC):
                     response_format=response_format,
                     tools=tools,
                     tool_choice=tool_choice or self._tool_choice,
+                    run_response=run_response,
                 )
 
                 # Populate assistant message from stream data
                 if stream_data.response_content:
                     assistant_message.content = stream_data.response_content
-                if stream_data.response_thinking:
-                    assistant_message.thinking = stream_data.response_thinking
-                if stream_data.response_redacted_thinking:
-                    assistant_message.redacted_thinking = stream_data.response_redacted_thinking
+                if stream_data.response_reasoning_content:
+                    assistant_message.reasoning_content = stream_data.response_reasoning_content
+                if stream_data.response_redacted_reasoning_content:
+                    assistant_message.redacted_reasoning_content = stream_data.response_redacted_reasoning_content
                 if stream_data.response_provider_data:
                     assistant_message.provider_data = stream_data.response_provider_data
                 if stream_data.response_citations:
@@ -828,6 +761,10 @@ class Model(ABC):
                 else:
                     self.format_function_call_results(messages=messages, function_call_results=function_call_results)
 
+                # Handle function call media
+                if any(msg.images or msg.videos or msg.audio for msg in function_call_results):
+                    self._handle_function_call_media(messages=messages, function_call_results=function_call_results)
+
                 for function_call_result in function_call_results:
                     function_call_result.log(metrics=True)
 
@@ -863,23 +800,28 @@ class Model(ABC):
         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]:
         """
         Process a streaming response from the model.
         """
-        assistant_message.metrics.start_timer()
         async for response_delta in self.ainvoke_stream(
             messages=messages,
+            assistant_message=assistant_message,
             response_format=response_format,
             tools=tools,
             tool_choice=tool_choice or self._tool_choice,
+            run_response=run_response,
         ):  # type: ignore
-            model_response_delta = self.parse_provider_response_delta(response_delta)
             for model_response in self._populate_stream_data_and_assistant_message(
-                stream_data=stream_data, assistant_message=assistant_message, model_response_delta=model_response_delta
+                stream_data=stream_data,
+                assistant_message=assistant_message,
+                model_response_delta=response_delta,
             ):
                 yield model_response
-        assistant_message.metrics.stop_timer()
+
+        # Populate the assistant message
+        self._populate_assistant_message(assistant_message=assistant_message, provider_response=model_response)
 
     async def aresponse_stream(
         self,
@@ -890,7 +832,8 @@ class Model(ABC):
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
         tool_call_limit: Optional[int] = None,
         stream_model_response: bool = True,
-    ) -> AsyncIterator[Union[ModelResponse, RunResponseEvent, TeamRunResponseEvent]]:
+        run_response: Optional[RunOutput] = None,
+    ) -> AsyncIterator[Union[ModelResponse, RunOutputEvent, TeamRunOutputEvent]]:
         """
         Generate an asynchronous streaming response from the model.
         """
@@ -914,16 +857,17 @@ class Model(ABC):
                     response_format=response_format,
                     tools=tools,
                     tool_choice=tool_choice or self._tool_choice,
+                    run_response=run_response,
                 ):
                     yield response
 
                 # Populate assistant message from stream data
                 if stream_data.response_content:
                     assistant_message.content = stream_data.response_content
-                if stream_data.response_thinking:
-                    assistant_message.thinking = stream_data.response_thinking
-                if stream_data.response_redacted_thinking:
-                    assistant_message.redacted_thinking = stream_data.response_redacted_thinking
+                if stream_data.response_reasoning_content:
+                    assistant_message.reasoning_content = stream_data.response_reasoning_content
+                if stream_data.response_redacted_reasoning_content:
+                    assistant_message.redacted_reasoning_content = stream_data.response_redacted_reasoning_content
                 if stream_data.response_provider_data:
                     assistant_message.provider_data = stream_data.response_provider_data
                 if stream_data.response_audio:
@@ -940,6 +884,7 @@ class Model(ABC):
                     response_format=response_format,
                     tools=tools,
                     tool_choice=tool_choice or self._tool_choice,
+                    run_response=run_response,
                 )
                 yield model_response
 
@@ -975,6 +920,10 @@ class Model(ABC):
                 else:
                     self.format_function_call_results(messages=messages, function_call_results=function_call_results)
 
+                # Handle function call media
+                if any(msg.images or msg.videos or msg.audio for msg in function_call_results):
+                    self._handle_function_call_media(messages=messages, function_call_results=function_call_results)
+
                 for function_call_result in function_call_results:
                     function_call_result.log(metrics=True)
 
@@ -1006,11 +955,6 @@ class Model(ABC):
         self, stream_data: MessageData, assistant_message: Message, model_response_delta: ModelResponse
     ) -> Iterator[ModelResponse]:
         """Update the stream data and assistant message with the model response."""
-
-        # Update metrics
-        if not assistant_message.metrics.time_to_first_token:
-            assistant_message.metrics.set_time_to_first_token()
-
         # Add role to assistant message
         if model_response_delta.role is not None:
             assistant_message.role = model_response_delta.role
@@ -1021,16 +965,12 @@ class Model(ABC):
             stream_data.response_content += model_response_delta.content
             should_yield = True
 
-        if model_response_delta.thinking is not None:
-            stream_data.response_thinking += model_response_delta.thinking
-            should_yield = True
-
         if model_response_delta.reasoning_content is not None:
-            stream_data.response_thinking += model_response_delta.reasoning_content
+            stream_data.response_reasoning_content += model_response_delta.reasoning_content
             should_yield = True
 
-        if model_response_delta.redacted_thinking is not None:
-            stream_data.response_redacted_thinking += model_response_delta.redacted_thinking
+        if model_response_delta.redacted_reasoning_content is not None:
+            stream_data.response_redacted_reasoning_content += model_response_delta.redacted_reasoning_content
             should_yield = True
 
         if model_response_delta.citations is not None:
@@ -1049,40 +989,45 @@ class Model(ABC):
             stream_data.response_tool_calls.extend(model_response_delta.tool_calls)
             should_yield = True
 
-        if model_response_delta.audio is not None:
+        if model_response_delta.audio is not None and isinstance(model_response_delta.audio, Audio):
             if stream_data.response_audio is None:
-                stream_data.response_audio = AudioResponse(id=str(uuid4()), content="", transcript="")
+                stream_data.response_audio = Audio(id=str(uuid4()), content="", transcript="")
+
+            from typing import cast
+
+            audio_response = cast(Audio, model_response_delta.audio)
 
             # Update the stream data with audio information
-            if model_response_delta.audio.id is not None:
-                stream_data.response_audio.id = model_response_delta.audio.id  # type: ignore
-            if model_response_delta.audio.content is not None:
-                stream_data.response_audio.content += model_response_delta.audio.content  # type: ignore
-            if model_response_delta.audio.transcript is not None:
-                stream_data.response_audio.transcript += model_response_delta.audio.transcript  # type: ignore
-            if model_response_delta.audio.expires_at is not None:
-                stream_data.response_audio.expires_at = model_response_delta.audio.expires_at
-            if model_response_delta.audio.mime_type is not None:
-                stream_data.response_audio.mime_type = model_response_delta.audio.mime_type
-            stream_data.response_audio.sample_rate = model_response_delta.audio.sample_rate
-            stream_data.response_audio.channels = model_response_delta.audio.channels
+            if audio_response.id is not None:
+                stream_data.response_audio.id = audio_response.id  # type: ignore
+            if audio_response.content is not None:
+                stream_data.response_audio.content += audio_response.content  # type: ignore
+            if audio_response.transcript is not None:
+                stream_data.response_audio.transcript += audio_response.transcript  # type: ignore
+            if audio_response.expires_at is not None:
+                stream_data.response_audio.expires_at = audio_response.expires_at
+            if audio_response.mime_type is not None:
+                stream_data.response_audio.mime_type = audio_response.mime_type
+            stream_data.response_audio.sample_rate = audio_response.sample_rate
+            stream_data.response_audio.channels = audio_response.channels
 
             should_yield = True
 
-        if model_response_delta.image:
+        if model_response_delta.images:
             if stream_data.response_image is None:
-                stream_data.response_image = model_response_delta.image
+                stream_data.response_image = model_response_delta.images[-1]
+            should_yield = True
+
+        if model_response_delta.videos:
+            if stream_data.response_video is None:
+                stream_data.response_video = model_response_delta.videos[-1]
+            should_yield = True
 
         if model_response_delta.extra is not None:
             if stream_data.extra is None:
                 stream_data.extra = {}
             stream_data.extra.update(model_response_delta.extra)
 
-        if model_response_delta.response_usage is not None:
-            _add_usage_metrics_to_assistant_message(
-                assistant_message=assistant_message, response_usage=model_response_delta.response_usage
-            )
-
         if should_yield:
             yield model_response_delta
 
@@ -1142,11 +1087,24 @@ class Model(ABC):
         success: bool,
         output: Optional[Union[List[Any], str]] = None,
         timer: Optional[Timer] = None,
+        function_execution_result: Optional[FunctionExecutionResult] = None,
     ) -> Message:
         """Create a function call result message."""
         kwargs = {}
         if timer is not None:
-            kwargs["metrics"] = MessageMetrics(time=timer.elapsed)
+            kwargs["metrics"] = Metrics(duration=timer.elapsed)
+
+        # Include media artifacts from function execution result in the tool message
+        images = None
+        videos = None
+        audios = None
+
+        if success and function_execution_result:
+            # With unified classes, no conversion needed - use directly
+            images = function_execution_result.images
+            videos = function_execution_result.videos
+            audios = function_execution_result.audios
+
         return Message(
             role=self.tool_message_role,
             content=output if success else function_call.error,
@@ -1155,7 +1113,10 @@ class Model(ABC):
             tool_args=function_call.arguments,
             tool_call_error=not success,
             stop_after_tool_call=function_call.function.stop_after_tool_call,
-            **kwargs,
+            images=images,
+            videos=videos,
+            audio=audios,
+            **kwargs,  # type: ignore
         )
 
     def create_tool_call_limit_error_result(self, function_call: FunctionCall) -> Message:
@@ -1172,8 +1133,8 @@ class Model(ABC):
         self,
         function_call: FunctionCall,
         function_call_results: List[Message],
-        additional_messages: Optional[List[Message]] = None,
-    ) -> Iterator[Union[ModelResponse, RunResponseEvent, TeamRunResponseEvent]]:
+        additional_input: Optional[List[Message]] = None,
+    ) -> Iterator[Union[ModelResponse, RunOutputEvent, TeamRunOutputEvent]]:
         # Start function call
         function_call_timer = Timer()
         function_call_timer.start()
@@ -1196,7 +1157,7 @@ class Model(ABC):
             function_execution_result = function_call.execute()
         except AgentRunException as a_exc:
             # Update additional messages from function call
-            _handle_agent_exception(a_exc, additional_messages)
+            _handle_agent_exception(a_exc, additional_input)
             # Set function call success to False if an exception occurred
         except Exception as e:
             log_error(f"Error executing function {function_call.function.name}: {e}")
@@ -1210,14 +1171,14 @@ class Model(ABC):
         # Process function call output
         function_call_output: str = ""
 
-        if isinstance(function_call.result, (GeneratorType, collections.abc.Iterator)):
-            for item in function_call.result:
+        if isinstance(function_execution_result.result, (GeneratorType, collections.abc.Iterator)):
+            for item in function_execution_result.result:
                 # This function yields agent/team run events
-                if isinstance(item, tuple(get_args(RunResponseEvent))) or isinstance(
-                    item, tuple(get_args(TeamRunResponseEvent))
+                if isinstance(item, tuple(get_args(RunOutputEvent))) or isinstance(
+                    item, tuple(get_args(TeamRunOutputEvent))
                 ):
                     # We only capture content events
-                    if isinstance(item, RunResponseContentEvent) or isinstance(item, TeamRunResponseContentEvent):
+                    if isinstance(item, RunContentEvent) or isinstance(item, TeamRunContentEvent):
                         if item.content is not None and isinstance(item.content, BaseModel):
                             function_call_output += item.content.model_dump_json()
                         else:
@@ -1227,6 +1188,9 @@ class Model(ABC):
                         if function_call.function.show_result:
                             yield ModelResponse(content=item.content)
 
+                        if isinstance(item, CustomEvent):
+                            function_call_output += str(item)
+
                     # Yield the event itself to bubble it up
                     yield item
 
@@ -1235,13 +1199,33 @@ class Model(ABC):
                     if function_call.function.show_result:
                         yield ModelResponse(content=str(item))
         else:
-            function_call_output = str(function_call.result)
+            from agno.tools.function import ToolResult
+
+            if isinstance(function_execution_result.result, ToolResult):
+                # Extract content and media from ToolResult
+                tool_result = function_execution_result.result
+                function_call_output = tool_result.content
+
+                # Transfer media from ToolResult to FunctionExecutionResult
+                if tool_result.images:
+                    function_execution_result.images = tool_result.images
+                if tool_result.videos:
+                    function_execution_result.videos = tool_result.videos
+                if tool_result.audios:
+                    function_execution_result.audios = tool_result.audios
+            else:
+                function_call_output = str(function_execution_result.result) if function_execution_result.result else ""
+
             if function_call.function.show_result:
                 yield ModelResponse(content=function_call_output)
 
         # Create and yield function call result
         function_call_result = self.create_function_call_result(
-            function_call, success=function_call_success, output=function_call_output, timer=function_call_timer
+            function_call,
+            success=function_call_success,
+            output=function_call_output,
+            timer=function_call_timer,
+            function_execution_result=function_execution_result,
         )
         yield ModelResponse(
             content=f"{function_call.get_call_str()} completed in {function_call_timer.elapsed:.4f}s.",
@@ -1257,6 +1241,11 @@ class Model(ABC):
                 )
             ],
             event=ModelResponseEvent.tool_call_completed.value,
+            updated_session_state=function_execution_result.updated_session_state,
+            # Add media artifacts from function execution
+            images=function_execution_result.images,
+            videos=function_execution_result.videos,
+            audios=function_execution_result.audios,
         )
 
         # Add function call to function call results
@@ -1266,13 +1255,13 @@ class Model(ABC):
         self,
         function_calls: List[FunctionCall],
         function_call_results: List[Message],
-        additional_messages: Optional[List[Message]] = None,
+        additional_input: Optional[List[Message]] = None,
         current_function_call_count: int = 0,
         function_call_limit: Optional[int] = None,
-    ) -> Iterator[Union[ModelResponse, RunResponseEvent, TeamRunResponseEvent]]:
+    ) -> Iterator[Union[ModelResponse, RunOutputEvent, TeamRunOutputEvent]]:
         # Additional messages from function calls that will be added to the function call results
-        if additional_messages is None:
-            additional_messages = []
+        if additional_input is None:
+            additional_input = []
 
         for fc in function_calls:
             if function_call_limit is not None:
@@ -1358,17 +1347,17 @@ class Model(ABC):
                 continue
 
             yield from self.run_function_call(
-                function_call=fc, function_call_results=function_call_results, additional_messages=additional_messages
+                function_call=fc, function_call_results=function_call_results, additional_input=additional_input
             )
 
         # Add any additional messages at the end
-        if additional_messages:
-            function_call_results.extend(additional_messages)
+        if additional_input:
+            function_call_results.extend(additional_input)
 
     async def arun_function_call(
         self,
         function_call: FunctionCall,
-    ) -> Tuple[Union[bool, AgentRunException], Timer, FunctionCall]:
+    ) -> Tuple[Union[bool, AgentRunException], Timer, FunctionCall, FunctionExecutionResult]:
         """Run a single function call and return its success status, timer, and the FunctionCall object."""
         from inspect import isasyncgenfunction, iscoroutine, iscoroutinefunction
 
@@ -1402,20 +1391,20 @@ class Model(ABC):
             raise e
 
         function_call_timer.stop()
-        return success, function_call_timer, function_call
+        return success, function_call_timer, function_call, result
 
     async def arun_function_calls(
         self,
         function_calls: List[FunctionCall],
         function_call_results: List[Message],
-        additional_messages: Optional[List[Message]] = None,
+        additional_input: Optional[List[Message]] = None,
         current_function_call_count: int = 0,
         function_call_limit: Optional[int] = None,
         skip_pause_check: bool = False,
-    ) -> AsyncIterator[Union[ModelResponse, RunResponseEvent, TeamRunResponseEvent]]:
+    ) -> AsyncIterator[Union[ModelResponse, RunOutputEvent, TeamRunOutputEvent]]:
         # Additional messages from function calls that will be added to the function call results
-        if additional_messages is None:
-            additional_messages = []
+        if additional_input is None:
+            additional_input = []
 
         function_calls_to_run = []
         for fc in function_calls:
@@ -1548,33 +1537,35 @@ class Model(ABC):
                 raise result
 
             # Unpack result
-            function_call_success, function_call_timer, fc = result
+            function_call_success, function_call_timer, function_call, function_execution_result = result
+
+            updated_session_state = function_execution_result.updated_session_state
 
             # Handle AgentRunException
             if isinstance(function_call_success, AgentRunException):
                 a_exc = function_call_success
                 # Update additional messages from function call
-                _handle_agent_exception(a_exc, additional_messages)
+                _handle_agent_exception(a_exc, additional_input)
                 # Set function call success to False if an exception occurred
                 function_call_success = False
 
             # Process function call output
             function_call_output: str = ""
-            if isinstance(fc.result, (GeneratorType, collections.abc.Iterator)):
-                for item in fc.result:
+            if isinstance(function_call.result, (GeneratorType, collections.abc.Iterator)):
+                for item in function_call.result:
                     # This function yields agent/team run events
-                    if isinstance(item, tuple(get_args(RunResponseEvent))) or isinstance(
-                        item, tuple(get_args(TeamRunResponseEvent))
+                    if isinstance(item, tuple(get_args(RunOutputEvent))) or isinstance(
+                        item, tuple(get_args(TeamRunOutputEvent))
                     ):
                         # We only capture content events
-                        if isinstance(item, RunResponseContentEvent) or isinstance(item, TeamRunResponseContentEvent):
+                        if isinstance(item, RunContentEvent) or isinstance(item, TeamRunContentEvent):
                             if item.content is not None and isinstance(item.content, BaseModel):
                                 function_call_output += item.content.model_dump_json()
                             else:
                                 # Capture output
                                 function_call_output += item.content or ""
 
-                            if fc.function.show_result:
+                            if function_call.function.show_result:
                                 yield ModelResponse(content=item.content)
                                 continue
 
@@ -1582,43 +1573,66 @@ class Model(ABC):
                         yield item
                     else:
                         function_call_output += str(item)
-                        if fc.function.show_result:
+                        if function_call.function.show_result:
                             yield ModelResponse(content=str(item))
-            elif isinstance(fc.result, (AsyncGeneratorType, collections.abc.AsyncIterator)):
-                async for item in fc.result:
+            elif isinstance(function_call.result, (AsyncGeneratorType, collections.abc.AsyncIterator)):
+                async for item in function_call.result:
                     # This function yields agent/team run events
-                    if isinstance(item, tuple(get_args(RunResponseEvent))) or isinstance(
-                        item, tuple(get_args(TeamRunResponseEvent))
+                    if isinstance(item, tuple(get_args(RunOutputEvent))) or isinstance(
+                        item, tuple(get_args(TeamRunOutputEvent))
                     ):
                         # We only capture content events
-                        if isinstance(item, RunResponseContentEvent) or isinstance(item, TeamRunResponseContentEvent):
+                        if isinstance(item, RunContentEvent) or isinstance(item, TeamRunContentEvent):
                             if item.content is not None and isinstance(item.content, BaseModel):
                                 function_call_output += item.content.model_dump_json()
                             else:
                                 # Capture output
                                 function_call_output += item.content or ""
 
-                            if fc.function.show_result:
+                            if function_call.function.show_result:
                                 yield ModelResponse(content=item.content)
                                 continue
 
+                            if isinstance(item, CustomEvent):
+                                function_call_output += str(item)
+
                         # Yield the event itself to bubble it up
                         yield item
+
+                    # Yield custom events emitted by the tool
                     else:
                         function_call_output += str(item)
-                        if fc.function.show_result:
+                        if function_call.function.show_result:
                             yield ModelResponse(content=str(item))
             else:
-                function_call_output = str(fc.result)
-                if fc.function.show_result:
+                from agno.tools.function import ToolResult
+
+                if isinstance(function_execution_result.result, ToolResult):
+                    tool_result = function_execution_result.result
+                    function_call_output = tool_result.content
+
+                    if tool_result.images:
+                        function_execution_result.images = tool_result.images
+                    if tool_result.videos:
+                        function_execution_result.videos = tool_result.videos
+                    if tool_result.audios:
+                        function_execution_result.audios = tool_result.audios
+                else:
+                    function_call_output = str(function_call.result)
+
+                if function_call.function.show_result:
                     yield ModelResponse(content=function_call_output)
 
             # Create and yield function call result
             function_call_result = self.create_function_call_result(
-                fc, success=function_call_success, output=function_call_output, timer=function_call_timer
+                function_call,
+                success=function_call_success,
+                output=function_call_output,
+                timer=function_call_timer,
+                function_execution_result=function_execution_result,
             )
             yield ModelResponse(
-                content=f"{fc.get_call_str()} completed in {function_call_timer.elapsed:.4f}s.",
+                content=f"{function_call.get_call_str()} completed in {function_call_timer.elapsed:.4f}s.",
                 tool_executions=[
                     ToolExecution(
                         tool_call_id=function_call_result.tool_call_id,
@@ -1631,14 +1645,18 @@ class Model(ABC):
                     )
                 ],
                 event=ModelResponseEvent.tool_call_completed.value,
+                updated_session_state=updated_session_state,
+                images=function_execution_result.images,
+                videos=function_execution_result.videos,
+                audios=function_execution_result.audios,
             )
 
             # Add function call result to function call results
             function_call_results.append(function_call_result)
 
         # Add any additional messages at the end
-        if additional_messages:
-            function_call_results.extend(additional_messages)
+        if additional_input:
+            function_call_results.extend(additional_input)
 
     def _prepare_function_calls(
         self,
@@ -1669,6 +1687,44 @@ class Model(ABC):
         if len(function_call_results) > 0:
             messages.extend(function_call_results)
 
+    def _handle_function_call_media(self, messages: List[Message], function_call_results: List[Message]) -> None:
+        """
+        Handle media artifacts from function calls by adding follow-up user messages for generated media if needed.
+        """
+        if not function_call_results:
+            return
+
+        # Collect all media artifacts from function calls
+        all_images: List[Image] = []
+        all_videos: List[Video] = []
+        all_audio: List[Audio] = []
+
+        for result_message in function_call_results:
+            if result_message.images:
+                all_images.extend(result_message.images)
+                # Remove images from tool message to avoid errors on the LLMs
+                result_message.images = None
+
+            if result_message.videos:
+                all_videos.extend(result_message.videos)
+                result_message.videos = None
+
+            if result_message.audio:
+                all_audio.extend(result_message.audio)
+                result_message.audio = None
+
+        # If we have media artifacts, add a follow-up "user" message instead of a "tool"
+        # message with the media artifacts which throws error for some models
+        if all_images or all_videos or all_audio:
+            media_message = Message(
+                role="user",
+                content="Take note of the following content",
+                images=all_images if all_images else None,
+                videos=all_videos if all_videos else None,
+                audio=all_audio if all_audio else None,
+            )
+            messages.append(media_message)
+
     def get_system_message_for_model(self, tools: Optional[List[Any]] = None) -> Optional[str]:
         return self.system_prompt