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
2026-06-16 21:04:09 +08:00 · 2025-11-05 14:25:24 +09:00
parent 0c862e97a6
commit 35a8565495
51 changed files with 7677 additions and 163 deletions
@@ -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