def mock_agent_output_fn_stateful(
    state: Dict[str, Any], output: str
) -> Tuple[AgentChatResponse, bool]:
    """Mock agent output function (for StatefulFnComponent)."""
    return mock_agent_output_fn(state["task"], state["step_state"], output)
Mock agent output function (for StatefulFnComponent).
llama-index-core/tests/agent/custom/test_query_pipeline.py