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

agno/workflow/parallel.py

@@ -7,6 +7,7 @@ from uuid import uuid4
 
 from agno.models.metrics import Metrics
 from agno.run.agent import RunOutputEvent
+from agno.run.base import RunContext
 from agno.run.team import TeamRunOutputEvent
 from agno.run.workflow import (
     ParallelExecutionCompletedEvent,
@@ -200,6 +201,7 @@ class Parallel:
         user_id: Optional[str] = None,
         workflow_run_response: Optional[WorkflowRunOutput] = None,
         store_executor_outputs: bool = True,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         workflow_session: Optional[WorkflowSession] = None,
         add_workflow_history_to_steps: Optional[bool] = False,
@@ -214,10 +216,14 @@ class Parallel:
         # Create individual session_state copies for each step to prevent race conditions
         session_state_copies = []
         for _ in range(len(self.steps)):
-            if session_state is not None:
-                session_state_copies.append(deepcopy(session_state))
+            # If using run context, no need to deepcopy the state. We want the direct reference.
+            if run_context is not None and run_context.session_state is not None:
+                session_state_copies.append(run_context.session_state)
             else:
-                session_state_copies.append({})
+                if session_state is not None:
+                    session_state_copies.append(deepcopy(session_state))
+                else:
+                    session_state_copies.append({})
 
         def execute_step_with_index(step_with_index):
             """Execute a single step and preserve its original index"""
@@ -235,6 +241,7 @@ class Parallel:
                     workflow_session=workflow_session,
                     add_workflow_history_to_steps=add_workflow_history_to_steps,
                     num_history_runs=num_history_runs,
+                    run_context=run_context,
                     session_state=step_session_state,
                 )  # type: ignore[union-attr]
                 return idx, step_result, step_session_state
@@ -288,7 +295,7 @@ class Parallel:
                         )
                     )
 
-        if session_state is not None:
+        if run_context is None and session_state is not None:
             merge_parallel_session_states(session_state, modified_session_states)
 
         # Sort by original index to preserve order
@@ -322,6 +329,7 @@ class Parallel:
         workflow_run_response: Optional[WorkflowRunOutput] = None,
         step_index: Optional[Union[int, tuple]] = None,
         store_executor_outputs: bool = True,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         parent_step_id: Optional[str] = None,
         workflow_session: Optional[WorkflowSession] = None,
@@ -338,10 +346,14 @@ class Parallel:
         # Create individual session_state copies for each step to prevent race conditions
         session_state_copies = []
         for _ in range(len(self.steps)):
-            if session_state is not None:
-                session_state_copies.append(deepcopy(session_state))
+            # If using run context, no need to deepcopy the state. We want the direct reference.
+            if run_context is not None and run_context.session_state is not None:
+                session_state_copies.append(run_context.session_state)
             else:
-                session_state_copies.append({})
+                if session_state is not None:
+                    session_state_copies.append(deepcopy(session_state))
+                else:
+                    session_state_copies.append({})
 
         # Considering both stream_events and stream_intermediate_steps (deprecated)
         stream_events = stream_events or stream_intermediate_steps
@@ -468,7 +480,7 @@ class Parallel:
                     logger.error(f"Future completion error: {e}")
 
         # Merge all session_state changes back into the original session_state
-        if session_state is not None:
+        if run_context is None and session_state is not None:
             merge_parallel_session_states(session_state, modified_session_states)
 
         # Flatten step_results - handle steps that return List[StepOutput] (like Condition/Loop)
@@ -509,6 +521,7 @@ class Parallel:
         user_id: Optional[str] = None,
         workflow_run_response: Optional[WorkflowRunOutput] = None,
         store_executor_outputs: bool = True,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         workflow_session: Optional[WorkflowSession] = None,
         add_workflow_history_to_steps: Optional[bool] = False,
@@ -523,10 +536,14 @@ class Parallel:
         # Create individual session_state copies for each step to prevent race conditions
         session_state_copies = []
         for _ in range(len(self.steps)):
-            if session_state is not None:
-                session_state_copies.append(deepcopy(session_state))
+            # If using run context, no need to deepcopy the state. We want the direct reference.
+            if run_context is not None and run_context.session_state is not None:
+                session_state_copies.append(run_context.session_state)
             else:
-                session_state_copies.append({})
+                if session_state is not None:
+                    session_state_copies.append(deepcopy(session_state))
+                else:
+                    session_state_copies.append({})
 
         async def execute_step_async_with_index(step_with_index):
             """Execute a single step asynchronously and preserve its original index"""
@@ -598,7 +615,7 @@ class Parallel:
                 log_debug(f"Parallel step {step_name} completed")
 
         # Smart merge all session_state changes back into the original session_state
-        if session_state is not None:
+        if run_context is None and session_state is not None:
             merge_parallel_session_states(session_state, modified_session_states)
 
         # Sort by original index to preserve order
@@ -632,6 +649,7 @@ class Parallel:
         workflow_run_response: Optional[WorkflowRunOutput] = None,
         step_index: Optional[Union[int, tuple]] = None,
         store_executor_outputs: bool = True,
+        run_context: Optional[RunContext] = None,
         session_state: Optional[Dict[str, Any]] = None,
         parent_step_id: Optional[str] = None,
         workflow_session: Optional[WorkflowSession] = None,
@@ -648,10 +666,14 @@ class Parallel:
         # Create individual session_state copies for each step to prevent race conditions
         session_state_copies = []
         for _ in range(len(self.steps)):
-            if session_state is not None:
-                session_state_copies.append(deepcopy(session_state))
+            # If using run context, no need to deepcopy the state. We want the direct reference.
+            if run_context is not None and run_context.session_state is not None:
+                session_state_copies.append(run_context.session_state)
             else:
-                session_state_copies.append({})
+                if session_state is not None:
+                    session_state_copies.append(deepcopy(session_state))
+                else:
+                    session_state_copies.append({})
 
         # Considering both stream_events and stream_intermediate_steps (deprecated)
         stream_events = stream_events or stream_intermediate_steps
@@ -705,6 +727,7 @@ class Parallel:
                     step_index=sub_step_index,
                     store_executor_outputs=store_executor_outputs,
                     session_state=step_session_state,
+                    run_context=run_context,
                     parent_step_id=parallel_step_id,
                     workflow_session=workflow_session,
                     add_workflow_history_to_steps=add_workflow_history_to_steps,
@@ -766,7 +789,7 @@ class Parallel:
         await asyncio.gather(*tasks, return_exceptions=True)
 
         # Merge all session_state changes back into the original session_state
-        if session_state is not None:
+        if run_context is None and session_state is not None:
             merge_parallel_session_states(session_state, modified_session_states)
 
         # Flatten step_results - handle steps that return List[StepOutput] (like Condition/Loop)

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,