rasa-pro 3.14.0.dev5__py3-none-any.whl → 3.14.0.dev7__py3-none-any.whl

rasa/core/policies/flows/flow_executor.py

@@ -1,14 +1,25 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, Text
+import asyncio
+from typing import Any, Dict, List, Optional, Text, cast
 
 import structlog
 from jinja2 import Template
 from pypred import Predicate
-from structlog.contextvars import (
-    bound_contextvars,
+from structlog.contextvars import bound_contextvars
+
+from rasa.agents.agent_manager import AgentManager
+from rasa.agents.constants import (
+    A2A_AGENT_CONTEXT_ID_KEY,
+    AGENT_METADATA_AGENT_RESPONSE_KEY,
+    AGENT_METADATA_EXIT_IF_KEY,
+    AGENT_METADATA_TOOL_RESULTS_KEY,
+    MAX_AGENT_RETRY_DELAY_SECONDS,
 )
-
+from rasa.agents.core.types import AgentStatus, ProtocolType
+from rasa.agents.schemas import AgentInput, AgentOutput
+from rasa.agents.schemas.agent_input import AgentInputSlot
+from rasa.core.available_agents import AvailableAgents
 from rasa.core.available_endpoints import AvailableEndpoints
 from rasa.core.constants import ACTIVE_FLOW_METADATA_KEY, STEP_ID_METADATA_KEY
 from rasa.core.policies.flows.flow_exceptions import (
@@ -22,7 +33,8 @@ from rasa.core.policies.flows.flow_step_result import (
     FlowStepResult,
     PauseFlowReturnPrediction,
 )
-from rasa.dialogue_understanding.commands import CancelFlowCommand
+from rasa.core.policies.flows.mcp_tool_executor import call_mcp_tool
+from rasa.core.utils import get_slot_names_from_exit_conditions
 from rasa.dialogue_understanding.patterns.cancel import CancelPatternFlowStackFrame
 from rasa.dialogue_understanding.patterns.collect_information import (
     FLOW_PATTERN_COLLECT_INFORMATION,
@@ -49,17 +61,30 @@ from rasa.dialogue_understanding.stack.frames import (
     UserFlowStackFrame,
 )
 from rasa.dialogue_understanding.stack.frames.flow_stack_frame import (
+    AgentStackFrame,
+    AgentState,
     FlowStackFrameType,
 )
 from rasa.dialogue_understanding.stack.utils import (
-    top_user_flow_frame,
+    user_frames_on_the_stack,
 )
+from rasa.dialogue_understanding.utils import assemble_options_string
+from rasa.shared.agents.utils import get_protocol_type
 from rasa.shared.constants import RASA_PATTERN_HUMAN_HANDOFF
 from rasa.shared.core.constants import (
+    ACTION_AGENT_REQUEST_USER_INPUT_NAME,
     ACTION_LISTEN_NAME,
+    ACTION_METADATA_MESSAGE_KEY,
+    ACTION_METADATA_TEXT_KEY,
+    ACTION_SEND_TEXT_NAME,
+    FLOW_HASHES_SLOT,
     SILENCE_TIMEOUT_SLOT,
 )
 from rasa.shared.core.events import (
+    AgentCancelled,
+    AgentCompleted,
+    AgentResumed,
+    AgentStarted,
     Event,
     FlowCompleted,
     FlowResumed,
@@ -67,11 +92,7 @@ from rasa.shared.core.events import (
     SlotSet,
 )
 from rasa.shared.core.flows import FlowsList
-from rasa.shared.core.flows.flow import (
-    END_STEP,
-    Flow,
-    FlowStep,
-)
+from rasa.shared.core.flows.flow import END_STEP, Flow, FlowStep
 from rasa.shared.core.flows.flow_step_links import (
     ElseFlowStepLink,
     IfFlowStepLink,
@@ -88,15 +109,19 @@ from rasa.shared.core.flows.steps import (
     SetSlotsFlowStep,
 )
 from rasa.shared.core.flows.steps.constants import START_STEP
-from rasa.shared.core.slots import Slot, SlotRejection
-from rasa.shared.core.trackers import (
-    DialogueStateTracker,
-)
+from rasa.shared.core.slots import CategoricalSlot, Slot, SlotRejection
+from rasa.shared.core.trackers import DialogueStateTracker
+from rasa.shared.utils.llm import tracker_as_readable_transcript
 
 structlogger = structlog.get_logger()
 
 MAX_NUMBER_OF_STEPS = 250
 
+MAX_AGENT_RETRIES = 3
+
+# Slots that should not be forwarded to sub-agents via AgentInput
+SLOTS_EXCLUDED_FOR_AGENT = [FLOW_HASHES_SLOT]
+
 
 def render_template_variables(text: str, context: Dict[Text, Any]) -> str:
     """Replace context variables in a text."""
@@ -148,6 +173,13 @@ def select_next_step_id(
     tracker: DialogueStateTracker,
 ) -> Optional[Text]:
     """Selects the next step id based on the current step."""
+    # if the current step is a call step to an agent, and we already have an
+    # AgentStackFrame on top of the stack, we need to return the current
+    # step id again in order to loop back to the agent.
+    agent_stack_frame = tracker.stack.top()
+    if agent_stack_frame and isinstance(agent_stack_frame, AgentStackFrame):
+        return current.id
+
     next_step = current.next
     if len(next_step.links) == 1 and isinstance(next_step.links[0], StaticFlowStepLink):
         return next_step.links[0].target
@@ -234,38 +266,58 @@ def trigger_pattern_continue_interrupted(
     stack: DialogueStack,
     flows: FlowsList,
     tracker: DialogueStateTracker,
-) -> List[Event]:
+) -> None:
     """Trigger the pattern to continue an interrupted flow if needed."""
-    events: List[Event] = []
-
-    # get previously started user flow that will be continued
-    interrupted_user_flow_frame = top_user_flow_frame(stack)
-    interrupted_user_flow_step = (
-        interrupted_user_flow_frame.step(flows) if interrupted_user_flow_frame else None
-    )
-    interrupted_user_flow = (
-        interrupted_user_flow_frame.flow(flows) if interrupted_user_flow_frame else None
-    )
-
+    # only trigger the pattern if the current frame is a user flow frame
+    # with a frame type of interrupt
     if (
-        isinstance(current_frame, UserFlowStackFrame)
-        and interrupted_user_flow_step is not None
-        and interrupted_user_flow is not None
-        and current_frame.frame_type == FlowStackFrameType.INTERRUPT
-        and not is_step_end_of_flow(interrupted_user_flow_step)
+        not isinstance(current_frame, UserFlowStackFrame)
+        or current_frame.frame_type != FlowStackFrameType.INTERRUPT
     ):
-        stack.push(
-            ContinueInterruptedPatternFlowStackFrame(
-                previous_flow_name=interrupted_user_flow.readable_name(
-                    language=tracker.current_language
-                ),
-            )
+        return None
+
+    # get all previously interrupted user flows
+    interrupted_user_flow_stack_frames = user_frames_on_the_stack(stack)
+
+    interrupted_user_flows_to_continue: List[UserFlowStackFrame] = []
+    # check if interrupted user flows can be continued
+    # i.e. the flow is not at the end of the flow
+    for frame in interrupted_user_flow_stack_frames:
+        interrupted_user_flow_step = frame.step(flows)
+        interrupted_user_flow = frame.flow(flows)
+        if (
+            interrupted_user_flow_step is not None
+            and interrupted_user_flow is not None
+            and not is_step_end_of_flow(interrupted_user_flow_step)
+        ):
+            interrupted_user_flows_to_continue.append(frame)
+
+    # if there are no interrupted user flows to continue,
+    # we don't need to trigger the pattern
+    if len(interrupted_user_flows_to_continue) == 0:
+        return None
+
+    # get the flow names and ids of the interrupted flows
+    # and assemble the options string
+    flow_names: List[str] = []
+    flow_ids: List[str] = []
+    for frame in interrupted_user_flows_to_continue:
+        flow_names.append(
+            frame.flow(flows).readable_name(language=tracker.current_language)
         )
-        events.append(
-            FlowResumed(interrupted_user_flow.id, interrupted_user_flow_step.id)
+        flow_ids.append(frame.flow_id)
+    options_string = assemble_options_string(flow_names)
+
+    # trigger the pattern to continue the interrupted flows
+    stack.push(
+        ContinueInterruptedPatternFlowStackFrame(
+            interrupted_flow_names=flow_names,
+            interrupted_flow_ids=flow_ids,
+            interrupted_flow_options=options_string,
         )
+    )
 
-    return events
+    return None
 
 
 def trigger_pattern_completed(
@@ -359,8 +411,11 @@ def reset_scoped_slots(
     return events
 
 
-def advance_flows(
-    tracker: DialogueStateTracker, available_actions: List[str], flows: FlowsList
+async def advance_flows(
+    tracker: DialogueStateTracker,
+    available_actions: List[str],
+    flows: FlowsList,
+    slots: List[Slot],
 ) -> FlowActionPrediction:
     """Advance the current flows until the next action.
 
@@ -368,6 +423,7 @@ def advance_flows(
         tracker: The tracker to get the next action for.
         available_actions: The actions that are available in the domain.
         flows: All flows.
+        slots: The slots that are available in the domain.
 
     Returns:
     The predicted action and the events to run.
@@ -377,13 +433,16 @@ def advance_flows(
         # if there are no flows, there is nothing to do
         return FlowActionPrediction(None, 0.0)
 
-    return advance_flows_until_next_action(tracker, available_actions, flows)
+    return await advance_flows_until_next_action(
+        tracker, available_actions, flows, slots
+    )
 
 
-def advance_flows_until_next_action(
+async def advance_flows_until_next_action(
     tracker: DialogueStateTracker,
     available_actions: List[str],
     flows: FlowsList,
+    slots: List[Slot],
 ) -> FlowActionPrediction:
     """Advance the flow and select the next action to execute.
 
@@ -441,7 +500,7 @@ def advance_flows_until_next_action(
 
             with bound_contextvars(step_id=next_step.id):
                 step_stack = tracker.stack
-                step_result = run_step(
+                step_result = await run_step(
                     next_step,
                     current_flow,
                     step_stack,
@@ -449,6 +508,7 @@ def advance_flows_until_next_action(
                     available_actions,
                     flows,
                     previous_step_id,
+                    slots,
                 )
                 new_events = step_result.events
                 if (
@@ -477,10 +537,9 @@ def advance_flows_until_next_action(
         # make sure we really return all events that got created during the
         # step execution of all steps (not only the last one)
         prediction.events = gathered_events
-        prediction.metadata = {
-            ACTIVE_FLOW_METADATA_KEY: tracker.active_flow,
-            STEP_ID_METADATA_KEY: tracker.current_step_id,
-        }
+        prediction.metadata = prediction.metadata or {}
+        prediction.metadata[ACTIVE_FLOW_METADATA_KEY] = tracker.active_flow
+        prediction.metadata[STEP_ID_METADATA_KEY] = tracker.current_step_id
         return prediction
     else:
         structlogger.warning("flow.step.execution.no_action")
@@ -522,6 +581,8 @@ def validate_collect_step(
 
 def cancel_flow_and_push_internal_error(stack: DialogueStack, flow_name: str) -> None:
     """Cancel the top user flow and push the internal error pattern."""
+    from rasa.dialogue_understanding.commands import CancelFlowCommand
+
     top_frame = stack.top()
 
     if isinstance(top_frame, BaseFlowStackFrame):
@@ -550,7 +611,7 @@ def attach_stack_metadata_to_events(
         event.metadata[ACTIVE_FLOW_METADATA_KEY] = flow_id
 
 
-def run_step(
+async def run_step(
     step: FlowStep,
     flow: Flow,
     stack: DialogueStack,
@@ -558,6 +619,7 @@ def run_step(
     available_actions: List[str],
     flows: FlowsList,
     previous_step_id: str,
+    slots: List[Slot],
 ) -> FlowStepResult:
     """Run a single step of a flow.
 
@@ -576,6 +638,7 @@ def run_step(
         available_actions: The actions that are available in the domain.
         flows: All flows.
         previous_step_id: The ID of the previous step.
+        slots: The slots that are available in the domain.
 
     Returns:
     A result of running the step describing where to transition to.
@@ -615,7 +678,7 @@ def run_step(
         return _run_link_step(initial_events, stack, step)
 
     elif isinstance(step, CallFlowStep):
-        return _run_call_step(initial_events, stack, step)
+        return await _run_call_step(initial_events, stack, step, tracker, slots)
 
     elif isinstance(step, SetSlotsFlowStep):
         return _run_set_slot_step(initial_events, step)
@@ -667,12 +730,10 @@ def _run_end_step(
     structlogger.debug("flow.step.run.flow_end")
     current_frame = stack.pop()
     trigger_pattern_completed(current_frame, stack, flows)
-    resumed_events = trigger_pattern_continue_interrupted(
-        current_frame, stack, flows, tracker
-    )
+    trigger_pattern_continue_interrupted(current_frame, stack, flows, tracker)
     reset_events: List[Event] = reset_scoped_slots(current_frame, flow, tracker)
     return ContinueFlowWithNextStep(
-        events=initial_events + reset_events + resumed_events, has_flow_ended=True
+        events=initial_events + reset_events, has_flow_ended=True
     )
 
 
@@ -684,17 +745,26 @@ def _run_set_slot_step(
     return ContinueFlowWithNextStep(events=initial_events + slot_events)
 
 
-def _run_call_step(
-    initial_events: List[Event], stack: DialogueStack, step: CallFlowStep
+async def _run_call_step(
+    initial_events: List[Event],
+    stack: DialogueStack,
+    step: CallFlowStep,
+    tracker: DialogueStateTracker,
+    slots: List[Slot],
 ) -> FlowStepResult:
     structlogger.debug("flow.step.run.call")
-    stack.push(
-        UserFlowStackFrame(
-            flow_id=step.call,
-            frame_type=FlowStackFrameType.CALL,
-        ),
-    )
-    return ContinueFlowWithNextStep(events=initial_events)
+    if step.is_calling_mcp_tool():
+        return await call_mcp_tool(initial_events, stack, step, tracker)
+    elif step.is_calling_agent():
+        return await run_agent(initial_events, stack, step, tracker, slots)
+    else:
+        stack.push(
+            UserFlowStackFrame(
+                flow_id=step.call,
+                frame_type=FlowStackFrameType.CALL,
+            ),
+        )
+        return ContinueFlowWithNextStep(events=initial_events)
 
 
 def _run_link_step(
@@ -740,14 +810,7 @@ def _run_action_step(
             # do not log about non-existing validation actions of collect steps
             utter_action_name = render_template_variables("{{context.utter}}", context)
             if utter_action_name not in available_actions:
-                structlogger.warning(
-                    "flow.step.run.action.unknown",
-                    action=action_name,
-                    event_info=(
-                        f"The action '{action_name}' is not defined in the domain but "
-                        f"getting triggered by the flow '{step.flow_id}'."
-                    ),
-                )
+                structlogger.warning("flow.step.run.action.unknown", action=action_name)
         return ContinueFlowWithNextStep(events=initial_events)
 
 
@@ -839,3 +902,386 @@ def _append_global_silence_timeout_event(
                 AvailableEndpoints.get_instance().interaction_handling.global_silence_timeout,
             )
         )
+
+
+def _reset_slots_covered_by_exit_if(
+    exit_conditions: List[str], tracker: DialogueStateTracker
+) -> None:
+    """Reset the slots covered by the exit_if condition."""
+    reset_slot_names = get_slot_names_from_exit_conditions(exit_conditions)
+    for slot_name in reset_slot_names:
+        if tracker.slots.get(slot_name) is not None:
+            tracker.update(SlotSet(slot_name, None))
+
+
+async def run_agent(
+    initial_events: List[Event],
+    stack: DialogueStack,
+    step: CallFlowStep,
+    tracker: DialogueStateTracker,
+    slots: List[Slot],
+) -> FlowStepResult:
+    """Run an agent call step."""
+    structlogger.debug(
+        "flow.step.run_agent", agent_id=step.call, step_id=step.id, flow_id=step.flow_id
+    )
+
+    final_events = initial_events
+    agent_stack_frame = tracker.stack.find_agent_stack_frame_by_agent(
+        agent_id=step.call
+    )
+
+    if (
+        agent_stack_frame
+        and agent_stack_frame == stack.top()
+        and agent_stack_frame.state == AgentState.INTERRUPTED
+    ):
+        structlogger.debug(
+            "flow.step.run_agent.resume_interrupted_agent",
+            agent_id=step.call,
+            step_id=step.id,
+            flow_id=step.flow_id,
+        )
+        # The agent was previously interrupted when waiting for user input.
+        # Now we're back to the agent execution step and need to output the last message
+        # from the agent (user input request) again and wait for user input
+        cast(AgentStackFrame, stack.top()).state = AgentState.WAITING_FOR_INPUT
+        tracker.update_stack(stack)
+        utterance = (
+            agent_stack_frame.metadata.get(AGENT_METADATA_AGENT_RESPONSE_KEY, "")
+            if agent_stack_frame.metadata
+            else ""
+        )
+        final_events.append(AgentResumed(agent_id=step.call, flow_id=step.flow_id))
+        return PauseFlowReturnPrediction(
+            _create_agent_request_user_input_prediction(utterance, final_events)
+        )
+
+    agent_input_metadata = (
+        agent_stack_frame.metadata
+        if agent_stack_frame and agent_stack_frame.metadata
+        else {}
+    )
+    _update_agent_input_metadata_with_events(
+        agent_input_metadata, step.call, step.flow_id, tracker
+    )
+    if step.exit_if:
+        # TODO: this is a temporary fix to reset the slots covered by the exit_if
+        if (
+            agent_stack_frame
+            and agent_stack_frame.frame_id == f"restart_agent_{step.call}"
+        ):
+            # when restarting an agent, we need to reset the slots covered by the
+            # exit_if condition so that the agent can run again.
+            _reset_slots_covered_by_exit_if(step.exit_if, tracker)
+        agent_input_metadata[AGENT_METADATA_EXIT_IF_KEY] = step.exit_if
+    agent_input = AgentInput(
+        id=step.call,
+        user_message=tracker.latest_message.text or ""
+        if tracker.latest_message
+        else "",
+        slots=_prepare_slots_for_agent(tracker.current_slot_values(), slots),
+        conversation_history=tracker_as_readable_transcript(tracker),
+        events=tracker.current_state().get("events") or [],
+        metadata=agent_input_metadata,
+    )
+
+    final_events.append(AgentStarted(step.call, step.flow_id))
+
+    protocol_type = get_protocol_type(step, AvailableAgents.get_agent_config(step.call))
+    # send the input to the agent and wait for a response
+    structlogger.debug(
+        "flow.step.run_agent.agent_input",
+        agent_name=step.call,
+        step_id=step.id,
+        flow_id=step.flow_id,
+        agent_input=agent_input,
+    )
+    output: AgentOutput = await _call_agent_with_retry(
+        agent_name=step.call,
+        protocol_type=protocol_type,
+        agent_input=agent_input,
+        max_retries=MAX_AGENT_RETRIES,
+    )
+    structlogger.debug(
+        "flow.step.run_agent.agent_response",
+        agent_name=step.call,
+        step_id=step.id,
+        flow_id=step.flow_id,
+        agent_response=output,
+    )
+
+    # add the set slot events returned by the agent to the list of final events
+    if output.events:
+        final_events.extend(output.events)
+
+    if output.status == AgentStatus.INPUT_REQUIRED:
+        output.metadata = output.metadata or {}
+        output.metadata[AGENT_METADATA_AGENT_RESPONSE_KEY] = (
+            output.response_message or ""
+        )
+        output.metadata[AGENT_METADATA_TOOL_RESULTS_KEY] = output.tool_results or []
+        _update_agent_events(final_events, output.metadata)
+
+        top_stack_frame = stack.top()
+        # update the agent stack frame if it is already on the stack
+        # otherwise push a new one
+        if isinstance(top_stack_frame, AgentStackFrame):
+            top_stack_frame.state = AgentState.WAITING_FOR_INPUT
+            top_stack_frame.metadata = output.metadata
+            top_stack_frame.step_id = step.id
+            top_stack_frame.agent_id = step.call
+            top_stack_frame.flow_id = step.flow_id
+        else:
+            stack.push(
+                AgentStackFrame(
+                    flow_id=step.flow_id,
+                    agent_id=step.call,
+                    state=AgentState.WAITING_FOR_INPUT,
+                    step_id=step.id,
+                    metadata=output.metadata,
+                )
+            )
+
+        action_prediction = _create_agent_request_user_input_prediction(
+            output.response_message, final_events
+        )
+        return PauseFlowReturnPrediction(action_prediction)
+    elif output.status == AgentStatus.COMPLETED:
+        output.metadata = output.metadata or {}
+        _update_agent_events(final_events, output.metadata)
+        structlogger.debug(
+            "flow.step.run_agent.completed",
+            agent_name=step.call,
+            step_id=step.id,
+            flow_id=step.flow_id,
+        )
+        remove_agent_stack_frame(stack, step.call)
+        agent_completed_event = AgentCompleted(agent_id=step.call, flow_id=step.flow_id)
+        final_events.append(agent_completed_event)
+        if output.response_message:
+            # for open-ended agents we want to utter the last agent message
+
+            return PauseFlowReturnPrediction(
+                _create_send_text_prediction(output.response_message, final_events)
+            )
+        else:
+            return ContinueFlowWithNextStep(events=final_events)
+    elif output.status == AgentStatus.FATAL_ERROR:
+        output.metadata = output.metadata or {}
+        _update_agent_events(final_events, output.metadata)
+        # the agent failed, trigger pattern_internal_error
+        structlogger.error(
+            "flow.step.run_agent.fatal_error",
+            agent_name=step.call,
+            step_id=step.id,
+            flow_id=step.flow_id,
+            error_message=output.error_message,
+        )
+        remove_agent_stack_frame(stack, step.call)
+        final_events.append(
+            AgentCancelled(
+                agent_id=step.call, flow_id=step.flow_id, reason=output.error_message
+            )
+        )
+        stack.push(InternalErrorPatternFlowStackFrame())
+        return ContinueFlowWithNextStep(events=final_events)
+    else:
+        output.metadata = output.metadata or {}
+        _update_agent_events(final_events, output.metadata)
+        structlogger.error(
+            "flow.step.run_agent.unknown_status",
+            agent_name=step.call,
+            step_id=step.id,
+            flow_id=step.flow_id,
+            status=output.status,
+        )
+        remove_agent_stack_frame(stack, step.call)
+        final_events.append(AgentCancelled(agent_id=step.call, flow_id=step.flow_id))
+        stack.push(InternalErrorPatternFlowStackFrame())
+        return ContinueFlowWithNextStep(events=final_events)
+
+
+def remove_agent_stack_frame(stack: DialogueStack, agent_id: str) -> None:
+    """Finishes the agentic loop by popping the agent stack frame from the
+    provided `stack`. The `tracker.stack` is NOT modified.
+    """
+    agent_stack_frame = stack.find_agent_stack_frame_by_agent(agent_id)
+    if not agent_stack_frame:
+        return
+
+    while removed_frame := stack.pop():
+        structlogger.debug(
+            "flow_executor.remove_agent_stack_frame",
+            removed_frame=removed_frame,
+        )
+        if removed_frame == agent_stack_frame:
+            break
+
+
+def _create_action_prediction(
+    action_name: str, message: Optional[str], events: Optional[List[Event]]
+) -> FlowActionPrediction:
+    """Create a prediction for an action with a text message."""
+    action_metadata = {
+        ACTION_METADATA_MESSAGE_KEY: {
+            ACTION_METADATA_TEXT_KEY: message,
+        }
+    }
+    return FlowActionPrediction(
+        action_name,
+        1.0,
+        events=events if events else [],
+        metadata=action_metadata,
+    )
+
+
+def _create_agent_request_user_input_prediction(
+    message: Optional[str], events: Optional[List[Event]]
+) -> FlowActionPrediction:
+    """Create a prediction for requesting user input from the agent
+    and waiting for it.
+    """
+    return _create_action_prediction(
+        ACTION_AGENT_REQUEST_USER_INPUT_NAME, message, events
+    )
+
+
+def _create_send_text_prediction(
+    message: Optional[str], events: Optional[List[Event]]
+) -> FlowActionPrediction:
+    """Create a prediction for sending a text message to the user."""
+    return _create_action_prediction(ACTION_SEND_TEXT_NAME, message, events)
+
+
+async def _call_agent_with_retry(
+    agent_name: str,
+    protocol_type: ProtocolType,
+    agent_input: AgentInput,
+    max_retries: int,
+) -> AgentOutput:
+    """Call an agent with retries in case of recoverable errors."""
+    for attempt in range(max_retries):
+        if attempt > 0:
+            structlogger.debug(
+                "flow_executor.call_agent_with_retry.retrying",
+                agent_name=agent_name,
+                attempt=attempt + 1,
+                num_retries=max_retries,
+            )
+        try:
+            agent_response: AgentOutput = await AgentManager().run_agent(
+                agent_name=agent_name, protocol_type=protocol_type, context=agent_input
+            )
+        except Exception as e:
+            # We don't have a vaild agent response at this time to act based
+            # on the agent status, so we return a fatal error.
+            structlogger.error(
+                "flow_executor.call_agent_with_retry.exception",
+                agent_name=agent_name,
+                error_message=str(e),
+            )
+            return AgentOutput(
+                id=agent_name,
+                status=AgentStatus.FATAL_ERROR,
+                error_message=str(e),
+            )
+
+        if agent_response.status != AgentStatus.RECOVERABLE_ERROR:
+            return agent_response
+
+        structlogger.warning(
+            "flow_executor.call_agent_with_retry.recoverable_error",
+            agent_name=agent_name,
+            attempt=attempt + 1,
+            num_retries=max_retries,
+            error_message=agent_response.error_message,
+        )
+        if attempt < max_retries - 1:
+            # exponential backoff - wait longer with each retry
+            # 1 second, 2 seconds, 4 seconds, etc.
+            await asyncio.sleep(min(2**attempt, MAX_AGENT_RETRY_DELAY_SECONDS))
+
+    # we exhausted all retries, return fatal error
+    structlogger.warning(
+        "flow_executor.call_agent_with_retry.exhausted_retries",
+        agent_name=agent_name,
+        num_retries=max_retries,
+    )
+    return AgentOutput(
+        id=agent_name,
+        status=AgentStatus.FATAL_ERROR,
+        error_message="Exhausted all retries for agent call.",
+    )
+
+
+def _prepare_slots_for_agent(
+    slot_values: Dict[str, Any], slot_definitions: List[Slot]
+) -> List[AgentInputSlot]:
+    """Prepare the slots for the agent.
+
+    Filter out slots that should not be forwarded to agents.
+    Add the slot type and allowed values to the slot dictionary.
+
+    Args:
+        slot_values: The full slot dictionary from the tracker.
+        slot_definitions: The slot definitions from the domain.
+
+    Returns:
+        A list of slots containing the name, current value, type, and allowed values.
+    """
+
+    def _get_slot_definition(slot_name: str) -> Optional[Slot]:
+        for slot in slot_definitions:
+            if slot.name == slot_name:
+                return slot
+        return None
+
+    filtered_slots: List[AgentInputSlot] = []
+    for key, value in slot_values.items():
+        if key in SLOTS_EXCLUDED_FOR_AGENT:
+            continue
+        slot_definition = _get_slot_definition(key)
+        if slot_definition:
+            filtered_slots.append(
+                AgentInputSlot(
+                    name=key,
+                    value=value,
+                    type=slot_definition.type_name if slot_definition else "any",
+                    allowed_values=slot_definition.values
+                    if isinstance(slot_definition, CategoricalSlot)
+                    else None,
+                )
+            )
+
+    return filtered_slots
+
+
+def _update_agent_events(events: List[Event], metadata: Dict[str, Any]) -> None:
+    """Update the agent events based on the agent output metadata if needed."""
+    if A2A_AGENT_CONTEXT_ID_KEY in metadata:
+        # If the context ID is present, we need to store it in the AgentStarted
+        # event, so that it can be re-used later in case the agent is restarted.
+        for event in events:
+            if isinstance(event, AgentStarted):
+                event.context_id = metadata[A2A_AGENT_CONTEXT_ID_KEY]
+
+
+def _update_agent_input_metadata_with_events(
+    metadata: Dict[str, Any], agent_id: str, flow_id: str, tracker: DialogueStateTracker
+) -> None:
+    """Update the agent input metadata with the events."""
+    agent_started_events = [
+        event
+        for event in tracker.events
+        if type(event) == AgentStarted
+        and event.agent_id == agent_id
+        and event.flow_id == flow_id
+    ]
+    if agent_started_events:
+        # If we have context ID from the previous agent run, we want to
+        # include it in the metadata so that the agent can continue the same
+        # context.
+        agent_started_event = agent_started_events[-1]
+        if agent_started_event.context_id:
+            metadata[A2A_AGENT_CONTEXT_ID_KEY] = agent_started_event.context_id