mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: fix(ag-ui): Execute tools with approval_mode, fix shared state, code cleanup (#3079)
* fix(ag-ui): execute tools after approval in human-in-the-loop flow * Fix shared state bug * Bug fix finalized * Refactoring to clean up code * Code cleanup * More fixes * More code cleanup * Add version detection in __init__.py to ruff ignore list
This commit is contained in:
Evan Mattson
committed by
GitHub
Unverified
GitHub
parent
50d34aec91
commit
88968da0bd
@@ -630,3 +630,179 @@ async def test_suppressed_summary_with_document_state():
|
||||
# Should contain some reference to the document
|
||||
full_text = "".join(e.delta for e in text_events)
|
||||
assert "written" in full_text.lower() or "document" in full_text.lower()
|
||||
|
||||
|
||||
async def test_function_approval_mode_executes_tool():
|
||||
"""Test that function approval with approval_mode='always_require' sends the correct messages."""
|
||||
from agent_framework import FunctionResultContent, ai_function
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
messages_received: list[Any] = []
|
||||
|
||||
@ai_function(
|
||||
name="get_datetime",
|
||||
description="Get the current date and time",
|
||||
approval_mode="always_require",
|
||||
)
|
||||
def get_datetime() -> str:
|
||||
return "2025/12/01 12:00:00"
|
||||
|
||||
async def stream_fn(
|
||||
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
|
||||
) -> AsyncIterator[ChatResponseUpdate]:
|
||||
# Capture the messages received by the chat client
|
||||
messages_received.clear()
|
||||
messages_received.extend(messages)
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Processing completed")])
|
||||
|
||||
agent = ChatAgent(
|
||||
name="test_agent",
|
||||
instructions="Test",
|
||||
chat_client=StreamingChatClientStub(stream_fn),
|
||||
tools=[get_datetime],
|
||||
)
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
# Simulate the conversation history with:
|
||||
# 1. User message asking for time
|
||||
# 2. Assistant message with the function call that needs approval
|
||||
# 3. Tool approval message from user
|
||||
tool_result: dict[str, Any] = {"accepted": True}
|
||||
input_data: dict[str, Any] = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What time is it?",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_get_datetime_123",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_datetime",
|
||||
"arguments": "{}",
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(tool_result),
|
||||
"toolCallId": "call_get_datetime_123",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
events: list[Any] = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Verify the run completed successfully
|
||||
run_started = [e for e in events if e.type == "RUN_STARTED"]
|
||||
run_finished = [e for e in events if e.type == "RUN_FINISHED"]
|
||||
assert len(run_started) == 1
|
||||
assert len(run_finished) == 1
|
||||
|
||||
# Verify that a FunctionResultContent was created and sent to the agent
|
||||
# Approved tool calls are resolved before the model run.
|
||||
tool_result_found = False
|
||||
for msg in messages_received:
|
||||
for content in msg.contents:
|
||||
if isinstance(content, FunctionResultContent):
|
||||
tool_result_found = True
|
||||
assert content.call_id == "call_get_datetime_123"
|
||||
assert content.result == "2025/12/01 12:00:00"
|
||||
break
|
||||
|
||||
assert tool_result_found, (
|
||||
"FunctionResultContent should be included in messages sent to agent. "
|
||||
"This is required for the model to see the approved tool execution result."
|
||||
)
|
||||
|
||||
|
||||
async def test_function_approval_mode_rejection():
|
||||
"""Test that function approval rejection creates a rejection response."""
|
||||
from agent_framework import FunctionResultContent, ai_function
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
messages_received: list[Any] = []
|
||||
|
||||
@ai_function(
|
||||
name="delete_all_data",
|
||||
description="Delete all user data",
|
||||
approval_mode="always_require",
|
||||
)
|
||||
def delete_all_data() -> str:
|
||||
return "All data deleted"
|
||||
|
||||
async def stream_fn(
|
||||
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
|
||||
) -> AsyncIterator[ChatResponseUpdate]:
|
||||
# Capture the messages received by the chat client
|
||||
messages_received.clear()
|
||||
messages_received.extend(messages)
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Operation cancelled")])
|
||||
|
||||
agent = ChatAgent(
|
||||
name="test_agent",
|
||||
instructions="Test",
|
||||
chat_client=StreamingChatClientStub(stream_fn),
|
||||
tools=[delete_all_data],
|
||||
)
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
# Simulate rejection
|
||||
tool_result: dict[str, Any] = {"accepted": False}
|
||||
input_data: dict[str, Any] = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Delete all my data",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_delete_123",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "delete_all_data",
|
||||
"arguments": "{}",
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(tool_result),
|
||||
"toolCallId": "call_delete_123",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
events: list[Any] = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Verify the run completed
|
||||
run_finished = [e for e in events if e.type == "RUN_FINISHED"]
|
||||
assert len(run_finished) == 1
|
||||
|
||||
# Verify that a FunctionResultContent with rejection payload was created
|
||||
rejection_found = False
|
||||
for msg in messages_received:
|
||||
for content in msg.contents:
|
||||
if isinstance(content, FunctionResultContent):
|
||||
rejection_found = True
|
||||
assert content.call_id == "call_delete_123"
|
||||
assert content.result == "Error: Tool call invocation was rejected by user."
|
||||
break
|
||||
|
||||
assert rejection_found, (
|
||||
"FunctionResultContent with rejection details should be included in messages sent to agent. "
|
||||
"This tells the model that the tool was rejected."
|
||||
)
|
||||
|
||||
@@ -52,8 +52,8 @@ async def test_tool_call_flow():
|
||||
update2 = AgentRunResponseUpdate(contents=[tool_result])
|
||||
events2 = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Should have: ToolCallEndEvent, ToolCallResultEvent, MessagesSnapshotEvent
|
||||
assert len(events2) == 3
|
||||
# Should have: ToolCallEndEvent, ToolCallResultEvent
|
||||
assert len(events2) == 2
|
||||
assert isinstance(events2[0], ToolCallEndEvent)
|
||||
assert isinstance(events2[1], ToolCallResultEvent)
|
||||
|
||||
|
||||
@@ -231,7 +231,12 @@ async def test_function_approval_request_basic():
|
||||
"""Test FunctionApprovalRequestContent conversion."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
# Set require_confirmation=False to test just the function_approval_request event
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
require_confirmation=False,
|
||||
)
|
||||
|
||||
func_call = FunctionCallContent(
|
||||
call_id="call_123",
|
||||
@@ -284,14 +289,12 @@ async def test_empty_predict_state_config():
|
||||
assert "STATE_DELTA" not in event_types
|
||||
assert "STATE_SNAPSHOT" not in event_types
|
||||
|
||||
# Should have: ToolCallStart, ToolCallArgs, ToolCallEnd, ToolCallResult, MessagesSnapshot
|
||||
# MessagesSnapshotEvent is emitted after tool results to track the conversation
|
||||
# Should have: ToolCallStart, ToolCallArgs, ToolCallEnd, ToolCallResult
|
||||
assert event_types == [
|
||||
"TOOL_CALL_START",
|
||||
"TOOL_CALL_ARGS",
|
||||
"TOOL_CALL_END",
|
||||
"TOOL_CALL_RESULT",
|
||||
"MESSAGES_SNAPSHOT",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ from agent_framework import (
|
||||
from agent_framework._clients import BaseChatClient
|
||||
from agent_framework._types import ChatResponse, ChatResponseUpdate
|
||||
|
||||
from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
|
||||
from agent_framework_ag_ui._orchestrators import ExecutionContext
|
||||
|
||||
StreamFn = Callable[..., AsyncIterator[ChatResponseUpdate]]
|
||||
@@ -134,5 +135,9 @@ class StubAgent(AgentProtocol):
|
||||
class TestExecutionContext(ExecutionContext):
|
||||
"""ExecutionContext helper that allows setting messages for tests."""
|
||||
|
||||
def set_messages(self, messages: list[ChatMessage]) -> None:
|
||||
self._messages = messages
|
||||
def set_messages(self, messages: list[ChatMessage], *, normalize: bool = True) -> None:
|
||||
if normalize:
|
||||
self._messages = _deduplicate_messages(_sanitize_tool_history(messages))
|
||||
else:
|
||||
self._messages = messages
|
||||
self._snapshot_messages = None
|
||||
|
||||
@@ -10,9 +10,11 @@ from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
async def test_function_approval_request_emission():
|
||||
"""Test that CustomEvent is emitted for FunctionApprovalRequestContent."""
|
||||
# Set require_confirmation=False to test just the function_approval_request event
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
require_confirmation=False,
|
||||
)
|
||||
|
||||
# Create approval request
|
||||
@@ -47,11 +49,65 @@ async def test_function_approval_request_emission():
|
||||
assert event.value["function_call"]["arguments"]["subject"] == "Test"
|
||||
|
||||
|
||||
async def test_multiple_approval_requests():
|
||||
"""Test handling multiple approval requests in one update."""
|
||||
async def test_function_approval_request_with_confirm_changes():
|
||||
"""Test that confirm_changes is also emitted when require_confirmation=True."""
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
require_confirmation=True,
|
||||
)
|
||||
|
||||
func_call = FunctionCallContent(
|
||||
call_id="call_456",
|
||||
name="delete_file",
|
||||
arguments={"path": "/tmp/test.txt"},
|
||||
)
|
||||
approval_request = FunctionApprovalRequestContent(
|
||||
id="approval_002",
|
||||
function_call=func_call,
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[approval_request])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should emit: ToolCallEndEvent, CustomEvent, and confirm_changes (Start, Args, End) = 5 events
|
||||
assert len(events) == 5
|
||||
|
||||
# Check ToolCallEndEvent
|
||||
assert events[0].type == "TOOL_CALL_END"
|
||||
assert events[0].tool_call_id == "call_456"
|
||||
|
||||
# Check function_approval_request CustomEvent
|
||||
assert events[1].type == "CUSTOM"
|
||||
assert events[1].name == "function_approval_request"
|
||||
|
||||
# Check confirm_changes tool call events
|
||||
assert events[2].type == "TOOL_CALL_START"
|
||||
assert events[2].tool_call_name == "confirm_changes"
|
||||
assert events[3].type == "TOOL_CALL_ARGS"
|
||||
# Verify confirm_changes includes function info for Dojo UI
|
||||
import json
|
||||
|
||||
args = json.loads(events[3].delta)
|
||||
assert args["function_name"] == "delete_file"
|
||||
assert args["function_call_id"] == "call_456"
|
||||
assert args["function_arguments"] == {"path": "/tmp/test.txt"}
|
||||
assert args["steps"] == [
|
||||
{
|
||||
"description": "Execute delete_file",
|
||||
"status": "enabled",
|
||||
}
|
||||
]
|
||||
assert events[4].type == "TOOL_CALL_END"
|
||||
|
||||
|
||||
async def test_multiple_approval_requests():
|
||||
"""Test handling multiple approval requests in one update."""
|
||||
# Set require_confirmation=False to simplify the test
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
require_confirmation=False,
|
||||
)
|
||||
|
||||
func_call_1 = FunctionCallContent(
|
||||
@@ -94,3 +150,32 @@ async def test_multiple_approval_requests():
|
||||
assert events[3].type == "CUSTOM"
|
||||
assert events[3].name == "function_approval_request"
|
||||
assert events[3].value["id"] == "approval_2"
|
||||
|
||||
|
||||
async def test_function_approval_request_sets_stop_flag():
|
||||
"""Test that function approval request sets should_stop_after_confirm flag.
|
||||
|
||||
This ensures the orchestrator stops the run after emitting the approval request,
|
||||
allowing the UI to send back an approval response.
|
||||
"""
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
)
|
||||
|
||||
assert bridge.should_stop_after_confirm is False
|
||||
|
||||
func_call = FunctionCallContent(
|
||||
call_id="call_stop_test",
|
||||
name="get_datetime",
|
||||
arguments={},
|
||||
)
|
||||
approval_request = FunctionApprovalRequestContent(
|
||||
id="approval_stop_test",
|
||||
function_call=func_call,
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[approval_request])
|
||||
await bridge.from_agent_run_update(update)
|
||||
|
||||
assert bridge.should_stop_after_confirm is True
|
||||
|
||||
@@ -2,12 +2,15 @@
|
||||
|
||||
"""Tests for message adapters."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, Role, TextContent
|
||||
|
||||
from agent_framework_ag_ui._message_adapters import (
|
||||
agent_framework_messages_to_agui,
|
||||
agui_messages_to_agent_framework,
|
||||
agui_messages_to_snapshot_format,
|
||||
extract_text_from_contents,
|
||||
)
|
||||
|
||||
@@ -43,6 +46,32 @@ def test_agent_framework_to_agui_basic(sample_agent_framework_message):
|
||||
assert messages[0]["id"] == "msg-123"
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_normalizes_dict_roles():
|
||||
"""Dict inputs normalize unknown roles for UI compatibility."""
|
||||
messages = [
|
||||
{"role": "developer", "content": "policy"},
|
||||
{"role": "weird_role", "content": "payload"},
|
||||
]
|
||||
|
||||
converted = agent_framework_messages_to_agui(messages)
|
||||
|
||||
assert converted[0]["role"] == "system"
|
||||
assert converted[1]["role"] == "user"
|
||||
|
||||
|
||||
def test_agui_snapshot_format_normalizes_roles():
|
||||
"""Snapshot normalization coerces roles into supported AG-UI values."""
|
||||
messages = [
|
||||
{"role": "Developer", "content": "policy"},
|
||||
{"role": "unknown", "content": "payload"},
|
||||
]
|
||||
|
||||
normalized = agui_messages_to_snapshot_format(messages)
|
||||
|
||||
assert normalized[0]["role"] == "system"
|
||||
assert normalized[1]["role"] == "user"
|
||||
|
||||
|
||||
def test_agui_tool_result_to_agent_framework():
|
||||
"""Test converting AG-UI tool result message to Agent Framework."""
|
||||
tool_result_message = {
|
||||
@@ -68,6 +97,237 @@ def test_agui_tool_result_to_agent_framework():
|
||||
assert message.additional_properties.get("tool_call_id") == "call_123"
|
||||
|
||||
|
||||
def test_agui_tool_approval_updates_tool_call_arguments():
|
||||
"""Tool approval updates matching tool call arguments for snapshots and agent context."""
|
||||
messages_input = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_123",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "generate_task_steps",
|
||||
"arguments": {
|
||||
"steps": [
|
||||
{"description": "Boil water", "status": "enabled"},
|
||||
{"description": "Brew coffee", "status": "enabled"},
|
||||
{"description": "Serve coffee", "status": "enabled"},
|
||||
]
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
"id": "msg_1",
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(
|
||||
{
|
||||
"accepted": True,
|
||||
"steps": [
|
||||
{"description": "Boil water", "status": "enabled"},
|
||||
{"description": "Serve coffee", "status": "enabled"},
|
||||
],
|
||||
}
|
||||
),
|
||||
"toolCallId": "call_123",
|
||||
"id": "msg_2",
|
||||
},
|
||||
]
|
||||
|
||||
messages = agui_messages_to_agent_framework(messages_input)
|
||||
|
||||
assert len(messages) == 2
|
||||
assistant_msg = messages[0]
|
||||
func_call = next(content for content in assistant_msg.contents if isinstance(content, FunctionCallContent))
|
||||
assert func_call.arguments == {
|
||||
"steps": [
|
||||
{"description": "Boil water", "status": "enabled"},
|
||||
{"description": "Brew coffee", "status": "disabled"},
|
||||
{"description": "Serve coffee", "status": "enabled"},
|
||||
]
|
||||
}
|
||||
assert messages_input[0]["tool_calls"][0]["function"]["arguments"] == {
|
||||
"steps": [
|
||||
{"description": "Boil water", "status": "enabled"},
|
||||
{"description": "Brew coffee", "status": "disabled"},
|
||||
{"description": "Serve coffee", "status": "enabled"},
|
||||
]
|
||||
}
|
||||
|
||||
from agent_framework import FunctionApprovalResponseContent
|
||||
|
||||
approval_msg = messages[1]
|
||||
approval_content = next(
|
||||
content for content in approval_msg.contents if isinstance(content, FunctionApprovalResponseContent)
|
||||
)
|
||||
assert approval_content.function_call.parse_arguments() == {
|
||||
"steps": [
|
||||
{"description": "Boil water", "status": "enabled"},
|
||||
{"description": "Serve coffee", "status": "enabled"},
|
||||
]
|
||||
}
|
||||
assert approval_content.additional_properties is not None
|
||||
assert approval_content.additional_properties.get("ag_ui_state_args") == {
|
||||
"steps": [
|
||||
{"description": "Boil water", "status": "enabled"},
|
||||
{"description": "Brew coffee", "status": "disabled"},
|
||||
{"description": "Serve coffee", "status": "enabled"},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def test_agui_tool_approval_from_confirm_changes_maps_to_function_call():
|
||||
"""Confirm_changes approvals map back to the original tool call when metadata is present."""
|
||||
messages_input = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_tool",
|
||||
"type": "function",
|
||||
"function": {"name": "get_datetime", "arguments": {}},
|
||||
},
|
||||
{
|
||||
"id": "call_confirm",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "confirm_changes",
|
||||
"arguments": {"function_call_id": "call_tool"},
|
||||
},
|
||||
},
|
||||
],
|
||||
"id": "msg_1",
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps({"accepted": True, "function_call_id": "call_tool"}),
|
||||
"toolCallId": "call_confirm",
|
||||
"id": "msg_2",
|
||||
},
|
||||
]
|
||||
|
||||
messages = agui_messages_to_agent_framework(messages_input)
|
||||
|
||||
from agent_framework import FunctionApprovalResponseContent
|
||||
|
||||
approval_msg = messages[1]
|
||||
approval_content = next(
|
||||
content for content in approval_msg.contents if isinstance(content, FunctionApprovalResponseContent)
|
||||
)
|
||||
|
||||
assert approval_content.function_call.call_id == "call_tool"
|
||||
assert approval_content.function_call.name == "get_datetime"
|
||||
assert approval_content.function_call.parse_arguments() == {}
|
||||
assert messages_input[0]["tool_calls"][0]["function"]["arguments"] == {}
|
||||
|
||||
|
||||
def test_agui_tool_approval_from_confirm_changes_falls_back_to_sibling_call():
|
||||
"""Confirm_changes approvals map to the only sibling tool call when metadata is missing."""
|
||||
messages_input = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_tool",
|
||||
"type": "function",
|
||||
"function": {"name": "get_datetime", "arguments": {}},
|
||||
},
|
||||
{
|
||||
"id": "call_confirm",
|
||||
"type": "function",
|
||||
"function": {"name": "confirm_changes", "arguments": {}},
|
||||
},
|
||||
],
|
||||
"id": "msg_1",
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(
|
||||
{
|
||||
"accepted": True,
|
||||
"steps": [{"description": "Approve get_datetime", "status": "enabled"}],
|
||||
}
|
||||
),
|
||||
"toolCallId": "call_confirm",
|
||||
"id": "msg_2",
|
||||
},
|
||||
]
|
||||
|
||||
messages = agui_messages_to_agent_framework(messages_input)
|
||||
|
||||
from agent_framework import FunctionApprovalResponseContent
|
||||
|
||||
approval_msg = messages[1]
|
||||
approval_content = next(
|
||||
content for content in approval_msg.contents if isinstance(content, FunctionApprovalResponseContent)
|
||||
)
|
||||
|
||||
assert approval_content.function_call.call_id == "call_tool"
|
||||
assert approval_content.function_call.name == "get_datetime"
|
||||
assert approval_content.function_call.parse_arguments() == {}
|
||||
assert messages_input[0]["tool_calls"][0]["function"]["arguments"] == {}
|
||||
|
||||
|
||||
def test_agui_tool_approval_from_generate_task_steps_maps_to_function_call():
|
||||
"""Approval tool payloads map to the referenced function call when function_call_id is present."""
|
||||
messages_input = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_tool",
|
||||
"type": "function",
|
||||
"function": {"name": "get_datetime", "arguments": {}},
|
||||
},
|
||||
{
|
||||
"id": "call_steps",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "generate_task_steps",
|
||||
"arguments": {
|
||||
"function_name": "get_datetime",
|
||||
"function_call_id": "call_tool",
|
||||
"function_arguments": {},
|
||||
"steps": [{"description": "Execute get_datetime", "status": "enabled"}],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
"id": "msg_1",
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(
|
||||
{
|
||||
"accepted": True,
|
||||
"steps": [{"description": "Execute get_datetime", "status": "enabled"}],
|
||||
}
|
||||
),
|
||||
"toolCallId": "call_steps",
|
||||
"id": "msg_2",
|
||||
},
|
||||
]
|
||||
|
||||
messages = agui_messages_to_agent_framework(messages_input)
|
||||
|
||||
from agent_framework import FunctionApprovalResponseContent
|
||||
|
||||
approval_msg = messages[1]
|
||||
approval_content = next(
|
||||
content for content in approval_msg.contents if isinstance(content, FunctionApprovalResponseContent)
|
||||
)
|
||||
|
||||
assert approval_content.function_call.call_id == "call_tool"
|
||||
assert approval_content.function_call.name == "get_datetime"
|
||||
assert approval_content.function_call.parse_arguments() == {}
|
||||
|
||||
|
||||
def test_agui_multiple_messages_to_agent_framework():
|
||||
"""Test converting multiple AG-UI messages."""
|
||||
messages_input = [
|
||||
|
||||
@@ -2,10 +2,7 @@
|
||||
|
||||
from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent
|
||||
|
||||
from agent_framework_ag_ui._orchestration._message_hygiene import (
|
||||
deduplicate_messages,
|
||||
sanitize_tool_history,
|
||||
)
|
||||
from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
|
||||
|
||||
|
||||
def test_sanitize_tool_history_injects_confirm_changes_result() -> None:
|
||||
@@ -26,7 +23,7 @@ def test_sanitize_tool_history_injects_confirm_changes_result() -> None:
|
||||
),
|
||||
]
|
||||
|
||||
sanitized = sanitize_tool_history(messages)
|
||||
sanitized = _sanitize_tool_history(messages)
|
||||
|
||||
tool_messages = [
|
||||
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
|
||||
@@ -48,6 +45,6 @@ def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
|
||||
),
|
||||
]
|
||||
|
||||
deduped = deduplicate_messages(messages)
|
||||
deduped = _deduplicate_messages(messages)
|
||||
assert len(deduped) == 1
|
||||
assert deduped[0].contents[0].result == "result data"
|
||||
|
||||
@@ -42,6 +42,29 @@ class DummyAgent:
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="ok")], role="assistant")
|
||||
|
||||
|
||||
class RecordingAgent:
|
||||
"""Agent stub that captures messages passed to run_stream."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.chat_options = SimpleNamespace(tools=[], response_format=None)
|
||||
self.tools: list[Any] = []
|
||||
self.chat_client = SimpleNamespace(
|
||||
function_invocation_configuration=FunctionInvocationConfiguration(),
|
||||
)
|
||||
self.seen_messages: list[Any] | None = None
|
||||
|
||||
async def run_stream(
|
||||
self,
|
||||
messages: list[Any],
|
||||
*,
|
||||
thread: Any,
|
||||
tools: list[Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncGenerator[AgentRunResponseUpdate, None]:
|
||||
self.seen_messages = messages
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="ok")], role="assistant")
|
||||
|
||||
|
||||
async def test_default_orchestrator_merges_client_tools() -> None:
|
||||
"""Client tool declarations are merged with server tools before running agent."""
|
||||
|
||||
@@ -151,3 +174,104 @@ async def test_default_orchestrator_with_snake_case_ids() -> None:
|
||||
last_event = events[-1]
|
||||
assert last_event.run_id == "test-snakecase-runid"
|
||||
assert last_event.thread_id == "test-snakecase-threadid"
|
||||
|
||||
|
||||
async def test_state_context_injected_when_tool_call_state_mismatch() -> None:
|
||||
"""State context should be injected when current state differs from tool call args."""
|
||||
|
||||
agent = RecordingAgent()
|
||||
orchestrator = DefaultOrchestrator()
|
||||
|
||||
tool_recipe = {"title": "Salad", "special_preferences": []}
|
||||
current_recipe = {"title": "Salad", "special_preferences": ["Vegetarian"]}
|
||||
|
||||
input_data = {
|
||||
"state": {"recipe": current_recipe},
|
||||
"messages": [
|
||||
{"role": "system", "content": "Instructions"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "update_recipe", "arguments": {"recipe": tool_recipe}},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "What are the dietary preferences?"},
|
||||
],
|
||||
}
|
||||
|
||||
context = ExecutionContext(
|
||||
input_data=input_data,
|
||||
agent=agent,
|
||||
config=AgentConfig(
|
||||
state_schema={"recipe": {"type": "object"}},
|
||||
predict_state_config={"recipe": {"tool": "update_recipe", "tool_argument": "recipe"}},
|
||||
require_confirmation=False,
|
||||
),
|
||||
)
|
||||
|
||||
async for _event in orchestrator.run(context):
|
||||
pass
|
||||
|
||||
assert agent.seen_messages is not None
|
||||
state_messages = []
|
||||
for msg in agent.seen_messages:
|
||||
role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
|
||||
if role_value != "system":
|
||||
continue
|
||||
for content in msg.contents or []:
|
||||
if isinstance(content, TextContent) and content.text.startswith("Current state of the application:"):
|
||||
state_messages.append(content.text)
|
||||
assert state_messages
|
||||
assert "Vegetarian" in state_messages[0]
|
||||
|
||||
|
||||
async def test_state_context_not_injected_when_tool_call_matches_state() -> None:
|
||||
"""State context should be skipped when tool call args match current state."""
|
||||
|
||||
agent = RecordingAgent()
|
||||
orchestrator = DefaultOrchestrator()
|
||||
|
||||
input_data = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "Instructions"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "update_recipe", "arguments": {"recipe": {}}},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "What are the dietary preferences?"},
|
||||
],
|
||||
}
|
||||
|
||||
context = ExecutionContext(
|
||||
input_data=input_data,
|
||||
agent=agent,
|
||||
config=AgentConfig(
|
||||
state_schema={"recipe": {"type": "object"}},
|
||||
predict_state_config={"recipe": {"tool": "update_recipe", "tool_argument": "recipe"}},
|
||||
require_confirmation=False,
|
||||
),
|
||||
)
|
||||
|
||||
async for _event in orchestrator.run(context):
|
||||
pass
|
||||
|
||||
assert agent.seen_messages is not None
|
||||
state_messages = []
|
||||
for msg in agent.seen_messages:
|
||||
role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
|
||||
if role_value != "system":
|
||||
continue
|
||||
for content in msg.contents or []:
|
||||
if isinstance(content, TextContent) and content.text.startswith("Current state of the application:"):
|
||||
state_messages.append(content.text)
|
||||
assert not state_messages
|
||||
|
||||
@@ -62,7 +62,7 @@ async def test_human_in_the_loop_json_decode_error() -> None:
|
||||
agent=agent,
|
||||
config=AgentConfig(),
|
||||
)
|
||||
context.set_messages(messages)
|
||||
context.set_messages(messages, normalize=False)
|
||||
|
||||
assert orchestrator.can_handle(context)
|
||||
|
||||
@@ -385,8 +385,8 @@ async def test_state_context_injection() -> None:
|
||||
assert "banana" in system_messages[0].contents[0].text
|
||||
|
||||
|
||||
async def test_no_state_context_injection_with_tool_calls() -> None:
|
||||
"""Test state context is NOT injected if conversation has tool calls."""
|
||||
async def test_state_context_injection_with_tool_calls_and_input_state() -> None:
|
||||
"""Test state context is injected when state is provided, even with tool calls."""
|
||||
from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent
|
||||
|
||||
messages = [
|
||||
@@ -420,13 +420,13 @@ async def test_no_state_context_injection_with_tool_calls() -> None:
|
||||
async for event in orchestrator.run(context):
|
||||
events.append(event)
|
||||
|
||||
# Should NOT inject state context system message since conversation has tool calls
|
||||
# Should inject state context system message because input state is provided
|
||||
system_messages = [
|
||||
msg
|
||||
for msg in agent.messages_received
|
||||
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "system"
|
||||
]
|
||||
assert len(system_messages) == 0
|
||||
assert len(system_messages) == 1
|
||||
|
||||
|
||||
async def test_structured_output_processing() -> None:
|
||||
@@ -685,6 +685,54 @@ async def test_confirm_changes_with_invalid_json_fallback() -> None:
|
||||
assert len(user_messages) == 1
|
||||
|
||||
|
||||
async def test_confirm_changes_closes_active_message_before_finish() -> None:
|
||||
"""Confirm-changes flow closes any active text message before run finishes."""
|
||||
from ag_ui.core import TextMessageEndEvent, TextMessageStartEvent
|
||||
from agent_framework import FunctionCallContent, FunctionResultContent
|
||||
|
||||
updates = [
|
||||
AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(
|
||||
name="write_document_local",
|
||||
call_id="call_1",
|
||||
arguments='{"document": "Draft"}',
|
||||
)
|
||||
]
|
||||
),
|
||||
AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_1", result="Done")]),
|
||||
]
|
||||
|
||||
orchestrator = DefaultOrchestrator()
|
||||
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Start"}]}
|
||||
agent = StubAgent(
|
||||
chat_options=DEFAULT_CHAT_OPTIONS,
|
||||
updates=updates,
|
||||
)
|
||||
context = TestExecutionContext(
|
||||
input_data=input_data,
|
||||
agent=agent,
|
||||
config=AgentConfig(
|
||||
predict_state_config={"document": {"tool": "write_document_local", "tool_argument": "document"}},
|
||||
require_confirmation=True,
|
||||
),
|
||||
)
|
||||
|
||||
events: list[Any] = []
|
||||
async for event in orchestrator.run(context):
|
||||
events.append(event)
|
||||
|
||||
start_events = [e for e in events if isinstance(e, TextMessageStartEvent)]
|
||||
end_events = [e for e in events if isinstance(e, TextMessageEndEvent)]
|
||||
assert len(start_events) == 1
|
||||
assert len(end_events) == 1
|
||||
assert end_events[0].message_id == start_events[0].message_id
|
||||
|
||||
end_index = events.index(end_events[0])
|
||||
finished_index = events.index([e for e in events if e.type == "RUN_FINISHED"][0])
|
||||
assert end_index < finished_index
|
||||
|
||||
|
||||
async def test_tool_result_kept_when_call_id_matches() -> None:
|
||||
"""Test tool result is kept when call_id matches pending tool calls."""
|
||||
from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent
|
||||
|
||||
Reference in New Issue
Block a user