agno 2.0.1__py3-none-any.whl → 2.3.0__py3-none-any.whl

agno/workflow/step.py

@@ -1,15 +1,21 @@
 import inspect
+import warnings
 from copy import copy
 from dataclasses import dataclass
-from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterator, List, Optional, Union
+from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterator, List, Optional, Union, cast
 from uuid import uuid4
 
 from pydantic import BaseModel
+from typing_extensions import TypeGuard
 
 from agno.agent import Agent
 from agno.media import Audio, Image, Video
+from agno.models.message import Message
 from agno.models.metrics import Metrics
-from agno.run.agent import RunOutput
+from agno.run import RunContext
+from agno.run.agent import RunContentEvent, RunOutput
+from agno.run.base import BaseRunOutputEvent
+from agno.run.team import RunContentEvent as TeamRunContentEvent
 from agno.run.team import TeamRunOutput
 from agno.run.workflow import (
     StepCompletedEvent,
@@ -17,8 +23,11 @@ from agno.run.workflow import (
     WorkflowRunOutput,
     WorkflowRunOutputEvent,
 )
+from agno.session.agent import AgentSession
+from agno.session.team import TeamSession
+from agno.session.workflow import WorkflowSession
 from agno.team import Team
-from agno.utils.log import log_debug, logger, use_agent_logger, use_team_logger, use_workflow_logger
+from agno.utils.log import log_debug, log_warning, logger, use_agent_logger, use_team_logger, use_workflow_logger
 from agno.utils.merge_dict import merge_dictionaries
 from agno.workflow.types import StepInput, StepOutput, StepType
 
@@ -60,6 +69,9 @@ class Step:
     # If False, only warn about missing inputs
     strict_input_validation: bool = False
 
+    add_workflow_history: Optional[bool] = None
+    num_history_runs: int = 3
+
     _retry_count: int = 0
 
     def __init__(
@@ -74,6 +86,8 @@ class Step:
         timeout_seconds: Optional[int] = None,
         skip_on_failure: bool = False,
         strict_input_validation: bool = False,
+        add_workflow_history: Optional[bool] = None,
+        num_history_runs: int = 3,
     ):
         # Auto-detect name for function executors if not provided
         if name is None and executor is not None:
@@ -93,6 +107,8 @@ class Step:
         self.timeout_seconds = timeout_seconds
         self.skip_on_failure = skip_on_failure
         self.strict_input_validation = strict_input_validation
+        self.add_workflow_history = add_workflow_history
+        self.num_history_runs = num_history_runs
         self.step_id = step_id
 
         if step_id is None:
@@ -164,14 +180,55 @@ class Step:
             return response.metrics
         return None
 
+    def _call_custom_function(
+        self,
+        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():
+            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):
+            return func(step_input, **kwargs)
+        else:
+            return await func(step_input, **kwargs)
+
     def execute(
         self,
         step_input: StepInput,
         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,
+        add_workflow_history_to_steps: Optional[bool] = False,
+        num_history_runs: int = 3,
     ) -> StepOutput:
         """Execute the step with StepInput, returning final StepOutput (non-streaming)"""
         log_debug(f"Executing step: {self.name}")
@@ -179,46 +236,89 @@ class Step:
         if step_input.previous_step_outputs:
             step_input.previous_step_content = step_input.get_last_step_content()
 
+        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 = 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):
             try:
                 response: Union[RunOutput, TeamRunOutput, StepOutput]
                 if self._executor_type == "function":
-                    if inspect.iscoroutinefunction(self.active_executor) or inspect.isasyncgenfunction(
-                        self.active_executor
-                    ):
+                    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 inspect.isgeneratorfunction(self.active_executor):
+                    if _is_generator_function(self.active_executor):
                         content = ""
                         final_response = None
                         try:
-                            for chunk in self.active_executor(step_input):  # type: ignore
-                                if (
-                                    hasattr(chunk, "content")
-                                    and chunk.content is not None
-                                    and isinstance(chunk.content, str)
-                                ):
-                                    content += chunk.content
+                            for chunk in self._call_custom_function(
+                                self.active_executor,
+                                step_input,
+                                session_state_copy,  # type: ignore[arg-type]
+                                run_context,
+                            ):  # type: ignore
+                                if isinstance(chunk, (BaseRunOutputEvent)):
+                                    if (
+                                        isinstance(chunk, (RunContentEvent, TeamRunContentEvent))
+                                        and chunk.content is not None
+                                    ):
+                                        # Its a regular chunk of content
+                                        if isinstance(chunk.content, str):
+                                            content += chunk.content
+                                        # Its the BaseModel object, set it as the content. Replace any previous content.
+                                        # There should be no previous str content at this point
+                                        elif isinstance(chunk.content, BaseModel):
+                                            content = chunk.content  # type: ignore[assignment]
+                                        else:
+                                            # Case when parse_response is False and the content is a dict
+                                            content += str(chunk.content)
+                                elif isinstance(chunk, (RunOutput, TeamRunOutput)):
+                                    # This is the final response from the agent/team
+                                    content = chunk.content  # type: ignore[assignment]
+                                # If the chunk is a StepOutput, use it as the final response
+                                elif isinstance(chunk, StepOutput):
+                                    final_response = chunk
+                                    break
+                                # Non Agent/Team data structure that was yielded
                                 else:
                                     content += str(chunk)
-                                if isinstance(chunk, StepOutput):
-                                    final_response = chunk
 
                         except StopIteration as e:
                             if hasattr(e, "value") and isinstance(e.value, StepOutput):
                                 final_response = e.value
 
+                        # Merge session_state changes back
+                        if run_context is None and session_state is not None:
+                            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
+                        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 run_context is None and session_state is not None:
+                            merge_dictionaries(session_state, session_state_copy)
 
                         # If function returns StepOutput, use it directly
                         if isinstance(result, StepOutput):
                             response = result
+                        elif isinstance(result, (RunOutput, TeamRunOutput)):
+                            response = StepOutput(content=result.content)
                         else:
                             response = StepOutput(content=str(result))
                 else:
@@ -248,9 +348,22 @@ class Step:
                         if isinstance(self.active_executor, Team):
                             kwargs["store_member_responses"] = True
 
-                        session_state_copy = copy(session_state)
+                        num_history_runs = self.num_history_runs if self.num_history_runs else num_history_runs
+
+                        use_history = (
+                            self.add_workflow_history
+                            if self.add_workflow_history is not None
+                            else add_workflow_history_to_steps
+                        )
+
+                        final_message = message
+                        if use_history and workflow_session:
+                            history_messages = workflow_session.get_workflow_history_context(num_runs=num_history_runs)
+                            if history_messages:
+                                final_message = f"{history_messages}{message}"
+
                         response = self.active_executor.run(  # type: ignore
-                            input=message,  # type: ignore
+                            input=final_message,  # type: ignore
                             images=images,
                             videos=videos,
                             audio=audios,
@@ -258,11 +371,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,
                         )
 
                         # Update workflow session state
-                        merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        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
@@ -291,25 +406,97 @@ 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":
+            return False
+
+        try:
+            sig = inspect.signature(self.active_executor)  # type: ignore
+            return "session_state" in sig.parameters
+        except Exception:
+            return False
+
+    def _enrich_event_with_context(
+        self,
+        event: Any,
+        workflow_run_response: Optional["WorkflowRunOutput"] = None,
+        step_index: Optional[Union[int, tuple]] = None,
+    ) -> Any:
+        """Enrich event with step and workflow context information"""
+        if workflow_run_response is None:
+            return event
+        if hasattr(event, "workflow_id"):
+            event.workflow_id = workflow_run_response.workflow_id
+        if hasattr(event, "workflow_run_id"):
+            event.workflow_run_id = workflow_run_response.run_id
+        if hasattr(event, "step_id"):
+            event.step_id = self.step_id
+        if hasattr(event, "step_name") and self.name is not None:
+            if getattr(event, "step_name", None) is None:
+                event.step_name = self.name
+        # Only set step_index if it's not already set (preserve parallel.py's tuples)
+        if hasattr(event, "step_index") and step_index is not None:
+            if event.step_index is None:
+                event.step_index = step_index
+
+        return event
+
     def execute_stream(
         self,
         step_input: StepInput,
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
+        stream_events: bool = False,
         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,
         parent_step_id: Optional[str] = None,
+        workflow_session: Optional["WorkflowSession"] = None,
+        add_workflow_history_to_steps: Optional[bool] = False,
+        num_history_runs: int = 3,
     ) -> Iterator[Union[WorkflowRunOutputEvent, StepOutput]]:
         """Execute the step with event-driven streaming support"""
 
         if step_input.previous_step_outputs:
             step_input.previous_step_content = step_input.get_last_step_content()
 
+        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 = 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)
+        if stream_intermediate_steps is not None:
+            warnings.warn(
+                "The 'stream_intermediate_steps' parameter is deprecated and will be removed in future versions. Use 'stream_events' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        stream_events = stream_events or stream_intermediate_steps
+
         # Emit StepStartedEvent
-        if stream_intermediate_steps and workflow_run_response:
+        if stream_events and workflow_run_response:
             yield StepStartedEvent(
                 run_id=workflow_run_response.run_id or "",
                 workflow_name=workflow_run_response.workflow_name or "",
@@ -329,30 +516,48 @@ class Step:
 
                 if self._executor_type == "function":
                     log_debug(f"Executing function executor for step: {self.name}")
-
-                    if inspect.iscoroutinefunction(self.active_executor) or inspect.isasyncgenfunction(
-                        self.active_executor
-                    ):
+                    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 inspect.isgeneratorfunction(self.active_executor):
+                    if _is_generator_function(self.active_executor):
                         log_debug("Function returned iterable, streaming events")
                         content = ""
                         try:
-                            for event in self.active_executor(step_input):  # type: ignore
-                                if (
-                                    hasattr(event, "content")
-                                    and event.content is not None
-                                    and isinstance(event.content, str)
-                                ):
-                                    content += event.content
-                                else:
-                                    content += str(event)
-                                if isinstance(event, StepOutput):
+                            iterator = self._call_custom_function(
+                                self.active_executor,
+                                step_input,
+                                session_state_copy,
+                                run_context,
+                            )
+                            for event in iterator:  # type: ignore
+                                if isinstance(event, (BaseRunOutputEvent)):
+                                    if (
+                                        isinstance(event, (RunContentEvent, TeamRunContentEvent))
+                                        and event.content is not None
+                                    ):
+                                        if isinstance(event.content, str):
+                                            content += event.content
+                                        elif isinstance(event.content, BaseModel):
+                                            content = event.content  # type: ignore[assignment]
+                                        else:
+                                            content = str(event.content)
+                                    # Only yield executor events if stream_executor_events is True
+                                    if stream_executor_events:
+                                        enriched_event = self._enrich_event_with_context(
+                                            event, workflow_run_response, step_index
+                                        )
+                                        yield enriched_event  # type: ignore[misc]
+                                elif isinstance(event, (RunOutput, TeamRunOutput)):
+                                    content = event.content  # type: ignore[assignment]
+                                elif isinstance(event, StepOutput):
                                     final_response = event
                                     break
                                 else:
-                                    yield event  # type: ignore[misc]
+                                    content += str(event)
+
+                            # Merge session_state changes back
+                            if run_context is None and session_state is not None:
+                                merge_dictionaries(session_state, session_state_copy)
+
                             if not final_response:
                                 final_response = StepOutput(content=content)
                         except StopIteration as e:
@@ -360,9 +565,21 @@ class Step:
                                 final_response = e.value
 
                     else:
-                        result = self.active_executor(step_input)  # 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 run_context is None and session_state is not None:
+                            merge_dictionaries(session_state, session_state_copy)
+
                         if isinstance(result, StepOutput):
                             final_response = result
+                        elif isinstance(result, (RunOutput, TeamRunOutput)):
+                            final_response = StepOutput(content=result.content)
                         else:
                             final_response = StepOutput(content=str(result))
                         log_debug("Function returned non-iterable, created StepOutput")
@@ -392,9 +609,22 @@ class Step:
                         if isinstance(self.active_executor, Team):
                             kwargs["store_member_responses"] = True
 
-                        session_state_copy = copy(session_state)
+                        num_history_runs = self.num_history_runs if self.num_history_runs else num_history_runs
+
+                        use_history = (
+                            self.add_workflow_history
+                            if self.add_workflow_history is not None
+                            else add_workflow_history_to_steps
+                        )
+
+                        final_message = message
+                        if use_history and workflow_session:
+                            history_messages = workflow_session.get_workflow_history_context(num_runs=num_history_runs)
+                            if history_messages:
+                                final_message = f"{history_messages}{message}"
+
                         response_stream = self.active_executor.run(  # type: ignore[call-overload, misc]
-                            input=message,
+                            input=final_message,
                             images=images,
                             videos=videos,
                             audio=audios,
@@ -403,16 +633,9 @@ class Step:
                             user_id=user_id,
                             session_state=session_state_copy,  # Send a copy to the executor
                             stream=True,
-                            stream_intermediate_steps=stream_intermediate_steps,
-                            # Pass workflow context directly via kwargs
-                            workflow_context={
-                                "workflow_id": workflow_run_response.workflow_id if workflow_run_response else None,
-                                "workflow_run_id": workflow_run_response.run_id if workflow_run_response else None,
-                                "step_id": self.step_id,
-                                "step_name": self.name,
-                                "step_index": step_index,
-                            },
+                            stream_events=stream_events,
                             yield_run_response=True,
+                            run_context=run_context,
                             **kwargs,
                         )
 
@@ -421,15 +644,21 @@ class Step:
                             if isinstance(event, RunOutput) or isinstance(event, TeamRunOutput):
                                 active_executor_run_response = event
                                 break
-                            yield event  # type: ignore[misc]
+                            # Only yield executor events if stream_executor_events is True
+                            if stream_executor_events:
+                                enriched_event = self._enrich_event_with_context(
+                                    event, workflow_run_response, step_index
+                                )
+                                yield enriched_event  # type: ignore[misc]
 
                         # Update workflow session state
-                        merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        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
 
-                        final_response = self._process_step_output(active_executor_run_response)  # type: ignore
+                        final_response = active_executor_run_response  # type: ignore
 
                     else:
                         raise ValueError(f"Unsupported executor type: {self._executor_type}")
@@ -443,10 +672,11 @@ class Step:
                 use_workflow_logger()
 
                 # Yield the step output
+                final_response = self._process_step_output(final_response)
                 yield final_response
 
                 # Emit StepCompletedEvent
-                if stream_intermediate_steps and workflow_run_response:
+                if stream_events and workflow_run_response:
                     yield StepCompletedEvent(
                         run_id=workflow_run_response.run_id or "",
                         workflow_name=workflow_run_response.workflow_name or "",
@@ -484,8 +714,12 @@ 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,
+        add_workflow_history_to_steps: Optional[bool] = False,
+        num_history_runs: int = 3,
     ) -> StepOutput:
         """Execute the step with StepInput, returning final StepOutput (non-streaming)"""
         logger.info(f"Executing async step (non-streaming): {self.name}")
@@ -494,61 +728,118 @@ class Step:
         if step_input.previous_step_outputs:
             step_input.previous_step_content = step_input.get_last_step_content()
 
+        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 = 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):
             try:
                 if self._executor_type == "function":
-                    import inspect
-
-                    if inspect.isgeneratorfunction(self.active_executor) or inspect.isasyncgenfunction(
+                    if _is_generator_function(self.active_executor) or _is_async_generator_function(
                         self.active_executor
                     ):
                         content = ""
                         final_response = None
                         try:
-                            if inspect.isgeneratorfunction(self.active_executor):
-                                for chunk in self.active_executor(step_input):  # type: ignore
-                                    if (
-                                        hasattr(chunk, "content")
-                                        and chunk.content is not None
-                                        and isinstance(chunk.content, str)
-                                    ):
-                                        content += chunk.content
-                                    else:
-                                        content += str(chunk)
-                                    if isinstance(chunk, StepOutput):
-                                        final_response = chunk
-                            else:
-                                if inspect.isasyncgenfunction(self.active_executor):
-                                    async for chunk in self.active_executor(step_input):  # type: ignore
+                            if _is_generator_function(self.active_executor):
+                                iterator = self._call_custom_function(
+                                    self.active_executor,
+                                    step_input,
+                                    session_state_copy,
+                                    run_context,
+                                )
+                                for chunk in iterator:  # type: ignore
+                                    if isinstance(chunk, (BaseRunOutputEvent)):
                                         if (
-                                            hasattr(chunk, "content")
+                                            isinstance(chunk, (RunContentEvent, TeamRunContentEvent))
                                             and chunk.content is not None
-                                            and isinstance(chunk.content, str)
                                         ):
-                                            content += chunk.content
+                                            if isinstance(chunk.content, str):
+                                                content += chunk.content
+                                            elif isinstance(chunk.content, BaseModel):
+                                                content = chunk.content  # type: ignore[assignment]
+                                            else:
+                                                content = str(chunk.content)
+                                    elif isinstance(chunk, (RunOutput, TeamRunOutput)):
+                                        content = chunk.content  # type: ignore[assignment]
+                                    elif isinstance(chunk, StepOutput):
+                                        final_response = chunk
+                                        break
+                                    else:
+                                        content += str(chunk)
+
+                            else:
+                                if _is_async_generator_function(self.active_executor):
+                                    iterator = await self._acall_custom_function(
+                                        self.active_executor,
+                                        step_input,
+                                        session_state_copy,
+                                        run_context,
+                                    )
+                                    async for chunk in iterator:  # type: ignore
+                                        if isinstance(chunk, (BaseRunOutputEvent)):
+                                            if (
+                                                isinstance(chunk, (RunContentEvent, TeamRunContentEvent))
+                                                and chunk.content is not None
+                                            ):
+                                                if isinstance(chunk.content, str):
+                                                    content += chunk.content
+                                                elif isinstance(chunk.content, BaseModel):
+                                                    content = chunk.content  # type: ignore[assignment]
+                                                else:
+                                                    content = str(chunk.content)
+                                        elif isinstance(chunk, (RunOutput, TeamRunOutput)):
+                                            content = chunk.content  # type: ignore[assignment]
+                                        elif isinstance(chunk, StepOutput):
+                                            final_response = chunk
+                                            break
                                         else:
                                             content += str(chunk)
-                                        if isinstance(chunk, StepOutput):
-                                            final_response = chunk
 
                         except StopIteration as e:
                             if hasattr(e, "value") and isinstance(e.value, StepOutput):
                                 final_response = e.value
 
+                        # Merge session_state changes back
+                        if run_context is None and session_state is not None:
+                            merge_dictionaries(session_state, session_state_copy)
+
                         if final_response is not None:
                             response = final_response
                         else:
                             response = StepOutput(content=content)
                     else:
-                        if inspect.iscoroutinefunction(self.active_executor):
-                            result = await self.active_executor(step_input)  # type: ignore
+                        if _is_async_callable(self.active_executor):
+                            result = await self._acall_custom_function(
+                                self.active_executor,
+                                step_input,
+                                session_state_copy,
+                                run_context,
+                            )
                         else:
-                            result = self.active_executor(step_input)  # 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 run_context is None and session_state is not None:
+                            merge_dictionaries(session_state, session_state_copy)
 
                         # If function returns StepOutput, use it directly
                         if isinstance(result, StepOutput):
                             response = result
+                        elif isinstance(result, (RunOutput, TeamRunOutput)):
+                            response = StepOutput(content=result.content)
                         else:
                             response = StepOutput(content=str(result))
 
@@ -579,9 +870,22 @@ class Step:
                         if isinstance(self.active_executor, Team):
                             kwargs["store_member_responses"] = True
 
-                        session_state_copy = copy(session_state)
+                        num_history_runs = self.num_history_runs if self.num_history_runs else num_history_runs
+
+                        use_history = (
+                            self.add_workflow_history
+                            if self.add_workflow_history is not None
+                            else add_workflow_history_to_steps
+                        )
+
+                        final_message = message
+                        if use_history and workflow_session:
+                            history_messages = workflow_session.get_workflow_history_context(num_runs=num_history_runs)
+                            if history_messages:
+                                final_message = f"{history_messages}{message}"
+
                         response = await self.active_executor.arun(  # type: ignore
-                            input=message,  # type: ignore
+                            input=final_message,  # type: ignore
                             images=images,
                             videos=videos,
                             audio=audios,
@@ -589,11 +893,13 @@ class Step:
                             session_id=session_id,
                             user_id=user_id,
                             session_state=session_state_copy,
+                            run_context=run_context,
                             **kwargs,
                         )
 
                         # Update workflow session state
-                        merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        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
@@ -627,19 +933,44 @@ class Step:
         step_input: StepInput,
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
+        stream_events: bool = False,
         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,
         parent_step_id: Optional[str] = None,
+        workflow_session: Optional["WorkflowSession"] = None,
+        add_workflow_history_to_steps: Optional[bool] = False,
+        num_history_runs: int = 3,
     ) -> AsyncIterator[Union[WorkflowRunOutputEvent, StepOutput]]:
         """Execute the step with event-driven streaming support"""
 
         if step_input.previous_step_outputs:
             step_input.previous_step_content = step_input.get_last_step_content()
 
-        if stream_intermediate_steps and workflow_run_response:
+        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 = 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)
+        if stream_intermediate_steps is not None:
+            warnings.warn(
+                "The 'stream_intermediate_steps' parameter is deprecated and will be removed in future versions. Use 'stream_events' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        stream_events = stream_events or stream_intermediate_steps
+
+        if stream_events and workflow_run_response:
             # Emit StepStartedEvent
             yield StepStartedEvent(
                 run_id=workflow_run_response.run_id or "",
@@ -660,61 +991,117 @@ class Step:
 
                 if self._executor_type == "function":
                     log_debug(f"Executing async function executor for step: {self.name}")
-                    import inspect
 
                     # Check if the function is an async generator
-                    if inspect.isasyncgenfunction(self.active_executor):
+                    if _is_async_generator_function(self.active_executor):
                         content = ""
                         # It's an async generator - iterate over it
-                        async for event in self.active_executor(step_input):  # type: ignore
-                            if (
-                                hasattr(event, "content")
-                                and event.content is not None
-                                and isinstance(event.content, str)
-                            ):
-                                content += event.content
-                            else:
-                                content += str(event)
-                            if isinstance(event, StepOutput):
+                        iterator = await self._acall_custom_function(
+                            self.active_executor,
+                            step_input,
+                            session_state_copy,
+                            run_context,
+                        )
+                        async for event in iterator:  # type: ignore
+                            if isinstance(event, (BaseRunOutputEvent)):
+                                if (
+                                    isinstance(event, (RunContentEvent, TeamRunContentEvent))
+                                    and event.content is not None
+                                ):
+                                    if isinstance(event.content, str):
+                                        content += event.content
+                                    elif isinstance(event.content, BaseModel):
+                                        content = event.content  # type: ignore[assignment]
+                                    else:
+                                        content = str(event.content)
+
+                                # Only yield executor events if stream_executor_events is True
+                                if stream_executor_events:
+                                    enriched_event = self._enrich_event_with_context(
+                                        event, workflow_run_response, step_index
+                                    )
+                                    yield enriched_event  # type: ignore[misc]
+                            elif isinstance(event, (RunOutput, TeamRunOutput)):
+                                content = event.content  # type: ignore[assignment]
+                            elif isinstance(event, StepOutput):
                                 final_response = event
                                 break
                             else:
-                                yield event  # type: ignore[misc]
+                                content += str(event)
                         if not final_response:
                             final_response = StepOutput(content=content)
-                    elif inspect.iscoroutinefunction(self.active_executor):
+                    elif _is_async_callable(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,
+                            run_context,
+                        )
                         if isinstance(result, StepOutput):
                             final_response = result
+                        elif isinstance(result, (RunOutput, TeamRunOutput)):
+                            final_response = StepOutput(content=result.content)
                         else:
                             final_response = StepOutput(content=str(result))
-                    elif inspect.isgeneratorfunction(self.active_executor):
+                    elif _is_generator_function(self.active_executor):
                         content = ""
                         # It's a regular generator function - iterate over it
-                        for event in self.active_executor(step_input):  # type: ignore
-                            if (
-                                hasattr(event, "content")
-                                and event.content is not None
-                                and isinstance(event.content, str)
-                            ):
-                                content += event.content
-                            else:
-                                content += str(event)
-                            if isinstance(event, StepOutput):
+                        iterator = self._call_custom_function(
+                            self.active_executor,
+                            step_input,
+                            session_state_copy,
+                            run_context,
+                        )
+                        for event in iterator:  # type: ignore
+                            if isinstance(event, (BaseRunOutputEvent)):
+                                if (
+                                    isinstance(event, (RunContentEvent, TeamRunContentEvent))
+                                    and event.content is not None
+                                ):
+                                    if isinstance(event.content, str):
+                                        content += event.content
+                                    elif isinstance(event.content, BaseModel):
+                                        content = event.content  # type: ignore[assignment]
+                                    else:
+                                        content = str(event.content)
+
+                                # Only yield executor events if stream_executor_events is True
+                                if stream_executor_events:
+                                    enriched_event = self._enrich_event_with_context(
+                                        event, workflow_run_response, step_index
+                                    )
+                                    yield enriched_event  # type: ignore[misc]
+                            elif isinstance(event, (RunOutput, TeamRunOutput)):
+                                content = event.content  # type: ignore[assignment]
+                            elif isinstance(event, StepOutput):
                                 final_response = event
                                 break
                             else:
-                                yield event  # type: ignore[misc]
+                                if isinstance(content, str):
+                                    content += str(event)
+                                else:
+                                    content = str(event)
                         if not final_response:
                             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,  # type: ignore[arg-type]
+                            step_input,
+                            session_state_copy,
+                            run_context,
+                        )
                         if isinstance(result, StepOutput):
                             final_response = result
+                        elif isinstance(result, (RunOutput, TeamRunOutput)):
+                            final_response = StepOutput(content=result.content)
                         else:
                             final_response = StepOutput(content=str(result))
+
+                    # Merge session_state changes back
+                    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
                     message = self._prepare_message(
@@ -741,9 +1128,22 @@ class Step:
                         if isinstance(self.active_executor, Team):
                             kwargs["store_member_responses"] = True
 
-                        session_state_copy = copy(session_state)
+                        num_history_runs = self.num_history_runs if self.num_history_runs else num_history_runs
+
+                        use_history = (
+                            self.add_workflow_history
+                            if self.add_workflow_history is not None
+                            else add_workflow_history_to_steps
+                        )
+
+                        final_message = message
+                        if use_history and workflow_session:
+                            history_messages = workflow_session.get_workflow_history_context(num_runs=num_history_runs)
+                            if history_messages:
+                                final_message = f"{history_messages}{message}"
+
                         response_stream = self.active_executor.arun(  # type: ignore
-                            input=message,
+                            input=final_message,
                             images=images,
                             videos=videos,
                             audio=audios,
@@ -752,34 +1152,32 @@ class Step:
                             user_id=user_id,
                             session_state=session_state_copy,
                             stream=True,
-                            stream_intermediate_steps=stream_intermediate_steps,
-                            # Pass workflow context directly via kwargs
-                            workflow_context={
-                                "workflow_id": workflow_run_response.workflow_id if workflow_run_response else None,
-                                "workflow_run_id": workflow_run_response.run_id if workflow_run_response else None,
-                                "step_id": self.step_id,
-                                "step_name": self.name,
-                                "step_index": step_index,
-                            },
+                            stream_events=stream_events,
+                            run_context=run_context,
                             yield_run_response=True,
                             **kwargs,
                         )
 
                         active_executor_run_response = None
                         async for event in response_stream:
-                            log_debug(f"Received async event from agent: {type(event).__name__}")
                             if isinstance(event, RunOutput) or isinstance(event, TeamRunOutput):
                                 active_executor_run_response = event
                                 break
-                            yield event  # type: ignore[misc]
+                            # Only yield executor events if stream_executor_events is True
+                            if stream_executor_events:
+                                enriched_event = self._enrich_event_with_context(
+                                    event, workflow_run_response, step_index
+                                )
+                                yield enriched_event  # type: ignore[misc]
 
                         # Update workflow session state
-                        merge_dictionaries(session_state, session_state_copy)  # type: ignore
+                        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
 
-                        final_response = self._process_step_output(active_executor_run_response)  # type: ignore
+                        final_response = active_executor_run_response  # type: ignore
                     else:
                         raise ValueError(f"Unsupported executor type: {self._executor_type}")
 
@@ -791,9 +1189,10 @@ class Step:
                 use_workflow_logger()
 
                 # Yield the final response
+                final_response = self._process_step_output(final_response)
                 yield final_response
 
-                if stream_intermediate_steps and workflow_run_response:
+                if stream_events and workflow_run_response:
                     # Emit StepCompletedEvent
                     yield StepCompletedEvent(
                         run_id=workflow_run_response.run_id or "",
@@ -826,6 +1225,83 @@ class Step:
 
         return
 
+    def get_chat_history(self, session_id: str, last_n_runs: Optional[int] = None) -> List[Message]:
+        """Return the step's Agent or Team chat history for the given session.
+
+        Args:
+            session_id: The session ID to get the chat history for. If not provided, the current cached session ID is used.
+            last_n_runs: Number of recent runs to include. If None, all runs will be considered.
+
+        Returns:
+            List[Message]: The step's Agent or Team chat history for the given session.
+        """
+        session: Union[AgentSession, TeamSession, WorkflowSession, None] = None
+
+        if self.agent:
+            session = self.agent.get_session(session_id=session_id)
+            if not session:
+                log_warning("Session not found")
+                return []
+
+            if not isinstance(session, WorkflowSession):
+                raise ValueError("The provided session is not a WorkflowSession")
+
+            session = cast(WorkflowSession, session)
+            return session.get_messages(last_n_runs=last_n_runs, agent_id=self.agent.id)
+
+        elif self.team:
+            session = self.team.get_session(session_id=session_id)
+            if not session:
+                log_warning("Session not found")
+                return []
+
+            if not isinstance(session, WorkflowSession):
+                raise ValueError("The provided session is not a WorkflowSession")
+
+            session = cast(WorkflowSession, session)
+            return session.get_messages(last_n_runs=last_n_runs, team_id=self.team.id)
+
+        return []
+
+    async def aget_chat_history(
+        self, session_id: Optional[str] = None, last_n_runs: Optional[int] = None
+    ) -> List[Message]:
+        """Return the step's Agent or Team chat history for the given session.
+
+        Args:
+            session_id: The session ID to get the chat history for. If not provided, the current cached session ID is used.
+            last_n_runs: Number of recent runs to include. If None, all runs will be considered.
+
+        Returns:
+            List[Message]: The step's Agent or Team chat history for the given session.
+        """
+        session: Union[AgentSession, TeamSession, WorkflowSession, None] = None
+
+        if self.agent:
+            session = await self.agent.aget_session(session_id=session_id)
+            if not session:
+                log_warning("Session not found")
+                return []
+
+            if not isinstance(session, WorkflowSession):
+                raise ValueError("The provided session is not a WorkflowSession")
+
+            session = cast(WorkflowSession, session)
+            return session.get_messages(last_n_runs=last_n_runs, agent_id=self.agent.id)
+
+        elif self.team:
+            session = await self.team.aget_session(session_id=session_id)
+            if not session:
+                log_warning("Session not found")
+                return []
+
+            if not isinstance(session, WorkflowSession):
+                raise ValueError("The provided session is not a WorkflowSession")
+
+            return session.get_messages(last_n_runs=last_n_runs, team_id=self.team.id)
+
+        return []
+
     def _store_executor_response(
         self, workflow_run_response: "WorkflowRunOutput", executor_run_response: Union[RunOutput, TeamRunOutput]
     ) -> None:
@@ -835,6 +1311,14 @@ class Step:
             executor_run_response.parent_run_id = workflow_run_response.run_id
             executor_run_response.workflow_step_id = self.step_id
 
+            # Scrub the executor response based on the executor's storage flags before storing
+            if (
+                not self.active_executor.store_media
+                or not self.active_executor.store_tool_messages
+                or not self.active_executor.store_history_messages
+            ):  # type: ignore
+                self.active_executor._scrub_run_output_for_storage(executor_run_response)  # type: ignore
+
             # Get the raw response from the step's active executor
             raw_response = executor_run_response
             if raw_response and isinstance(raw_response, (RunOutput, TeamRunOutput)):
@@ -1021,3 +1505,28 @@ class Step:
                 continue
 
         return videos
+
+
+def _is_async_callable(obj: Any) -> TypeGuard[Callable[..., Any]]:
+    """Checks if obj is an async callable (coroutine function or callable with async __call__)"""
+    return inspect.iscoroutinefunction(obj) or (callable(obj) and inspect.iscoroutinefunction(obj.__call__))
+
+
+def _is_generator_function(obj: Any) -> TypeGuard[Callable[..., Any]]:
+    """Checks if obj is a generator function, including callable class instances with generator __call__ methods"""
+    if inspect.isgeneratorfunction(obj):
+        return True
+    # Check if it's a callable class instance with a generator __call__ method
+    if callable(obj) and hasattr(obj, "__call__"):
+        return inspect.isgeneratorfunction(obj.__call__)
+    return False
+
+
+def _is_async_generator_function(obj: Any) -> TypeGuard[Callable[..., Any]]:
+    """Checks if obj is an async generator function, including callable class instances"""
+    if inspect.isasyncgenfunction(obj):
+        return True
+    # Check if it's a callable class instance with an async generator __call__ method
+    if callable(obj) and hasattr(obj, "__call__"):
+        return inspect.isasyncgenfunction(obj.__call__)
+    return False