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

* 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

python/packages/ag-ui/tests/test_structured_output.py

@@ -3,12 +3,18 @@
 """Tests for structured output handling in _agent.py."""
 
 import json
+import sys
+from collections.abc import AsyncIterator, MutableSequence
+from pathlib import Path
 from typing import Any
 
-from agent_framework import ChatAgent, ChatOptions, 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, stream_from_updates
+
 
 class RecipeOutput(BaseModel):
     """Test Pydantic model for recipe output."""
@@ -34,14 +40,14 @@ async def test_structured_output_with_recipe():
     """Test structured output processing with recipe state."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Simulate structured output
-            yield ChatResponseUpdate(
-                contents=[TextContent(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
-            )
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(
+            contents=[TextContent(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
+        )
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -51,7 +57,7 @@ async def test_structured_output_with_recipe():
 
     input_data = {"messages": [{"role": "user", "content": "Make pasta"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -72,17 +78,18 @@ async def test_structured_output_with_steps():
     """Test structured output processing with steps state."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            steps_data = {
-                "steps": [
-                    {"id": "1", "description": "Step 1", "status": "pending"},
-                    {"id": "2", "description": "Step 2", "status": "pending"},
-                ]
-            }
-            yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(steps_data))])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        steps_data = {
+            "steps": [
+                {"id": "1", "description": "Step 1", "status": "pending"},
+                {"id": "2", "description": "Step 2", "status": "pending"},
+            ]
+        }
+        yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(steps_data))])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=StepsOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -92,7 +99,7 @@ async def test_structured_output_with_steps():
 
     input_data = {"messages": [{"role": "user", "content": "Do steps"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -111,12 +118,13 @@ async def test_structured_output_with_no_schema_match():
     """Test structured output when response fields don't match state_schema keys."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Response has "data" field but schema expects "result" field
-            yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}}')])
+    updates = [
+        ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}}')]),
+    ]
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(
+        name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_from_updates(updates))
+    )
     agent.chat_options = ChatOptions(response_format=GenericOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -126,7 +134,7 @@ async def test_structured_output_with_no_schema_match():
 
     input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -146,11 +154,12 @@ async def test_structured_output_without_schema():
         data: dict[str, Any]
         info: str
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}, "info": "processed"}')])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}, "info": "processed"}')])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=DataOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -160,7 +169,7 @@ async def test_structured_output_without_schema():
 
     input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -177,18 +186,20 @@ async def test_no_structured_output_when_no_response_format():
     """Test that structured output path is skipped when no response_format."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Regular text")])
+    updates = [ChatResponseUpdate(contents=[TextContent(text="Regular text")])]
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(
+        name="test",
+        instructions="Test",
+        chat_client=StreamingChatClientStub(stream_from_updates(updates)),
+    )
     # No response_format set
 
     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)
 
@@ -202,12 +213,13 @@ async def test_structured_output_with_message_field():
     """Test structured output that includes a message field."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
-            yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(output_data))])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
+        yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(output_data))])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -217,7 +229,7 @@ async def test_structured_output_with_message_field():
 
     input_data = {"messages": [{"role": "user", "content": "Make salad"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -236,20 +248,20 @@ async def test_empty_updates_no_structured_processing():
     """Test that empty updates don't trigger structured output processing."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Return nothing
-            if False:
-                yield
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        if False:
+            yield ChatResponseUpdate(contents=[])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {"messages": [{"role": "user", "content": "Test"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)