agno 2.1.0__py3-none-any.whl → 2.1.2__py3-none-any.whl

agno/workflow/parallel.py

@@ -1,5 +1,6 @@
 import asyncio
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterator, List, Optional, Union
 from uuid import uuid4
@@ -13,6 +14,7 @@ from agno.run.workflow import (
     WorkflowRunOutput,
     WorkflowRunOutputEvent,
 )
+from agno.utils.merge_dict import merge_parallel_session_states
 from agno.utils.log import log_debug, logger
 from agno.workflow.condition import Condition
 from agno.workflow.step import Step
@@ -205,9 +207,20 @@ class Parallel:
 
         self._prepare_steps()
 
+        # 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))
+            else:
+                session_state_copies.append({})
+
         def execute_step_with_index(step_with_index):
             """Execute a single step and preserve its original index"""
             idx, step = step_with_index
+            # Use the individual session_state copy for this step
+            step_session_state = session_state_copies[idx]
+            
             try:
                 step_result = step.execute(
                     step_input,
@@ -215,9 +228,9 @@ class Parallel:
                     user_id=user_id,
                     workflow_run_response=workflow_run_response,
                     store_executor_outputs=store_executor_outputs,
-                    session_state=session_state,
+                    session_state=step_session_state,
                 )  # type: ignore[union-attr]
-                return idx, step_result
+                return idx, step_result, step_session_state
             except Exception as exc:
                 parallel_step_name = getattr(step, "name", f"step_{idx}")
                 logger.error(f"Parallel step {parallel_step_name} failed: {exc}")
@@ -229,6 +242,7 @@ class Parallel:
                         success=False,
                         error=str(exc),
                     ),
+                    step_session_state,
                 )
 
         # Use index to preserve order
@@ -241,12 +255,14 @@ class Parallel:
                 for indexed_step in indexed_steps
             }
 
-            # Collect results
+            # Collect results and modified session_state copies
             results_with_indices = []
+            modified_session_states = []
             for future in as_completed(future_to_index):
                 try:
-                    index, result = future.result()
+                    index, result, modified_session_state = future.result()
                     results_with_indices.append((index, result))
+                    modified_session_states.append(modified_session_state)
                     step_name = getattr(self.steps[index], "name", f"step_{index}")
                     log_debug(f"Parallel step {step_name} completed")
                 except Exception as e:
@@ -265,6 +281,9 @@ class Parallel:
                         )
                     )
 
+        if session_state is not None:
+            merge_parallel_session_states(session_state, modified_session_states)
+
         # Sort by original index to preserve order
         results_with_indices.sort(key=lambda x: x[0])
         results = [result for _, result in results_with_indices]
@@ -304,6 +323,14 @@ class Parallel:
 
         self._prepare_steps()
 
+        # 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))
+            else:
+                session_state_copies.append({})
+
         if stream_intermediate_steps and workflow_run_response:
             # Yield parallel step started event
             yield ParallelExecutionStartedEvent(
@@ -321,6 +348,9 @@ class Parallel:
         def execute_step_stream_with_index(step_with_index):
             """Execute a single step with streaming and preserve its original index"""
             idx, step = step_with_index
+            # Use the individual session_state copy for this step
+            step_session_state = session_state_copies[idx]
+            
             try:
                 step_events = []
 
@@ -342,11 +372,11 @@ class Parallel:
                     workflow_run_response=workflow_run_response,
                     step_index=sub_step_index,
                     store_executor_outputs=store_executor_outputs,
-                    session_state=session_state,
+                    session_state=step_session_state,
                     parent_step_id=parallel_step_id,
                 ):
                     step_events.append(event)
-                return idx, step_events
+                return idx, step_events, step_session_state
             except Exception as exc:
                 parallel_step_name = getattr(step, "name", f"step_{idx}")
                 logger.error(f"Parallel step {parallel_step_name} streaming failed: {exc}")
@@ -360,12 +390,14 @@ class Parallel:
                             error=str(exc),
                         )
                     ],
+                    step_session_state,
                 )
 
         # Use index to preserve order
         indexed_steps = list(enumerate(self.steps))
         all_events_with_indices = []
         step_results = []
+        modified_session_states = []
 
         with ThreadPoolExecutor(max_workers=len(self.steps)) as executor:
             # Submit all tasks with their original indices
@@ -374,11 +406,12 @@ class Parallel:
                 for indexed_step in indexed_steps
             }
 
-            # Collect results as they complete
+            # Collect results and modified session_state copies
             for future in as_completed(future_to_index):
                 try:
-                    index, events = future.result()
+                    index, events, modified_session_state = future.result()
                     all_events_with_indices.append((index, events))
+                    modified_session_states.append(modified_session_state)
 
                     # Extract StepOutput from events for the final result
                     step_outputs = [event for event in events if isinstance(event, StepOutput)]
@@ -400,6 +433,10 @@ class Parallel:
                     all_events_with_indices.append((index, [error_event]))
                     step_results.append(error_event)
 
+        # Merge all session_state changes back into the original session_state
+        if session_state is not None:
+            merge_parallel_session_states(session_state, modified_session_states)
+
         # Sort events by original index to preserve order
         all_events_with_indices.sort(key=lambda x: x[0])
 
@@ -456,9 +493,20 @@ class Parallel:
 
         self._prepare_steps()
 
+        # 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))
+            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"""
             idx, step = step_with_index
+            # Use the individual session_state copy for this step
+            step_session_state = session_state_copies[idx]
+            
             try:
                 inner_step_result = await step.aexecute(
                     step_input,
@@ -466,9 +514,9 @@ class Parallel:
                     user_id=user_id,
                     workflow_run_response=workflow_run_response,
                     store_executor_outputs=store_executor_outputs,
-                    session_state=session_state,
+                    session_state=step_session_state,
                 )  # type: ignore[union-attr]
-                return idx, inner_step_result
+                return idx, inner_step_result, step_session_state
             except Exception as exc:
                 parallel_step_name = getattr(step, "name", f"step_{idx}")
                 logger.error(f"Parallel step {parallel_step_name} failed: {exc}")
@@ -480,6 +528,7 @@ class Parallel:
                         success=False,
                         error=str(exc),
                     ),
+                    step_session_state,
                 )
 
         # Use index to preserve order
@@ -493,6 +542,7 @@ class Parallel:
 
         # Process results and handle exceptions, preserving order
         processed_results_with_indices = []
+        modified_session_states = []
         for i, result in enumerate(results_with_indices):
             if isinstance(result, Exception):
                 step_name = getattr(self.steps[i], "name", f"step_{i}")
@@ -508,12 +558,19 @@ class Parallel:
                         ),
                     )
                 )
+                # Still collect the session state copy for failed steps
+                modified_session_states.append(session_state_copies[i])
             else:
-                index, step_result = result  # type: ignore[misc]
+                index, step_result, modified_session_state = result  # type: ignore[misc]
                 processed_results_with_indices.append((index, step_result))
+                modified_session_states.append(modified_session_state)
                 step_name = getattr(self.steps[index], "name", f"step_{index}")
                 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:
+            merge_parallel_session_states(session_state, modified_session_states)
+
         # Sort by original index to preserve order
         processed_results_with_indices.sort(key=lambda x: x[0])
         results = [result for _, result in processed_results_with_indices]
@@ -553,6 +610,14 @@ class Parallel:
 
         self._prepare_steps()
 
+        # 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))
+            else:
+                session_state_copies.append({})
+
         if stream_intermediate_steps and workflow_run_response:
             # Yield parallel step started event
             yield ParallelExecutionStartedEvent(
@@ -570,6 +635,9 @@ class Parallel:
         async def execute_step_stream_async_with_index(step_with_index):
             """Execute a single step with async streaming and preserve its original index"""
             idx, step = step_with_index
+            # Use the individual session_state copy for this step
+            step_session_state = session_state_copies[idx]
+            
             try:
                 step_events = []
 
@@ -591,11 +659,11 @@ class Parallel:
                     workflow_run_response=workflow_run_response,
                     step_index=sub_step_index,
                     store_executor_outputs=store_executor_outputs,
-                    session_state=session_state,
+                    session_state=step_session_state,
                     parent_step_id=parallel_step_id,
                 ):  # type: ignore[union-attr]
                     step_events.append(event)
-                return idx, step_events
+                return idx, step_events, step_session_state
             except Exception as e:
                 parallel_step_name = getattr(step, "name", f"step_{idx}")
                 logger.error(f"Parallel step {parallel_step_name} async streaming failed: {e}")
@@ -609,12 +677,14 @@ class Parallel:
                             error=str(e),
                         )
                     ],
+                    step_session_state,
                 )
 
         # Use index to preserve order
         indexed_steps = list(enumerate(self.steps))
         all_events_with_indices = []
         step_results = []
+        modified_session_states = []
 
         # Create tasks for all steps with their indices
         tasks = [execute_step_stream_async_with_index(indexed_step) for indexed_step in indexed_steps]
@@ -635,9 +705,11 @@ class Parallel:
                 )
                 all_events_with_indices.append((i, [error_event]))
                 step_results.append(error_event)
+                modified_session_states.append(session_state_copies[i])
             else:
-                index, events = result  # type: ignore[misc]
+                index, events, modified_session_state = result  # type: ignore[misc]
                 all_events_with_indices.append((index, events))
+                modified_session_states.append(modified_session_state)
 
                 # Extract StepOutput from events for the final result
                 step_outputs = [event for event in events if isinstance(event, StepOutput)]
@@ -647,6 +719,10 @@ class Parallel:
                 step_name = getattr(self.steps[index], "name", f"step_{index}")
                 log_debug(f"Parallel step {step_name} async streaming completed")
 
+        # Merge all session_state changes back into the original session_state
+        if session_state is not None:
+            merge_parallel_session_states(session_state, modified_session_states)
+
         # Sort events by original index to preserve order
         all_events_with_indices.sort(key=lambda x: x[0])
 

agno/workflow/step.py

@@ -211,6 +211,8 @@ class Step:
         if step_input.previous_step_outputs:
             step_input.previous_step_content = step_input.get_last_step_content()
 
+        session_state_copy = copy(session_state) if session_state is not None else {}
+
         # Execute with retries
         for attempt in range(self.max_retries + 1):
             try:
@@ -223,7 +225,6 @@ class Step:
                     if inspect.isgeneratorfunction(self.active_executor):
                         content = ""
                         final_response = None
-                        session_state_copy = copy(session_state) if session_state else None
                         try:
                             for chunk in self._call_custom_function(
                                 self.active_executor, step_input, session_state_copy
@@ -244,7 +245,7 @@ class Step:
                                 final_response = e.value
 
                         # Merge session_state changes back
-                        if session_state_copy and session_state:
+                        if session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         if final_response is not None:
@@ -253,11 +254,10 @@ class Step:
                             response = StepOutput(content=content)
                     else:
                         # Execute function with signature inspection for session_state support
-                        session_state_copy = copy(session_state) if session_state else None
                         result = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
 
                         # Merge session_state changes back
-                        if session_state_copy and session_state:
+                        if session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         # If function returns StepOutput, use it directly
@@ -292,7 +292,6 @@ class Step:
                         if isinstance(self.active_executor, Team):
                             kwargs["store_member_responses"] = True
 
-                        session_state_copy = copy(session_state)
                         response = self.active_executor.run(  # type: ignore
                             input=message,  # type: ignore
                             images=images,
@@ -305,8 +304,9 @@ class Step:
                             **kwargs,
                         )
 
-                        # Update workflow session state
-                        merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        if session_state is not None:
+                            # Update workflow session state
+                            merge_dictionaries(session_state, session_state_copy)  # type: ignore
 
                         if store_executor_outputs and workflow_run_response is not None:
                             self._store_executor_response(workflow_run_response, response)  # type: ignore
@@ -365,6 +365,9 @@ class Step:
         if step_input.previous_step_outputs:
             step_input.previous_step_content = step_input.get_last_step_content()
 
+        # Create session_state copy once to avoid duplication
+        session_state_copy = copy(session_state) if session_state is not None else {}
+
         # Emit StepStartedEvent
         if stream_intermediate_steps and workflow_run_response:
             yield StepStartedEvent(
@@ -395,7 +398,6 @@ class Step:
                     if inspect.isgeneratorfunction(self.active_executor):
                         log_debug("Function returned iterable, streaming events")
                         content = ""
-                        session_state_copy = copy(session_state) if session_state else None
                         try:
                             iterator = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
                             for event in iterator:  # type: ignore
@@ -414,7 +416,7 @@ class Step:
                                     yield event  # type: ignore[misc]
 
                             # Merge session_state changes back
-                            if session_state_copy and session_state:
+                            if session_state is not None:
                                 merge_dictionaries(session_state, session_state_copy)
 
                             if not final_response:
@@ -424,11 +426,10 @@ class Step:
                                 final_response = e.value
 
                     else:
-                        session_state_copy = copy(session_state) if session_state else None
                         result = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
 
                         # Merge session_state changes back
-                        if session_state_copy and session_state:
+                        if session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         if isinstance(result, StepOutput):
@@ -462,7 +463,6 @@ class Step:
                         if isinstance(self.active_executor, Team):
                             kwargs["store_member_responses"] = True
 
-                        session_state_copy = copy(session_state)
                         response_stream = self.active_executor.run(  # type: ignore[call-overload, misc]
                             input=message,
                             images=images,
@@ -493,8 +493,9 @@ class Step:
                                 break
                             yield event  # type: ignore[misc]
 
-                        # Update workflow session state
-                        merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        if session_state is not None:
+                            # Update workflow session state
+                            merge_dictionaries(session_state, session_state_copy)  # type: ignore
 
                         if store_executor_outputs and workflow_run_response is not None:
                             self._store_executor_response(workflow_run_response, active_executor_run_response)  # type: ignore
@@ -565,6 +566,9 @@ class Step:
         if step_input.previous_step_outputs:
             step_input.previous_step_content = step_input.get_last_step_content()
 
+        # Create session_state copy once to avoid duplication
+        session_state_copy = copy(session_state) if session_state is not None else {}
+
         # Execute with retries
         for attempt in range(self.max_retries + 1):
             try:
@@ -576,7 +580,6 @@ class Step:
                     ):
                         content = ""
                         final_response = None
-                        session_state_copy = copy(session_state) if session_state else None
                         try:
                             if inspect.isgeneratorfunction(self.active_executor):
                                 iterator = self._call_custom_function(
@@ -615,7 +618,7 @@ class Step:
                                 final_response = e.value
 
                         # Merge session_state changes back
-                        if session_state_copy and session_state:
+                        if session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         if final_response is not None:
@@ -623,7 +626,6 @@ class Step:
                         else:
                             response = StepOutput(content=content)
                     else:
-                        session_state_copy = copy(session_state) if session_state else None
                         if inspect.iscoroutinefunction(self.active_executor):
                             result = await self._acall_custom_function(
                                 self.active_executor, step_input, session_state_copy
@@ -632,7 +634,7 @@ class Step:
                             result = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
 
                         # Merge session_state changes back
-                        if session_state_copy and session_state:
+                        if session_state is not None:
                             merge_dictionaries(session_state, session_state_copy)
 
                         # If function returns StepOutput, use it directly
@@ -668,7 +670,6 @@ class Step:
                         if isinstance(self.active_executor, Team):
                             kwargs["store_member_responses"] = True
 
-                        session_state_copy = copy(session_state)
                         response = await self.active_executor.arun(  # type: ignore
                             input=message,  # type: ignore
                             images=images,
@@ -681,8 +682,9 @@ class Step:
                             **kwargs,
                         )
 
-                        # Update workflow session state
-                        merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        if session_state is not None:
+                            # Update workflow session state
+                            merge_dictionaries(session_state, session_state_copy)  # type: ignore
 
                         if store_executor_outputs and workflow_run_response is not None:
                             self._store_executor_response(workflow_run_response, response)  # type: ignore
@@ -728,6 +730,9 @@ class Step:
         if step_input.previous_step_outputs:
             step_input.previous_step_content = step_input.get_last_step_content()
 
+        # Create session_state copy once to avoid duplication
+        session_state_copy = copy(session_state) if session_state is not None else {}
+
         if stream_intermediate_steps and workflow_run_response:
             # Emit StepStartedEvent
             yield StepStartedEvent(
@@ -751,8 +756,6 @@ class Step:
                     log_debug(f"Executing async function executor for step: {self.name}")
                     import inspect
 
-                    session_state_copy = copy(session_state) if session_state else None
-
                     # Check if the function is an async generator
                     if inspect.isasyncgenfunction(self.active_executor):
                         content = ""
@@ -812,7 +815,7 @@ class Step:
                             final_response = StepOutput(content=str(result))
 
                     # Merge session_state changes back
-                    if session_state_copy and session_state:
+                    if session_state is not None:
                         merge_dictionaries(session_state, session_state_copy)
                 else:
                     # For agents and teams, prepare message with context
@@ -840,7 +843,6 @@ class Step:
                         if isinstance(self.active_executor, Team):
                             kwargs["store_member_responses"] = True
 
-                        session_state_copy = copy(session_state)
                         response_stream = self.active_executor.arun(  # type: ignore
                             input=message,
                             images=images,
@@ -871,8 +873,9 @@ class Step:
                                 break
                             yield event  # type: ignore[misc]
 
-                        # Update workflow session state
-                        merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        if session_state is not None:
+                            # Update workflow session state
+                            merge_dictionaries(session_state, session_state_copy)  # type: ignore
 
                         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/types.py

@@ -1,3 +1,4 @@
+import json
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
@@ -8,6 +9,7 @@ from pydantic import BaseModel
 from agno.media import Audio, File, Image, Video
 from agno.models.metrics import Metrics
 from agno.utils.log import log_warning
+from agno.utils.serialize import json_serializer
 
 
 @dataclass
@@ -443,9 +445,7 @@ class WebSocketHandler:
             else:
                 data = {"type": "message", "content": str(event)}
 
-            import json
-
-            await self.websocket.send_text(self.format_sse_event(json.dumps(data)))
+            await self.websocket.send_text(self.format_sse_event(json.dumps(data, default=json_serializer)))
 
         except Exception as e:
             log_warning(f"Failed to handle WebSocket event: {e}")
@@ -466,9 +466,7 @@ class WebSocketHandler:
             return
 
         try:
-            import json
-
-            await self.websocket.send_text(self.format_sse_event(json.dumps(data)))
+            await self.websocket.send_text(self.format_sse_event(json.dumps(data, default=json_serializer)))
         except Exception as e:
             log_warning(f"Failed to send WebSocket dict: {e}")
 

agno/workflow/workflow.py

@@ -875,10 +875,10 @@ class Workflow:
             return func(**call_kwargs)
         except TypeError as e:
             # If signature inspection fails, fall back to original method
-            logger.warning(
-                f"Async function signature inspection failed: {e}. Falling back to original calling convention."
+            logger.error(
+                f"Function signature inspection failed: {e}. Falling back to original calling convention."
             )
-            return func(**call_kwargs)
+            return func(**kwargs)
 
     def _execute(
         self,
@@ -2158,6 +2158,7 @@ class Workflow:
                     additional_data=additional_data,
                     user_id=user_id,
                     session_id=session_id,
+                    session_state=session_state,
                     audio=audio,
                     images=images,
                     videos=videos,
@@ -2176,6 +2177,7 @@ class Workflow:
                     additional_data=additional_data,
                     user_id=user_id,
                     session_id=session_id,
+                    session_state=session_state,
                     audio=audio,
                     images=images,
                     videos=videos,