mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Evan Mattson
6d7690e485
* fix(ag-ui): properly handle json serialize with handoff workflows as agent * Other improvements around handling non-serializable objects
106 lines
3.9 KiB
Python
106 lines
3.9 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
from ag_ui.core import CustomEvent, EventType
|
|
from agent_framework import ChatMessage
|
|
|
|
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
|
|
from agent_framework_ag_ui._orchestration._state_manager import StateManager
|
|
|
|
|
|
def test_state_manager_initializes_defaults_and_snapshot() -> None:
|
|
state_manager = StateManager(
|
|
state_schema={"items": {"type": "array"}, "metadata": {"type": "object"}},
|
|
predict_state_config=None,
|
|
require_confirmation=True,
|
|
)
|
|
current_state = state_manager.initialize({"metadata": {"a": 1}})
|
|
bridge = AgentFrameworkEventBridge(run_id="run", thread_id="thread", current_state=current_state)
|
|
|
|
snapshot_event = state_manager.initial_snapshot_event(bridge)
|
|
assert snapshot_event is not None
|
|
assert snapshot_event.snapshot["items"] == []
|
|
assert snapshot_event.snapshot["metadata"] == {"a": 1}
|
|
|
|
|
|
def test_state_manager_predict_state_event_shape() -> None:
|
|
state_manager = StateManager(
|
|
state_schema=None,
|
|
predict_state_config={"doc": {"tool": "write_document_local", "tool_argument": "document"}},
|
|
require_confirmation=True,
|
|
)
|
|
predict_event = state_manager.predict_state_event()
|
|
assert isinstance(predict_event, CustomEvent)
|
|
assert predict_event.type == EventType.CUSTOM
|
|
assert predict_event.name == "PredictState"
|
|
assert predict_event.value[0]["state_key"] == "doc"
|
|
|
|
|
|
def test_state_context_only_when_new_user_turn() -> None:
|
|
state_manager = StateManager(
|
|
state_schema={"items": {"type": "array"}},
|
|
predict_state_config=None,
|
|
require_confirmation=True,
|
|
)
|
|
state_manager.initialize({"items": [1]})
|
|
|
|
assert state_manager.state_context_message(is_new_user_turn=False, conversation_has_tool_calls=False) is None
|
|
|
|
message = state_manager.state_context_message(is_new_user_turn=True, conversation_has_tool_calls=False)
|
|
assert isinstance(message, ChatMessage)
|
|
assert message.contents[0].type == "text"
|
|
assert "Current state of the application" in message.contents[0].text
|
|
|
|
|
|
def test_state_manager_with_dataclass_in_state() -> None:
|
|
"""Test that state containing dataclasses can be serialized without crashing.
|
|
|
|
This test ensures the fix for JSON serialization errors when state
|
|
contains dataclass or other non-JSON-serializable objects.
|
|
"""
|
|
from dataclasses import dataclass
|
|
|
|
@dataclass
|
|
class UserData:
|
|
name: str
|
|
age: int
|
|
|
|
state_manager = StateManager(
|
|
state_schema={"user": {"type": "object"}},
|
|
predict_state_config=None,
|
|
require_confirmation=True,
|
|
)
|
|
# Initialize with a dataclass object in the state
|
|
state_manager.initialize({"user": UserData(name="Alice", age=30)})
|
|
|
|
# This should NOT raise TypeError when generating the context message
|
|
message = state_manager.state_context_message(is_new_user_turn=True, conversation_has_tool_calls=False)
|
|
|
|
assert message is not None
|
|
assert isinstance(message, ChatMessage)
|
|
# The dataclass should be serialized to JSON in the message
|
|
assert "Alice" in message.contents[0].text
|
|
assert "30" in message.contents[0].text
|
|
|
|
|
|
def test_state_manager_with_pydantic_in_state() -> None:
|
|
"""Test that state containing Pydantic models can be serialized without crashing."""
|
|
from pydantic import BaseModel
|
|
|
|
class UserModel(BaseModel):
|
|
email: str
|
|
active: bool
|
|
|
|
state_manager = StateManager(
|
|
state_schema={"user": {"type": "object"}},
|
|
predict_state_config=None,
|
|
require_confirmation=True,
|
|
)
|
|
# Initialize with a Pydantic model in the state
|
|
state_manager.initialize({"user": UserModel(email="test@example.com", active=True)})
|
|
|
|
# This should NOT raise TypeError
|
|
message = state_manager.state_context_message(is_new_user_turn=True, conversation_has_tool_calls=False)
|
|
|
|
assert message is not None
|
|
assert "test@example.com" in message.contents[0].text
|