mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Dmytro Struk
@@ -1,317 +0,0 @@
|
||||
"""Tests for AGUIChatClient."""
|
||||
|
||||
import json
|
||||
|
||||
from agent_framework import ChatMessage, ChatOptions, FunctionCallContent, Role, ai_function
|
||||
|
||||
from agent_framework_ag_ui._client import AGUIChatClient, ServerFunctionCallContent
|
||||
|
||||
|
||||
class TestAGUIChatClient:
|
||||
"""Test suite for AGUIChatClient."""
|
||||
|
||||
async def test_client_initialization(self) -> None:
|
||||
"""Test client initialization."""
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
|
||||
assert client._http_service is not None
|
||||
assert client._http_service.endpoint.startswith("http://localhost:8888")
|
||||
|
||||
async def test_client_context_manager(self) -> None:
|
||||
"""Test client as async context manager."""
|
||||
async with AGUIChatClient(endpoint="http://localhost:8888/") as client:
|
||||
assert client is not None
|
||||
|
||||
async def test_extract_state_from_messages_no_state(self) -> None:
|
||||
"""Test state extraction when no state is present."""
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
messages = [
|
||||
ChatMessage(role="user", text="Hello"),
|
||||
ChatMessage(role="assistant", text="Hi there"),
|
||||
]
|
||||
|
||||
result_messages, state = client._extract_state_from_messages(messages)
|
||||
|
||||
assert result_messages == messages
|
||||
assert state is None
|
||||
|
||||
async def test_extract_state_from_messages_with_state(self) -> None:
|
||||
"""Test state extraction from last message."""
|
||||
import base64
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
|
||||
state_data = {"key": "value", "count": 42}
|
||||
state_json = json.dumps(state_data)
|
||||
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
|
||||
|
||||
from agent_framework import DataContent
|
||||
|
||||
messages = [
|
||||
ChatMessage(role="user", text="Hello"),
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[DataContent(uri=f"data:application/json;base64,{state_b64}")],
|
||||
),
|
||||
]
|
||||
|
||||
result_messages, state = client._extract_state_from_messages(messages)
|
||||
|
||||
assert len(result_messages) == 1
|
||||
assert result_messages[0].text == "Hello"
|
||||
assert state == state_data
|
||||
|
||||
async def test_extract_state_invalid_json(self) -> None:
|
||||
"""Test state extraction with invalid JSON."""
|
||||
import base64
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
|
||||
invalid_json = "not valid json"
|
||||
state_b64 = base64.b64encode(invalid_json.encode("utf-8")).decode("utf-8")
|
||||
|
||||
from agent_framework import DataContent
|
||||
|
||||
messages = [
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[DataContent(uri=f"data:application/json;base64,{state_b64}")],
|
||||
),
|
||||
]
|
||||
|
||||
result_messages, state = client._extract_state_from_messages(messages)
|
||||
|
||||
assert result_messages == messages
|
||||
assert state is None
|
||||
|
||||
async def test_convert_messages_to_agui_format(self) -> None:
|
||||
"""Test message conversion to AG-UI format."""
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
messages = [
|
||||
ChatMessage(role=Role.USER, text="What is the weather?"),
|
||||
ChatMessage(role=Role.ASSISTANT, text="Let me check.", message_id="msg_123"),
|
||||
]
|
||||
|
||||
agui_messages = client._convert_messages_to_agui_format(messages)
|
||||
|
||||
assert len(agui_messages) == 2
|
||||
assert agui_messages[0]["role"] == "user"
|
||||
assert agui_messages[0]["content"] == "What is the weather?"
|
||||
assert agui_messages[1]["role"] == "assistant"
|
||||
assert agui_messages[1]["content"] == "Let me check."
|
||||
assert agui_messages[1]["id"] == "msg_123"
|
||||
|
||||
async def test_get_thread_id_from_metadata(self) -> None:
|
||||
"""Test thread ID extraction from metadata."""
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
chat_options = ChatOptions(metadata={"thread_id": "existing_thread_123"})
|
||||
|
||||
thread_id = client._get_thread_id(chat_options)
|
||||
|
||||
assert thread_id == "existing_thread_123"
|
||||
|
||||
async def test_get_thread_id_generation(self) -> None:
|
||||
"""Test automatic thread ID generation."""
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
chat_options = ChatOptions()
|
||||
|
||||
thread_id = client._get_thread_id(chat_options)
|
||||
|
||||
assert thread_id.startswith("thread_")
|
||||
assert len(thread_id) > 7
|
||||
|
||||
async def test_get_streaming_response(self, monkeypatch) -> None:
|
||||
"""Test streaming response method."""
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " world"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args, **kwargs):
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [ChatMessage(role="user", text="Test message")]
|
||||
chat_options = ChatOptions()
|
||||
|
||||
updates = []
|
||||
async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
|
||||
updates.append(update)
|
||||
|
||||
assert len(updates) == 4
|
||||
assert updates[0].additional_properties["thread_id"] == "thread_1"
|
||||
assert updates[1].contents[0].text == "Hello"
|
||||
assert updates[2].contents[0].text == " world"
|
||||
|
||||
async def test_get_response_non_streaming(self, monkeypatch) -> None:
|
||||
"""Test non-streaming response method."""
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Complete response"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args, **kwargs):
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [ChatMessage(role="user", text="Test message")]
|
||||
chat_options = ChatOptions()
|
||||
|
||||
response = await client._inner_get_response(messages=messages, chat_options=chat_options)
|
||||
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert "Complete response" in response.text
|
||||
|
||||
async def test_tool_handling(self, monkeypatch) -> None:
|
||||
"""Test that client tool metadata is sent to server.
|
||||
|
||||
Client tool metadata (name, description, schema) is sent to server for planning.
|
||||
When server requests a client function, @use_function_invocation decorator
|
||||
intercepts and executes it locally. This matches .NET AG-UI implementation.
|
||||
"""
|
||||
from agent_framework import ai_function
|
||||
|
||||
@ai_function
|
||||
def test_tool(param: str) -> str:
|
||||
"""Test tool."""
|
||||
return "result"
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args, **kwargs):
|
||||
# Client tool metadata should be sent to server
|
||||
tools = kwargs.get("tools")
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
assert tools[0]["name"] == "test_tool"
|
||||
assert tools[0]["description"] == "Test tool."
|
||||
assert "parameters" in tools[0]
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [ChatMessage(role="user", text="Test with tools")]
|
||||
chat_options = ChatOptions(tools=[test_tool])
|
||||
|
||||
response = await client._inner_get_response(messages=messages, chat_options=chat_options)
|
||||
|
||||
assert response is not None
|
||||
|
||||
async def test_server_tool_calls_unwrapped_after_invocation(self, monkeypatch) -> None:
|
||||
"""Ensure server-side tool calls are exposed as FunctionCallContent after processing."""
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_time_zone"},
|
||||
{"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"location": "Seattle"}'},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args, **kwargs):
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [ChatMessage(role="user", text="Test server tool execution")]
|
||||
chat_options = ChatOptions()
|
||||
|
||||
updates = []
|
||||
async for update in client.get_streaming_response(messages, chat_options=chat_options):
|
||||
updates.append(update)
|
||||
|
||||
function_calls = [
|
||||
content for update in updates for content in update.contents if isinstance(content, FunctionCallContent)
|
||||
]
|
||||
assert function_calls
|
||||
assert function_calls[0].name == "get_time_zone"
|
||||
assert not any(
|
||||
isinstance(content, ServerFunctionCallContent) for update in updates for content in update.contents
|
||||
)
|
||||
|
||||
async def test_server_tool_calls_not_executed_locally(self, monkeypatch) -> None:
|
||||
"""Server tools should not trigger local function invocation even when client tools exist."""
|
||||
|
||||
@ai_function
|
||||
def client_tool() -> str:
|
||||
"""Client tool stub."""
|
||||
return "client"
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_time_zone"},
|
||||
{"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"location": "Seattle"}'},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args, **kwargs):
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
async def fake_auto_invoke(*args, **kwargs):
|
||||
function_call = kwargs.get("function_call_content") or args[0]
|
||||
raise AssertionError(f"Unexpected local execution of server tool: {getattr(function_call, 'name', '?')}")
|
||||
|
||||
monkeypatch.setattr("agent_framework._tools._auto_invoke_function", fake_auto_invoke)
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [ChatMessage(role="user", text="Test server tool execution")]
|
||||
chat_options = ChatOptions(tool_choice="auto", tools=[client_tool])
|
||||
|
||||
async for _ in client.get_streaming_response(messages, chat_options=chat_options):
|
||||
pass
|
||||
|
||||
async def test_state_transmission(self, monkeypatch) -> None:
|
||||
"""Test state is properly transmitted to server."""
|
||||
import base64
|
||||
|
||||
state_data = {"user_id": "123", "session": "abc"}
|
||||
state_json = json.dumps(state_data)
|
||||
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
|
||||
|
||||
from agent_framework import DataContent
|
||||
|
||||
messages = [
|
||||
ChatMessage(role="user", text="Hello"),
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[DataContent(uri=f"data:application/json;base64,{state_b64}")],
|
||||
),
|
||||
]
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args, **kwargs):
|
||||
assert kwargs.get("state") == state_data
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
|
||||
|
||||
chat_options = ChatOptions()
|
||||
|
||||
response = await client._inner_get_response(messages=messages, chat_options=chat_options)
|
||||
|
||||
assert response is not None
|
||||
@@ -1,287 +0,0 @@
|
||||
"""Tests for AG-UI event converter."""
|
||||
|
||||
from agent_framework import FinishReason, Role
|
||||
|
||||
from agent_framework_ag_ui._event_converters import AGUIEventConverter
|
||||
|
||||
|
||||
class TestAGUIEventConverter:
|
||||
"""Test suite for AGUIEventConverter."""
|
||||
|
||||
def test_run_started_event(self) -> None:
|
||||
"""Test conversion of RUN_STARTED event."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "RUN_STARTED",
|
||||
"threadId": "thread_123",
|
||||
"runId": "run_456",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.role == Role.ASSISTANT
|
||||
assert update.additional_properties["thread_id"] == "thread_123"
|
||||
assert update.additional_properties["run_id"] == "run_456"
|
||||
assert converter.thread_id == "thread_123"
|
||||
assert converter.run_id == "run_456"
|
||||
|
||||
def test_text_message_start_event(self) -> None:
|
||||
"""Test conversion of TEXT_MESSAGE_START event."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "TEXT_MESSAGE_START",
|
||||
"messageId": "msg_789",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.role == Role.ASSISTANT
|
||||
assert update.message_id == "msg_789"
|
||||
assert converter.current_message_id == "msg_789"
|
||||
|
||||
def test_text_message_content_event(self) -> None:
|
||||
"""Test conversion of TEXT_MESSAGE_CONTENT event."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "TEXT_MESSAGE_CONTENT",
|
||||
"messageId": "msg_1",
|
||||
"delta": "Hello",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.role == Role.ASSISTANT
|
||||
assert update.message_id == "msg_1"
|
||||
assert len(update.contents) == 1
|
||||
assert update.contents[0].text == "Hello"
|
||||
|
||||
def test_text_message_streaming(self) -> None:
|
||||
"""Test streaming text across multiple TEXT_MESSAGE_CONTENT events."""
|
||||
converter = AGUIEventConverter()
|
||||
events = [
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " world"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "!"},
|
||||
]
|
||||
|
||||
updates = [converter.convert_event(event) for event in events]
|
||||
|
||||
assert all(update is not None for update in updates)
|
||||
assert all(update.message_id == "msg_1" for update in updates)
|
||||
assert updates[0].contents[0].text == "Hello"
|
||||
assert updates[1].contents[0].text == " world"
|
||||
assert updates[2].contents[0].text == "!"
|
||||
|
||||
def test_text_message_end_event(self) -> None:
|
||||
"""Test conversion of TEXT_MESSAGE_END event."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "TEXT_MESSAGE_END",
|
||||
"messageId": "msg_1",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is None
|
||||
|
||||
def test_tool_call_start_event(self) -> None:
|
||||
"""Test conversion of TOOL_CALL_START event."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "TOOL_CALL_START",
|
||||
"toolCallId": "call_123",
|
||||
"toolName": "get_weather",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.role == Role.ASSISTANT
|
||||
assert len(update.contents) == 1
|
||||
assert update.contents[0].call_id == "call_123"
|
||||
assert update.contents[0].name == "get_weather"
|
||||
assert update.contents[0].arguments == ""
|
||||
assert converter.current_tool_call_id == "call_123"
|
||||
assert converter.current_tool_name == "get_weather"
|
||||
|
||||
def test_tool_call_start_with_tool_call_name(self) -> None:
|
||||
"""Ensure TOOL_CALL_START with toolCallName still sets the tool name."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "TOOL_CALL_START",
|
||||
"toolCallId": "call_abc",
|
||||
"toolCallName": "get_weather",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.contents[0].name == "get_weather"
|
||||
assert converter.current_tool_name == "get_weather"
|
||||
|
||||
def test_tool_call_start_with_tool_call_name_snake_case(self) -> None:
|
||||
"""Support tool_call_name snake_case field for backwards compatibility."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "TOOL_CALL_START",
|
||||
"toolCallId": "call_snake",
|
||||
"tool_call_name": "get_weather",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.contents[0].name == "get_weather"
|
||||
assert converter.current_tool_name == "get_weather"
|
||||
|
||||
def test_tool_call_args_streaming(self) -> None:
|
||||
"""Test streaming tool arguments across multiple TOOL_CALL_ARGS events."""
|
||||
converter = AGUIEventConverter()
|
||||
converter.current_tool_call_id = "call_123"
|
||||
converter.current_tool_name = "search"
|
||||
|
||||
events = [
|
||||
{"type": "TOOL_CALL_ARGS", "delta": '{"query": "'},
|
||||
{"type": "TOOL_CALL_ARGS", "delta": 'latest news"}'},
|
||||
]
|
||||
|
||||
updates = [converter.convert_event(event) for event in events]
|
||||
|
||||
assert all(update is not None for update in updates)
|
||||
assert updates[0].contents[0].arguments == '{"query": "'
|
||||
assert updates[1].contents[0].arguments == 'latest news"}'
|
||||
assert converter.accumulated_tool_args == '{"query": "latest news"}'
|
||||
|
||||
def test_tool_call_end_event(self) -> None:
|
||||
"""Test conversion of TOOL_CALL_END event."""
|
||||
converter = AGUIEventConverter()
|
||||
converter.accumulated_tool_args = '{"location": "Seattle"}'
|
||||
|
||||
event = {
|
||||
"type": "TOOL_CALL_END",
|
||||
"toolCallId": "call_123",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is None
|
||||
assert converter.accumulated_tool_args == ""
|
||||
|
||||
def test_tool_call_result_event(self) -> None:
|
||||
"""Test conversion of TOOL_CALL_RESULT event."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "TOOL_CALL_RESULT",
|
||||
"toolCallId": "call_123",
|
||||
"result": {"temperature": 22, "condition": "sunny"},
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.role == Role.TOOL
|
||||
assert len(update.contents) == 1
|
||||
assert update.contents[0].call_id == "call_123"
|
||||
assert update.contents[0].result == {"temperature": 22, "condition": "sunny"}
|
||||
|
||||
def test_run_finished_event(self) -> None:
|
||||
"""Test conversion of RUN_FINISHED event."""
|
||||
converter = AGUIEventConverter()
|
||||
converter.thread_id = "thread_123"
|
||||
converter.run_id = "run_456"
|
||||
|
||||
event = {
|
||||
"type": "RUN_FINISHED",
|
||||
"threadId": "thread_123",
|
||||
"runId": "run_456",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.role == Role.ASSISTANT
|
||||
assert update.finish_reason == FinishReason.STOP
|
||||
assert update.additional_properties["thread_id"] == "thread_123"
|
||||
assert update.additional_properties["run_id"] == "run_456"
|
||||
|
||||
def test_run_error_event(self) -> None:
|
||||
"""Test conversion of RUN_ERROR event."""
|
||||
converter = AGUIEventConverter()
|
||||
converter.thread_id = "thread_123"
|
||||
converter.run_id = "run_456"
|
||||
|
||||
event = {
|
||||
"type": "RUN_ERROR",
|
||||
"message": "Connection timeout",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is not None
|
||||
assert update.role == Role.ASSISTANT
|
||||
assert update.finish_reason == FinishReason.CONTENT_FILTER
|
||||
assert len(update.contents) == 1
|
||||
assert update.contents[0].message == "Connection timeout"
|
||||
assert update.contents[0].error_code == "RUN_ERROR"
|
||||
|
||||
def test_unknown_event_type(self) -> None:
|
||||
"""Test handling of unknown event types."""
|
||||
converter = AGUIEventConverter()
|
||||
event = {
|
||||
"type": "UNKNOWN_EVENT",
|
||||
"data": "some data",
|
||||
}
|
||||
|
||||
update = converter.convert_event(event)
|
||||
|
||||
assert update is None
|
||||
|
||||
def test_full_conversation_flow(self) -> None:
|
||||
"""Test complete conversation flow with multiple event types."""
|
||||
converter = AGUIEventConverter()
|
||||
|
||||
events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TEXT_MESSAGE_START", "messageId": "msg_1"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "I'll check"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " the weather."},
|
||||
{"type": "TEXT_MESSAGE_END", "messageId": "msg_1"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_weather"},
|
||||
{"type": "TOOL_CALL_ARGS", "delta": '{"location": "Seattle"}'},
|
||||
{"type": "TOOL_CALL_END", "toolCallId": "call_1"},
|
||||
{"type": "TOOL_CALL_RESULT", "toolCallId": "call_1", "result": "Sunny, 72°F"},
|
||||
{"type": "TEXT_MESSAGE_START", "messageId": "msg_2"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_2", "delta": "It's sunny!"},
|
||||
{"type": "TEXT_MESSAGE_END", "messageId": "msg_2"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
updates = [converter.convert_event(event) for event in events]
|
||||
non_none_updates = [u for u in updates if u is not None]
|
||||
|
||||
assert len(non_none_updates) == 10
|
||||
assert converter.thread_id == "thread_1"
|
||||
assert converter.run_id == "run_1"
|
||||
|
||||
def test_multiple_tool_calls(self) -> None:
|
||||
"""Test handling multiple tool calls in sequence."""
|
||||
converter = AGUIEventConverter()
|
||||
|
||||
events = [
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "search"},
|
||||
{"type": "TOOL_CALL_ARGS", "delta": '{"query": "weather"}'},
|
||||
{"type": "TOOL_CALL_END", "toolCallId": "call_1"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "call_2", "toolName": "fetch"},
|
||||
{"type": "TOOL_CALL_ARGS", "delta": '{"url": "http://api.weather.com"}'},
|
||||
{"type": "TOOL_CALL_END", "toolCallId": "call_2"},
|
||||
]
|
||||
|
||||
updates = [converter.convert_event(event) for event in events]
|
||||
non_none_updates = [u for u in updates if u is not None]
|
||||
|
||||
assert len(non_none_updates) == 4
|
||||
assert non_none_updates[0].contents[0].name == "search"
|
||||
assert non_none_updates[2].contents[0].name == "fetch"
|
||||
@@ -1,238 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for AGUIHttpService."""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from agent_framework_ag_ui._http_service import AGUIHttpService
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_http_client():
|
||||
"""Create a mock httpx.AsyncClient."""
|
||||
client = AsyncMock(spec=httpx.AsyncClient)
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_events():
|
||||
"""Sample AG-UI events for testing."""
|
||||
return [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_123", "runId": "run_456"},
|
||||
{"type": "TEXT_MESSAGE_START", "messageId": "msg_1", "role": "assistant"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " world"},
|
||||
{"type": "TEXT_MESSAGE_END", "messageId": "msg_1"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_123", "runId": "run_456"},
|
||||
]
|
||||
|
||||
|
||||
def create_sse_response(events: list[dict]) -> str:
|
||||
"""Create SSE formatted response from events."""
|
||||
lines = []
|
||||
for event in events:
|
||||
lines.append(f"data: {json.dumps(event)}\n")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
async def test_http_service_initialization():
|
||||
"""Test AGUIHttpService initialization."""
|
||||
# Test with default client
|
||||
service = AGUIHttpService("http://localhost:8888/")
|
||||
assert service.endpoint == "http://localhost:8888"
|
||||
assert service._owns_client is True
|
||||
assert isinstance(service.http_client, httpx.AsyncClient)
|
||||
await service.close()
|
||||
|
||||
# Test with custom client
|
||||
custom_client = httpx.AsyncClient()
|
||||
service = AGUIHttpService("http://localhost:8888/", http_client=custom_client)
|
||||
assert service._owns_client is False
|
||||
assert service.http_client is custom_client
|
||||
# Shouldn't close the custom client
|
||||
await service.close()
|
||||
await custom_client.aclose()
|
||||
|
||||
|
||||
async def test_http_service_strips_trailing_slash():
|
||||
"""Test that endpoint trailing slash is stripped."""
|
||||
service = AGUIHttpService("http://localhost:8888/")
|
||||
assert service.endpoint == "http://localhost:8888"
|
||||
await service.close()
|
||||
|
||||
|
||||
async def test_post_run_successful_streaming(mock_http_client, sample_events):
|
||||
"""Test successful streaming of events."""
|
||||
|
||||
# Create async generator for lines
|
||||
async def mock_aiter_lines():
|
||||
sse_data = create_sse_response(sample_events)
|
||||
for line in sse_data.split("\n"):
|
||||
if line:
|
||||
yield line
|
||||
|
||||
# Create mock response
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
# aiter_lines is called as a method, so it should return a new generator each time
|
||||
mock_response.aiter_lines = mock_aiter_lines
|
||||
|
||||
# Setup mock streaming context manager
|
||||
mock_stream_context = AsyncMock()
|
||||
mock_stream_context.__aenter__.return_value = mock_response
|
||||
mock_stream_context.__aexit__.return_value = None
|
||||
mock_http_client.stream.return_value = mock_stream_context
|
||||
|
||||
service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
|
||||
|
||||
events = []
|
||||
async for event in service.post_run(
|
||||
thread_id="thread_123", run_id="run_456", messages=[{"role": "user", "content": "Hello"}]
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == len(sample_events)
|
||||
assert events[0]["type"] == "RUN_STARTED"
|
||||
assert events[-1]["type"] == "RUN_FINISHED"
|
||||
|
||||
# Verify request was made correctly
|
||||
mock_http_client.stream.assert_called_once()
|
||||
call_args = mock_http_client.stream.call_args
|
||||
assert call_args.args[0] == "POST"
|
||||
assert call_args.args[1] == "http://localhost:8888"
|
||||
assert call_args.kwargs["headers"] == {"Accept": "text/event-stream"}
|
||||
|
||||
|
||||
async def test_post_run_with_state_and_tools(mock_http_client):
|
||||
"""Test posting run with state and tools."""
|
||||
|
||||
async def mock_aiter_lines():
|
||||
return
|
||||
yield # Make it an async generator
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.aiter_lines = mock_aiter_lines
|
||||
|
||||
mock_stream_context = AsyncMock()
|
||||
mock_stream_context.__aenter__.return_value = mock_response
|
||||
mock_stream_context.__aexit__.return_value = None
|
||||
mock_http_client.stream.return_value = mock_stream_context
|
||||
|
||||
service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
|
||||
|
||||
state = {"user_context": {"name": "Alice"}}
|
||||
tools = [{"type": "function", "function": {"name": "test_tool"}}]
|
||||
|
||||
async for _ in service.post_run(thread_id="thread_123", run_id="run_456", messages=[], state=state, tools=tools):
|
||||
pass
|
||||
|
||||
# Verify state and tools were included in request
|
||||
call_args = mock_http_client.stream.call_args
|
||||
request_data = call_args.kwargs["json"]
|
||||
assert request_data["state"] == state
|
||||
assert request_data["tools"] == tools
|
||||
|
||||
|
||||
async def test_post_run_http_error(mock_http_client):
|
||||
"""Test handling of HTTP errors."""
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 500
|
||||
mock_response.text = "Internal Server Error"
|
||||
|
||||
def raise_http_error():
|
||||
raise httpx.HTTPStatusError("Server error", request=Mock(), response=mock_response)
|
||||
|
||||
mock_response_async = AsyncMock()
|
||||
mock_response_async.raise_for_status = raise_http_error
|
||||
|
||||
mock_stream_context = AsyncMock()
|
||||
mock_stream_context.__aenter__.return_value = mock_response_async
|
||||
mock_stream_context.__aexit__.return_value = None
|
||||
mock_http_client.stream.return_value = mock_stream_context
|
||||
|
||||
service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
|
||||
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
async for _ in service.post_run(thread_id="thread_123", run_id="run_456", messages=[]):
|
||||
pass
|
||||
|
||||
|
||||
async def test_post_run_invalid_json(mock_http_client):
|
||||
"""Test handling of invalid JSON in SSE stream."""
|
||||
invalid_sse = "data: {invalid json}\n\ndata: " + json.dumps({"type": "RUN_FINISHED"}) + "\n"
|
||||
|
||||
async def mock_aiter_lines():
|
||||
for line in invalid_sse.split("\n"):
|
||||
if line:
|
||||
yield line
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.aiter_lines = mock_aiter_lines
|
||||
|
||||
mock_stream_context = AsyncMock()
|
||||
mock_stream_context.__aenter__.return_value = mock_response
|
||||
mock_stream_context.__aexit__.return_value = None
|
||||
mock_http_client.stream.return_value = mock_stream_context
|
||||
|
||||
service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
|
||||
|
||||
events = []
|
||||
async for event in service.post_run(thread_id="thread_123", run_id="run_456", messages=[]):
|
||||
events.append(event)
|
||||
|
||||
# Should skip invalid JSON and continue with valid events
|
||||
assert len(events) == 1
|
||||
assert events[0]["type"] == "RUN_FINISHED"
|
||||
|
||||
|
||||
async def test_context_manager():
|
||||
"""Test context manager functionality."""
|
||||
async with AGUIHttpService("http://localhost:8888/") as service:
|
||||
assert service.http_client is not None
|
||||
assert service._owns_client is True
|
||||
|
||||
# Client should be closed after exiting context
|
||||
|
||||
|
||||
async def test_context_manager_with_external_client():
|
||||
"""Test context manager doesn't close external client."""
|
||||
external_client = httpx.AsyncClient()
|
||||
|
||||
async with AGUIHttpService("http://localhost:8888/", http_client=external_client) as service:
|
||||
assert service.http_client is external_client
|
||||
assert service._owns_client is False
|
||||
|
||||
# External client should still be open
|
||||
# (caller's responsibility to close)
|
||||
await external_client.aclose()
|
||||
|
||||
|
||||
async def test_post_run_empty_response(mock_http_client):
|
||||
"""Test handling of empty response stream."""
|
||||
|
||||
async def mock_aiter_lines():
|
||||
return
|
||||
yield # Make it an async generator
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.aiter_lines = mock_aiter_lines
|
||||
|
||||
mock_stream_context = AsyncMock()
|
||||
mock_stream_context.__aenter__.return_value = mock_response
|
||||
mock_stream_context.__aexit__.return_value = None
|
||||
mock_http_client.stream.return_value = mock_stream_context
|
||||
|
||||
service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
|
||||
|
||||
events = []
|
||||
async for event in service.post_run(thread_id="thread_123", run_id="run_456", messages=[]):
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 0
|
||||
@@ -63,9 +63,10 @@ def test_agui_tool_result_to_agent_framework():
|
||||
assert isinstance(message.contents[0], TextContent)
|
||||
assert message.contents[0].text == '{"accepted": true, "steps": []}'
|
||||
|
||||
assert message.additional_properties is not None
|
||||
assert message.additional_properties.get("is_tool_result") is True
|
||||
assert message.additional_properties.get("tool_call_id") == "call_123"
|
||||
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():
|
||||
@@ -158,36 +159,6 @@ def test_agui_message_without_id():
|
||||
assert messages[0].message_id is None
|
||||
|
||||
|
||||
def test_agui_with_tool_calls_to_agent_framework():
|
||||
"""Assistant message with tool_calls is converted to FunctionCallContent."""
|
||||
agui_msg = {
|
||||
"role": "assistant",
|
||||
"content": "Calling tool",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call-123",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": {"location": "Seattle"}},
|
||||
}
|
||||
],
|
||||
"id": "msg-789",
|
||||
}
|
||||
|
||||
messages = agui_messages_to_agent_framework([agui_msg])
|
||||
|
||||
assert len(messages) == 1
|
||||
msg = messages[0]
|
||||
assert msg.role == Role.ASSISTANT
|
||||
assert msg.message_id == "msg-789"
|
||||
# First content is text, second is the function call
|
||||
assert isinstance(msg.contents[0], TextContent)
|
||||
assert msg.contents[0].text == "Calling tool"
|
||||
assert isinstance(msg.contents[1], FunctionCallContent)
|
||||
assert msg.contents[1].call_id == "call-123"
|
||||
assert msg.contents[1].name == "get_weather"
|
||||
assert msg.contents[1].arguments == {"location": "Seattle"}
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_with_tool_calls():
|
||||
"""Test converting Agent Framework message with tool calls to AG-UI."""
|
||||
msg = ChatMessage(
|
||||
@@ -227,15 +198,13 @@ def test_agent_framework_to_agui_multiple_text_contents():
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_no_message_id():
|
||||
"""Test message without message_id - should auto-generate 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" in messages[0] # ID should be auto-generated
|
||||
assert messages[0]["id"] # ID should not be empty
|
||||
assert len(messages[0]["id"]) > 0 # ID should be a valid string
|
||||
assert "id" not in messages[0]
|
||||
|
||||
|
||||
def test_agent_framework_to_agui_system_role():
|
||||
|
||||
@@ -1,82 +0,0 @@
|
||||
"""Tests for AG-UI orchestrators."""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentRunResponseUpdate, TextContent, ai_function
|
||||
from agent_framework._tools import FunctionInvocationConfiguration
|
||||
|
||||
from agent_framework_ag_ui._agent import AgentConfig
|
||||
from agent_framework_ag_ui._orchestrators import DefaultOrchestrator, ExecutionContext
|
||||
|
||||
|
||||
@ai_function
|
||||
def server_tool() -> str:
|
||||
"""Server-executable tool."""
|
||||
return "server"
|
||||
|
||||
|
||||
class DummyAgent:
|
||||
"""Minimal agent stub to capture run_stream parameters."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.chat_options = SimpleNamespace(tools=[server_tool], response_format=None)
|
||||
self.tools = [server_tool]
|
||||
self.chat_client = SimpleNamespace(
|
||||
function_invocation_configuration=FunctionInvocationConfiguration(),
|
||||
)
|
||||
self.seen_tools: list[Any] | None = None
|
||||
|
||||
async def run_stream(
|
||||
self,
|
||||
messages: list[Any],
|
||||
*,
|
||||
thread: Any,
|
||||
tools: list[Any] | None = None,
|
||||
) -> AsyncGenerator[AgentRunResponseUpdate, None]:
|
||||
self.seen_tools = tools
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="ok")], role="assistant")
|
||||
|
||||
|
||||
async def test_default_orchestrator_merges_client_tools() -> None:
|
||||
"""Client tool declarations are merged with server tools before running agent."""
|
||||
|
||||
agent = DummyAgent()
|
||||
orchestrator = DefaultOrchestrator()
|
||||
|
||||
input_data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "Hello"}],
|
||||
}
|
||||
],
|
||||
"tools": [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Client weather lookup.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"location": {"type": "string"}},
|
||||
"required": ["location"],
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
context = ExecutionContext(
|
||||
input_data=input_data,
|
||||
agent=agent,
|
||||
config=AgentConfig(),
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in orchestrator.run(context):
|
||||
events.append(event)
|
||||
|
||||
assert agent.seen_tools is not None
|
||||
tool_names = [getattr(tool, "name", "?") for tool in agent.seen_tools]
|
||||
assert "server_tool" in tool_names
|
||||
assert "get_weather" in tool_names
|
||||
assert agent.chat_client.function_invocation_configuration.additional_tools
|
||||
@@ -197,109 +197,3 @@ def test_make_json_safe_fallback():
|
||||
result = make_json_safe(obj)
|
||||
# Objects with __dict__ return their __dict__ dict
|
||||
assert isinstance(result, dict)
|
||||
|
||||
|
||||
def test_convert_tools_to_agui_format_with_ai_function():
|
||||
"""Test converting AIFunction to AG-UI format."""
|
||||
from agent_framework import ai_function
|
||||
|
||||
from agent_framework_ag_ui._utils import convert_tools_to_agui_format
|
||||
|
||||
@ai_function
|
||||
def test_func(param: str, count: int = 5) -> str:
|
||||
"""Test function."""
|
||||
return f"{param} {count}"
|
||||
|
||||
result = convert_tools_to_agui_format([test_func])
|
||||
|
||||
assert result is not None
|
||||
assert len(result) == 1
|
||||
assert result[0]["name"] == "test_func"
|
||||
assert result[0]["description"] == "Test function."
|
||||
assert "parameters" in result[0]
|
||||
assert "properties" in result[0]["parameters"]
|
||||
|
||||
|
||||
def test_convert_tools_to_agui_format_with_callable():
|
||||
"""Test converting plain callable to AG-UI format."""
|
||||
from agent_framework_ag_ui._utils import convert_tools_to_agui_format
|
||||
|
||||
def plain_func(x: int) -> int:
|
||||
"""A plain function."""
|
||||
return x * 2
|
||||
|
||||
result = convert_tools_to_agui_format([plain_func])
|
||||
|
||||
assert result is not None
|
||||
assert len(result) == 1
|
||||
assert result[0]["name"] == "plain_func"
|
||||
assert result[0]["description"] == "A plain function."
|
||||
assert "parameters" in result[0]
|
||||
|
||||
|
||||
def test_convert_tools_to_agui_format_with_dict():
|
||||
"""Test converting dict tool to AG-UI format."""
|
||||
from agent_framework_ag_ui._utils import convert_tools_to_agui_format
|
||||
|
||||
tool_dict = {
|
||||
"name": "custom_tool",
|
||||
"description": "Custom tool",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
|
||||
result = convert_tools_to_agui_format([tool_dict])
|
||||
|
||||
assert result is not None
|
||||
assert len(result) == 1
|
||||
assert result[0] == tool_dict
|
||||
|
||||
|
||||
def test_convert_tools_to_agui_format_with_none():
|
||||
"""Test converting None tools."""
|
||||
from agent_framework_ag_ui._utils import convert_tools_to_agui_format
|
||||
|
||||
result = convert_tools_to_agui_format(None)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_convert_tools_to_agui_format_with_single_tool():
|
||||
"""Test converting single tool (not in list)."""
|
||||
from agent_framework import ai_function
|
||||
|
||||
from agent_framework_ag_ui._utils import convert_tools_to_agui_format
|
||||
|
||||
@ai_function
|
||||
def single_tool(arg: str) -> str:
|
||||
"""Single tool."""
|
||||
return arg
|
||||
|
||||
result = convert_tools_to_agui_format(single_tool)
|
||||
|
||||
assert result is not None
|
||||
assert len(result) == 1
|
||||
assert result[0]["name"] == "single_tool"
|
||||
|
||||
|
||||
def test_convert_tools_to_agui_format_with_multiple_tools():
|
||||
"""Test converting multiple tools."""
|
||||
from agent_framework import ai_function
|
||||
|
||||
from agent_framework_ag_ui._utils import convert_tools_to_agui_format
|
||||
|
||||
@ai_function
|
||||
def tool1(x: int) -> int:
|
||||
"""Tool 1."""
|
||||
return x
|
||||
|
||||
@ai_function
|
||||
def tool2(y: str) -> str:
|
||||
"""Tool 2."""
|
||||
return y
|
||||
|
||||
result = convert_tools_to_agui_format([tool1, tool2])
|
||||
|
||||
assert result is not None
|
||||
assert len(result) == 2
|
||||
assert result[0]["name"] == "tool1"
|
||||
assert result[1]["name"] == "tool2"
|
||||
|
||||
Reference in New Issue
Block a user