mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Evan Mattson
8cf8b0f995
* 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
52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
from ag_ui.core import CustomEvent, EventType
|
|
from agent_framework import ChatMessage, TextContent
|
|
|
|
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 isinstance(message.contents[0], TextContent)
|
|
assert "Current state of the application" in message.contents[0].text
|