mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: (ag-ui): Add Workflow Support, Harden Streaming Semantics, and add Dynamic Handoff Demo (#3911)
* fix Workflow.as_agent() streaming regression in ag-ui * Address PR feedback * workflows wip * wip * wip * Workflow AG-UI demo * Fixes for handoff workflow demo * Fixes to workflows support in AG-UI * Fixes * Add headers to some demo files * Fix comment * Fixes for store * Make _input_schema lazy-loaded * fix mypy * revert session change to handoff only for now --------- Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>
This commit is contained in:
Evan Mattson
committed by
GitHub
Unverified
GitHub
parent
b1c7c7c844
commit
d8b9409e96
@@ -41,7 +41,7 @@ from agent_framework import Agent, SupportsAgentRun
|
||||
from agent_framework._middleware import FunctionInvocationContext, FunctionMiddleware
|
||||
from agent_framework._sessions import AgentSession
|
||||
from agent_framework._tools import FunctionTool, tool
|
||||
from agent_framework._types import AgentResponse, AgentResponseUpdate, Message
|
||||
from agent_framework._types import AgentResponse, AgentResponseUpdate, Content, Message
|
||||
from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
|
||||
from agent_framework._workflows._agent_utils import resolve_agent_id
|
||||
from agent_framework._workflows._checkpoint import CheckpointStorage
|
||||
@@ -267,6 +267,88 @@ class HandoffAgentExecutor(AgentExecutor):
|
||||
|
||||
return cloned_agent
|
||||
|
||||
def _persist_pending_approval_function_calls(self) -> None:
|
||||
"""Persist pending approval function calls for stateless provider resumes.
|
||||
|
||||
Handoff workflows force ``store=False`` and replay conversation state from ``_full_conversation``.
|
||||
When a run pauses on function approval, ``AgentExecutor`` returns ``None`` and the assistant
|
||||
function-call message is not returned as an ``AgentResponse``. Without persisting that call, the
|
||||
next turn may submit only a function result, which responses-style APIs reject.
|
||||
"""
|
||||
pending_calls: list[Content] = []
|
||||
for request in self._pending_agent_requests.values():
|
||||
if request.type != "function_approval_request":
|
||||
continue
|
||||
function_call = getattr(request, "function_call", None)
|
||||
if isinstance(function_call, Content) and function_call.type == "function_call":
|
||||
pending_calls.append(function_call)
|
||||
|
||||
if not pending_calls:
|
||||
return
|
||||
|
||||
self._full_conversation.append(
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=pending_calls,
|
||||
author_name=self._agent.name,
|
||||
)
|
||||
)
|
||||
|
||||
def _persist_missing_approved_function_results(
|
||||
self,
|
||||
*,
|
||||
runtime_tool_messages: list[Message],
|
||||
response_messages: list[Message],
|
||||
) -> None:
|
||||
"""Persist fallback function_result entries for approved calls when missing.
|
||||
|
||||
In approval resumes, function invocation can execute approved tools without
|
||||
always surfacing those tool outputs in the returned ``AgentResponse.messages``.
|
||||
For stateless handoff replays, we must keep call/output pairs balanced.
|
||||
"""
|
||||
candidate_results: dict[str, Content] = {}
|
||||
for message in runtime_tool_messages:
|
||||
for content in message.contents:
|
||||
if content.type == "function_result":
|
||||
call_id = getattr(content, "call_id", None)
|
||||
if isinstance(call_id, str) and call_id:
|
||||
candidate_results[call_id] = content
|
||||
continue
|
||||
|
||||
if content.type != "function_approval_response" or not content.approved:
|
||||
continue
|
||||
|
||||
function_call = getattr(content, "function_call", None)
|
||||
call_id = getattr(function_call, "call_id", None) or getattr(content, "id", None)
|
||||
if isinstance(call_id, str) and call_id and call_id not in candidate_results:
|
||||
# Fallback content for approved calls when runtime messages do not include
|
||||
# a concrete function_result payload.
|
||||
candidate_results[call_id] = Content.from_function_result(
|
||||
call_id=call_id,
|
||||
result='{"status":"approved"}',
|
||||
)
|
||||
|
||||
if not candidate_results:
|
||||
return
|
||||
|
||||
observed_result_call_ids: set[str] = set()
|
||||
for message in [*self._full_conversation, *response_messages]:
|
||||
for content in message.contents:
|
||||
if content.type == "function_result" and isinstance(content.call_id, str) and content.call_id:
|
||||
observed_result_call_ids.add(content.call_id)
|
||||
|
||||
missing_call_ids = sorted(set(candidate_results.keys()) - observed_result_call_ids)
|
||||
if not missing_call_ids:
|
||||
return
|
||||
|
||||
self._full_conversation.append(
|
||||
Message(
|
||||
role="tool",
|
||||
contents=[candidate_results[call_id] for call_id in missing_call_ids],
|
||||
author_name=self._agent.name,
|
||||
)
|
||||
)
|
||||
|
||||
def _clone_chat_agent(self, agent: Agent) -> Agent:
|
||||
"""Produce a deep copy of the Agent while preserving runtime configuration."""
|
||||
options = agent.default_options
|
||||
@@ -287,6 +369,10 @@ class HandoffAgentExecutor(AgentExecutor):
|
||||
# Disable parallel tool calls to prevent the agent from invoking multiple handoff tools at once.
|
||||
cloned_options: dict[str, Any] = {
|
||||
"allow_multiple_tool_calls": False,
|
||||
# Handoff workflows already manage full conversation context explicitly
|
||||
# across executors. Keep provider-side conversation storage disabled to
|
||||
# avoid stale tool-call state (Responses API previous_response chains).
|
||||
"store": False,
|
||||
"frequency_penalty": options.get("frequency_penalty"),
|
||||
"instructions": options.get("instructions"),
|
||||
"logit_bias": dict(logit_bias) if logit_bias else None,
|
||||
@@ -297,7 +383,6 @@ class HandoffAgentExecutor(AgentExecutor):
|
||||
"response_format": options.get("response_format"),
|
||||
"seed": options.get("seed"),
|
||||
"stop": options.get("stop"),
|
||||
"store": options.get("store"),
|
||||
"temperature": options.get("temperature"),
|
||||
"tool_choice": options.get("tool_choice"),
|
||||
"tools": all_tools if all_tools else None,
|
||||
@@ -366,14 +451,44 @@ class HandoffAgentExecutor(AgentExecutor):
|
||||
self, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate]
|
||||
) -> None:
|
||||
"""Override to support handoff."""
|
||||
incoming_messages = list(self._cache)
|
||||
cleaned_incoming_messages = clean_conversation_for_handoff(incoming_messages)
|
||||
runtime_tool_messages = [
|
||||
message
|
||||
for message in incoming_messages
|
||||
if any(
|
||||
content.type
|
||||
in {
|
||||
"function_result",
|
||||
"function_approval_response",
|
||||
}
|
||||
for content in message.contents
|
||||
)
|
||||
or message.role == "tool"
|
||||
]
|
||||
|
||||
# When the full conversation is empty, it means this is the first run.
|
||||
# Broadcast the initial cache to all other agents. Subsequent runs won't
|
||||
# need this since responses are broadcast after each agent run and user input.
|
||||
if self._is_start_agent and not self._full_conversation:
|
||||
await self._broadcast_messages(self._cache.copy(), cast(WorkflowContext[AgentExecutorRequest], ctx))
|
||||
await self._broadcast_messages(cleaned_incoming_messages, cast(WorkflowContext[AgentExecutorRequest], ctx))
|
||||
|
||||
# Append the cache to the full conversation history
|
||||
self._full_conversation.extend(self._cache)
|
||||
# Persist only cleaned chat history between turns to avoid replaying stale tool calls.
|
||||
self._full_conversation.extend(cleaned_incoming_messages)
|
||||
|
||||
# Always run with full conversation context for request_info resumes.
|
||||
# Keep runtime tool-control messages for this run only (e.g., approval responses).
|
||||
self._cache = list(self._full_conversation)
|
||||
self._cache.extend(runtime_tool_messages)
|
||||
|
||||
# Handoff workflows are orchestrator-stateful and provider-stateless by design.
|
||||
# If an existing session still has a service conversation id, clear it to avoid
|
||||
# replaying stale unresolved tool calls across resumed turns.
|
||||
if (
|
||||
cast(Agent, self._agent).default_options.get("store") is False
|
||||
and self._session.service_session_id is not None
|
||||
):
|
||||
self._session.service_session_id = None
|
||||
|
||||
# Check termination condition before running the agent
|
||||
if await self._check_terminate_and_yield(cast(WorkflowContext[Never, list[Message]], ctx)):
|
||||
@@ -392,17 +507,26 @@ class HandoffAgentExecutor(AgentExecutor):
|
||||
|
||||
# A function approval request is issued by the base AgentExecutor
|
||||
if response is None:
|
||||
if cast(Agent, self._agent).default_options.get("store") is False:
|
||||
self._persist_pending_approval_function_calls()
|
||||
# Agent did not complete (e.g., waiting for user input); do not emit response
|
||||
logger.debug("AgentExecutor %s: Agent did not complete, awaiting user input", self.id)
|
||||
return
|
||||
|
||||
# Remove function call related content from the agent response for full conversation history
|
||||
# Remove function call related content from the agent response for broadcast.
|
||||
# This prevents replaying stale tool artifacts to other agents.
|
||||
cleaned_response = clean_conversation_for_handoff(response.messages)
|
||||
# Append the agent response to the full conversation history. This list removes
|
||||
# function call related content such that the result stays consistent regardless
|
||||
# of which agent yields the final output.
|
||||
self._full_conversation.extend(cleaned_response)
|
||||
# Broadcast the cleaned response to all other agents
|
||||
|
||||
# For internal tracking, preserve the full response (including function_calls)
|
||||
# in _full_conversation so that Azure OpenAI can match function_calls with
|
||||
# function_results when the workflow resumes after user approvals.
|
||||
self._full_conversation.extend(response.messages)
|
||||
self._persist_missing_approved_function_results(
|
||||
runtime_tool_messages=runtime_tool_messages,
|
||||
response_messages=response.messages,
|
||||
)
|
||||
|
||||
# Broadcast only the cleaned response to other agents (without function_calls/results)
|
||||
await self._broadcast_messages(cleaned_response, cast(WorkflowContext[AgentExecutorRequest], ctx))
|
||||
|
||||
# Check if a handoff was requested
|
||||
@@ -422,6 +546,12 @@ class HandoffAgentExecutor(AgentExecutor):
|
||||
self._autonomous_mode_turns = 0 # Reset autonomous mode turn counter on handoff
|
||||
return
|
||||
|
||||
# Re-evaluate termination after appending and broadcasting this response.
|
||||
# Without this check, workflows that become terminal due to the latest assistant
|
||||
# message would still emit request_info and require an unnecessary extra resume.
|
||||
if await self._check_terminate_and_yield(cast(WorkflowContext[Never, list[Message]], ctx)):
|
||||
return
|
||||
|
||||
# Handle case where no handoff was requested
|
||||
if self._autonomous_mode and self._autonomous_mode_turns < self._autonomous_mode_turn_limit:
|
||||
# In autonomous mode, continue running the agent until a handoff is requested
|
||||
@@ -497,22 +627,20 @@ class HandoffAgentExecutor(AgentExecutor):
|
||||
last_message = response.messages[-1]
|
||||
for content in last_message.contents:
|
||||
if content.type == "function_result":
|
||||
if not content.result:
|
||||
continue
|
||||
|
||||
parsed_result: dict[str, Any] | None = None
|
||||
if isinstance(content.result, dict):
|
||||
parsed_result = content.result
|
||||
elif isinstance(content.result, str):
|
||||
payload = content.result
|
||||
parsed_payload: dict[str, Any] | None = None
|
||||
if isinstance(payload, dict):
|
||||
parsed_payload = payload
|
||||
elif isinstance(payload, str):
|
||||
try:
|
||||
loaded_result = json.loads(content.result)
|
||||
maybe_payload = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(loaded_result, dict):
|
||||
parsed_result = loaded_result
|
||||
maybe_payload = None
|
||||
if isinstance(maybe_payload, dict):
|
||||
parsed_payload = maybe_payload
|
||||
|
||||
if parsed_result is not None:
|
||||
handoff_target = parsed_result.get(HANDOFF_FUNCTION_RESULT_KEY)
|
||||
if parsed_payload:
|
||||
handoff_target = parsed_payload.get(HANDOFF_FUNCTION_RESULT_KEY)
|
||||
if isinstance(handoff_target, str):
|
||||
return handoff_target
|
||||
else:
|
||||
|
||||
+20
-44
@@ -14,57 +14,33 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def clean_conversation_for_handoff(conversation: list[Message]) -> list[Message]:
|
||||
"""Remove tool-related content from conversation for clean handoffs.
|
||||
"""Keep only plain text chat history for handoff routing.
|
||||
|
||||
During handoffs, tool calls can cause API errors because:
|
||||
1. Assistant messages with tool_calls must be followed by tool responses
|
||||
2. Tool response messages must follow an assistant message with tool_calls
|
||||
Handoff executors must not replay prior tool-control artifacts (function calls,
|
||||
tool outputs, approval payloads) into future model turns, or providers may reject
|
||||
the next request due to unmatched tool-call state.
|
||||
|
||||
This creates a cleaned copy removing ALL tool-related content.
|
||||
|
||||
Removes:
|
||||
- function_approval_request and function_call from assistant messages
|
||||
- Tool response messages (role="tool")
|
||||
- Messages with only tool calls and no text
|
||||
|
||||
Preserves:
|
||||
- User messages
|
||||
- Assistant messages with text content
|
||||
|
||||
Args:
|
||||
conversation: Original conversation with potential tool content
|
||||
|
||||
Returns:
|
||||
Cleaned conversation safe for handoff routing
|
||||
This helper builds a text-only copy of the conversation:
|
||||
- Drops all non-text content from every message.
|
||||
- Drops messages with no remaining text content.
|
||||
- Preserves original roles and author names for retained text messages.
|
||||
"""
|
||||
cleaned: list[Message] = []
|
||||
for msg in conversation:
|
||||
# Skip tool response messages entirely
|
||||
if msg.role == "tool":
|
||||
# Keep only plain text history for handoff routing. Tool-control content
|
||||
# (function_call/function_result/approval payloads) is runtime-only and
|
||||
# must not be replayed in future model turns.
|
||||
text_parts = [content.text for content in msg.contents if content.type == "text" and content.text]
|
||||
if not text_parts:
|
||||
continue
|
||||
|
||||
# Check for tool-related content
|
||||
has_tool_content = False
|
||||
if msg.contents:
|
||||
has_tool_content = any(
|
||||
content.type in ("function_approval_request", "function_call") for content in msg.contents
|
||||
)
|
||||
|
||||
# If no tool content, keep original
|
||||
if not has_tool_content:
|
||||
cleaned.append(msg)
|
||||
continue
|
||||
|
||||
# Has tool content - only keep if it also has text
|
||||
if msg.text and msg.text.strip():
|
||||
# Create fresh text-only message while preserving additional_properties
|
||||
msg_copy = Message(
|
||||
role=msg.role,
|
||||
text=msg.text,
|
||||
author_name=msg.author_name,
|
||||
additional_properties=dict(msg.additional_properties) if msg.additional_properties else None,
|
||||
)
|
||||
cleaned.append(msg_copy)
|
||||
msg_copy = Message(
|
||||
role=msg.role,
|
||||
text=" ".join(text_parts),
|
||||
author_name=msg.author_name,
|
||||
additional_properties=dict(msg.additional_properties) if msg.additional_properties else None,
|
||||
)
|
||||
cleaned.append(msg_copy)
|
||||
|
||||
return cleaned
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import re
|
||||
from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
|
||||
from typing import Any, cast
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
@@ -15,18 +16,21 @@ from agent_framework import (
|
||||
ResponseStream,
|
||||
WorkflowEvent,
|
||||
resolve_agent_id,
|
||||
tool,
|
||||
)
|
||||
from agent_framework._clients import BaseChatClient
|
||||
from agent_framework._middleware import ChatMiddlewareLayer, FunctionInvocationContext, MiddlewareTermination
|
||||
from agent_framework._tools import FunctionInvocationLayer, FunctionTool, tool
|
||||
from agent_framework._tools import FunctionInvocationLayer, FunctionTool
|
||||
from agent_framework.orchestrations import HandoffAgentUserRequest, HandoffBuilder
|
||||
|
||||
from agent_framework_orchestrations._handoff import (
|
||||
HANDOFF_FUNCTION_RESULT_KEY,
|
||||
HandoffAgentExecutor,
|
||||
HandoffConfiguration,
|
||||
_AutoHandoffMiddleware, # pyright: ignore[reportPrivateUsage]
|
||||
get_handoff_tool_name,
|
||||
)
|
||||
from agent_framework_orchestrations._orchestrator_helpers import clean_conversation_for_handoff
|
||||
|
||||
|
||||
class MockChatClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
||||
@@ -130,6 +134,67 @@ class MockHandoffAgent(Agent):
|
||||
super().__init__(client=MockChatClient(name=name, handoff_to=handoff_to), name=name, id=name)
|
||||
|
||||
|
||||
class ContextAwareRefundClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
||||
"""Mock client that expects prior user context to remain available on resume."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
ChatMiddlewareLayer.__init__(self)
|
||||
FunctionInvocationLayer.__init__(self)
|
||||
BaseChatClient.__init__(self)
|
||||
self._call_index = 0
|
||||
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
del kwargs
|
||||
del options
|
||||
|
||||
contents = self._next_contents(messages)
|
||||
if stream:
|
||||
return self._build_streaming_response(contents)
|
||||
|
||||
async def _get() -> ChatResponse:
|
||||
return ChatResponse(messages=[Message(role="assistant", contents=contents)], response_id="context-aware")
|
||||
|
||||
return _get()
|
||||
|
||||
def _build_streaming_response(self, contents: list[Content]) -> ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
|
||||
|
||||
def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
|
||||
return ChatResponse.from_updates(updates)
|
||||
|
||||
return ResponseStream(_stream(), finalizer=_finalize)
|
||||
|
||||
def _next_contents(self, messages: Sequence[Message]) -> list[Content]:
|
||||
user_text = " ".join(message.text or "" for message in messages if message.role == "user")
|
||||
order_match = re.search(r"\b(\d{4,12})\b", user_text)
|
||||
order_id = order_match.group(1) if order_match else None
|
||||
asks_refund = any(token in user_text.lower() for token in ("broken", "damaged", "refund", "cracked"))
|
||||
|
||||
if self._call_index == 0:
|
||||
reply = "Refund Agent: Please share your order number."
|
||||
elif self._call_index == 1:
|
||||
if order_id:
|
||||
reply = f"Refund Agent: Thanks, I found order {order_id}. Why do you need the refund?"
|
||||
else:
|
||||
reply = "Refund Agent: I still need your order number."
|
||||
else:
|
||||
if order_id and asks_refund:
|
||||
reply = f"Refund Agent: Got it for order {order_id}. I can proceed with your refund."
|
||||
else:
|
||||
reply = "Refund Agent: I still need your order number."
|
||||
|
||||
self._call_index += 1
|
||||
return [Content.from_text(text=reply)]
|
||||
|
||||
|
||||
async def _drain(stream: AsyncIterable[WorkflowEvent]) -> list[WorkflowEvent]:
|
||||
return [event async for event in stream]
|
||||
|
||||
@@ -168,6 +233,567 @@ async def test_handoff():
|
||||
assert request.source_executor_id == escalation.name
|
||||
|
||||
|
||||
def _latest_request_info_event(events: list[WorkflowEvent]) -> WorkflowEvent[Any]:
|
||||
request_events = [event for event in events if event.type == "request_info"]
|
||||
assert request_events
|
||||
request_event = request_events[-1]
|
||||
assert isinstance(request_event.data, HandoffAgentUserRequest)
|
||||
return request_event
|
||||
|
||||
|
||||
def _request_text(event: WorkflowEvent[Any]) -> str:
|
||||
request_payload = cast(HandoffAgentUserRequest, event.data)
|
||||
messages = request_payload.agent_response.messages
|
||||
assert messages
|
||||
return messages[-1].text or ""
|
||||
|
||||
|
||||
async def test_resume_keeps_prior_user_context_for_same_agent() -> None:
|
||||
"""Ensure same-agent request_info resumes retain prior turn context."""
|
||||
refund_agent = Agent(
|
||||
id="refund_agent",
|
||||
name="refund_agent",
|
||||
client=ContextAwareRefundClient(),
|
||||
)
|
||||
workflow = (
|
||||
HandoffBuilder(participants=[refund_agent], termination_condition=lambda _: False)
|
||||
.with_start_agent(refund_agent)
|
||||
.build()
|
||||
)
|
||||
|
||||
first_events = await _drain(workflow.run("My order arrived damaged.", stream=True))
|
||||
first_request = _latest_request_info_event(first_events)
|
||||
assert "order number" in _request_text(first_request).lower()
|
||||
|
||||
second_events = await _drain(
|
||||
workflow.run(
|
||||
stream=True,
|
||||
responses={first_request.request_id: [Message(role="user", text="Order 2939393")]},
|
||||
)
|
||||
)
|
||||
second_request = _latest_request_info_event(second_events)
|
||||
second_text = _request_text(second_request).lower()
|
||||
assert "order 2939393" in second_text
|
||||
assert "order number" not in second_text
|
||||
|
||||
third_events = await _drain(
|
||||
workflow.run(
|
||||
stream=True,
|
||||
responses={second_request.request_id: [Message(role="user", text="It arrived broken and unusable.")]},
|
||||
)
|
||||
)
|
||||
third_request = _latest_request_info_event(third_events)
|
||||
third_text = _request_text(third_request).lower()
|
||||
assert "order 2939393" in third_text
|
||||
assert "order number" not in third_text
|
||||
|
||||
|
||||
async def test_tool_approval_responses_are_not_replayed_from_history() -> None:
|
||||
"""Ensure persisted history does not re-execute previously approved tool calls."""
|
||||
execution_count = 0
|
||||
|
||||
@tool(name="submit_refund_counted", approval_mode="always_require")
|
||||
def submit_refund_counted() -> str:
|
||||
nonlocal execution_count
|
||||
execution_count += 1
|
||||
return "ok"
|
||||
|
||||
class ApprovalReplayClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
||||
def __init__(self) -> None:
|
||||
ChatMiddlewareLayer.__init__(self)
|
||||
FunctionInvocationLayer.__init__(self)
|
||||
BaseChatClient.__init__(self)
|
||||
self._call_index = 0
|
||||
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
del messages
|
||||
del options
|
||||
del kwargs
|
||||
|
||||
if self._call_index == 0:
|
||||
contents = [
|
||||
Content.from_function_call(
|
||||
call_id="refund-call-1",
|
||||
name="submit_refund_counted",
|
||||
arguments={},
|
||||
)
|
||||
]
|
||||
elif self._call_index == 1:
|
||||
contents = [Content.from_text(text="Refund approved and recorded.")]
|
||||
else:
|
||||
contents = [Content.from_text(text="No additional tool work needed.")]
|
||||
self._call_index += 1
|
||||
|
||||
if stream:
|
||||
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
|
||||
|
||||
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
|
||||
|
||||
async def _get() -> ChatResponse:
|
||||
return ChatResponse(
|
||||
messages=[Message(role="assistant", contents=contents)],
|
||||
response_id="approval-replay",
|
||||
)
|
||||
|
||||
return _get()
|
||||
|
||||
agent = Agent(
|
||||
id="refund_agent",
|
||||
name="refund_agent",
|
||||
client=ApprovalReplayClient(),
|
||||
tools=[submit_refund_counted],
|
||||
)
|
||||
workflow = (
|
||||
HandoffBuilder(participants=[agent], termination_condition=lambda _: False).with_start_agent(agent).build()
|
||||
)
|
||||
|
||||
first_events = await _drain(workflow.run("start", stream=True))
|
||||
first_requests = [event for event in first_events if event.type == "request_info"]
|
||||
assert first_requests
|
||||
first_request = first_requests[-1]
|
||||
assert isinstance(first_request.data, Content)
|
||||
approval_response = first_request.data.to_function_approval_response(approved=True)
|
||||
|
||||
second_events = await _drain(workflow.run(stream=True, responses={first_request.request_id: approval_response}))
|
||||
second_request = _latest_request_info_event(second_events)
|
||||
|
||||
await _drain(
|
||||
workflow.run(
|
||||
stream=True,
|
||||
responses={second_request.request_id: [Message(role="user", text="Thanks, what's next?")]},
|
||||
)
|
||||
)
|
||||
|
||||
assert execution_count == 1
|
||||
|
||||
|
||||
async def test_handoff_resume_preserves_approval_function_call_for_stateless_runs() -> None:
|
||||
"""Approval resume turns must replay matching function calls when store=False."""
|
||||
|
||||
@tool(name="submit_refund", approval_mode="always_require")
|
||||
def submit_refund() -> str:
|
||||
return "ok"
|
||||
|
||||
class StrictStatelessApprovalClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
||||
def __init__(self) -> None:
|
||||
ChatMiddlewareLayer.__init__(self)
|
||||
FunctionInvocationLayer.__init__(self)
|
||||
BaseChatClient.__init__(self)
|
||||
self._call_index = 0
|
||||
self.resume_validated = False
|
||||
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
del options
|
||||
del kwargs
|
||||
|
||||
if self._call_index == 0:
|
||||
contents = [
|
||||
Content.from_function_call(
|
||||
call_id="refund-call-1",
|
||||
name="submit_refund",
|
||||
arguments={},
|
||||
)
|
||||
]
|
||||
else:
|
||||
function_call_ids = {
|
||||
content.call_id
|
||||
for message in messages
|
||||
for content in message.contents
|
||||
if content.type == "function_call" and content.call_id
|
||||
}
|
||||
function_result_ids = {
|
||||
content.call_id
|
||||
for message in messages
|
||||
for content in message.contents
|
||||
if content.type == "function_result" and content.call_id
|
||||
}
|
||||
missing_call_ids = sorted(function_result_ids - function_call_ids)
|
||||
if missing_call_ids:
|
||||
raise AssertionError(
|
||||
f"No tool call found for function call output with call_id {missing_call_ids[0]}."
|
||||
)
|
||||
self.resume_validated = True
|
||||
contents = [Content.from_text(text="Refund submitted.")]
|
||||
|
||||
self._call_index += 1
|
||||
|
||||
if stream:
|
||||
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
|
||||
|
||||
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
|
||||
|
||||
async def _get() -> ChatResponse:
|
||||
return ChatResponse(
|
||||
messages=[Message(role="assistant", contents=contents)],
|
||||
response_id="strict-stateless",
|
||||
)
|
||||
|
||||
return _get()
|
||||
|
||||
client = StrictStatelessApprovalClient()
|
||||
agent = Agent(
|
||||
id="refund_agent",
|
||||
name="refund_agent",
|
||||
client=client,
|
||||
tools=[submit_refund],
|
||||
)
|
||||
workflow = (
|
||||
HandoffBuilder(participants=[agent], termination_condition=lambda _: False).with_start_agent(agent).build()
|
||||
)
|
||||
|
||||
first_events = await _drain(workflow.run("start", stream=True))
|
||||
approval_requests = [
|
||||
event for event in first_events if event.type == "request_info" and isinstance(event.data, Content)
|
||||
]
|
||||
assert approval_requests
|
||||
first_request = approval_requests[0]
|
||||
|
||||
approval_response = first_request.data.to_function_approval_response(True)
|
||||
await _drain(workflow.run(stream=True, responses={first_request.request_id: approval_response}))
|
||||
|
||||
assert client.resume_validated is True
|
||||
|
||||
|
||||
async def test_handoff_replay_serializes_handoff_function_results() -> None:
|
||||
"""Returning to the same agent must not replay dict tool outputs."""
|
||||
|
||||
class ReplaySafeHandoffClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
||||
def __init__(self, name: str, handoff_sequence: list[str | None]) -> None:
|
||||
ChatMiddlewareLayer.__init__(self)
|
||||
FunctionInvocationLayer.__init__(self)
|
||||
BaseChatClient.__init__(self)
|
||||
self._name = name
|
||||
self._handoff_sequence = handoff_sequence
|
||||
self._call_index = 0
|
||||
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
del options
|
||||
del kwargs
|
||||
|
||||
for message in messages:
|
||||
for content in message.contents:
|
||||
if content.type == "function_result" and isinstance(content.result, dict):
|
||||
raise AssertionError("Expected replayed function_result payloads to be JSON strings.")
|
||||
|
||||
handoff_to = (
|
||||
self._handoff_sequence[self._call_index] if self._call_index < len(self._handoff_sequence) else None
|
||||
)
|
||||
call_id = f"{self._name}-handoff-{self._call_index}" if handoff_to else None
|
||||
contents = _build_reply_contents(self._name, handoff_to, call_id)
|
||||
self._call_index += 1
|
||||
|
||||
if stream:
|
||||
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
|
||||
|
||||
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
|
||||
|
||||
async def _get() -> ChatResponse:
|
||||
return ChatResponse(messages=[Message(role="assistant", contents=contents)], response_id="replay-safe")
|
||||
|
||||
return _get()
|
||||
|
||||
triage = Agent(
|
||||
id="triage",
|
||||
name="triage",
|
||||
client=ReplaySafeHandoffClient(name="triage", handoff_sequence=["specialist", None]),
|
||||
)
|
||||
specialist = Agent(
|
||||
id="specialist",
|
||||
name="specialist",
|
||||
client=ReplaySafeHandoffClient(name="specialist", handoff_sequence=["triage"]),
|
||||
)
|
||||
|
||||
workflow = (
|
||||
HandoffBuilder(participants=[triage, specialist], termination_condition=lambda _: False)
|
||||
.with_start_agent(triage)
|
||||
.build()
|
||||
)
|
||||
|
||||
events = await _drain(workflow.run("start", stream=True))
|
||||
requests = [event for event in events if event.type == "request_info"]
|
||||
assert requests
|
||||
assert requests[-1].source_executor_id == triage.name
|
||||
|
||||
|
||||
async def test_handoff_resume_preserves_approved_tool_output_for_stateless_runs() -> None:
|
||||
"""Approved calls must keep function_call/function_result pairs for later replays."""
|
||||
submit_call_id = "call_submit_refund_approved"
|
||||
|
||||
@tool(name="submit_refund", approval_mode="always_require")
|
||||
def submit_refund() -> str:
|
||||
return "submitted"
|
||||
|
||||
class RefundReplayClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
||||
def __init__(self) -> None:
|
||||
ChatMiddlewareLayer.__init__(self)
|
||||
FunctionInvocationLayer.__init__(self)
|
||||
BaseChatClient.__init__(self)
|
||||
self._call_index = 0
|
||||
self.resume_validated = False
|
||||
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
del options
|
||||
del kwargs
|
||||
|
||||
if self._call_index == 0:
|
||||
contents = [Content.from_function_call(call_id=submit_call_id, name="submit_refund", arguments={})]
|
||||
elif self._call_index == 1:
|
||||
contents = _build_reply_contents("refund_agent", "order_agent", "refund-order-handoff-1")
|
||||
else:
|
||||
function_call_ids = {
|
||||
content.call_id
|
||||
for message in messages
|
||||
for content in message.contents
|
||||
if content.type == "function_call" and content.call_id
|
||||
}
|
||||
function_result_ids = {
|
||||
content.call_id
|
||||
for message in messages
|
||||
for content in message.contents
|
||||
if content.type == "function_result" and content.call_id
|
||||
}
|
||||
if submit_call_id in function_call_ids and submit_call_id not in function_result_ids:
|
||||
raise AssertionError(f"No tool output found for function call {submit_call_id}.")
|
||||
self.resume_validated = True
|
||||
contents = [Content.from_text(text="Refund agent resumed.")]
|
||||
|
||||
self._call_index += 1
|
||||
|
||||
if stream:
|
||||
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
|
||||
|
||||
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
|
||||
|
||||
async def _get() -> ChatResponse:
|
||||
return ChatResponse(
|
||||
messages=[Message(role="assistant", contents=contents)],
|
||||
response_id="refund-replay",
|
||||
)
|
||||
|
||||
return _get()
|
||||
|
||||
class OrderReplayClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
||||
def __init__(self) -> None:
|
||||
ChatMiddlewareLayer.__init__(self)
|
||||
FunctionInvocationLayer.__init__(self)
|
||||
BaseChatClient.__init__(self)
|
||||
self._call_index = 0
|
||||
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
del messages
|
||||
del options
|
||||
del kwargs
|
||||
|
||||
if self._call_index == 0:
|
||||
contents = [Content.from_text(text="Would you like a replacement or a refund?")]
|
||||
else:
|
||||
contents = _build_reply_contents("order_agent", "refund_agent", "order-refund-handoff-1")
|
||||
self._call_index += 1
|
||||
|
||||
if stream:
|
||||
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
|
||||
|
||||
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
|
||||
|
||||
async def _get() -> ChatResponse:
|
||||
return ChatResponse(messages=[Message(role="assistant", contents=contents)], response_id="order-replay")
|
||||
|
||||
return _get()
|
||||
|
||||
refund_client = RefundReplayClient()
|
||||
refund_agent = Agent(
|
||||
id="refund_agent",
|
||||
name="refund_agent",
|
||||
client=refund_client,
|
||||
tools=[submit_refund],
|
||||
)
|
||||
order_agent = Agent(
|
||||
id="order_agent",
|
||||
name="order_agent",
|
||||
client=OrderReplayClient(),
|
||||
)
|
||||
workflow = (
|
||||
HandoffBuilder(participants=[refund_agent, order_agent], termination_condition=lambda _: False)
|
||||
.with_start_agent(refund_agent)
|
||||
.build()
|
||||
)
|
||||
|
||||
first_events = await _drain(workflow.run("start", stream=True))
|
||||
approval_requests = [
|
||||
event for event in first_events if event.type == "request_info" and isinstance(event.data, Content)
|
||||
]
|
||||
assert approval_requests
|
||||
approval_request = approval_requests[-1]
|
||||
approval_response = approval_request.data.to_function_approval_response(True)
|
||||
|
||||
second_events = await _drain(workflow.run(stream=True, responses={approval_request.request_id: approval_response}))
|
||||
order_request = _latest_request_info_event(second_events)
|
||||
assert order_request.source_executor_id == order_agent.name
|
||||
|
||||
await _drain(
|
||||
workflow.run(
|
||||
stream=True,
|
||||
responses={order_request.request_id: [Message(role="user", text="Please continue with refund.")]},
|
||||
)
|
||||
)
|
||||
|
||||
assert refund_client.resume_validated is True
|
||||
|
||||
|
||||
def test_handoff_clone_disables_provider_side_storage() -> None:
|
||||
"""Handoff executors should force store=False to avoid stale provider call state."""
|
||||
triage = MockHandoffAgent(name="triage")
|
||||
workflow = HandoffBuilder(participants=[triage]).with_start_agent(triage).build()
|
||||
|
||||
executor = workflow.executors[resolve_agent_id(triage)]
|
||||
assert isinstance(executor, HandoffAgentExecutor)
|
||||
assert executor._agent.default_options.get("store") is False
|
||||
|
||||
|
||||
async def test_handoff_clears_stale_service_session_id_before_run() -> None:
|
||||
"""Stale service session IDs must be dropped before each handoff agent turn."""
|
||||
triage = MockHandoffAgent(name="triage", handoff_to="specialist")
|
||||
specialist = MockHandoffAgent(name="specialist")
|
||||
workflow = HandoffBuilder(participants=[triage, specialist]).with_start_agent(triage).build()
|
||||
|
||||
triage_executor = workflow.executors[resolve_agent_id(triage)]
|
||||
assert isinstance(triage_executor, HandoffAgentExecutor)
|
||||
triage_executor._session.service_session_id = "resp_stale_value"
|
||||
|
||||
await _drain(workflow.run("My order is damaged", stream=True))
|
||||
|
||||
assert triage_executor._session.service_session_id is None
|
||||
|
||||
|
||||
def test_clean_conversation_for_handoff_keeps_text_only_history() -> None:
|
||||
"""Tool-control messages must be excluded from persisted handoff history."""
|
||||
function_call = Content.from_function_call(
|
||||
call_id="handoff-call-1",
|
||||
name="handoff_to_refund_agent",
|
||||
arguments={"context": "route to refund"},
|
||||
)
|
||||
approval_response = Content.from_function_approval_response(
|
||||
approved=True,
|
||||
id="approval-1",
|
||||
function_call=function_call,
|
||||
)
|
||||
|
||||
conversation = [
|
||||
Message(role="user", text="My order arrived damaged."),
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
function_call,
|
||||
Content.from_text(text="Triage Agent: Routing you to Refund."),
|
||||
],
|
||||
),
|
||||
Message(role="tool", contents=[Content.from_function_result(call_id="handoff-call-1", result="ok")]),
|
||||
Message(role="user", contents=[approval_response]),
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[Content.from_function_call(call_id="handoff-call-2", name="handoff_to_order_agent")],
|
||||
),
|
||||
]
|
||||
|
||||
cleaned = clean_conversation_for_handoff(conversation)
|
||||
assert [message.role for message in cleaned] == ["user", "assistant"]
|
||||
assert [message.text for message in cleaned] == [
|
||||
"My order arrived damaged.",
|
||||
"Triage Agent: Routing you to Refund.",
|
||||
]
|
||||
|
||||
|
||||
def test_persist_missing_approved_function_results_handles_runtime_and_fallback_outputs() -> None:
|
||||
"""Persisted history should retain approved call outputs across runtime shapes."""
|
||||
agent = MockHandoffAgent(name="triage")
|
||||
executor = HandoffAgentExecutor(agent, handoffs=[])
|
||||
|
||||
call_with_runtime_result = "call-runtime-result"
|
||||
call_with_approval_only = "call-approval-only"
|
||||
|
||||
executor._full_conversation = [
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id=call_with_runtime_result, name="submit_refund", arguments={}),
|
||||
Content.from_function_call(call_id=call_with_approval_only, name="submit_refund", arguments={}),
|
||||
],
|
||||
)
|
||||
]
|
||||
|
||||
approval_response = Content.from_function_approval_response(
|
||||
approved=True,
|
||||
id=call_with_approval_only,
|
||||
function_call=Content.from_function_call(call_id=call_with_approval_only, name="submit_refund", arguments={}),
|
||||
)
|
||||
runtime_messages = [
|
||||
Message(
|
||||
role="tool",
|
||||
contents=[Content.from_function_result(call_id=call_with_runtime_result, result='{"submitted":true}')],
|
||||
),
|
||||
Message(role="user", contents=[approval_response]),
|
||||
]
|
||||
|
||||
executor._persist_missing_approved_function_results(runtime_tool_messages=runtime_messages, response_messages=[])
|
||||
|
||||
persisted_tool_messages = [message for message in executor._full_conversation if message.role == "tool"]
|
||||
assert persisted_tool_messages
|
||||
persisted_results = [
|
||||
content
|
||||
for message in persisted_tool_messages
|
||||
for content in message.contents
|
||||
if content.type == "function_result" and content.call_id
|
||||
]
|
||||
result_by_call_id = {content.call_id: content.result for content in persisted_results}
|
||||
assert result_by_call_id[call_with_runtime_result] == '{"submitted":true}'
|
||||
assert result_by_call_id[call_with_approval_only] == '{"status":"approved"}'
|
||||
|
||||
|
||||
async def test_autonomous_mode_yields_output_without_user_request():
|
||||
"""Ensure autonomous interaction mode yields output without requesting user input."""
|
||||
triage = MockHandoffAgent(name="triage", handoff_to="specialist")
|
||||
@@ -278,6 +904,61 @@ async def test_handoff_async_termination_condition() -> None:
|
||||
assert termination_call_count > 0
|
||||
|
||||
|
||||
async def test_handoff_terminates_without_request_info_when_latest_response_meets_condition() -> None:
|
||||
"""Termination triggered by the latest assistant response should not emit request_info."""
|
||||
|
||||
class FinalizingClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
||||
def __init__(self) -> None:
|
||||
ChatMiddlewareLayer.__init__(self)
|
||||
FunctionInvocationLayer.__init__(self)
|
||||
BaseChatClient.__init__(self)
|
||||
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
del messages, options, kwargs
|
||||
contents = [Content.from_text(text="Replacement request submitted. Case complete.")]
|
||||
|
||||
if stream:
|
||||
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
|
||||
|
||||
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
|
||||
|
||||
async def _get() -> ChatResponse:
|
||||
return ChatResponse(messages=[Message(role="assistant", contents=contents)], response_id="finalizing")
|
||||
|
||||
return _get()
|
||||
|
||||
agent = Agent(id="order_agent", name="order_agent", client=FinalizingClient())
|
||||
workflow = (
|
||||
HandoffBuilder(
|
||||
participants=[agent],
|
||||
termination_condition=lambda conv: any(
|
||||
message.role == "assistant" and "case complete." in (message.text or "").lower() for message in conv
|
||||
),
|
||||
)
|
||||
.with_start_agent(agent)
|
||||
.build()
|
||||
)
|
||||
|
||||
events = await _drain(workflow.run("ship replacement", stream=True))
|
||||
|
||||
requests = [event for event in events if event.type == "request_info"]
|
||||
assert not requests
|
||||
|
||||
outputs = [event for event in events if event.type == "output"]
|
||||
assert outputs
|
||||
conversation_outputs = [event for event in outputs if isinstance(event.data, list)]
|
||||
assert len(conversation_outputs) == 1
|
||||
|
||||
|
||||
async def test_tool_choice_preserved_from_agent_config():
|
||||
"""Verify that agent-level tool_choice configuration is preserved and not overridden."""
|
||||
# Create a mock chat client that records the tool_choice used
|
||||
|
||||
Reference in New Issue
Block a user