agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Refactor ag-ui to clean up some patterns (#2363)

Browse Source 
* Refactor ag-ui to clean up some patterns

* Mypy fixes

* Fix imports, typing, tests, logging.

* Fix test import error

* Fix imports again

* Fix thread handling
This commit is contained in:
Evan Mattson
2025-11-27 11:13:03 +09:00
committed by GitHub
Unverified
parent 6c624319db
commit 8cf8b0f995
26 changed files with 1887 additions and 1415 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/ag-ui/tests/test_agent_wrapper_comprehensive.py
+182 -132
View File
@@ -3,21 +3,30 @@
"""Comprehensive tests for AgentFrameworkAgent (_agent.py)."""
import json
import sys
from collections.abc import AsyncIterator, MutableSequence
from pathlib import Path
from typing import Any
import pytest
from agent_framework import ChatAgent, TextContent
from agent_framework import ChatAgent, ChatMessage, ChatOptions, TextContent
from agent_framework._types import ChatResponseUpdate
from pydantic import BaseModel
sys.path.insert(0, str(Path(__file__).parent))
from test_helpers_ag_ui import StreamingChatClientStub
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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
assert wrapper.name == "test_agent"
@@ -30,12 +39,13 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
state_schema = {"document": {"type": "string"}}
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"document": {"type": "string"}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
assert wrapper.config.state_schema == state_schema
@@ -45,31 +55,56 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
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_agent_initialization_with_pydantic_state_schema():
"""Test agent initialization when state_schema is provided as Pydantic model/class."""
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
class MyState(BaseModel):
document: str
tags: list[str] = []
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper_class_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState)
wrapper_instance_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState(document="hi"))
expected_properties = MyState.model_json_schema().get("properties", {})
assert wrapper_class_schema.config.state_schema == expected_properties
assert wrapper_instance_schema.config.state_schema == expected_properties
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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -83,11 +118,12 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
predict_config = {
"document": {"tool": "write_doc", "tool_argument": "content"},
"summary": {"tool": "summarize", "tool_argument": "text"},
@@ -96,7 +132,7 @@ async def test_predict_state_custom_event_emission():
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -114,11 +150,12 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
state_schema = {"document": {"type": "string"}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
@@ -127,7 +164,7 @@ async def test_initial_state_snapshot_with_schema():
"state": {"document": "Initial content"},
}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -143,17 +180,18 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
state_schema = {"recipe": {"type": "object", "properties": {}}}
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"recipe": {"type": "object", "properties": {}}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -169,17 +207,18 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
state_schema = {"steps": {"type": "array", "items": {}}}
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"steps": {"type": "array", "items": {}}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -195,16 +234,17 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -216,11 +256,12 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Document updated")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(
agent=agent,
state_schema={"document": {"type": "string"}},
@@ -228,8 +269,8 @@ async def test_tool_result_confirm_changes_accepted():
)
# Simulate tool result message with acceptance
tool_result = {"accepted": True, "steps": []}
input_data = {
tool_result: dict[str, Any] = {"accepted": True, "steps": []}
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool", # Tool result from UI
@@ -240,7 +281,7 @@ async def test_tool_result_confirm_changes_accepted():
"state": {"document": "Updated content"},
}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -262,16 +303,17 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result message with rejection
tool_result = {"accepted": False, "steps": []}
input_data = {
tool_result: dict[str, Any] = {"accepted": False, "steps": []}
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
@@ -281,7 +323,7 @@ async def test_tool_result_confirm_changes_rejected():
],
}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -295,22 +337,23 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result with multiple steps
tool_result = {
tool_result: dict[str, Any] = {
"accepted": True,
"steps": [
{"id": "step1", "description": "Send email", "status": "enabled"},
{"id": "step2", "description": "Create calendar event", "status": "enabled"},
],
}
input_data = {
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
@@ -320,7 +363,7 @@ async def test_tool_result_function_approval_accepted():
],
}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -340,19 +383,20 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result rejection with steps
tool_result = {
tool_result: dict[str, Any] = {
"accepted": False,
"steps": [{"id": "step1", "description": "Send email", "status": "disabled"}],
}
input_data = {
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
@@ -362,7 +406,7 @@ async def test_tool_result_function_approval_rejected():
],
}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -376,17 +420,16 @@ 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 = {}
thread_metadata: dict[str, Any] = {}
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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if chat_options.metadata:
thread_metadata.update(chat_options.metadata)
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {
@@ -395,28 +438,28 @@ async def test_thread_metadata_tracking():
"run_id": "test_run_456",
}
events = []
events: list[Any] = []
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
assert thread_metadata.get("ag_ui_thread_id") == "test_thread_123"
assert thread_metadata.get("ag_ui_run_id") == "test_run_456"
async def test_state_context_injection():
"""Test that current state is injected into thread metadata."""
from agent_framework.ag_ui import AgentFrameworkAgent
thread_metadata = {}
thread_metadata: dict[str, Any] = {}
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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if chat_options.metadata:
thread_metadata.update(chat_options.metadata)
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(
agent=agent,
state_schema={"document": {"type": "string"}},
@@ -427,27 +470,31 @@ async def test_state_context_injection():
"state": {"document": "Test content"},
}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
# State should be injected - this is validated by agent execution flow
current_state = thread_metadata.get("current_state")
if isinstance(current_state, str):
current_state = json.loads(current_state)
assert current_state == {"document": "Test content"}
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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": []}
input_data: dict[str, Any] = {"messages": []}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -461,16 +508,17 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello world")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -488,19 +536,20 @@ 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")
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if False:
yield ChatResponseUpdate(contents=[])
raise RuntimeError("Simulated failure")
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=FailingChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
with pytest.raises(RuntimeError, match="Simulated failure"):
async for event in wrapper.run_agent(input_data):
async for _ in wrapper.run_agent(input_data):
pass
@@ -508,18 +557,18 @@ async def test_json_decode_error_in_tool_result():
"""Test handling of orphaned tool result - should be sanitized out."""
from agent_framework.ag_ui import AgentFrameworkAgent
class MockChatClient:
async def get_streaming_response(self, messages, chat_options, **kwargs):
# Should not be called since orphaned tool result is dropped
if False:
yield
raise AssertionError("ChatClient should not be called with orphaned tool result")
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if False:
yield ChatResponseUpdate(contents=[])
raise AssertionError("ChatClient should not be called with orphaned tool result")
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Send invalid JSON as tool result without preceding tool call
input_data = {
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
@@ -529,7 +578,7 @@ async def test_json_decode_error_in_tool_result():
],
}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
@@ -545,11 +594,12 @@ 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")])
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Response")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(
agent=agent,
state_schema={"document": {"type": "string"}},
@@ -558,8 +608,8 @@ async def test_suppressed_summary_with_document_state():
)
# Simulate confirmation with document state
tool_result = {"accepted": True, "steps": []}
input_data = {
tool_result: dict[str, Any] = {"accepted": True, "steps": []}
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
@@ -570,7 +620,7 @@ async def test_suppressed_summary_with_document_state():
"state": {"document": "This is the beginning of a document. It contains important information."},
}
events = []
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)