agno 2.0.2__py3-none-any.whl → 2.0.4__py3-none-any.whl

agno/workflow/step.py

@@ -164,6 +164,38 @@ class Step:
             return response.metrics
         return None
 
+    def _call_custom_function(
+        self,
+        func: Callable,
+        step_input: StepInput,
+        session_state: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """Call custom function with session_state support if the function accepts it"""
+        if session_state is not None and self._function_has_session_state_param():
+            return func(step_input, session_state)
+        else:
+            return func(step_input)
+
+    async def _acall_custom_function(
+        self,
+        func: Callable,
+        step_input: StepInput,
+        session_state: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """Call custom async function with session_state support if the function accepts it"""
+        import inspect
+
+        if inspect.isasyncgenfunction(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)
+        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)
+
     def execute(
         self,
         step_input: StepInput,
@@ -191,8 +223,11 @@ 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.active_executor(step_input):  # type: ignore
+                            for chunk in self._call_custom_function(
+                                self.active_executor, step_input, session_state_copy
+                            ):  # type: ignore
                                 if (
                                     hasattr(chunk, "content")
                                     and chunk.content is not None
@@ -208,13 +243,22 @@ class Step:
                             if hasattr(e, "value") and isinstance(e.value, StepOutput):
                                 final_response = e.value
 
+                        # Merge session_state changes back
+                        if session_state_copy and session_state:
+                            merge_dictionaries(session_state, session_state_copy)
+
                         if final_response is not None:
                             response = final_response
                         else:
                             response = StepOutput(content=content)
                     else:
-                        # Execute function directly with StepInput
-                        result = self.active_executor(step_input)  # type: ignore
+                        # 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:
+                            merge_dictionaries(session_state, session_state_copy)
 
                         # If function returns StepOutput, use it directly
                         if isinstance(result, StepOutput):
@@ -291,6 +335,19 @@ class Step:
 
         return StepOutput(content=f"Step {self.name} failed but skipped", success=False)
 
+    def _function_has_session_state_param(self) -> bool:
+        """Check if the custom function has a session_state parameter"""
+        if self._executor_type != "function":
+            return False
+
+        try:
+            from inspect import signature
+
+            sig = signature(self.active_executor)  # type: ignore
+            return "session_state" in sig.parameters
+        except Exception:
+            return False
+
     def execute_stream(
         self,
         step_input: StepInput,
@@ -338,8 +395,10 @@ 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:
-                            for event in self.active_executor(step_input):  # type: ignore
+                            iterator = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                            for event in iterator:  # type: ignore
                                 if (
                                     hasattr(event, "content")
                                     and event.content is not None
@@ -353,6 +412,11 @@ class Step:
                                     break
                                 else:
                                     yield event  # type: ignore[misc]
+
+                            # Merge session_state changes back
+                            if session_state_copy and session_state:
+                                merge_dictionaries(session_state, session_state_copy)
+
                             if not final_response:
                                 final_response = StepOutput(content=content)
                         except StopIteration as e:
@@ -360,7 +424,13 @@ class Step:
                                 final_response = e.value
 
                     else:
-                        result = self.active_executor(step_input)  # type: ignore
+                        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:
+                            merge_dictionaries(session_state, session_state_copy)
+
                         if isinstance(result, StepOutput):
                             final_response = result
                         else:
@@ -505,9 +575,13 @@ class Step:
                     ):
                         content = ""
                         final_response = None
+                        session_state_copy = copy(session_state) if session_state else None
                         try:
                             if inspect.isgeneratorfunction(self.active_executor):
-                                for chunk in self.active_executor(step_input):  # type: ignore
+                                iterator = self._call_custom_function(
+                                    self.active_executor, step_input, session_state_copy
+                                )  # type: ignore
+                                for chunk in iterator:  # type: ignore
                                     if (
                                         hasattr(chunk, "content")
                                         and chunk.content is not None
@@ -520,7 +594,10 @@ class Step:
                                         final_response = chunk
                             else:
                                 if inspect.isasyncgenfunction(self.active_executor):
-                                    async for chunk in self.active_executor(step_input):  # type: ignore
+                                    iterator = await self._acall_custom_function(
+                                        self.active_executor, step_input, session_state_copy
+                                    )  # type: ignore
+                                    async for chunk in iterator:  # type: ignore
                                         if (
                                             hasattr(chunk, "content")
                                             and chunk.content is not None
@@ -536,15 +613,26 @@ class Step:
                             if hasattr(e, "value") and isinstance(e.value, StepOutput):
                                 final_response = e.value
 
+                        # Merge session_state changes back
+                        if session_state_copy and session_state:
+                            merge_dictionaries(session_state, session_state_copy)
+
                         if final_response is not None:
                             response = final_response
                         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.active_executor(step_input)  # type: ignore
+                            result = await self._acall_custom_function(
+                                self.active_executor, step_input, session_state_copy
+                            )  # type: ignore
                         else:
-                            result = self.active_executor(step_input)  # type: ignore
+                            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:
+                            merge_dictionaries(session_state, session_state_copy)
 
                         # If function returns StepOutput, use it directly
                         if isinstance(result, StepOutput):
@@ -662,11 +750,16 @@ 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 = ""
                         # It's an async generator - iterate over it
-                        async for event in self.active_executor(step_input):  # type: ignore
+                        iterator = await self._acall_custom_function(
+                            self.active_executor, step_input, session_state_copy
+                        )  # type: ignore
+                        async for event in iterator:  # type: ignore
                             if (
                                 hasattr(event, "content")
                                 and event.content is not None
@@ -684,7 +777,7 @@ class Step:
                             final_response = StepOutput(content=content)
                     elif inspect.iscoroutinefunction(self.active_executor):
                         # It's a regular async function - await it
-                        result = await self.active_executor(step_input)  # type: ignore
+                        result = await self._acall_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
                         if isinstance(result, StepOutput):
                             final_response = result
                         else:
@@ -692,7 +785,8 @@ class Step:
                     elif inspect.isgeneratorfunction(self.active_executor):
                         content = ""
                         # It's a regular generator function - iterate over it
-                        for event in self.active_executor(step_input):  # type: ignore
+                        iterator = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
+                        for event in iterator:  # type: ignore
                             if (
                                 hasattr(event, "content")
                                 and event.content is not None
@@ -710,11 +804,15 @@ class Step:
                             final_response = StepOutput(content=content)
                     else:
                         # It's a regular function - call it directly
-                        result = self.active_executor(step_input)  # type: ignore
+                        result = self._call_custom_function(self.active_executor, step_input, session_state_copy)  # type: ignore
                         if isinstance(result, StepOutput):
                             final_response = result
                         else:
                             final_response = StepOutput(content=str(result))
+
+                    # Merge session_state changes back
+                    if session_state_copy and session_state:
+                        merge_dictionaries(session_state, session_state_copy)
                 else:
                     # For agents and teams, prepare message with context
                     message = self._prepare_message(

agno/workflow/workflow.py

@@ -51,6 +51,7 @@ from agno.run.workflow import (
 )
 from agno.session.workflow import WorkflowSession
 from agno.team.team import Team
+from agno.utils.common import is_typed_dict, validate_typed_dict
 from agno.utils.log import (
     log_debug,
     log_warning,
@@ -217,14 +218,18 @@ class Workflow:
 
     def _validate_input(
         self, input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel, List[Message]]]
-    ) -> Optional[BaseModel]:
+    ) -> Optional[Union[str, List, Dict, Message, BaseModel]]:
         """Parse and validate input against input_schema if provided"""
         if self.input_schema is None:
-            return None
+            return input  # Return input unchanged if no schema is set
 
         if input is None:
             raise ValueError("Input required when input_schema is set")
 
+        # Handle Message objects - extract content
+        if isinstance(input, Message):
+            input = input.content  # type: ignore
+
         # If input is a string, convert it to a dict
         if isinstance(input, str):
             import json
@@ -238,8 +243,6 @@ class Workflow:
         if isinstance(input, BaseModel):
             if isinstance(input, self.input_schema):
                 try:
-                    # Re-validate to catch any field validation errors
-                    input.model_validate(input.model_dump())
                     return input
                 except Exception as e:
                     raise ValueError(f"BaseModel validation failed: {str(e)}")
@@ -250,8 +253,13 @@ class Workflow:
         # Case 2: Message is a dict
         elif isinstance(input, dict):
             try:
-                validated_model = self.input_schema(**input)
-                return validated_model
+                # Check if the schema is a TypedDict
+                if is_typed_dict(self.input_schema):
+                    validated_dict = validate_typed_dict(input, self.input_schema)
+                    return validated_dict
+                else:
+                    validated_model = self.input_schema(**input)
+                    return validated_model
             except Exception as e:
                 raise ValueError(f"Failed to parse dict into {self.input_schema.__name__}: {str(e)}")
 
@@ -1924,10 +1932,7 @@ class Workflow:
     ) -> Union[WorkflowRunOutput, Iterator[WorkflowRunOutputEvent]]:
         """Execute the workflow synchronously with optional streaming"""
 
-        validated_input = self._validate_input(input)
-        if validated_input is not None:
-            input = validated_input
-
+        input = self._validate_input(input)
         if background:
             raise RuntimeError("Background execution is not supported for sync run()")
 
@@ -2059,9 +2064,7 @@ class Workflow:
     ) -> Union[WorkflowRunOutput, AsyncIterator[WorkflowRunOutputEvent]]:
         """Execute the workflow synchronously with optional streaming"""
 
-        validated_input = self._validate_input(input)
-        if validated_input is not None:
-            input = validated_input
+        input = self._validate_input(input)
 
         websocket_handler = None
         if websocket:

{agno-2.0.2.dist-info → agno-2.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agno
-Version: 2.0.2
+Version: 2.0.4
 Summary: Agno: a lightweight library for building Multi-Agent Systems
 Author-email: Ashpreet Bedi <ashpreet@agno.com>
 Project-URL: homepage, https://agno.com