mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: AG-UI protocol support (#1826)
* Add AG-UI integration * Fix tests. PR feedback * Cleanup * PR Feedback * Improve README and getting started experience * Fix links
This commit is contained in:
Evan Mattson
committed by
GitHub
Unverified
GitHub
parent
0c862e97a6
commit
35a8565495
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
@@ -0,0 +1,577 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Comprehensive tests for AgentFrameworkAgent (_agent.py)."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from agent_framework import ChatAgent, TextContent
|
||||
from agent_framework._types import ChatResponseUpdate
|
||||
|
||||
|
||||
async def test_agent_initialization_basic():
|
||||
"""Test basic agent initialization without state schema."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
assert wrapper.name == "test_agent"
|
||||
assert wrapper.agent == agent
|
||||
assert wrapper.config.state_schema == {}
|
||||
assert wrapper.config.predict_state_config == {}
|
||||
|
||||
|
||||
async def test_agent_initialization_with_state_schema():
|
||||
"""Test agent initialization with state_schema."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
state_schema = {"document": {"type": "string"}}
|
||||
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
|
||||
|
||||
assert wrapper.config.state_schema == state_schema
|
||||
|
||||
|
||||
async def test_agent_initialization_with_predict_state_config():
|
||||
"""Test agent initialization with predict_state_config."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
|
||||
wrapper = AgentFrameworkAgent(agent=agent, predict_state_config=predict_config)
|
||||
|
||||
assert wrapper.config.predict_state_config == predict_config
|
||||
|
||||
|
||||
async def test_run_started_event_emission():
|
||||
"""Test RunStartedEvent is emitted at start of run."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# First event should be RunStartedEvent
|
||||
assert events[0].type == "RUN_STARTED"
|
||||
assert events[0].run_id is not None
|
||||
assert events[0].thread_id is not None
|
||||
|
||||
|
||||
async def test_predict_state_custom_event_emission():
|
||||
"""Test PredictState CustomEvent is emitted when predict_state_config is present."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
predict_config = {
|
||||
"document": {"tool": "write_doc", "tool_argument": "content"},
|
||||
"summary": {"tool": "summarize", "tool_argument": "text"},
|
||||
}
|
||||
wrapper = AgentFrameworkAgent(agent=agent, predict_state_config=predict_config)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Find PredictState event
|
||||
predict_events = [e for e in events if e.type == "CUSTOM" and e.name == "PredictState"]
|
||||
assert len(predict_events) == 1
|
||||
|
||||
predict_value = predict_events[0].value
|
||||
assert len(predict_value) == 2
|
||||
assert {"state_key": "document", "tool": "write_doc", "tool_argument": "content"} in predict_value
|
||||
assert {"state_key": "summary", "tool": "summarize", "tool_argument": "text"} in predict_value
|
||||
|
||||
|
||||
async def test_initial_state_snapshot_with_schema():
|
||||
"""Test initial StateSnapshotEvent emission when state_schema present."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
state_schema = {"document": {"type": "string"}}
|
||||
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
|
||||
|
||||
input_data = {
|
||||
"messages": [{"role": "user", "content": "Hi"}],
|
||||
"state": {"document": "Initial content"},
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Find StateSnapshotEvent
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) >= 1
|
||||
|
||||
# First snapshot should have initial state
|
||||
assert snapshot_events[0].snapshot == {"document": "Initial content"}
|
||||
|
||||
|
||||
async def test_state_initialization_object_type():
|
||||
"""Test state initialization with object type in schema."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
state_schema = {"recipe": {"type": "object", "properties": {}}}
|
||||
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Find StateSnapshotEvent
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) >= 1
|
||||
|
||||
# Should initialize as empty object
|
||||
assert snapshot_events[0].snapshot == {"recipe": {}}
|
||||
|
||||
|
||||
async def test_state_initialization_array_type():
|
||||
"""Test state initialization with array type in schema."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
state_schema = {"steps": {"type": "array", "items": {}}}
|
||||
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Find StateSnapshotEvent
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) >= 1
|
||||
|
||||
# Should initialize as empty array
|
||||
assert snapshot_events[0].snapshot == {"steps": []}
|
||||
|
||||
|
||||
async def test_run_finished_event_emission():
|
||||
"""Test RunFinishedEvent is emitted at end of run."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Last event should be RunFinishedEvent
|
||||
assert events[-1].type == "RUN_FINISHED"
|
||||
|
||||
|
||||
async def test_tool_result_confirm_changes_accepted():
|
||||
"""Test confirm_changes tool result handling when accepted."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Document updated")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema={"document": {"type": "string"}},
|
||||
predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}},
|
||||
)
|
||||
|
||||
# Simulate tool result message with acceptance
|
||||
tool_result = {"accepted": True, "steps": []}
|
||||
input_data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "tool", # Tool result from UI
|
||||
"content": json.dumps(tool_result),
|
||||
"toolCallId": "confirm_call_123",
|
||||
}
|
||||
],
|
||||
"state": {"document": "Updated content"},
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit text message confirming acceptance
|
||||
text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert len(text_content_events) > 0
|
||||
# Should contain confirmation message mentioning the state key or generic confirmation
|
||||
confirmation_found = any(
|
||||
"document" in e.delta.lower()
|
||||
or "confirm" in e.delta.lower()
|
||||
or "applied" in e.delta.lower()
|
||||
or "changes" in e.delta.lower()
|
||||
for e in text_content_events
|
||||
)
|
||||
assert confirmation_found, f"No confirmation in deltas: {[e.delta for e in text_content_events]}"
|
||||
|
||||
|
||||
async def test_tool_result_confirm_changes_rejected():
|
||||
"""Test confirm_changes tool result handling when rejected."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
# Simulate tool result message with rejection
|
||||
tool_result = {"accepted": False, "steps": []}
|
||||
input_data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(tool_result),
|
||||
"toolCallId": "confirm_call_123",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit text message asking what to change
|
||||
text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert len(text_content_events) > 0
|
||||
assert any("what would you like me to change" in e.delta.lower() for e in text_content_events)
|
||||
|
||||
|
||||
async def test_tool_result_function_approval_accepted():
|
||||
"""Test function approval tool result when steps are accepted."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
# Simulate tool result with multiple steps
|
||||
tool_result = {
|
||||
"accepted": True,
|
||||
"steps": [
|
||||
{"id": "step1", "description": "Send email", "status": "enabled"},
|
||||
{"id": "step2", "description": "Create calendar event", "status": "enabled"},
|
||||
],
|
||||
}
|
||||
input_data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(tool_result),
|
||||
"toolCallId": "approval_call_123",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should list enabled steps
|
||||
text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert len(text_content_events) > 0
|
||||
|
||||
# Concatenate all text content
|
||||
full_text = "".join(e.delta for e in text_content_events)
|
||||
assert "executing" in full_text.lower()
|
||||
assert "2 approved steps" in full_text.lower()
|
||||
assert "send email" in full_text.lower()
|
||||
assert "create calendar event" in full_text.lower()
|
||||
|
||||
|
||||
async def test_tool_result_function_approval_rejected():
|
||||
"""Test function approval tool result when rejected."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
# Simulate tool result rejection with steps
|
||||
tool_result = {
|
||||
"accepted": False,
|
||||
"steps": [{"id": "step1", "description": "Send email", "status": "disabled"}],
|
||||
}
|
||||
input_data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(tool_result),
|
||||
"toolCallId": "approval_call_123",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should ask what to change about the plan
|
||||
text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert len(text_content_events) > 0
|
||||
assert any("what would you like me to change about the plan" in e.delta.lower() for e in text_content_events)
|
||||
|
||||
|
||||
async def test_thread_metadata_tracking():
|
||||
"""Test that thread metadata includes ag_ui_thread_id and ag_ui_run_id."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
thread_metadata = {}
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
# Capture thread metadata from kwargs
|
||||
nonlocal thread_metadata
|
||||
if "thread" in kwargs:
|
||||
thread_metadata = kwargs["thread"].metadata
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
input_data = {
|
||||
"messages": [{"role": "user", "content": "Hi"}],
|
||||
"thread_id": "test_thread_123",
|
||||
"run_id": "test_run_456",
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Check thread metadata was set
|
||||
# Note: This test may need adjustment based on actual thread passing mechanism
|
||||
|
||||
|
||||
async def test_state_context_injection():
|
||||
"""Test that current state is injected into thread metadata."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
thread_metadata = {}
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
# Track if state context message was added
|
||||
nonlocal thread_metadata
|
||||
# In actual implementation, thread is passed and state is in metadata
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema={"document": {"type": "string"}},
|
||||
)
|
||||
|
||||
input_data = {
|
||||
"messages": [{"role": "user", "content": "Hi"}],
|
||||
"state": {"document": "Test content"},
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# State should be injected - this is validated by agent execution flow
|
||||
|
||||
|
||||
async def test_no_messages_provided():
|
||||
"""Test handling when no messages are provided."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
input_data = {"messages": []}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit RunStartedEvent and RunFinishedEvent only
|
||||
assert len(events) == 2
|
||||
assert events[0].type == "RUN_STARTED"
|
||||
assert events[-1].type == "RUN_FINISHED"
|
||||
|
||||
|
||||
async def test_message_end_event_emission():
|
||||
"""Test TextMessageEndEvent is emitted for assistant messages."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello world")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should have TextMessageEndEvent before RunFinishedEvent
|
||||
end_events = [e for e in events if e.type == "TEXT_MESSAGE_END"]
|
||||
assert len(end_events) == 1
|
||||
|
||||
# EndEvent should come before FinishedEvent
|
||||
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_error_handling_with_exception():
|
||||
"""Test that exceptions during agent execution are re-raised."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class FailingChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
if False:
|
||||
yield
|
||||
raise RuntimeError("Simulated failure")
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=FailingChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
with pytest.raises(RuntimeError, match="Simulated failure"):
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
pass
|
||||
|
||||
|
||||
async def test_json_decode_error_in_tool_result():
|
||||
"""Test handling of JSONDecodeError when parsing tool result."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Fallback response")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
# Send invalid JSON as tool result
|
||||
input_data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "invalid json {not valid}",
|
||||
"toolCallId": "call_123",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should fall through to normal agent processing
|
||||
text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert len(text_events) > 0
|
||||
assert text_events[0].delta == "Fallback response"
|
||||
|
||||
|
||||
async def test_suppressed_summary_with_document_state():
|
||||
"""Test suppressed summary uses document state for confirmation message."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Response")])
|
||||
|
||||
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
|
||||
wrapper = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema={"document": {"type": "string"}},
|
||||
predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}},
|
||||
confirmation_strategy=DocumentWriterConfirmationStrategy(),
|
||||
)
|
||||
|
||||
# Simulate confirmation with document state
|
||||
tool_result = {"accepted": True, "steps": []}
|
||||
input_data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": json.dumps(tool_result),
|
||||
"toolCallId": "confirm_123",
|
||||
}
|
||||
],
|
||||
"state": {"document": "This is the beginning of a document. It contains important information."},
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should generate fallback summary from document state
|
||||
text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert len(text_events) > 0
|
||||
# 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()
|
||||
@@ -0,0 +1,124 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for backend tool rendering."""
|
||||
|
||||
from ag_ui.core import (
|
||||
TextMessageContentEvent,
|
||||
TextMessageStartEvent,
|
||||
ToolCallArgsEvent,
|
||||
ToolCallEndEvent,
|
||||
ToolCallResultEvent,
|
||||
ToolCallStartEvent,
|
||||
)
|
||||
from agent_framework import AgentRunResponseUpdate, FunctionCallContent, FunctionResultContent, TextContent
|
||||
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
|
||||
async def test_tool_call_flow():
|
||||
"""Test complete tool call flow: call -> args -> end -> result."""
|
||||
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
|
||||
|
||||
# Step 1: Tool call starts
|
||||
tool_call = FunctionCallContent(
|
||||
call_id="weather-123",
|
||||
name="get_weather",
|
||||
arguments={"location": "Seattle"},
|
||||
)
|
||||
|
||||
update1 = AgentRunResponseUpdate(contents=[tool_call])
|
||||
events1 = await bridge.from_agent_run_update(update1)
|
||||
|
||||
# Should have: ToolCallStartEvent, ToolCallArgsEvent
|
||||
assert len(events1) == 2
|
||||
assert isinstance(events1[0], ToolCallStartEvent)
|
||||
assert isinstance(events1[1], ToolCallArgsEvent)
|
||||
|
||||
start_event = events1[0]
|
||||
assert start_event.tool_call_id == "weather-123"
|
||||
assert start_event.tool_call_name == "get_weather"
|
||||
|
||||
args_event = events1[1]
|
||||
assert "Seattle" in args_event.delta
|
||||
|
||||
# Step 2: Tool result comes back
|
||||
tool_result = FunctionResultContent(
|
||||
call_id="weather-123",
|
||||
result="Weather in Seattle: Rainy, 52°F",
|
||||
)
|
||||
|
||||
update2 = AgentRunResponseUpdate(contents=[tool_result])
|
||||
events2 = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Should have: ToolCallEndEvent, ToolCallResultEvent, MessagesSnapshotEvent
|
||||
assert len(events2) == 3
|
||||
assert isinstance(events2[0], ToolCallEndEvent)
|
||||
assert isinstance(events2[1], ToolCallResultEvent)
|
||||
|
||||
end_event = events2[0]
|
||||
assert end_event.tool_call_id == "weather-123"
|
||||
|
||||
result_event = events2[1]
|
||||
assert result_event.tool_call_id == "weather-123"
|
||||
assert "Seattle" in result_event.content
|
||||
assert "Rainy" in result_event.content
|
||||
|
||||
|
||||
async def test_text_with_tool_call():
|
||||
"""Test agent response with both text and tool calls."""
|
||||
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
|
||||
|
||||
# Agent says something then calls a tool
|
||||
text_content = TextContent(text="Let me check the weather for you.")
|
||||
tool_call = FunctionCallContent(
|
||||
call_id="weather-456",
|
||||
name="get_forecast",
|
||||
arguments={"location": "San Francisco", "days": 3},
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[text_content, tool_call])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should have: TextMessageStart, TextMessageContent, ToolCallStart, ToolCallArgs
|
||||
assert len(events) == 4
|
||||
|
||||
assert isinstance(events[0], TextMessageStartEvent)
|
||||
assert isinstance(events[1], TextMessageContentEvent)
|
||||
assert isinstance(events[2], ToolCallStartEvent)
|
||||
assert isinstance(events[3], ToolCallArgsEvent)
|
||||
|
||||
text_event = events[1]
|
||||
assert "check the weather" in text_event.delta
|
||||
|
||||
tool_start = events[2]
|
||||
assert tool_start.tool_call_name == "get_forecast"
|
||||
|
||||
|
||||
async def test_multiple_tool_results():
|
||||
"""Test handling multiple tool results in sequence."""
|
||||
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
|
||||
|
||||
# Multiple tool results
|
||||
results = [
|
||||
FunctionResultContent(call_id="tool-1", result="Result 1"),
|
||||
FunctionResultContent(call_id="tool-2", result="Result 2"),
|
||||
FunctionResultContent(call_id="tool-3", result="Result 3"),
|
||||
]
|
||||
|
||||
update = AgentRunResponseUpdate(contents=results)
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should have 3 pairs of ToolCallEndEvent + ToolCallResultEvent = 6 events
|
||||
assert len(events) == 6
|
||||
|
||||
# Verify the pattern: End, Result, End, Result, End, Result
|
||||
for i in range(3):
|
||||
end_idx = i * 2
|
||||
result_idx = i * 2 + 1
|
||||
|
||||
assert isinstance(events[end_idx], ToolCallEndEvent)
|
||||
assert isinstance(events[result_idx], ToolCallResultEvent)
|
||||
|
||||
assert events[end_idx].tool_call_id == f"tool-{i + 1}"
|
||||
assert events[result_idx].tool_call_id == f"tool-{i + 1}"
|
||||
assert f"Result {i + 1}" in events[result_idx].content
|
||||
@@ -0,0 +1,275 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Comprehensive tests for all confirmation strategies."""
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework_ag_ui._confirmation_strategies import (
|
||||
ConfirmationStrategy,
|
||||
DefaultConfirmationStrategy,
|
||||
DocumentWriterConfirmationStrategy,
|
||||
RecipeConfirmationStrategy,
|
||||
TaskPlannerConfirmationStrategy,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_steps():
|
||||
"""Sample steps for testing approval messages."""
|
||||
return [
|
||||
{"description": "Step 1: Do something", "status": "enabled"},
|
||||
{"description": "Step 2: Do another thing", "status": "enabled"},
|
||||
{"description": "Step 3: Disabled step", "status": "disabled"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def all_enabled_steps():
|
||||
"""All steps enabled."""
|
||||
return [
|
||||
{"description": "Task A", "status": "enabled"},
|
||||
{"description": "Task B", "status": "enabled"},
|
||||
{"description": "Task C", "status": "enabled"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def empty_steps():
|
||||
"""Empty steps list."""
|
||||
return []
|
||||
|
||||
|
||||
class TestDefaultConfirmationStrategy:
|
||||
"""Tests for DefaultConfirmationStrategy."""
|
||||
|
||||
def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
|
||||
strategy = DefaultConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(sample_steps)
|
||||
|
||||
assert "Executing 2 approved steps" in message
|
||||
assert "Step 1: Do something" in message
|
||||
assert "Step 2: Do another thing" in message
|
||||
assert "Step 3" not in message # Disabled step shouldn't appear
|
||||
assert "All steps completed successfully!" in message
|
||||
|
||||
def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
|
||||
strategy = DefaultConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(all_enabled_steps)
|
||||
|
||||
assert "Executing 3 approved steps" in message
|
||||
assert "Task A" in message
|
||||
assert "Task B" in message
|
||||
assert "Task C" in message
|
||||
|
||||
def test_on_approval_accepted_with_empty_steps(self, empty_steps):
|
||||
strategy = DefaultConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(empty_steps)
|
||||
|
||||
assert "Executing 0 approved steps" in message
|
||||
assert "All steps completed successfully!" in message
|
||||
|
||||
def test_on_approval_rejected(self, sample_steps):
|
||||
strategy = DefaultConfirmationStrategy()
|
||||
message = strategy.on_approval_rejected(sample_steps)
|
||||
|
||||
assert "No problem!" in message
|
||||
assert "What would you like me to change" in message
|
||||
|
||||
def test_on_state_confirmed(self):
|
||||
strategy = DefaultConfirmationStrategy()
|
||||
message = strategy.on_state_confirmed()
|
||||
|
||||
assert "Changes confirmed" in message
|
||||
assert "successfully" in message
|
||||
|
||||
def test_on_state_rejected(self):
|
||||
strategy = DefaultConfirmationStrategy()
|
||||
message = strategy.on_state_rejected()
|
||||
|
||||
assert "No problem!" in message
|
||||
assert "What would you like me to change" in message
|
||||
|
||||
|
||||
class TestTaskPlannerConfirmationStrategy:
|
||||
"""Tests for TaskPlannerConfirmationStrategy."""
|
||||
|
||||
def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
|
||||
strategy = TaskPlannerConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(sample_steps)
|
||||
|
||||
assert "Executing your requested tasks" in message
|
||||
assert "1. Step 1: Do something" in message
|
||||
assert "2. Step 2: Do another thing" in message
|
||||
assert "Step 3" not in message
|
||||
assert "All tasks completed successfully!" in message
|
||||
|
||||
def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
|
||||
strategy = TaskPlannerConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(all_enabled_steps)
|
||||
|
||||
assert "Executing your requested tasks" in message
|
||||
assert "1. Task A" in message
|
||||
assert "2. Task B" in message
|
||||
assert "3. Task C" in message
|
||||
|
||||
def test_on_approval_accepted_with_empty_steps(self, empty_steps):
|
||||
strategy = TaskPlannerConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(empty_steps)
|
||||
|
||||
assert "Executing your requested tasks" in message
|
||||
assert "All tasks completed successfully!" in message
|
||||
|
||||
def test_on_approval_rejected(self, sample_steps):
|
||||
strategy = TaskPlannerConfirmationStrategy()
|
||||
message = strategy.on_approval_rejected(sample_steps)
|
||||
|
||||
assert "No problem!" in message
|
||||
assert "revise the plan" in message
|
||||
|
||||
def test_on_state_confirmed(self):
|
||||
strategy = TaskPlannerConfirmationStrategy()
|
||||
message = strategy.on_state_confirmed()
|
||||
|
||||
assert "Tasks confirmed" in message
|
||||
assert "ready to execute" in message
|
||||
|
||||
def test_on_state_rejected(self):
|
||||
strategy = TaskPlannerConfirmationStrategy()
|
||||
message = strategy.on_state_rejected()
|
||||
|
||||
assert "No problem!" in message
|
||||
assert "adjust the task list" in message
|
||||
|
||||
|
||||
class TestRecipeConfirmationStrategy:
|
||||
"""Tests for RecipeConfirmationStrategy."""
|
||||
|
||||
def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
|
||||
strategy = RecipeConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(sample_steps)
|
||||
|
||||
assert "Updating your recipe" in message
|
||||
assert "1. Step 1: Do something" in message
|
||||
assert "2. Step 2: Do another thing" in message
|
||||
assert "Step 3" not in message
|
||||
assert "Recipe updated successfully!" in message
|
||||
|
||||
def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
|
||||
strategy = RecipeConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(all_enabled_steps)
|
||||
|
||||
assert "Updating your recipe" in message
|
||||
assert "1. Task A" in message
|
||||
assert "2. Task B" in message
|
||||
assert "3. Task C" in message
|
||||
|
||||
def test_on_approval_accepted_with_empty_steps(self, empty_steps):
|
||||
strategy = RecipeConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(empty_steps)
|
||||
|
||||
assert "Updating your recipe" in message
|
||||
assert "Recipe updated successfully!" in message
|
||||
|
||||
def test_on_approval_rejected(self, sample_steps):
|
||||
strategy = RecipeConfirmationStrategy()
|
||||
message = strategy.on_approval_rejected(sample_steps)
|
||||
|
||||
assert "No problem!" in message
|
||||
assert "ingredients or steps" in message
|
||||
|
||||
def test_on_state_confirmed(self):
|
||||
strategy = RecipeConfirmationStrategy()
|
||||
message = strategy.on_state_confirmed()
|
||||
|
||||
assert "Recipe changes applied" in message
|
||||
assert "successfully" in message
|
||||
|
||||
def test_on_state_rejected(self):
|
||||
strategy = RecipeConfirmationStrategy()
|
||||
message = strategy.on_state_rejected()
|
||||
|
||||
assert "No problem!" in message
|
||||
assert "adjust in the recipe" in message
|
||||
|
||||
|
||||
class TestDocumentWriterConfirmationStrategy:
|
||||
"""Tests for DocumentWriterConfirmationStrategy."""
|
||||
|
||||
def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
|
||||
strategy = DocumentWriterConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(sample_steps)
|
||||
|
||||
assert "Applying your edits" in message
|
||||
assert "1. Step 1: Do something" in message
|
||||
assert "2. Step 2: Do another thing" in message
|
||||
assert "Step 3" not in message
|
||||
assert "Document updated successfully!" in message
|
||||
|
||||
def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
|
||||
strategy = DocumentWriterConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(all_enabled_steps)
|
||||
|
||||
assert "Applying your edits" in message
|
||||
assert "1. Task A" in message
|
||||
assert "2. Task B" in message
|
||||
assert "3. Task C" in message
|
||||
|
||||
def test_on_approval_accepted_with_empty_steps(self, empty_steps):
|
||||
strategy = DocumentWriterConfirmationStrategy()
|
||||
message = strategy.on_approval_accepted(empty_steps)
|
||||
|
||||
assert "Applying your edits" in message
|
||||
assert "Document updated successfully!" in message
|
||||
|
||||
def test_on_approval_rejected(self, sample_steps):
|
||||
strategy = DocumentWriterConfirmationStrategy()
|
||||
message = strategy.on_approval_rejected(sample_steps)
|
||||
|
||||
assert "No problem!" in message
|
||||
assert "keep or modify" in message
|
||||
|
||||
def test_on_state_confirmed(self):
|
||||
strategy = DocumentWriterConfirmationStrategy()
|
||||
message = strategy.on_state_confirmed()
|
||||
|
||||
assert "Document edits applied!" in message
|
||||
|
||||
def test_on_state_rejected(self):
|
||||
strategy = DocumentWriterConfirmationStrategy()
|
||||
message = strategy.on_state_rejected()
|
||||
|
||||
assert "No problem!" in message
|
||||
assert "change about the document" in message
|
||||
|
||||
|
||||
class TestConfirmationStrategyInterface:
|
||||
"""Tests for ConfirmationStrategy abstract base class."""
|
||||
|
||||
def test_cannot_instantiate_abstract_class(self):
|
||||
"""Verify ConfirmationStrategy is abstract and cannot be instantiated."""
|
||||
with pytest.raises(TypeError):
|
||||
ConfirmationStrategy() # type: ignore
|
||||
|
||||
def test_all_strategies_implement_interface(self):
|
||||
"""Verify all concrete strategies implement the full interface."""
|
||||
strategies = [
|
||||
DefaultConfirmationStrategy(),
|
||||
TaskPlannerConfirmationStrategy(),
|
||||
RecipeConfirmationStrategy(),
|
||||
DocumentWriterConfirmationStrategy(),
|
||||
]
|
||||
|
||||
sample_steps = [{"description": "Test", "status": "enabled"}]
|
||||
|
||||
for strategy in strategies:
|
||||
# All should have these methods
|
||||
assert callable(strategy.on_approval_accepted)
|
||||
assert callable(strategy.on_approval_rejected)
|
||||
assert callable(strategy.on_state_confirmed)
|
||||
assert callable(strategy.on_state_rejected)
|
||||
|
||||
# All should return strings
|
||||
assert isinstance(strategy.on_approval_accepted(sample_steps), str)
|
||||
assert isinstance(strategy.on_approval_rejected(sample_steps), str)
|
||||
assert isinstance(strategy.on_state_confirmed(), str)
|
||||
assert isinstance(strategy.on_state_rejected(), str)
|
||||
@@ -0,0 +1,243 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for document writer predictive state flow with confirm_changes."""
|
||||
|
||||
from ag_ui.core import EventType
|
||||
from agent_framework import FunctionCallContent, FunctionResultContent, TextContent
|
||||
from agent_framework._types import AgentRunResponseUpdate
|
||||
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
|
||||
async def test_streaming_document_with_state_deltas():
|
||||
"""Test that streaming tool arguments emit progressive StateDeltaEvents."""
|
||||
predict_config = {
|
||||
"document": {"tool": "write_document_local", "tool_argument": "document"},
|
||||
}
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config=predict_config,
|
||||
)
|
||||
|
||||
# Simulate streaming tool call - first chunk with name
|
||||
tool_call_start = FunctionCallContent(
|
||||
call_id="call_123",
|
||||
name="write_document_local",
|
||||
arguments='{"document":"Once',
|
||||
)
|
||||
update1 = AgentRunResponseUpdate(contents=[tool_call_start])
|
||||
events1 = await bridge.from_agent_run_update(update1)
|
||||
|
||||
# Should have ToolCallStartEvent and ToolCallArgsEvent
|
||||
assert any(e.type == EventType.TOOL_CALL_START for e in events1)
|
||||
assert any(e.type == EventType.TOOL_CALL_ARGS for e in events1)
|
||||
|
||||
# Second chunk - incomplete JSON, should try partial extraction
|
||||
tool_call_chunk2 = FunctionCallContent(
|
||||
call_id="call_123",
|
||||
name=None, # Name only in first chunk
|
||||
arguments=" upon a time",
|
||||
)
|
||||
update2 = AgentRunResponseUpdate(contents=[tool_call_chunk2])
|
||||
events2 = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Should emit StateDeltaEvent with partial document
|
||||
state_deltas = [e for e in events2 if e.type == EventType.STATE_DELTA]
|
||||
assert len(state_deltas) >= 1
|
||||
|
||||
# Check JSON Patch format
|
||||
delta = state_deltas[0]
|
||||
assert isinstance(delta.delta, list)
|
||||
assert len(delta.delta) > 0
|
||||
assert delta.delta[0]["op"] == "replace"
|
||||
assert delta.delta[0]["path"] == "/document"
|
||||
assert "Once upon a time" in delta.delta[0]["value"]
|
||||
|
||||
|
||||
async def test_confirm_changes_emission():
|
||||
"""Test that confirm_changes tool call is emitted after predictive tool completion."""
|
||||
predict_config = {
|
||||
"document": {"tool": "write_document_local", "tool_argument": "document"},
|
||||
}
|
||||
|
||||
current_state = {}
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config=predict_config,
|
||||
current_state=current_state,
|
||||
)
|
||||
|
||||
# Set current tool name (simulating earlier tool call start)
|
||||
bridge.current_tool_call_name = "write_document_local"
|
||||
bridge.pending_state_updates = {"document": "A short story"}
|
||||
|
||||
# Tool result
|
||||
tool_result = FunctionResultContent(
|
||||
call_id="call_123",
|
||||
result="Document written.",
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[tool_result])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should have: ToolCallEndEvent, ToolCallResultEvent, StateSnapshotEvent, confirm_changes sequence
|
||||
assert any(e.type == EventType.TOOL_CALL_END for e in events)
|
||||
assert any(e.type == EventType.TOOL_CALL_RESULT for e in events)
|
||||
assert any(e.type == EventType.STATE_SNAPSHOT for e in events)
|
||||
|
||||
# Check for confirm_changes tool call
|
||||
confirm_starts = [
|
||||
e for e in events if e.type == EventType.TOOL_CALL_START and e.tool_call_name == "confirm_changes"
|
||||
]
|
||||
assert len(confirm_starts) == 1
|
||||
|
||||
confirm_args = [e for e in events if e.type == EventType.TOOL_CALL_ARGS and e.delta == "{}"]
|
||||
assert len(confirm_args) >= 1
|
||||
|
||||
confirm_ends = [e for e in events if e.type == EventType.TOOL_CALL_END]
|
||||
# At least 2: one for write_document_local, one for confirm_changes
|
||||
assert len(confirm_ends) >= 2
|
||||
|
||||
# Check that stop flag is set
|
||||
assert bridge.should_stop_after_confirm is True
|
||||
|
||||
|
||||
async def test_text_suppression_before_confirm():
|
||||
"""Test that text messages are suppressed when confirm_changes is pending."""
|
||||
predict_config = {
|
||||
"document": {"tool": "write_document_local", "tool_argument": "document"},
|
||||
}
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config=predict_config,
|
||||
)
|
||||
|
||||
# Set flag indicating we're waiting for confirmation
|
||||
bridge.should_stop_after_confirm = True
|
||||
|
||||
# Text content that should be suppressed
|
||||
text = TextContent(text="I have written a story about pirates.")
|
||||
update = AgentRunResponseUpdate(contents=[text])
|
||||
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should NOT emit TextMessageContentEvent
|
||||
text_events = [e for e in events if e.type == EventType.TEXT_MESSAGE_CONTENT]
|
||||
assert len(text_events) == 0
|
||||
|
||||
# But should save the text
|
||||
assert bridge.suppressed_summary == "I have written a story about pirates."
|
||||
|
||||
|
||||
async def test_no_confirm_for_non_predictive_tools():
|
||||
"""Test that confirm_changes is NOT emitted for regular tool calls."""
|
||||
predict_config = {
|
||||
"document": {"tool": "write_document_local", "tool_argument": "document"},
|
||||
}
|
||||
|
||||
current_state = {}
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config=predict_config,
|
||||
current_state=current_state,
|
||||
)
|
||||
|
||||
# Different tool (not in predict_state_config)
|
||||
bridge.current_tool_call_name = "get_weather"
|
||||
|
||||
tool_result = FunctionResultContent(
|
||||
call_id="call_456",
|
||||
result="Sunny, 72°F",
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[tool_result])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should NOT have confirm_changes
|
||||
confirm_starts = [
|
||||
e for e in events if e.type == EventType.TOOL_CALL_START and e.tool_call_name == "confirm_changes"
|
||||
]
|
||||
assert len(confirm_starts) == 0
|
||||
|
||||
# Stop flag should NOT be set
|
||||
assert bridge.should_stop_after_confirm is False
|
||||
|
||||
|
||||
async def test_state_delta_deduplication():
|
||||
"""Test that duplicate state values don't emit multiple StateDeltaEvents."""
|
||||
predict_config = {
|
||||
"document": {"tool": "write_document_local", "tool_argument": "document"},
|
||||
}
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config=predict_config,
|
||||
)
|
||||
|
||||
# First tool call with document
|
||||
tool_call1 = FunctionCallContent(
|
||||
call_id="call_1",
|
||||
name="write_document_local",
|
||||
arguments='{"document":"Same text"}',
|
||||
)
|
||||
update1 = AgentRunResponseUpdate(contents=[tool_call1])
|
||||
events1 = await bridge.from_agent_run_update(update1)
|
||||
|
||||
# Count state deltas
|
||||
state_deltas_1 = [e for e in events1 if e.type == EventType.STATE_DELTA]
|
||||
assert len(state_deltas_1) >= 1
|
||||
|
||||
# Second tool call with SAME document (shouldn't emit new delta)
|
||||
bridge.current_tool_call_name = "write_document_local"
|
||||
tool_call2 = FunctionCallContent(
|
||||
call_id="call_2",
|
||||
name=None,
|
||||
arguments='{"document":"Same text"}', # Identical content
|
||||
)
|
||||
update2 = AgentRunResponseUpdate(contents=[tool_call2])
|
||||
events2 = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Should NOT emit state delta (same value)
|
||||
state_deltas_2 = [e for e in events2 if e.type == EventType.STATE_DELTA]
|
||||
assert len(state_deltas_2) == 0
|
||||
|
||||
|
||||
async def test_predict_state_config_multiple_fields():
|
||||
"""Test predictive state with multiple state fields."""
|
||||
predict_config = {
|
||||
"title": {"tool": "create_post", "tool_argument": "title"},
|
||||
"content": {"tool": "create_post", "tool_argument": "body"},
|
||||
}
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config=predict_config,
|
||||
)
|
||||
|
||||
# Tool call with both fields
|
||||
tool_call = FunctionCallContent(
|
||||
call_id="call_999",
|
||||
name="create_post",
|
||||
arguments='{"title":"My Post","body":"Post content"}',
|
||||
)
|
||||
update = AgentRunResponseUpdate(contents=[tool_call])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should emit StateDeltaEvent for both fields
|
||||
state_deltas = [e for e in events if e.type == EventType.STATE_DELTA]
|
||||
assert len(state_deltas) >= 2
|
||||
|
||||
# Check both fields are present
|
||||
paths = [delta.delta[0]["path"] for delta in state_deltas]
|
||||
assert "/title" in paths
|
||||
assert "/content" in paths
|
||||
@@ -0,0 +1,242 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for FastAPI endpoint creation (_endpoint.py)."""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import ChatAgent, TextContent
|
||||
from agent_framework._types import ChatResponseUpdate
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from agent_framework_ag_ui._agent import AgentFrameworkAgent
|
||||
from agent_framework_ag_ui._endpoint import add_agent_framework_fastapi_endpoint
|
||||
|
||||
|
||||
class MockChatClient:
|
||||
"""Mock chat client for testing."""
|
||||
|
||||
def __init__(self, response_text: str = "Test response"):
|
||||
self.response_text = response_text
|
||||
|
||||
async def get_streaming_response(self, messages: list[Any], chat_options: Any, **kwargs: Any):
|
||||
"""Mock streaming response."""
|
||||
yield ChatResponseUpdate(contents=[TextContent(text=self.response_text)])
|
||||
|
||||
|
||||
async def test_add_endpoint_with_agent_protocol():
|
||||
"""Test adding endpoint with raw AgentProtocol."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/test-agent")
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post("/test-agent", json={"messages": [{"role": "user", "content": "Hello"}]})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
|
||||
|
||||
|
||||
async def test_add_endpoint_with_wrapped_agent():
|
||||
"""Test adding endpoint with pre-wrapped AgentFrameworkAgent."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
wrapped_agent = AgentFrameworkAgent(agent=agent, name="wrapped")
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/wrapped-agent")
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post("/wrapped-agent", json={"messages": [{"role": "user", "content": "Hello"}]})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
|
||||
|
||||
|
||||
async def test_endpoint_with_state_schema():
|
||||
"""Test endpoint with state_schema parameter."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
state_schema = {"document": {"type": "string"}}
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/stateful", state_schema=state_schema)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/stateful", json={"messages": [{"role": "user", "content": "Hello"}], "state": {"document": ""}}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
async def test_endpoint_with_predict_state_config():
|
||||
"""Test endpoint with predict_state_config parameter."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/predictive", predict_state_config=predict_config)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post("/predictive", json={"messages": [{"role": "user", "content": "Hello"}]})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
async def test_endpoint_request_logging():
|
||||
"""Test that endpoint logs request details."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/logged")
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/logged",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Test"}],
|
||||
"run_id": "run-123",
|
||||
"thread_id": "thread-456",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
async def test_endpoint_event_streaming():
|
||||
"""Test that endpoint streams events correctly."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient("Streamed response"))
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/stream")
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post("/stream", json={"messages": [{"role": "user", "content": "Hello"}]})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
content = response.content.decode("utf-8")
|
||||
lines = [line for line in content.split("\n") if line.strip()]
|
||||
|
||||
found_run_started = False
|
||||
found_text_content = False
|
||||
found_run_finished = False
|
||||
|
||||
for line in lines:
|
||||
if line.startswith("data: "):
|
||||
event_data = json.loads(line[6:])
|
||||
if event_data.get("type") == "RUN_STARTED":
|
||||
found_run_started = True
|
||||
elif event_data.get("type") == "TEXT_MESSAGE_CONTENT":
|
||||
found_text_content = True
|
||||
elif event_data.get("type") == "RUN_FINISHED":
|
||||
found_run_finished = True
|
||||
|
||||
assert found_run_started
|
||||
assert found_text_content
|
||||
assert found_run_finished
|
||||
|
||||
|
||||
async def test_endpoint_error_handling():
|
||||
"""Test endpoint error handling during request parsing."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/failing")
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Send invalid JSON to trigger parsing error before streaming
|
||||
response = client.post("/failing", data="invalid json", headers={"content-type": "application/json"})
|
||||
|
||||
# The exception handler catches it and returns JSON error
|
||||
assert response.status_code == 200
|
||||
content = json.loads(response.content)
|
||||
assert "error" in content
|
||||
assert "Expecting value" in content["error"]
|
||||
|
||||
|
||||
async def test_endpoint_multiple_paths():
|
||||
"""Test adding multiple endpoints with different paths."""
|
||||
app = FastAPI()
|
||||
agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=MockChatClient("Response 1"))
|
||||
agent2 = ChatAgent(name="agent2", instructions="Second agent", chat_client=MockChatClient("Response 2"))
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent1, path="/agent1")
|
||||
add_agent_framework_fastapi_endpoint(app, agent2, path="/agent2")
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response1 = client.post("/agent1", json={"messages": [{"role": "user", "content": "Hi"}]})
|
||||
response2 = client.post("/agent2", json={"messages": [{"role": "user", "content": "Hi"}]})
|
||||
|
||||
assert response1.status_code == 200
|
||||
assert response2.status_code == 200
|
||||
|
||||
|
||||
async def test_endpoint_default_path():
|
||||
"""Test endpoint with default path."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post("/", json={"messages": [{"role": "user", "content": "Hello"}]})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
async def test_endpoint_response_headers():
|
||||
"""Test that endpoint sets correct response headers."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/headers")
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post("/headers", json={"messages": [{"role": "user", "content": "Test"}]})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
|
||||
assert "cache-control" in response.headers
|
||||
assert response.headers["cache-control"] == "no-cache"
|
||||
|
||||
|
||||
async def test_endpoint_empty_messages():
|
||||
"""Test endpoint with empty messages list."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/empty")
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post("/empty", json={"messages": []})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
async def test_endpoint_complex_input():
|
||||
"""Test endpoint with complex input data."""
|
||||
app = FastAPI()
|
||||
agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
|
||||
|
||||
add_agent_framework_fastapi_endpoint(app, agent, path="/complex")
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/complex",
|
||||
json={
|
||||
"messages": [
|
||||
{"role": "user", "content": "First message", "id": "msg-1"},
|
||||
{"role": "assistant", "content": "Response", "id": "msg-2"},
|
||||
{"role": "user", "content": "Follow-up", "id": "msg-3"},
|
||||
],
|
||||
"run_id": "complex-run-123",
|
||||
"thread_id": "complex-thread-456",
|
||||
"state": {"custom_field": "value"},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
@@ -0,0 +1,659 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Comprehensive tests for AgentFrameworkEventBridge (_events.py)."""
|
||||
|
||||
import json
|
||||
|
||||
from agent_framework import (
|
||||
AgentRunResponseUpdate,
|
||||
FunctionApprovalRequestContent,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
TextContent,
|
||||
)
|
||||
|
||||
|
||||
async def test_basic_text_message_conversion():
|
||||
"""Test basic TextContent to AG-UI events."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[TextContent(text="Hello")])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
assert len(events) == 2
|
||||
assert events[0].type == "TEXT_MESSAGE_START"
|
||||
assert events[0].role == "assistant"
|
||||
assert events[1].type == "TEXT_MESSAGE_CONTENT"
|
||||
assert events[1].delta == "Hello"
|
||||
|
||||
|
||||
async def test_text_message_streaming():
|
||||
"""Test streaming TextContent with multiple chunks."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
update1 = AgentRunResponseUpdate(contents=[TextContent(text="Hello ")])
|
||||
update2 = AgentRunResponseUpdate(contents=[TextContent(text="world")])
|
||||
|
||||
events1 = await bridge.from_agent_run_update(update1)
|
||||
events2 = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# First update: START + CONTENT
|
||||
assert len(events1) == 2
|
||||
assert events1[0].type == "TEXT_MESSAGE_START"
|
||||
assert events1[1].delta == "Hello "
|
||||
|
||||
# Second update: just CONTENT (same message)
|
||||
assert len(events2) == 1
|
||||
assert events2[0].type == "TEXT_MESSAGE_CONTENT"
|
||||
assert events2[0].delta == "world"
|
||||
|
||||
# Both content events should have same message_id
|
||||
assert events1[1].message_id == events2[0].message_id
|
||||
|
||||
|
||||
async def test_skip_text_content_for_structured_outputs():
|
||||
"""Test that text content is skipped when skip_text_content=True."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread", skip_text_content=True)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[TextContent(text='{"result": "data"}')])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# No events should be emitted
|
||||
assert len(events) == 0
|
||||
|
||||
|
||||
async def test_tool_call_with_name():
|
||||
"""Test FunctionCallContent with name emits ToolCallStartEvent."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[FunctionCallContent(name="search_web", call_id="call_123")])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "TOOL_CALL_START"
|
||||
assert events[0].tool_call_name == "search_web"
|
||||
assert events[0].tool_call_id == "call_123"
|
||||
|
||||
|
||||
async def test_tool_call_streaming_args():
|
||||
"""Test streaming tool call arguments."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
# First chunk: name only
|
||||
update1 = AgentRunResponseUpdate(contents=[FunctionCallContent(name="search_web", call_id="call_123")])
|
||||
events1 = await bridge.from_agent_run_update(update1)
|
||||
|
||||
# Second chunk: arguments chunk 1 (name can be empty string for continuation)
|
||||
update2 = AgentRunResponseUpdate(
|
||||
contents=[FunctionCallContent(name="", call_id="call_123", arguments='{"query": "')]
|
||||
)
|
||||
events2 = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Third chunk: arguments chunk 2
|
||||
update3 = AgentRunResponseUpdate(contents=[FunctionCallContent(name="", call_id="call_123", arguments='AI"}')])
|
||||
events3 = await bridge.from_agent_run_update(update3)
|
||||
|
||||
# First update: ToolCallStartEvent
|
||||
assert len(events1) == 1
|
||||
assert events1[0].type == "TOOL_CALL_START"
|
||||
|
||||
# Second update: ToolCallArgsEvent
|
||||
assert len(events2) == 1
|
||||
assert events2[0].type == "TOOL_CALL_ARGS"
|
||||
assert events2[0].delta == '{"query": "'
|
||||
|
||||
# Third update: ToolCallArgsEvent
|
||||
assert len(events3) == 1
|
||||
assert events3[0].type == "TOOL_CALL_ARGS"
|
||||
assert events3[0].delta == 'AI"}'
|
||||
|
||||
# All should have same tool_call_id
|
||||
assert events1[0].tool_call_id == events2[0].tool_call_id == events3[0].tool_call_id
|
||||
|
||||
|
||||
async def test_tool_result_with_dict():
|
||||
"""Test FunctionResultContent with dict result."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
result_data = {"status": "success", "count": 42}
|
||||
update = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_123", result=result_data)])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should emit ToolCallEndEvent + ToolCallResultEvent
|
||||
assert len(events) == 2
|
||||
assert events[0].type == "TOOL_CALL_END"
|
||||
assert events[0].tool_call_id == "call_123"
|
||||
|
||||
assert events[1].type == "TOOL_CALL_RESULT"
|
||||
assert events[1].tool_call_id == "call_123"
|
||||
assert events[1].role == "tool"
|
||||
# Result should be JSON-serialized
|
||||
assert json.loads(events[1].content) == result_data
|
||||
|
||||
|
||||
async def test_tool_result_with_string():
|
||||
"""Test FunctionResultContent with string result."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_123", result="Search complete")])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
assert len(events) == 2
|
||||
assert events[0].type == "TOOL_CALL_END"
|
||||
assert events[1].type == "TOOL_CALL_RESULT"
|
||||
assert events[1].content == "Search complete"
|
||||
|
||||
|
||||
async def test_tool_result_with_none():
|
||||
"""Test FunctionResultContent with None result."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_123", result=None)])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
assert len(events) == 2
|
||||
assert events[0].type == "TOOL_CALL_END"
|
||||
assert events[1].type == "TOOL_CALL_RESULT"
|
||||
assert events[1].content == ""
|
||||
|
||||
|
||||
async def test_multiple_tool_results_in_sequence():
|
||||
"""Test multiple tool results processed sequentially."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
update = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionResultContent(call_id="call_1", result="Result 1"),
|
||||
FunctionResultContent(call_id="call_2", result="Result 2"),
|
||||
]
|
||||
)
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Each result emits: ToolCallEndEvent + ToolCallResultEvent = 4 events total
|
||||
assert len(events) == 4
|
||||
assert events[0].tool_call_id == "call_1"
|
||||
assert events[1].tool_call_id == "call_1"
|
||||
assert events[2].tool_call_id == "call_2"
|
||||
assert events[3].tool_call_id == "call_2"
|
||||
|
||||
|
||||
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")
|
||||
|
||||
func_call = FunctionCallContent(
|
||||
call_id="call_123",
|
||||
name="send_email",
|
||||
arguments={"to": "user@example.com", "subject": "Test"},
|
||||
)
|
||||
approval = FunctionApprovalRequestContent(
|
||||
id="approval_001",
|
||||
function_call=func_call,
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[approval])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should emit: ToolCallEndEvent + CustomEvent
|
||||
assert len(events) == 2
|
||||
|
||||
# First: ToolCallEndEvent to close the tool call
|
||||
assert events[0].type == "TOOL_CALL_END"
|
||||
assert events[0].tool_call_id == "call_123"
|
||||
|
||||
# Second: CustomEvent with approval details
|
||||
assert events[1].type == "CUSTOM"
|
||||
assert events[1].name == "function_approval_request"
|
||||
assert events[1].value["id"] == "approval_001"
|
||||
assert events[1].value["function_call"]["name"] == "send_email"
|
||||
|
||||
|
||||
async def test_empty_predict_state_config():
|
||||
"""Test behavior with no predictive state configuration."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={}, # Empty config
|
||||
)
|
||||
|
||||
# Tool call with arguments
|
||||
update = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(name="write_doc", call_id="call_1", arguments='{"content": "test"}'),
|
||||
FunctionResultContent(call_id="call_1", result="Done"),
|
||||
]
|
||||
)
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should NOT emit StateDeltaEvent or confirm_changes
|
||||
event_types = [e.type for e in events]
|
||||
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
|
||||
assert event_types == [
|
||||
"TOOL_CALL_START",
|
||||
"TOOL_CALL_ARGS",
|
||||
"TOOL_CALL_END",
|
||||
"TOOL_CALL_RESULT",
|
||||
"MESSAGES_SNAPSHOT",
|
||||
]
|
||||
|
||||
|
||||
async def test_tool_not_in_predict_state_config():
|
||||
"""Test tool that doesn't match any predict_state_config entry."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={
|
||||
"document": {"tool": "write_document", "tool_argument": "content"},
|
||||
},
|
||||
)
|
||||
|
||||
# Different tool name
|
||||
update = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(name="search_web", call_id="call_1", arguments='{"query": "AI"}'),
|
||||
FunctionResultContent(call_id="call_1", result="Results"),
|
||||
]
|
||||
)
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should NOT emit StateDeltaEvent or confirm_changes
|
||||
event_types = [e.type for e in events]
|
||||
assert "STATE_DELTA" not in event_types
|
||||
assert "STATE_SNAPSHOT" not in event_types
|
||||
|
||||
|
||||
async def test_state_management_tracking():
|
||||
"""Test current_state and pending_state_updates tracking."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
initial_state = {"document": ""}
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={
|
||||
"document": {"tool": "write_doc", "tool_argument": "content"},
|
||||
},
|
||||
current_state=initial_state,
|
||||
)
|
||||
|
||||
# Streaming tool call
|
||||
update1 = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(name="write_doc", call_id="call_1"),
|
||||
FunctionCallContent(name="", call_id="call_1", arguments='{"content": "Hello"}'),
|
||||
]
|
||||
)
|
||||
await bridge.from_agent_run_update(update1)
|
||||
|
||||
# Check pending_state_updates was populated
|
||||
assert "document" in bridge.pending_state_updates
|
||||
assert bridge.pending_state_updates["document"] == "Hello"
|
||||
|
||||
# Tool result should update current_state
|
||||
update2 = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_1", result="Done")])
|
||||
await bridge.from_agent_run_update(update2)
|
||||
|
||||
# current_state should be updated
|
||||
assert bridge.current_state["document"] == "Hello"
|
||||
|
||||
# pending_state_updates should be cleared
|
||||
assert len(bridge.pending_state_updates) == 0
|
||||
|
||||
|
||||
async def test_wildcard_tool_argument():
|
||||
"""Test tool_argument='*' uses all arguments as state value."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={
|
||||
"recipe": {"tool": "create_recipe", "tool_argument": "*"},
|
||||
},
|
||||
current_state={},
|
||||
)
|
||||
|
||||
# Complete tool call with dict arguments
|
||||
update = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(
|
||||
name="create_recipe",
|
||||
call_id="call_1",
|
||||
arguments={"title": "Pasta", "ingredients": ["pasta", "sauce"]},
|
||||
),
|
||||
FunctionResultContent(call_id="call_1", result="Created"),
|
||||
]
|
||||
)
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Find StateDeltaEvent
|
||||
delta_events = [e for e in events if e.type == "STATE_DELTA"]
|
||||
assert len(delta_events) > 0
|
||||
|
||||
# Value should be the entire arguments dict
|
||||
delta = delta_events[0].delta[0]
|
||||
assert delta["path"] == "/recipe"
|
||||
assert delta["value"] == {"title": "Pasta", "ingredients": ["pasta", "sauce"]}
|
||||
|
||||
|
||||
async def test_run_lifecycle_events():
|
||||
"""Test RunStartedEvent and RunFinishedEvent creation."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
started = bridge.create_run_started_event()
|
||||
assert started.type == "RUN_STARTED"
|
||||
assert started.run_id == "test_run"
|
||||
assert started.thread_id == "test_thread"
|
||||
|
||||
finished = bridge.create_run_finished_event(result={"status": "complete"})
|
||||
assert finished.type == "RUN_FINISHED"
|
||||
assert finished.run_id == "test_run"
|
||||
assert finished.thread_id == "test_thread"
|
||||
assert finished.result == {"status": "complete"}
|
||||
|
||||
|
||||
async def test_message_lifecycle_events():
|
||||
"""Test TextMessageStartEvent and TextMessageEndEvent creation."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
start = bridge.create_message_start_event("msg_123", role="assistant")
|
||||
assert start.type == "TEXT_MESSAGE_START"
|
||||
assert start.message_id == "msg_123"
|
||||
assert start.role == "assistant"
|
||||
|
||||
end = bridge.create_message_end_event("msg_123")
|
||||
assert end.type == "TEXT_MESSAGE_END"
|
||||
assert end.message_id == "msg_123"
|
||||
|
||||
|
||||
async def test_state_event_creation():
|
||||
"""Test StateSnapshotEvent and StateDeltaEvent creation helpers."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
# StateSnapshotEvent
|
||||
snapshot = bridge.create_state_snapshot_event({"document": "content"})
|
||||
assert snapshot.type == "STATE_SNAPSHOT"
|
||||
assert snapshot.snapshot == {"document": "content"}
|
||||
|
||||
# StateDeltaEvent with JSON Patch
|
||||
delta = bridge.create_state_delta_event([{"op": "replace", "path": "/document", "value": "new content"}])
|
||||
assert delta.type == "STATE_DELTA"
|
||||
assert len(delta.delta) == 1
|
||||
assert delta.delta[0]["op"] == "replace"
|
||||
assert delta.delta[0]["path"] == "/document"
|
||||
assert delta.delta[0]["value"] == "new content"
|
||||
|
||||
|
||||
async def test_state_snapshot_after_tool_result():
|
||||
"""Test StateSnapshotEvent emission after tool result with pending updates."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={
|
||||
"document": {"tool": "write_doc", "tool_argument": "content"},
|
||||
},
|
||||
current_state={"document": ""},
|
||||
)
|
||||
|
||||
# Tool call with streaming args
|
||||
update1 = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(name="write_doc", call_id="call_1"),
|
||||
FunctionCallContent(name="", call_id="call_1", arguments='{"content": "Test"}'),
|
||||
]
|
||||
)
|
||||
await bridge.from_agent_run_update(update1)
|
||||
|
||||
# Tool result should trigger StateSnapshotEvent
|
||||
update2 = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_1", result="Done")])
|
||||
events = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Should have: ToolCallEnd, ToolCallResult, StateSnapshot, ToolCallStart (confirm_changes), ToolCallArgs, ToolCallEnd
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) == 1
|
||||
assert snapshot_events[0].snapshot["document"] == "Test"
|
||||
|
||||
|
||||
async def test_message_id_persistence_across_chunks():
|
||||
"""Test that message_id persists across multiple text chunks."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
# First chunk
|
||||
update1 = AgentRunResponseUpdate(contents=[TextContent(text="Hello ")])
|
||||
events1 = await bridge.from_agent_run_update(update1)
|
||||
message_id = events1[0].message_id
|
||||
|
||||
# Second chunk
|
||||
update2 = AgentRunResponseUpdate(contents=[TextContent(text="world")])
|
||||
events2 = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Should use same message_id
|
||||
assert events2[0].message_id == message_id
|
||||
assert bridge.current_message_id == message_id
|
||||
|
||||
|
||||
async def test_tool_call_id_tracking():
|
||||
"""Test tool_call_id tracking across streaming chunks."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
# First chunk with name
|
||||
update1 = AgentRunResponseUpdate(contents=[FunctionCallContent(name="search", call_id="call_1")])
|
||||
await bridge.from_agent_run_update(update1)
|
||||
|
||||
assert bridge.current_tool_call_id == "call_1"
|
||||
assert bridge.current_tool_call_name == "search"
|
||||
|
||||
# Second chunk with args but no name
|
||||
update2 = AgentRunResponseUpdate(contents=[FunctionCallContent(name="", call_id="call_1", arguments='{"q":"AI"}')])
|
||||
events2 = await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Should still track same tool call
|
||||
assert bridge.current_tool_call_id == "call_1"
|
||||
assert events2[0].tool_call_id == "call_1"
|
||||
|
||||
|
||||
async def test_tool_name_reset_after_result():
|
||||
"""Test current_tool_call_name is reset after tool result."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={
|
||||
"document": {"tool": "write_doc", "tool_argument": "content"},
|
||||
},
|
||||
)
|
||||
|
||||
# Tool call
|
||||
update1 = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(name="write_doc", call_id="call_1"),
|
||||
FunctionCallContent(name="", call_id="call_1", arguments='{"content": "Test"}'),
|
||||
]
|
||||
)
|
||||
await bridge.from_agent_run_update(update1)
|
||||
|
||||
assert bridge.current_tool_call_name == "write_doc"
|
||||
|
||||
# Tool result with predictive state (should trigger confirm_changes and reset)
|
||||
update2 = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_1", result="Done")])
|
||||
await bridge.from_agent_run_update(update2)
|
||||
|
||||
# Tool name should be reset
|
||||
assert bridge.current_tool_call_name is None
|
||||
|
||||
|
||||
async def test_function_approval_with_wildcard_argument():
|
||||
"""Test function approval with wildcard * argument."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={
|
||||
"payload": {"tool": "submit", "tool_argument": "*"},
|
||||
},
|
||||
)
|
||||
|
||||
approval_content = FunctionApprovalRequestContent(
|
||||
id="approval_1",
|
||||
function_call=FunctionCallContent(
|
||||
name="submit", call_id="call_1", arguments='{"key1": "value1", "key2": "value2"}'
|
||||
),
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[approval_content])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should emit StateSnapshotEvent with entire parsed args as value
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) == 1
|
||||
assert snapshot_events[0].snapshot["payload"] == {"key1": "value1", "key2": "value2"}
|
||||
|
||||
|
||||
async def test_function_approval_missing_argument():
|
||||
"""Test function approval when specified argument is not in parsed args."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={
|
||||
"data": {"tool": "process", "tool_argument": "missing_field"},
|
||||
},
|
||||
)
|
||||
|
||||
approval_content = FunctionApprovalRequestContent(
|
||||
id="approval_1",
|
||||
function_call=FunctionCallContent(name="process", call_id="call_1", arguments='{"other_field": "value"}'),
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[approval_content])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should not emit StateSnapshotEvent since argument not found
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) == 0
|
||||
|
||||
|
||||
async def test_empty_predict_state_config_no_deltas():
|
||||
"""Test with empty predict_state_config (no predictive updates)."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread", predict_state_config={})
|
||||
|
||||
# Tool call with arguments
|
||||
update = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(name="search", call_id="call_1"),
|
||||
FunctionCallContent(name="", call_id="call_1", arguments='{"query": "test"}'),
|
||||
]
|
||||
)
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should not emit any StateDeltaEvents
|
||||
delta_events = [e for e in events if e.type == "STATE_DELTA"]
|
||||
assert len(delta_events) == 0
|
||||
|
||||
|
||||
async def test_tool_with_no_matching_config():
|
||||
"""Test tool call for tool not in predict_state_config."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}},
|
||||
)
|
||||
|
||||
# Tool call for different tool
|
||||
update = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(name="search_web", call_id="call_1"),
|
||||
FunctionCallContent(name="", call_id="call_1", arguments='{"query": "test"}'),
|
||||
]
|
||||
)
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should not emit StateDeltaEvents
|
||||
delta_events = [e for e in events if e.type == "STATE_DELTA"]
|
||||
assert len(delta_events) == 0
|
||||
|
||||
|
||||
async def test_tool_call_without_name_or_id():
|
||||
"""Test handling FunctionCallContent with no name and no call_id."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
|
||||
|
||||
# This should not crash but log an error
|
||||
update = AgentRunResponseUpdate(contents=[FunctionCallContent(name="", call_id="", arguments='{"arg": "val"}')])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should emit ToolCallArgsEvent with generated ID
|
||||
assert len(events) >= 1
|
||||
|
||||
|
||||
async def test_state_delta_count_logging():
|
||||
"""Test that state delta count increments and logs at intervals."""
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}},
|
||||
)
|
||||
|
||||
# Emit multiple state deltas with different content each time
|
||||
for i in range(15):
|
||||
update = AgentRunResponseUpdate(
|
||||
contents=[
|
||||
FunctionCallContent(name="", call_id="call_1", arguments=f'{{"text": "Content variation {i}"}}'),
|
||||
]
|
||||
)
|
||||
# Set the tool name to match config
|
||||
bridge.current_tool_call_name = "write"
|
||||
await bridge.from_agent_run_update(update)
|
||||
|
||||
# State delta count should have incremented (one per unique state update)
|
||||
assert bridge.state_delta_count >= 1
|
||||
@@ -0,0 +1,96 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for human in the loop (function approval requests)."""
|
||||
|
||||
from agent_framework import FunctionApprovalRequestContent, FunctionCallContent
|
||||
from agent_framework._types import AgentRunResponseUpdate
|
||||
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
|
||||
async def test_function_approval_request_emission():
|
||||
"""Test that CustomEvent is emitted for FunctionApprovalRequestContent."""
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
)
|
||||
|
||||
# Create approval request
|
||||
func_call = FunctionCallContent(
|
||||
call_id="call_123",
|
||||
name="send_email",
|
||||
arguments={"to": "user@example.com", "subject": "Test"},
|
||||
)
|
||||
approval_request = FunctionApprovalRequestContent(
|
||||
id="approval_001",
|
||||
function_call=func_call,
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[approval_request])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should emit ToolCallEndEvent + CustomEvent for approval request
|
||||
assert len(events) == 2
|
||||
|
||||
# First event: ToolCallEndEvent to close the tool call
|
||||
assert events[0].type == "TOOL_CALL_END"
|
||||
assert events[0].tool_call_id == "call_123"
|
||||
|
||||
# Second event: CustomEvent with approval details
|
||||
event = events[1]
|
||||
assert event.type == "CUSTOM"
|
||||
assert event.name == "function_approval_request"
|
||||
assert event.value["id"] == "approval_001"
|
||||
assert event.value["function_call"]["call_id"] == "call_123"
|
||||
assert event.value["function_call"]["name"] == "send_email"
|
||||
assert event.value["function_call"]["arguments"]["to"] == "user@example.com"
|
||||
assert event.value["function_call"]["arguments"]["subject"] == "Test"
|
||||
|
||||
|
||||
async def test_multiple_approval_requests():
|
||||
"""Test handling multiple approval requests in one update."""
|
||||
bridge = AgentFrameworkEventBridge(
|
||||
run_id="test_run",
|
||||
thread_id="test_thread",
|
||||
)
|
||||
|
||||
func_call_1 = FunctionCallContent(
|
||||
call_id="call_1",
|
||||
name="create_event",
|
||||
arguments={"title": "Meeting"},
|
||||
)
|
||||
approval_1 = FunctionApprovalRequestContent(
|
||||
id="approval_1",
|
||||
function_call=func_call_1,
|
||||
)
|
||||
|
||||
func_call_2 = FunctionCallContent(
|
||||
call_id="call_2",
|
||||
name="book_room",
|
||||
arguments={"room": "Conference A"},
|
||||
)
|
||||
approval_2 = FunctionApprovalRequestContent(
|
||||
id="approval_2",
|
||||
function_call=func_call_2,
|
||||
)
|
||||
|
||||
update = AgentRunResponseUpdate(contents=[approval_1, approval_2])
|
||||
events = await bridge.from_agent_run_update(update)
|
||||
|
||||
# Should emit ToolCallEndEvent + CustomEvent for each approval (4 events total)
|
||||
assert len(events) == 4
|
||||
|
||||
# Events should alternate: End, Custom, End, Custom
|
||||
assert events[0].type == "TOOL_CALL_END"
|
||||
assert events[0].tool_call_id == "call_1"
|
||||
|
||||
assert events[1].type == "CUSTOM"
|
||||
assert events[1].name == "function_approval_request"
|
||||
assert events[1].value["id"] == "approval_1"
|
||||
|
||||
assert events[2].type == "TOOL_CALL_END"
|
||||
assert events[2].tool_call_id == "call_2"
|
||||
|
||||
assert events[3].type == "CUSTOM"
|
||||
assert events[3].name == "function_approval_request"
|
||||
assert events[3].value["id"] == "approval_2"
|
||||
@@ -0,0 +1,249 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for message adapters."""
|
||||
|
||||
import pytest
|
||||
from agent_framework import ChatMessage, FunctionCallContent, Role, TextContent
|
||||
|
||||
from agent_framework_ag_ui._message_adapters import (
|
||||
agent_framework_messages_to_agui,
|
||||
agui_messages_to_agent_framework,
|
||||
extract_text_from_contents,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_agui_message():
|
||||
"""Create a sample AG-UI message."""
|
||||
return {"role": "user", "content": "Hello", "id": "msg-123"}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_agent_framework_message():
|
||||
"""Create a sample Agent Framework message."""
|
||||
return ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")], message_id="msg-123")
|
||||
|
||||
|
||||
def test_agui_to_agent_framework_basic(sample_agui_message):
|
||||
"""Test converting AG-UI message to Agent Framework."""
|
||||
messages = agui_messages_to_agent_framework([sample_agui_message])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == Role.USER
|
||||
assert messages[0].message_id == "msg-123"
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_basic(sample_agent_framework_message):
|
||||
"""Test converting Agent Framework message to AG-UI."""
|
||||
messages = agent_framework_messages_to_agui([sample_agent_framework_message])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["role"] == "user"
|
||||
assert messages[0]["content"] == "Hello"
|
||||
assert messages[0]["id"] == "msg-123"
|
||||
|
||||
|
||||
def test_agui_tool_result_to_agent_framework():
|
||||
"""Test converting AG-UI tool result message to Agent Framework."""
|
||||
tool_result_message = {
|
||||
"role": "tool",
|
||||
"content": '{"accepted": true, "steps": []}',
|
||||
"toolCallId": "call_123",
|
||||
"id": "msg_456",
|
||||
}
|
||||
|
||||
messages = agui_messages_to_agent_framework([tool_result_message])
|
||||
|
||||
assert len(messages) == 1
|
||||
message = messages[0]
|
||||
|
||||
assert message.role == Role.USER
|
||||
|
||||
assert len(message.contents) == 1
|
||||
assert isinstance(message.contents[0], TextContent)
|
||||
assert message.contents[0].text == '{"accepted": true, "steps": []}'
|
||||
|
||||
assert hasattr(message, "metadata")
|
||||
assert message.metadata is not None
|
||||
assert message.metadata.get("is_tool_result") is True
|
||||
assert message.metadata.get("tool_call_id") == "call_123"
|
||||
|
||||
|
||||
def test_agui_multiple_messages_to_agent_framework():
|
||||
"""Test converting multiple AG-UI messages."""
|
||||
messages_input = [
|
||||
{"role": "user", "content": "First message", "id": "msg-1"},
|
||||
{"role": "assistant", "content": "Second message", "id": "msg-2"},
|
||||
{"role": "user", "content": "Third message", "id": "msg-3"},
|
||||
]
|
||||
|
||||
messages = agui_messages_to_agent_framework(messages_input)
|
||||
|
||||
assert len(messages) == 3
|
||||
assert messages[0].role == Role.USER
|
||||
assert messages[1].role == Role.ASSISTANT
|
||||
assert messages[2].role == Role.USER
|
||||
|
||||
|
||||
def test_agui_empty_messages():
|
||||
"""Test handling of empty messages list."""
|
||||
messages = agui_messages_to_agent_framework([])
|
||||
assert len(messages) == 0
|
||||
|
||||
|
||||
def test_agui_function_approvals():
|
||||
"""Test converting function approvals from AG-UI to Agent Framework."""
|
||||
agui_msg = {
|
||||
"role": "user",
|
||||
"function_approvals": [
|
||||
{
|
||||
"call_id": "call-1",
|
||||
"name": "search",
|
||||
"arguments": {"query": "test"},
|
||||
"approved": True,
|
||||
"id": "approval-1",
|
||||
},
|
||||
{
|
||||
"call_id": "call-2",
|
||||
"name": "update",
|
||||
"arguments": {"value": 42},
|
||||
"approved": False,
|
||||
"id": "approval-2",
|
||||
},
|
||||
],
|
||||
"id": "msg-123",
|
||||
}
|
||||
|
||||
messages = agui_messages_to_agent_framework([agui_msg])
|
||||
|
||||
assert len(messages) == 1
|
||||
msg = messages[0]
|
||||
assert msg.role == Role.USER
|
||||
assert len(msg.contents) == 2
|
||||
|
||||
from agent_framework import FunctionApprovalResponseContent
|
||||
|
||||
assert isinstance(msg.contents[0], FunctionApprovalResponseContent)
|
||||
assert msg.contents[0].approved is True
|
||||
assert msg.contents[0].id == "approval-1"
|
||||
assert msg.contents[0].function_call.name == "search"
|
||||
assert msg.contents[0].function_call.call_id == "call-1"
|
||||
|
||||
assert isinstance(msg.contents[1], FunctionApprovalResponseContent)
|
||||
assert msg.contents[1].approved is False
|
||||
|
||||
|
||||
def test_agui_system_role():
|
||||
"""Test converting system role messages."""
|
||||
messages = agui_messages_to_agent_framework([{"role": "system", "content": "System prompt"}])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == Role.SYSTEM
|
||||
|
||||
|
||||
def test_agui_non_string_content():
|
||||
"""Test handling non-string content."""
|
||||
messages = agui_messages_to_agent_framework([{"role": "user", "content": {"nested": "object"}}])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert len(messages[0].contents) == 1
|
||||
assert isinstance(messages[0].contents[0], TextContent)
|
||||
assert "nested" in messages[0].contents[0].text
|
||||
|
||||
|
||||
def test_agui_message_without_id():
|
||||
"""Test message without ID field."""
|
||||
messages = agui_messages_to_agent_framework([{"role": "user", "content": "No ID"}])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0].message_id is None
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_with_tool_calls():
|
||||
"""Test converting Agent Framework message with tool calls to AG-UI."""
|
||||
msg = ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
contents=[
|
||||
TextContent(text="Calling tool"),
|
||||
FunctionCallContent(call_id="call-123", name="search", arguments={"query": "test"}),
|
||||
],
|
||||
message_id="msg-456",
|
||||
)
|
||||
|
||||
messages = agent_framework_messages_to_agui([msg])
|
||||
|
||||
assert len(messages) == 1
|
||||
agui_msg = messages[0]
|
||||
assert agui_msg["role"] == "assistant"
|
||||
assert agui_msg["content"] == "Calling tool"
|
||||
assert "tool_calls" in agui_msg
|
||||
assert len(agui_msg["tool_calls"]) == 1
|
||||
assert agui_msg["tool_calls"][0]["id"] == "call-123"
|
||||
assert agui_msg["tool_calls"][0]["type"] == "function"
|
||||
assert agui_msg["tool_calls"][0]["function"]["name"] == "search"
|
||||
assert agui_msg["tool_calls"][0]["function"]["arguments"] == {"query": "test"}
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_multiple_text_contents():
|
||||
"""Test concatenating multiple text contents."""
|
||||
msg = ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
contents=[TextContent(text="Part 1 "), TextContent(text="Part 2")],
|
||||
)
|
||||
|
||||
messages = agent_framework_messages_to_agui([msg])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["content"] == "Part 1 Part 2"
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_no_message_id():
|
||||
"""Test message without message_id."""
|
||||
msg = ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])
|
||||
|
||||
messages = agent_framework_messages_to_agui([msg])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert "id" not in messages[0]
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_system_role():
|
||||
"""Test system role conversion."""
|
||||
msg = ChatMessage(role=Role.SYSTEM, contents=[TextContent(text="System")])
|
||||
|
||||
messages = agent_framework_messages_to_agui([msg])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["role"] == "system"
|
||||
|
||||
|
||||
def test_extract_text_from_contents():
|
||||
"""Test extracting text from contents list."""
|
||||
contents = [TextContent(text="Hello "), TextContent(text="World")]
|
||||
|
||||
result = extract_text_from_contents(contents)
|
||||
|
||||
assert result == "Hello World"
|
||||
|
||||
|
||||
def test_extract_text_from_empty_contents():
|
||||
"""Test extracting text from empty contents."""
|
||||
result = extract_text_from_contents([])
|
||||
|
||||
assert result == ""
|
||||
|
||||
|
||||
class CustomTextContent:
|
||||
"""Custom content with text attribute."""
|
||||
|
||||
def __init__(self, text: str):
|
||||
self.text = text
|
||||
|
||||
|
||||
def test_extract_text_from_custom_contents():
|
||||
"""Test extracting text from custom content objects."""
|
||||
contents = [CustomTextContent(text="Custom "), TextContent(text="Mixed")]
|
||||
|
||||
result = extract_text_from_contents(contents)
|
||||
|
||||
assert result == "Custom Mixed"
|
||||
@@ -0,0 +1,109 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for shared state management."""
|
||||
|
||||
import pytest
|
||||
from ag_ui.core import StateSnapshotEvent
|
||||
from agent_framework import ChatAgent, TextContent
|
||||
from agent_framework._types import ChatResponseUpdate
|
||||
|
||||
from agent_framework_ag_ui._agent import AgentFrameworkAgent
|
||||
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent():
|
||||
"""Create a mock agent for testing."""
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Hello!")])
|
||||
|
||||
return ChatAgent(
|
||||
name="test_agent",
|
||||
instructions="Test agent",
|
||||
chat_client=MockChatClient(),
|
||||
)
|
||||
|
||||
|
||||
def test_state_snapshot_event():
|
||||
"""Test creating state snapshot events."""
|
||||
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
|
||||
|
||||
state = {
|
||||
"recipe": {
|
||||
"name": "Chocolate Chip Cookies",
|
||||
"ingredients": ["flour", "sugar", "chocolate chips"],
|
||||
"instructions": ["Mix ingredients", "Bake at 350°F"],
|
||||
"servings": 24,
|
||||
}
|
||||
}
|
||||
|
||||
event = bridge.create_state_snapshot_event(state)
|
||||
|
||||
assert isinstance(event, StateSnapshotEvent)
|
||||
assert event.snapshot == state
|
||||
assert event.snapshot["recipe"]["name"] == "Chocolate Chip Cookies"
|
||||
assert len(event.snapshot["recipe"]["ingredients"]) == 3
|
||||
|
||||
|
||||
def test_state_delta_event():
|
||||
"""Test creating state delta events using JSON Patch format."""
|
||||
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
|
||||
|
||||
# JSON Patch operations (RFC 6902)
|
||||
delta = [
|
||||
{"op": "add", "path": "/recipe/ingredients/-", "value": "vanilla extract"},
|
||||
{"op": "replace", "path": "/recipe/servings", "value": 30},
|
||||
]
|
||||
|
||||
event = bridge.create_state_delta_event(delta)
|
||||
|
||||
assert event.delta == delta
|
||||
assert len(event.delta) == 2
|
||||
assert event.delta[0]["op"] == "add"
|
||||
assert event.delta[1]["op"] == "replace"
|
||||
|
||||
|
||||
async def test_agent_with_initial_state(mock_agent):
|
||||
"""Test agent emits state snapshot when initial state provided."""
|
||||
state_schema = {"recipe": {"type": "object", "properties": {"name": {"type": "string"}}}}
|
||||
|
||||
agent = AgentFrameworkAgent(
|
||||
agent=mock_agent,
|
||||
state_schema=state_schema,
|
||||
)
|
||||
|
||||
initial_state = {"recipe": {"name": "Test Recipe"}}
|
||||
|
||||
input_data = {
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"state": initial_state,
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in agent.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should have RunStartedEvent, StateSnapshotEvent, RunFinishedEvent at minimum
|
||||
snapshot_events = [e for e in events if isinstance(e, StateSnapshotEvent)]
|
||||
assert len(snapshot_events) == 1
|
||||
assert snapshot_events[0].snapshot == initial_state
|
||||
|
||||
|
||||
async def test_agent_without_state_schema(mock_agent):
|
||||
"""Test agent doesn't emit state events without state schema."""
|
||||
agent = AgentFrameworkAgent(agent=mock_agent)
|
||||
|
||||
input_data = {
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"state": {"some": "state"},
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in agent.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should NOT have any StateSnapshotEvent
|
||||
snapshot_events = [e for e in events if isinstance(e, StateSnapshotEvent)]
|
||||
assert len(snapshot_events) == 0
|
||||
@@ -0,0 +1,257 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for structured output handling in _agent.py."""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import ChatAgent, ChatOptions, TextContent
|
||||
from agent_framework._types import ChatResponseUpdate
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class RecipeOutput(BaseModel):
|
||||
"""Test Pydantic model for recipe output."""
|
||||
|
||||
recipe: dict[str, Any]
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class StepsOutput(BaseModel):
|
||||
"""Test Pydantic model for steps output."""
|
||||
|
||||
steps: list[dict[str, Any]]
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class GenericOutput(BaseModel):
|
||||
"""Test Pydantic model for generic data."""
|
||||
|
||||
data: dict[str, Any]
|
||||
|
||||
|
||||
async def test_structured_output_with_recipe():
|
||||
"""Test structured output processing with recipe state."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
# Simulate structured output
|
||||
yield ChatResponseUpdate(
|
||||
contents=[TextContent(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
|
||||
)
|
||||
|
||||
agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
|
||||
agent.chat_options = ChatOptions(response_format=RecipeOutput)
|
||||
|
||||
wrapper = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema={"recipe": {"type": "object"}},
|
||||
)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Make pasta"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit StateSnapshotEvent with recipe
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) >= 1
|
||||
# Find snapshot with recipe
|
||||
recipe_snapshots = [e for e in snapshot_events if "recipe" in e.snapshot]
|
||||
assert len(recipe_snapshots) >= 1
|
||||
assert recipe_snapshots[0].snapshot["recipe"] == {"name": "Pasta"}
|
||||
|
||||
# Should also emit message as text
|
||||
text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert any("Here is your recipe" in e.delta for e in text_events)
|
||||
|
||||
|
||||
async def test_structured_output_with_steps():
|
||||
"""Test structured output processing with steps state."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
steps_data = {
|
||||
"steps": [
|
||||
{"id": "1", "description": "Step 1", "status": "pending"},
|
||||
{"id": "2", "description": "Step 2", "status": "pending"},
|
||||
]
|
||||
}
|
||||
yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(steps_data))])
|
||||
|
||||
agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
|
||||
agent.chat_options = ChatOptions(response_format=StepsOutput)
|
||||
|
||||
wrapper = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema={"steps": {"type": "array"}},
|
||||
)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Do steps"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit StateSnapshotEvent with steps
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) >= 1
|
||||
|
||||
# Snapshot should contain steps
|
||||
steps_snapshots = [e for e in snapshot_events if "steps" in e.snapshot]
|
||||
assert len(steps_snapshots) >= 1
|
||||
assert len(steps_snapshots[0].snapshot["steps"]) == 2
|
||||
assert steps_snapshots[0].snapshot["steps"][0]["id"] == "1"
|
||||
|
||||
|
||||
async def test_structured_output_with_no_schema_match():
|
||||
"""Test structured output when response fields don't match state_schema keys."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
# Response has "data" field but schema expects "result" field
|
||||
yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}}')])
|
||||
|
||||
agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
|
||||
agent.chat_options = ChatOptions(response_format=GenericOutput)
|
||||
|
||||
wrapper = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema={"result": {"type": "object"}}, # Schema expects "result", not "data"
|
||||
)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit StateSnapshotEvent but with no state updates since no schema fields match
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
# Initial state snapshot from state_schema initialization
|
||||
assert len(snapshot_events) >= 1
|
||||
|
||||
|
||||
async def test_structured_output_without_schema():
|
||||
"""Test structured output without state_schema treats all fields as state."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class DataOutput(BaseModel):
|
||||
"""Output with data and info fields."""
|
||||
|
||||
data: dict[str, Any]
|
||||
info: str
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}, "info": "processed"}')])
|
||||
|
||||
agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
|
||||
agent.chat_options = ChatOptions(response_format=DataOutput)
|
||||
|
||||
wrapper = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
# No state_schema - all non-message fields treated as state
|
||||
)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit StateSnapshotEvent with both data and info fields
|
||||
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
|
||||
assert len(snapshot_events) >= 1
|
||||
assert "data" in snapshot_events[0].snapshot
|
||||
assert "info" in snapshot_events[0].snapshot
|
||||
assert snapshot_events[0].snapshot["data"] == {"key": "value"}
|
||||
assert snapshot_events[0].snapshot["info"] == "processed"
|
||||
|
||||
|
||||
async def test_no_structured_output_when_no_response_format():
|
||||
"""Test that structured output path is skipped when no response_format."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="Regular text")])
|
||||
|
||||
agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
|
||||
# No response_format set
|
||||
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit text content normally
|
||||
text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert len(text_events) > 0
|
||||
assert text_events[0].delta == "Regular text"
|
||||
|
||||
|
||||
async def test_structured_output_with_message_field():
|
||||
"""Test structured output that includes a message field."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
|
||||
yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(output_data))])
|
||||
|
||||
agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
|
||||
agent.chat_options = ChatOptions(response_format=RecipeOutput)
|
||||
|
||||
wrapper = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema={"recipe": {"type": "object"}},
|
||||
)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Make salad"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should emit the message as text
|
||||
text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
|
||||
assert any("Fresh salad recipe ready" in e.delta for e in text_events)
|
||||
|
||||
# Should also have TextMessageStart and TextMessageEnd
|
||||
start_events = [e for e in events if e.type == "TEXT_MESSAGE_START"]
|
||||
end_events = [e for e in events if e.type == "TEXT_MESSAGE_END"]
|
||||
assert len(start_events) >= 1
|
||||
assert len(end_events) >= 1
|
||||
|
||||
|
||||
async def test_empty_updates_no_structured_processing():
|
||||
"""Test that empty updates don't trigger structured output processing."""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
class MockChatClient:
|
||||
async def get_streaming_response(self, messages, chat_options, **kwargs):
|
||||
# Return nothing
|
||||
if False:
|
||||
yield
|
||||
|
||||
agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
|
||||
agent.chat_options = ChatOptions(response_format=RecipeOutput)
|
||||
|
||||
wrapper = AgentFrameworkAgent(agent=agent)
|
||||
|
||||
input_data = {"messages": [{"role": "user", "content": "Test"}]}
|
||||
|
||||
events = []
|
||||
async for event in wrapper.run_agent(input_data):
|
||||
events.append(event)
|
||||
|
||||
# Should only have start and end events
|
||||
assert len(events) == 2 # RunStarted, RunFinished
|
||||
@@ -0,0 +1,145 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for type definitions in _types.py."""
|
||||
|
||||
from agent_framework_ag_ui._types import AgentState, PredictStateConfig, RunMetadata
|
||||
|
||||
|
||||
class TestPredictStateConfig:
|
||||
"""Test PredictStateConfig TypedDict."""
|
||||
|
||||
def test_predict_state_config_creation(self) -> None:
|
||||
"""Test creating a PredictStateConfig dict."""
|
||||
config: PredictStateConfig = {
|
||||
"state_key": "document",
|
||||
"tool": "write_document",
|
||||
"tool_argument": "content",
|
||||
}
|
||||
|
||||
assert config["state_key"] == "document"
|
||||
assert config["tool"] == "write_document"
|
||||
assert config["tool_argument"] == "content"
|
||||
|
||||
def test_predict_state_config_with_none_tool_argument(self) -> None:
|
||||
"""Test PredictStateConfig with None tool_argument."""
|
||||
config: PredictStateConfig = {
|
||||
"state_key": "status",
|
||||
"tool": "update_status",
|
||||
"tool_argument": None,
|
||||
}
|
||||
|
||||
assert config["state_key"] == "status"
|
||||
assert config["tool"] == "update_status"
|
||||
assert config["tool_argument"] is None
|
||||
|
||||
def test_predict_state_config_type_validation(self) -> None:
|
||||
"""Test that PredictStateConfig validates field types at runtime."""
|
||||
config: PredictStateConfig = {
|
||||
"state_key": "test",
|
||||
"tool": "test_tool",
|
||||
"tool_argument": "arg",
|
||||
}
|
||||
|
||||
assert isinstance(config["state_key"], str)
|
||||
assert isinstance(config["tool"], str)
|
||||
assert isinstance(config["tool_argument"], (str, type(None)))
|
||||
|
||||
|
||||
class TestRunMetadata:
|
||||
"""Test RunMetadata TypedDict."""
|
||||
|
||||
def test_run_metadata_creation(self) -> None:
|
||||
"""Test creating a RunMetadata dict."""
|
||||
metadata: RunMetadata = {
|
||||
"run_id": "run-123",
|
||||
"thread_id": "thread-456",
|
||||
"predict_state": [
|
||||
{
|
||||
"state_key": "document",
|
||||
"tool": "write_document",
|
||||
"tool_argument": "content",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
assert metadata["run_id"] == "run-123"
|
||||
assert metadata["thread_id"] == "thread-456"
|
||||
assert metadata["predict_state"] is not None
|
||||
assert len(metadata["predict_state"]) == 1
|
||||
assert metadata["predict_state"][0]["state_key"] == "document"
|
||||
|
||||
def test_run_metadata_with_none_predict_state(self) -> None:
|
||||
"""Test RunMetadata with None predict_state."""
|
||||
metadata: RunMetadata = {
|
||||
"run_id": "run-789",
|
||||
"thread_id": "thread-012",
|
||||
"predict_state": None,
|
||||
}
|
||||
|
||||
assert metadata["run_id"] == "run-789"
|
||||
assert metadata["thread_id"] == "thread-012"
|
||||
assert metadata["predict_state"] is None
|
||||
|
||||
def test_run_metadata_empty_predict_state(self) -> None:
|
||||
"""Test RunMetadata with empty predict_state list."""
|
||||
metadata: RunMetadata = {
|
||||
"run_id": "run-345",
|
||||
"thread_id": "thread-678",
|
||||
"predict_state": [],
|
||||
}
|
||||
|
||||
assert metadata["run_id"] == "run-345"
|
||||
assert metadata["thread_id"] == "thread-678"
|
||||
assert metadata["predict_state"] == []
|
||||
|
||||
|
||||
class TestAgentState:
|
||||
"""Test AgentState TypedDict."""
|
||||
|
||||
def test_agent_state_creation(self) -> None:
|
||||
"""Test creating an AgentState dict."""
|
||||
state: AgentState = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi there"},
|
||||
]
|
||||
}
|
||||
|
||||
assert state["messages"] is not None
|
||||
assert len(state["messages"]) == 2
|
||||
assert state["messages"][0]["role"] == "user"
|
||||
assert state["messages"][1]["role"] == "assistant"
|
||||
|
||||
def test_agent_state_with_none_messages(self) -> None:
|
||||
"""Test AgentState with None messages."""
|
||||
state: AgentState = {"messages": None}
|
||||
|
||||
assert state["messages"] is None
|
||||
|
||||
def test_agent_state_empty_messages(self) -> None:
|
||||
"""Test AgentState with empty messages list."""
|
||||
state: AgentState = {"messages": []}
|
||||
|
||||
assert state["messages"] == []
|
||||
|
||||
def test_agent_state_complex_messages(self) -> None:
|
||||
"""Test AgentState with complex message structures."""
|
||||
state: AgentState = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Test",
|
||||
"metadata": {"timestamp": "2025-10-30"},
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Response",
|
||||
"tool_calls": [{"name": "search", "args": {}}],
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
assert state["messages"] is not None
|
||||
assert len(state["messages"]) == 2
|
||||
assert "metadata" in state["messages"][0]
|
||||
assert "tool_calls" in state["messages"][1]
|
||||
@@ -0,0 +1,199 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for utilities."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime
|
||||
|
||||
from agent_framework_ag_ui._utils import generate_event_id, make_json_safe, merge_state
|
||||
|
||||
|
||||
def test_generate_event_id():
|
||||
"""Test event ID generation."""
|
||||
id1 = generate_event_id()
|
||||
id2 = generate_event_id()
|
||||
|
||||
assert id1 != id2
|
||||
assert isinstance(id1, str)
|
||||
assert len(id1) > 0
|
||||
|
||||
|
||||
def test_merge_state():
|
||||
"""Test state merging."""
|
||||
current = {"a": 1, "b": 2}
|
||||
update = {"b": 3, "c": 4}
|
||||
|
||||
result = merge_state(current, update)
|
||||
|
||||
assert result["a"] == 1
|
||||
assert result["b"] == 3
|
||||
assert result["c"] == 4
|
||||
|
||||
|
||||
def test_merge_state_empty_update():
|
||||
"""Test merging with empty update."""
|
||||
current = {"x": 10, "y": 20}
|
||||
update = {}
|
||||
|
||||
result = merge_state(current, update)
|
||||
|
||||
assert result == current
|
||||
assert result is not current
|
||||
|
||||
|
||||
def test_merge_state_empty_current():
|
||||
"""Test merging with empty current state."""
|
||||
current = {}
|
||||
update = {"a": 1, "b": 2}
|
||||
|
||||
result = merge_state(current, update)
|
||||
|
||||
assert result == update
|
||||
|
||||
|
||||
def test_merge_state_deep_copy():
|
||||
"""Test that merge_state creates a deep copy preventing mutation of original."""
|
||||
current = {"recipe": {"name": "Cake", "ingredients": ["flour", "sugar"]}}
|
||||
update = {"other": "value"}
|
||||
|
||||
result = merge_state(current, update)
|
||||
|
||||
result["recipe"]["ingredients"].append("eggs")
|
||||
|
||||
assert "eggs" not in current["recipe"]["ingredients"]
|
||||
assert current["recipe"]["ingredients"] == ["flour", "sugar"]
|
||||
assert result["recipe"]["ingredients"] == ["flour", "sugar", "eggs"]
|
||||
|
||||
|
||||
def test_make_json_safe_basic():
|
||||
"""Test JSON serialization of basic types."""
|
||||
assert make_json_safe("text") == "text"
|
||||
assert make_json_safe(123) == 123
|
||||
assert make_json_safe(None) is None
|
||||
assert make_json_safe(3.14) == 3.14
|
||||
assert make_json_safe(True) is True
|
||||
assert make_json_safe(False) is False
|
||||
|
||||
|
||||
def test_make_json_safe_datetime():
|
||||
"""Test datetime serialization."""
|
||||
dt = datetime(2025, 10, 30, 12, 30, 45)
|
||||
result = make_json_safe(dt)
|
||||
assert result == "2025-10-30T12:30:45"
|
||||
|
||||
|
||||
def test_make_json_safe_date():
|
||||
"""Test date serialization."""
|
||||
d = date(2025, 10, 30)
|
||||
result = make_json_safe(d)
|
||||
assert result == "2025-10-30"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SampleDataclass:
|
||||
"""Sample dataclass for testing."""
|
||||
|
||||
name: str
|
||||
value: int
|
||||
|
||||
|
||||
def test_make_json_safe_dataclass():
|
||||
"""Test dataclass serialization."""
|
||||
obj = SampleDataclass(name="test", value=42)
|
||||
result = make_json_safe(obj)
|
||||
assert result == {"name": "test", "value": 42}
|
||||
|
||||
|
||||
class ModelDumpObject:
|
||||
"""Object with model_dump method."""
|
||||
|
||||
def model_dump(self):
|
||||
return {"type": "model", "data": "dump"}
|
||||
|
||||
|
||||
def test_make_json_safe_model_dump():
|
||||
"""Test object with model_dump method."""
|
||||
obj = ModelDumpObject()
|
||||
result = make_json_safe(obj)
|
||||
assert result == {"type": "model", "data": "dump"}
|
||||
|
||||
|
||||
class DictObject:
|
||||
"""Object with dict method."""
|
||||
|
||||
def dict(self):
|
||||
return {"type": "dict", "method": "call"}
|
||||
|
||||
|
||||
def test_make_json_safe_dict_method():
|
||||
"""Test object with dict method."""
|
||||
obj = DictObject()
|
||||
result = make_json_safe(obj)
|
||||
assert result == {"type": "dict", "method": "call"}
|
||||
|
||||
|
||||
class CustomObject:
|
||||
"""Custom object with __dict__."""
|
||||
|
||||
def __init__(self):
|
||||
self.field1 = "value1"
|
||||
self.field2 = 123
|
||||
|
||||
|
||||
def test_make_json_safe_dict_attribute():
|
||||
"""Test object with __dict__ attribute."""
|
||||
obj = CustomObject()
|
||||
result = make_json_safe(obj)
|
||||
assert result == {"field1": "value1", "field2": 123}
|
||||
|
||||
|
||||
def test_make_json_safe_list():
|
||||
"""Test list serialization."""
|
||||
lst = [1, "text", None, {"key": "value"}]
|
||||
result = make_json_safe(lst)
|
||||
assert result == [1, "text", None, {"key": "value"}]
|
||||
|
||||
|
||||
def test_make_json_safe_tuple():
|
||||
"""Test tuple serialization."""
|
||||
tpl = (1, 2, 3)
|
||||
result = make_json_safe(tpl)
|
||||
assert result == [1, 2, 3]
|
||||
|
||||
|
||||
def test_make_json_safe_dict():
|
||||
"""Test dict serialization."""
|
||||
d = {"a": 1, "b": {"c": 2}}
|
||||
result = make_json_safe(d)
|
||||
assert result == {"a": 1, "b": {"c": 2}}
|
||||
|
||||
|
||||
def test_make_json_safe_nested():
|
||||
"""Test nested structure serialization."""
|
||||
obj = {
|
||||
"datetime": datetime(2025, 10, 30),
|
||||
"list": [1, 2, CustomObject()],
|
||||
"nested": {"value": SampleDataclass(name="nested", value=99)},
|
||||
}
|
||||
result = make_json_safe(obj)
|
||||
|
||||
assert result["datetime"] == "2025-10-30T00:00:00"
|
||||
assert result["list"][0] == 1
|
||||
assert result["list"][2] == {"field1": "value1", "field2": 123}
|
||||
assert result["nested"]["value"] == {"name": "nested", "value": 99}
|
||||
|
||||
|
||||
class UnserializableObject:
|
||||
"""Object that can't be serialized by standard methods."""
|
||||
|
||||
def __init__(self):
|
||||
# Add attribute to trigger __dict__ fallback path
|
||||
pass
|
||||
|
||||
|
||||
def test_make_json_safe_fallback():
|
||||
"""Test fallback to dict for objects with __dict__."""
|
||||
obj = UnserializableObject()
|
||||
result = make_json_safe(obj)
|
||||
# Objects with __dict__ return their __dict__ dict
|
||||
assert isinstance(result, dict)
|
||||
Reference in New Issue
Block a user