agno 2.2.8__py3-none-any.whl → 2.2.10__py3-none-any.whl

agno/workflow/router.py

@@ -4,6 +4,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 (
     RouterExecutionCompletedEvent,
@@ -171,6 +172,7 @@ class Router:
         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,
@@ -185,7 +187,10 @@ class Router:
         self._prepare_steps()
 
         # Route to appropriate steps
-        steps_to_execute = self._route_steps(step_input, session_state)
+        if run_context is not None and run_context.session_state is not None:
+            steps_to_execute = self._route_steps(step_input, session_state=run_context.session_state)
+        else:
+            steps_to_execute = self._route_steps(step_input, session_state=session_state)
         log_debug(f"Router {self.name}: Selected {len(steps_to_execute)} steps to execute")
 
         if not steps_to_execute:
@@ -209,6 +214,7 @@ class Router:
                     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,
@@ -266,6 +272,7 @@ class Router:
         step_input: StepInput,
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         stream_events: bool = False,
         stream_intermediate_steps: bool = False,
@@ -286,7 +293,10 @@ class Router:
         router_step_id = str(uuid4())
 
         # Route to appropriate steps
-        steps_to_execute = self._route_steps(step_input, session_state)
+        if run_context is not None and run_context.session_state is not None:
+            steps_to_execute = self._route_steps(step_input, session_state=run_context.session_state)
+        else:
+            steps_to_execute = self._route_steps(step_input, session_state=session_state)
         log_debug(f"Router {self.name}: Selected {len(steps_to_execute)} steps to execute")
 
         # Considering both stream_events and stream_intermediate_steps (deprecated)
@@ -341,6 +351,7 @@ class Router:
                     workflow_run_response=workflow_run_response,
                     step_index=step_index,
                     store_executor_outputs=store_executor_outputs,
+                    run_context=run_context,
                     session_state=session_state,
                     parent_step_id=router_step_id,
                     workflow_session=workflow_session,
@@ -425,6 +436,7 @@ class Router:
         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,
@@ -439,7 +451,10 @@ class Router:
         self._prepare_steps()
 
         # Route to appropriate steps
-        steps_to_execute = await self._aroute_steps(step_input, session_state)
+        if run_context is not None and run_context.session_state is not None:
+            steps_to_execute = await self._aroute_steps(step_input, session_state=run_context.session_state)
+        else:
+            steps_to_execute = await self._aroute_steps(step_input, session_state=session_state)
         log_debug(f"Router {self.name} selected: {len(steps_to_execute)} steps to execute")
 
         if not steps_to_execute:
@@ -464,6 +479,7 @@ class Router:
                     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,
@@ -523,6 +539,7 @@ class Router:
         step_input: StepInput,
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         stream_events: bool = False,
         stream_intermediate_steps: bool = False,
@@ -543,7 +560,10 @@ class Router:
         router_step_id = str(uuid4())
 
         # Route to appropriate steps
-        steps_to_execute = await self._aroute_steps(step_input, session_state)
+        if run_context is not None and run_context.session_state is not None:
+            steps_to_execute = await self._aroute_steps(step_input, session_state=run_context.session_state)
+        else:
+            steps_to_execute = await self._aroute_steps(step_input, session_state=session_state)
         log_debug(f"Router {self.name} selected: {len(steps_to_execute)} steps to execute")
 
         # Considering both stream_events and stream_intermediate_steps (deprecated)
@@ -600,6 +620,7 @@ class Router:
                     workflow_run_response=workflow_run_response,
                     step_index=step_index,
                     store_executor_outputs=store_executor_outputs,
+                    run_context=run_context,
                     session_state=session_state,
                     parent_step_id=router_step_id,
                     workflow_session=workflow_session,

agno/workflow/step.py

@@ -11,7 +11,9 @@ 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, RunOutput, RunContentEvent
+from agno.run.base import BaseRunOutputEvent
+from agno.run.team import RunCompletedEvent as TeamRunCompletedEvent, RunContentEvent as TeamRunContentEvent
 from agno.run.team import TeamRunOutput
 from agno.run.workflow import (
     StepCompletedEvent,
@@ -179,31 +181,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 +235,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 +257,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 +274,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 +283,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 +349,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 +384,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 +457,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 +488,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 +521,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 +531,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 +599,7 @@ class Step:
                             stream=True,
                             stream_events=stream_events,
                             yield_run_response=True,
+                            run_context=run_context,
                             **kwargs,
                         )
 
@@ -577,9 +613,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 +676,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 +692,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 +714,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 +732,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 +751,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 +761,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 +834,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 +878,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 +895,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 +934,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 +960,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 +973,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 +1002,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 +1067,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 +1082,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