agno 2.2.9__py3-none-any.whl → 2.2.11__py3-none-any.whl

agno/workflow/step.py

@@ -11,7 +11,10 @@ from agno.agent import Agent
 from agno.media import Audio, Image, Video
 from agno.models.metrics import Metrics
 from agno.run import RunContext
-from agno.run.agent import RunOutput
+from agno.run.agent import RunCompletedEvent, RunContentEvent, RunOutput
+from agno.run.base import BaseRunOutputEvent
+from agno.run.team import RunCompletedEvent as TeamRunCompletedEvent
+from agno.run.team import RunContentEvent as TeamRunContentEvent
 from agno.run.team import TeamRunOutput
 from agno.run.workflow import (
     StepCompletedEvent,
@@ -179,31 +182,37 @@ class Step:
         func: Callable,
         step_input: StepInput,
         session_state: Optional[Dict[str, Any]] = None,
+        run_context: Optional[RunContext] = None,
     ) -> Any:
         """Call custom function with session_state support if the function accepts it"""
+
+        kwargs: Dict[str, Any] = {}
+        if run_context is not None and self._function_has_run_context_param():
+            kwargs["run_context"] = run_context
         if session_state is not None and self._function_has_session_state_param():
-            return func(step_input, session_state)
-        else:
-            return func(step_input)
+            kwargs["session_state"] = session_state
+
+        return func(step_input, **kwargs)
 
     async def _acall_custom_function(
         self,
         func: Callable,
         step_input: StepInput,
         session_state: Optional[Dict[str, Any]] = None,
+        run_context: Optional[RunContext] = None,
     ) -> Any:
         """Call custom async function with session_state support if the function accepts it"""
 
+        kwargs: Dict[str, Any] = {}
+        if run_context is not None and self._function_has_run_context_param():
+            kwargs["run_context"] = run_context
+        if session_state is not None and self._function_has_session_state_param():
+            kwargs["session_state"] = session_state
+
         if _is_async_generator_function(func):
-            if session_state is not None and self._function_has_session_state_param():
-                return func(step_input, session_state)
-            else:
-                return func(step_input)
+            return func(step_input, **kwargs)
         else:
-            if session_state is not None and self._function_has_session_state_param():
-                return await func(step_input, session_state)
-            else:
-                return await func(step_input)
+            return await func(step_input, **kwargs)
 
     def execute(
         self,
@@ -227,8 +236,10 @@ class Step:
         if workflow_session:
             step_input.workflow_session = workflow_session
 
+        # Create session_state copy once to avoid duplication.
+        # Consider both run_context.session_state and session_state.
         if run_context is not None and run_context.session_state is not None:
-            session_state_copy = copy(run_context.session_state)
+            session_state_copy = run_context.session_state
         else:
             session_state_copy = copy(session_state) if session_state is not None else {}
 
@@ -247,13 +258,13 @@ class Step:
                                 self.active_executor,
                                 step_input,
                                 session_state_copy,  # type: ignore[arg-type]
+                                run_context,
                             ):  # type: ignore
-                                if (
-                                    hasattr(chunk, "content")
-                                    and chunk.content is not None
-                                    and isinstance(chunk.content, str)
-                                ):
-                                    content += chunk.content
+                                if isinstance(chunk, (BaseRunOutputEvent)):
+                                    if isinstance(chunk, (RunContentEvent, TeamRunContentEvent)):
+                                        content += chunk.content if chunk.content is not None else ""
+                                    elif isinstance(chunk, (RunCompletedEvent, TeamRunCompletedEvent)):
+                                        content = chunk.content if chunk.content is not None else ""
                                 else:
                                     content += str(chunk)
                                 if isinstance(chunk, StepOutput):
@@ -264,7 +275,7 @@ class Step:
                                 final_response = e.value
 
                         # Merge session_state changes back
-                        if session_state is not None:
+                        if run_context is None and session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         if final_response is not None:
@@ -273,10 +284,15 @@ class Step:
                             response = StepOutput(content=content)
                     else:
                         # Execute function with signature inspection for session_state support
-                        result = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                        result = self._call_custom_function(
+                            self.active_executor,  # type: ignore[arg-type]
+                            step_input,
+                            session_state_copy,
+                            run_context,
+                        )
 
                         # Merge session_state changes back
-                        if session_state is not None:
+                        if run_context is None and session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         # If function returns StepOutput, use it directly
@@ -334,12 +350,13 @@ class Step:
                             session_id=session_id,
                             user_id=user_id,
                             session_state=session_state_copy,  # Send a copy to the executor
+                            run_context=run_context,
                             **kwargs,
                         )
 
-                        if session_state is not None:
-                            # Update workflow session state
-                            merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        # Update workflow session state
+                        if run_context is None and session_state is not None:
+                            merge_dictionaries(session_state, session_state_copy)
 
                         if store_executor_outputs and workflow_run_response is not None:
                             self._store_executor_response(workflow_run_response, response)  # type: ignore
@@ -368,6 +385,17 @@ class Step:
 
         return StepOutput(content=f"Step {self.name} failed but skipped", success=False)
 
+    def _function_has_run_context_param(self) -> bool:
+        """Check if the custom function has a run_context parameter"""
+        if self._executor_type != "function":
+            return False
+
+        try:
+            sig = inspect.signature(self.active_executor)  # type: ignore
+            return "run_context" in sig.parameters
+        except Exception:
+            return False
+
     def _function_has_session_state_param(self) -> bool:
         """Check if the custom function has a session_state parameter"""
         if self._executor_type != "function":
@@ -430,9 +458,10 @@ class Step:
         if workflow_session:
             step_input.workflow_session = workflow_session
 
-        # Create session_state copy once to avoid duplication
+        # Create session_state copy once to avoid duplication.
+        # Consider both run_context.session_state and session_state.
         if run_context is not None and run_context.session_state is not None:
-            session_state_copy = copy(run_context.session_state)
+            session_state_copy = run_context.session_state
         else:
             session_state_copy = copy(session_state) if session_state is not None else {}
 
@@ -460,22 +489,24 @@ class Step:
 
                 if self._executor_type == "function":
                     log_debug(f"Executing function executor for step: {self.name}")
-
                     if _is_async_callable(self.active_executor) or _is_async_generator_function(self.active_executor):
                         raise ValueError("Cannot use async function with synchronous execution")
-
                     if _is_generator_function(self.active_executor):
                         log_debug("Function returned iterable, streaming events")
                         content = ""
                         try:
-                            iterator = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                            iterator = self._call_custom_function(
+                                self.active_executor,
+                                step_input,
+                                session_state_copy,
+                                run_context,
+                            )
                             for event in iterator:  # type: ignore
-                                if (
-                                    hasattr(event, "content")
-                                    and event.content is not None
-                                    and isinstance(event.content, str)
-                                ):
-                                    content += event.content
+                                if isinstance(event, (BaseRunOutputEvent)):
+                                    if isinstance(event, (RunContentEvent, TeamRunContentEvent)):
+                                        content += event.content if event.content is not None else ""
+                                    elif isinstance(event, (RunCompletedEvent, TeamRunCompletedEvent)):
+                                        content = event.content if event.content is not None else ""
                                 else:
                                     content += str(event)
                                 if isinstance(event, StepOutput):
@@ -491,7 +522,7 @@ class Step:
                                         yield enriched_event  # type: ignore[misc]
 
                             # Merge session_state changes back
-                            if session_state is not None:
+                            if run_context is None and session_state is not None:
                                 merge_dictionaries(session_state, session_state_copy)
 
                             if not final_response:
@@ -501,10 +532,15 @@ class Step:
                                 final_response = e.value
 
                     else:
-                        result = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                        result = self._call_custom_function(
+                            self.active_executor,  # type: ignore[arg-type]
+                            step_input,
+                            session_state_copy,
+                            run_context,
+                        )
 
                         # Merge session_state changes back
-                        if session_state is not None:
+                        if run_context is None and session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         if isinstance(result, StepOutput):
@@ -564,6 +600,7 @@ class Step:
                             stream=True,
                             stream_events=stream_events,
                             yield_run_response=True,
+                            run_context=run_context,
                             **kwargs,
                         )
 
@@ -577,9 +614,9 @@ class Step:
                             if stream_executor_events:
                                 yield enriched_event  # type: ignore[misc]
 
-                        if session_state is not None:
-                            # Update workflow session state
-                            merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        # Update workflow session state
+                        if run_context is None and session_state is not None:
+                            merge_dictionaries(session_state, session_state_copy)
 
                         if store_executor_outputs and workflow_run_response is not None:
                             self._store_executor_response(workflow_run_response, active_executor_run_response)  # type: ignore
@@ -640,6 +677,7 @@ class Step:
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
         workflow_run_response: Optional["WorkflowRunOutput"] = None,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         store_executor_outputs: bool = True,
         workflow_session: Optional["WorkflowSession"] = None,
@@ -655,8 +693,13 @@ class Step:
 
         if workflow_session:
             step_input.workflow_session = workflow_session
-        # Create session_state copy once to avoid duplication
-        session_state_copy = copy(session_state) if session_state is not None else {}
+
+        # Create session_state copy once to avoid duplication.
+        # Consider both run_context.session_state and session_state.
+        if run_context is not None and run_context.session_state is not None:
+            session_state_copy = run_context.session_state
+        else:
+            session_state_copy = copy(session_state) if session_state is not None else {}
 
         # Execute with retries
         for attempt in range(self.max_retries + 1):
@@ -672,15 +715,15 @@ class Step:
                                 iterator = self._call_custom_function(
                                     self.active_executor,
                                     step_input,
-                                    session_state_copy,  # type: ignore[arg-type]
-                                )  # type: ignore
+                                    session_state_copy,
+                                    run_context,
+                                )
                                 for chunk in iterator:  # type: ignore
-                                    if (
-                                        hasattr(chunk, "content")
-                                        and chunk.content is not None
-                                        and isinstance(chunk.content, str)
-                                    ):
-                                        content += chunk.content
+                                    if isinstance(chunk, (BaseRunOutputEvent)):
+                                        if isinstance(chunk, (RunContentEvent, TeamRunContentEvent)):
+                                            content += chunk.content if chunk.content is not None else ""
+                                        elif isinstance(chunk, (RunCompletedEvent, TeamRunCompletedEvent)):
+                                            content = chunk.content if chunk.content is not None else ""
                                     else:
                                         content += str(chunk)
                                     if isinstance(chunk, StepOutput):
@@ -690,15 +733,15 @@ class Step:
                                     iterator = await self._acall_custom_function(
                                         self.active_executor,
                                         step_input,
-                                        session_state_copy,  # type: ignore[arg-type]
-                                    )  # type: ignore
+                                        session_state_copy,
+                                        run_context,
+                                    )
                                     async for chunk in iterator:  # type: ignore
-                                        if (
-                                            hasattr(chunk, "content")
-                                            and chunk.content is not None
-                                            and isinstance(chunk.content, str)
-                                        ):
-                                            content += chunk.content
+                                        if isinstance(chunk, (BaseRunOutputEvent)):
+                                            if isinstance(chunk, (RunContentEvent, TeamRunContentEvent)):
+                                                content += chunk.content if chunk.content is not None else ""
+                                            elif isinstance(chunk, (RunCompletedEvent, TeamRunCompletedEvent)):
+                                                content = chunk.content if chunk.content is not None else ""
                                         else:
                                             content += str(chunk)
                                         if isinstance(chunk, StepOutput):
@@ -709,7 +752,7 @@ class Step:
                                 final_response = e.value
 
                         # Merge session_state changes back
-                        if session_state is not None:
+                        if run_context is None and session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         if final_response is not None:
@@ -719,13 +762,21 @@ class Step:
                     else:
                         if _is_async_callable(self.active_executor):
                             result = await self._acall_custom_function(
-                                self.active_executor, step_input, session_state_copy
-                            )  # type: ignore
+                                self.active_executor,
+                                step_input,
+                                session_state_copy,
+                                run_context,
+                            )
                         else:
-                            result = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                            result = self._call_custom_function(
+                                self.active_executor,  # type: ignore[arg-type]
+                                step_input,
+                                session_state_copy,
+                                run_context,
+                            )
 
                         # Merge session_state changes back
-                        if session_state is not None:
+                        if run_context is None and session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         # If function returns StepOutput, use it directly
@@ -784,12 +835,13 @@ class Step:
                             session_id=session_id,
                             user_id=user_id,
                             session_state=session_state_copy,
+                            run_context=run_context,
                             **kwargs,
                         )
 
-                        if session_state is not None:
-                            # Update workflow session state
-                            merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        # Update workflow session state
+                        if run_context is None and session_state is not None:
+                            merge_dictionaries(session_state, session_state_copy)
 
                         if store_executor_outputs and workflow_run_response is not None:
                             self._store_executor_response(workflow_run_response, response)  # type: ignore
@@ -827,6 +879,7 @@ class Step:
         stream_intermediate_steps: bool = False,
         stream_executor_events: bool = True,
         workflow_run_response: Optional["WorkflowRunOutput"] = None,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         step_index: Optional[Union[int, tuple]] = None,
         store_executor_outputs: bool = True,
@@ -843,8 +896,12 @@ class Step:
         if workflow_session:
             step_input.workflow_session = workflow_session
 
-        # Create session_state copy once to avoid duplication
-        session_state_copy = copy(session_state) if session_state is not None else {}
+        # Create session_state copy once to avoid duplication.
+        # Consider both run_context.session_state and session_state.
+        if run_context is not None and run_context.session_state is not None:
+            session_state_copy = run_context.session_state
+        else:
+            session_state_copy = copy(session_state) if session_state is not None else {}
 
         # Considering both stream_events and stream_intermediate_steps (deprecated)
         stream_events = stream_events or stream_intermediate_steps
@@ -878,15 +935,15 @@ class Step:
                         iterator = await self._acall_custom_function(
                             self.active_executor,
                             step_input,
-                            session_state_copy,  # type: ignore[arg-type]
-                        )  # type: ignore
+                            session_state_copy,
+                            run_context,
+                        )
                         async for event in iterator:  # type: ignore
-                            if (
-                                hasattr(event, "content")
-                                and event.content is not None
-                                and isinstance(event.content, str)
-                            ):
-                                content += event.content
+                            if isinstance(event, (BaseRunOutputEvent)):
+                                if isinstance(event, (RunContentEvent, TeamRunContentEvent)):
+                                    content += event.content if event.content is not None else ""
+                                elif isinstance(event, (RunCompletedEvent, TeamRunCompletedEvent)):
+                                    content = event.content if event.content is not None else ""
                             else:
                                 content += str(event)
                             if isinstance(event, StepOutput):
@@ -904,7 +961,12 @@ class Step:
                             final_response = StepOutput(content=content)
                     elif _is_async_callable(self.active_executor):
                         # It's a regular async function - await it
-                        result = await self._acall_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                        result = await self._acall_custom_function(
+                            self.active_executor,
+                            step_input,
+                            session_state_copy,
+                            run_context,
+                        )
                         if isinstance(result, StepOutput):
                             final_response = result
                         else:
@@ -912,14 +974,18 @@ class Step:
                     elif _is_generator_function(self.active_executor):
                         content = ""
                         # It's a regular generator function - iterate over it
-                        iterator = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                        iterator = self._call_custom_function(
+                            self.active_executor,
+                            step_input,
+                            session_state_copy,
+                            run_context,
+                        )
                         for event in iterator:  # type: ignore
-                            if (
-                                hasattr(event, "content")
-                                and event.content is not None
-                                and isinstance(event.content, str)
-                            ):
-                                content += event.content
+                            if isinstance(event, (BaseRunOutputEvent)):
+                                if isinstance(event, (RunContentEvent, TeamRunContentEvent)):
+                                    content += event.content if event.content is not None else ""
+                                elif isinstance(event, (RunCompletedEvent, TeamRunCompletedEvent)):
+                                    content = event.content if event.content is not None else ""
                             else:
                                 content += str(event)
                             if isinstance(event, StepOutput):
@@ -937,14 +1003,19 @@ class Step:
                             final_response = StepOutput(content=content)
                     else:
                         # It's a regular function - call it directly
-                        result = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                        result = self._call_custom_function(
+                            self.active_executor,  # type: ignore[arg-type]
+                            step_input,
+                            session_state_copy,
+                            run_context,
+                        )
                         if isinstance(result, StepOutput):
                             final_response = result
                         else:
                             final_response = StepOutput(content=str(result))
 
                     # Merge session_state changes back
-                    if session_state is not None:
+                    if run_context is None and session_state is not None:
                         merge_dictionaries(session_state, session_state_copy)
                 else:
                     # For agents and teams, prepare message with context
@@ -997,6 +1068,7 @@ class Step:
                             session_state=session_state_copy,
                             stream=True,
                             stream_events=stream_events,
+                            run_context=run_context,
                             yield_run_response=True,
                             **kwargs,
                         )
@@ -1011,9 +1083,9 @@ class Step:
                             if stream_executor_events:
                                 yield enriched_event  # type: ignore[misc]
 
-                        if session_state is not None:
-                            # Update workflow session state
-                            merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        # Update workflow session state
+                        if run_context is None and session_state is not None:
+                            merge_dictionaries(session_state, session_state_copy)
 
                         if store_executor_outputs and workflow_run_response is not None:
                             self._store_executor_response(workflow_run_response, active_executor_run_response)  # type: ignore

agno/workflow/steps.py

@@ -3,6 +3,7 @@ from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterator, List
 from uuid import uuid4
 
 from agno.run.agent import RunOutputEvent
+from agno.run.base import RunContext
 from agno.run.team import TeamRunOutputEvent
 from agno.run.workflow import (
     StepsExecutionCompletedEvent,
@@ -118,6 +119,7 @@ class Steps:
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
         workflow_run_response: Optional[WorkflowRunOutput] = None,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         store_executor_outputs: bool = True,
         workflow_session: Optional[WorkflowSession] = None,
@@ -151,6 +153,7 @@ class Steps:
                     user_id=user_id,
                     workflow_run_response=workflow_run_response,
                     store_executor_outputs=store_executor_outputs,
+                    run_context=run_context,
                     session_state=session_state,
                     workflow_session=workflow_session,
                     add_workflow_history_to_steps=add_workflow_history_to_steps,
@@ -204,6 +207,7 @@ class Steps:
         self,
         step_input: StepInput,
         workflow_run_response: WorkflowRunOutput,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
@@ -269,6 +273,7 @@ class Steps:
                     current_step_input,
                     session_id=session_id,
                     user_id=user_id,
+                    run_context=run_context,
                     session_state=session_state,
                     stream_events=stream_events,
                     stream_executor_events=stream_executor_events,
@@ -354,6 +359,7 @@ class Steps:
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
         workflow_run_response: Optional[WorkflowRunOutput] = None,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         store_executor_outputs: bool = True,
         workflow_session: Optional[WorkflowSession] = None,
@@ -387,6 +393,7 @@ class Steps:
                     user_id=user_id,
                     workflow_run_response=workflow_run_response,
                     store_executor_outputs=store_executor_outputs,
+                    run_context=run_context,
                     session_state=session_state,
                     workflow_session=workflow_session,
                     add_workflow_history_to_steps=add_workflow_history_to_steps,
@@ -439,6 +446,7 @@ class Steps:
         self,
         step_input: StepInput,
         workflow_run_response: WorkflowRunOutput,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
@@ -504,6 +512,7 @@ class Steps:
                     current_step_input,
                     session_id=session_id,
                     user_id=user_id,
+                    run_context=run_context,
                     session_state=session_state,
                     stream_events=stream_events,
                     stream_executor_events=stream_executor_events,

agno/workflow/types.py

@@ -17,6 +17,7 @@ from agno.utils.media import (
     reconstruct_videos,
 )
 from agno.utils.serialize import json_serializer
+from agno.utils.timer import Timer
 
 
 @dataclass
@@ -405,12 +406,18 @@ class WorkflowMetrics:
     """Complete metrics for a workflow execution"""
 
     steps: Dict[str, StepMetrics]
+    # Timer utility for tracking execution time
+    timer: Optional[Timer] = None
+    # Total workflow execution time
+    duration: Optional[float] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary"""
-        return {
+        result: Dict[str, Any] = {
             "steps": {name: step.to_dict() for name, step in self.steps.items()},
+            "duration": self.duration,
         }
+        return result
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "WorkflowMetrics":
@@ -419,8 +426,20 @@ class WorkflowMetrics:
 
         return cls(
             steps=steps,
+            duration=data.get("duration"),
         )
 
+    def start_timer(self):
+        if self.timer is None:
+            self.timer = Timer()
+        self.timer.start()
+
+    def stop_timer(self, set_duration: bool = True):
+        if self.timer is not None:
+            self.timer.stop()
+            if set_duration:
+                self.duration = self.timer.elapsed
+
 
 @dataclass
 class WebSocketHandler: