Python: DevUI - Internal Refactor, Conversations API support, and per… (#1235)

* Python: DevUI - Internal Refactor, Conversations API support, and performance improvements Comprehensive refactor of DevUI package including samples relocation, frontend reorganization, OpenAI Conversations API support, and critical performance and code quality improvements. Key Changes: Architecture & Organization - Moved DevUI samples to python/samples/getting_started/devui/ - Consolidated with other framework samples for better discoverability - Added .env.example files and comprehensive README - Restructured frontend components into feature-based folders (agent, workflow, gallery, layout) - Created new OpenAI-compliant message renderers (devui should render oai responses types primarily) New Features - Added _conversations.py (467 lines) - Full conversation storage abstraction, replaces the /threads endpoint to better match oai conversations api - Implements OpenAI Conversations API for thread management, Supports in-memory and extensible storage backends API Simplification - Use 'model' field as entity_id (agent/workflow name) instead of extra_body - Use standard OpenAI 'conversation' field for conversation context. Performance & Quality Improvements - Improved context management in MessageMapper with bounded memory (~500KB max) - Implemented hybrid LRU + cleanup approach to prevent unbounded memory growth - General QOL improvement - Eliminated ~150 lines of dead/duplicate code, Consolidated helper functions into _utils.py, Extracted magic numbers to module-level constants, Optimized conversation item lookups with index-based approach Testing - Added test_conversations.py (13 tests) - Added test_performance_fixes.py (9 tests) - Updated existing tests for code consolidation - 53 tests passing Impact: 76 files changed: +4,106 insertions, -2,373 deletions All linting and formatting checks passing. No breaking changes - backward compatible. Migration: Samples moved to python/samples/getting_started/devui/ * readme lint fixes * initial support for function approval and minor ui fixes
2026-06-16 21:04:09 +08:00 · 2025-10-08 12:34:30 -07:00
parent f5abbc67ae
commit c341ee7ed2
75 changed files with 4605 additions and 2646 deletions
@@ -10,6 +10,7 @@ import asyncio
 import contextlib
 import http.client
 import json
+import logging
 import threading
 import time
 from pathlib import Path
@@ -20,12 +21,20 @@ from openai import OpenAI

 from agent_framework_devui import DevServer

+logger = logging.getLogger(__name__)
+

 def start_server() -> tuple[str, Any]:
    """Start server with samples directory."""
-    # Get samples directory
+    # Get samples directory - updated path after samples were moved
    current_dir = Path(__file__).parent
-    samples_dir = current_dir.parent / "samples"
+    # Samples are now in python/samples/getting_started/devui
+    samples_dir = current_dir.parent.parent.parent / "samples" / "getting_started" / "devui"
+
+    if not samples_dir.exists():
+        raise RuntimeError(f"Samples directory not found: {samples_dir}")
+
+    logger.info(f"Using samples directory: {samples_dir}")

    # Create and start server with simplified parameters
    server = DevServer(
@@ -41,7 +50,7 @@ def start_server() -> tuple[str, Any]:
        app=app,
        host="127.0.0.1",
        port=8085,
-        log_level="info",  # More verbose to see tracing setup
+        # log_level="info",  # More verbose to see tracing setup
    )
    server_instance = uvicorn.Server(server_config)

@@ -80,10 +89,9 @@ def capture_agent_stream_with_tracing(client: OpenAI, agent_id: str, scenario: s

    try:
        stream = client.responses.create(
-            model="agent-framework",
+            model=agent_id,  # DevUI uses model field as entity_id
            input="Tell me about the weather in Tokyo. I want details.",
            stream=True,
-            extra_body={"entity_id": agent_id},
        )

        events = []
@@ -122,13 +130,12 @@ def capture_workflow_stream_with_tracing(

    try:
        stream = client.responses.create(
-            model="agent-framework",
+            model=workflow_id,  # DevUI uses model field as entity_id
            input=(
                "Process this spam detection workflow with multiple emails: "
                "'Buy now!', 'Hello mom', 'URGENT: Click here!'"
            ),
            stream=True,
-            extra_body={"entity_id": workflow_id},
        )

        events = []
@@ -161,70 +168,6 @@ def capture_workflow_stream_with_tracing(
        return [error_event]


-def capture_agent_with_bad_config(base_url: str, agent_id: str) -> list[dict[str, Any]]:
-    """Capture agent events with intentionally bad configuration to test error handling."""
-
-    # Test with invalid API key
-    bad_client = OpenAI(base_url=f"{base_url}/v1", api_key="invalid-api-key-123")
-
-    try:
-        return capture_agent_stream_with_tracing(bad_client, agent_id, "bad_api_key")
-    except Exception as e:
-        return [
-            {
-                "type": "error",
-                "scenario": "bad_api_key",
-                "error_message": str(e),
-                "error_type": type(e).__name__,
-                "timestamp": time.time(),
-            }
-        ]
-
-
-def capture_agent_with_wrong_model(base_url: str, agent_id: str) -> list[dict[str, Any]]:
-    """Capture agent events with wrong model name to test error handling."""
-
-    client = OpenAI(
-        base_url=f"{base_url}/v1",
-        api_key="dummy-key",  # Use the same key as success case
-    )
-
-    try:
-        stream = client.responses.create(
-            model="gpt-4-nonexistent-model",  # Wrong model name
-            input="Tell me about the weather in Tokyo. I want details.",
-            stream=True,
-            extra_body={"entity_id": agent_id},
-        )
-
-        events = []
-        for event in stream:
-            # Serialize the entire event object
-            try:
-                event_dict = json.loads(event.model_dump_json())
-            except Exception:
-                # Fallback to dict conversion if model_dump_json fails
-                event_dict = event.__dict__ if hasattr(event, "__dict__") else str(event)
-
-            events.append(event_dict)
-
-            if len(events) >= 200:
-                break
-
-        return events
-
-    except Exception as e:
-        return [
-            {
-                "type": "error",
-                "scenario": "wrong_model",
-                "error_message": str(e),
-                "error_type": type(e).__name__,
-                "timestamp": time.time(),
-            }
-        ]
-
-
 def main():
    """Main capture script - testing both success and failure scenarios."""

@@ -0,0 +1,335 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for conversation store implementation."""
+
+from typing import cast
+
+import pytest
+from openai.types.conversations import InputFileContent, InputImageContent, InputTextContent
+
+from agent_framework_devui._conversations import InMemoryConversationStore
+
+
+@pytest.mark.asyncio
+async def test_create_conversation():
+    """Test creating a conversation."""
+    store = InMemoryConversationStore()
+
+    conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    assert conversation.id.startswith("conv_")
+    assert conversation.object == "conversation"
+    assert conversation.metadata == {"agent_id": "test_agent"}
+
+
+@pytest.mark.asyncio
+async def test_get_conversation():
+    """Test retrieving a conversation."""
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    created = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Retrieve it
+    retrieved = store.get_conversation(created.id)
+
+    assert retrieved is not None
+    assert retrieved.id == created.id
+    assert retrieved.metadata == {"agent_id": "test_agent"}
+
+
+@pytest.mark.asyncio
+async def test_get_conversation_not_found():
+    """Test retrieving non-existent conversation."""
+    store = InMemoryConversationStore()
+
+    conversation = store.get_conversation("conv_nonexistent")
+
+    assert conversation is None
+
+
+@pytest.mark.asyncio
+async def test_update_conversation():
+    """Test updating conversation metadata."""
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    created = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Update metadata
+    updated = store.update_conversation(created.id, metadata={"agent_id": "new_agent", "session_id": "sess_123"})
+
+    assert updated.id == created.id
+    assert updated.metadata == {"agent_id": "new_agent", "session_id": "sess_123"}
+
+
+@pytest.mark.asyncio
+async def test_delete_conversation():
+    """Test deleting a conversation."""
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    created = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Delete it
+    result = store.delete_conversation(created.id)
+
+    assert result.id == created.id
+    assert result.deleted is True
+    assert result.object == "conversation.deleted"
+
+    # Verify it's gone
+    assert store.get_conversation(created.id) is None
+
+
+@pytest.mark.asyncio
+async def test_get_thread():
+    """Test getting underlying AgentThread."""
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Get thread
+    thread = store.get_thread(conversation.id)
+
+    assert thread is not None
+    # AgentThread should have message_store
+    assert hasattr(thread, "message_store")
+
+
+@pytest.mark.asyncio
+async def test_get_thread_not_found():
+    """Test getting thread for non-existent conversation."""
+    store = InMemoryConversationStore()
+
+    thread = store.get_thread("conv_nonexistent")
+
+    assert thread is None
+
+
+@pytest.mark.asyncio
+async def test_list_conversations_by_metadata():
+    """Test filtering conversations by metadata."""
+    store = InMemoryConversationStore()
+
+    # Create multiple conversations
+    _conv1 = store.create_conversation(metadata={"agent_id": "agent1"})
+    _conv2 = store.create_conversation(metadata={"agent_id": "agent2"})
+    conv3 = store.create_conversation(metadata={"agent_id": "agent1", "session_id": "sess_1"})
+
+    # Filter by agent_id
+    results = store.list_conversations_by_metadata({"agent_id": "agent1"})
+
+    assert len(results) == 2
+    assert all(cast(dict[str, str], c.metadata).get("agent_id") == "agent1" for c in results if c.metadata)
+
+    # Filter by agent_id and session_id
+    results = store.list_conversations_by_metadata({"agent_id": "agent1", "session_id": "sess_1"})
+
+    assert len(results) == 1
+    assert results[0].id == conv3.id
+
+
+@pytest.mark.asyncio
+async def test_add_items():
+    """Test adding items to conversation."""
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Add items
+    items = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}]
+
+    conv_items = await store.add_items(conversation.id, items=items)
+
+    assert len(conv_items) == 1
+    # Message is a ConversationItem type - check standard OpenAI fields
+    assert conv_items[0].type == "message"
+    assert conv_items[0].role == "user"
+    assert conv_items[0].status == "completed"
+    assert len(conv_items[0].content) == 1
+    assert conv_items[0].content[0].type == "text"
+    text_content = cast(InputTextContent, conv_items[0].content[0])
+    assert text_content.text == "Hello"
+
+
+@pytest.mark.asyncio
+async def test_list_items():
+    """Test listing conversation items."""
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Add items
+    items = [
+        {"role": "user", "content": [{"type": "text", "text": "Hello"}]},
+        {"role": "assistant", "content": [{"type": "text", "text": "Hi there"}]},
+    ]
+    await store.add_items(conversation.id, items=items)
+
+    # List items
+    retrieved_items, has_more = await store.list_items(conversation.id)
+
+    assert len(retrieved_items) >= 2  # At least the items we added
+    assert has_more is False
+
+
+@pytest.mark.asyncio
+async def test_list_items_pagination():
+    """Test pagination when listing items."""
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Add multiple items
+    items = [{"role": "user", "content": [{"type": "text", "text": f"Message {i}"}]} for i in range(5)]
+    await store.add_items(conversation.id, items=items)
+
+    # List with limit
+    retrieved_items, has_more = await store.list_items(conversation.id, limit=3)
+
+    assert len(retrieved_items) == 3
+    assert has_more is True
+
+
+@pytest.mark.asyncio
+async def test_list_items_converts_function_calls():
+    """Test that list_items properly converts function calls to ResponseFunctionToolCallItem."""
+    from agent_framework import ChatMessage, ChatMessageStore, Role
+
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Get the underlying thread and set up message store
+    thread = store.get_thread(conversation.id)
+    assert thread is not None
+
+    # Initialize message store if not present
+    if thread.message_store is None:
+        thread.message_store = ChatMessageStore()
+
+    # Simulate messages from agent execution with function calls
+    messages = [
+        ChatMessage(role=Role.USER, contents=[{"type": "text", "text": "What's the weather in SF?"}]),
+        ChatMessage(
+            role=Role.ASSISTANT,
+            contents=[
+                {
+                    "type": "function_call",
+                    "name": "get_weather",
+                    "arguments": '{"city": "San Francisco"}',
+                    "call_id": "call_test123",
+                }
+            ],
+        ),
+        ChatMessage(
+            role=Role.TOOL,
+            contents=[
+                {
+                    "type": "function_result",
+                    "call_id": "call_test123",
+                    "output": '{"temperature": 65, "condition": "sunny"}',
+                }
+            ],
+        ),
+        ChatMessage(role=Role.ASSISTANT, contents=[{"type": "text", "text": "The weather is sunny, 65°F"}]),
+    ]
+
+    # Add messages to thread
+    await thread.on_new_messages(messages)
+
+    # List conversation items
+    items, has_more = await store.list_items(conversation.id)
+
+    # Verify we got the right number and types of items
+    assert len(items) == 4, f"Expected 4 items, got {len(items)}"
+    assert has_more is False
+
+    # Check item types
+    assert items[0].type == "message", "First item should be a message"
+    assert items[0].role == "user"
+    assert len(items[0].content) == 1
+    text_content_0 = cast(InputTextContent, items[0].content[0])
+    assert text_content_0.text == "What's the weather in SF?"
+
+    assert items[1].type == "function_call", "Second item should be a function_call"
+    assert items[1].call_id == "call_test123"
+    assert items[1].name == "get_weather"
+    assert items[1].arguments == '{"city": "San Francisco"}'
+    assert items[1].status == "completed"
+
+    assert items[2].type == "function_call_output", "Third item should be a function_call_output"
+    assert items[2].call_id == "call_test123"
+    assert items[2].output == '{"temperature": 65, "condition": "sunny"}'
+    assert items[2].status == "completed"
+
+    assert items[3].type == "message", "Fourth item should be a message"
+    assert items[3].role == "assistant"
+    assert len(items[3].content) == 1
+    text_content_3 = cast(InputTextContent, items[3].content[0])
+    assert text_content_3.text == "The weather is sunny, 65°F"
+
+    # CRITICAL: Ensure no empty message items
+    for item in items:
+        if item.type == "message":
+            assert len(item.content) > 0, f"Message item {item.id} has empty content!"
+
+
+@pytest.mark.asyncio
+async def test_list_items_handles_images_and_files():
+    """Test that list_items properly converts data content (images/files) to OpenAI types."""
+    from agent_framework import ChatMessage, ChatMessageStore, Role
+
+    store = InMemoryConversationStore()
+
+    # Create conversation
+    conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
+
+    # Get the underlying thread
+    thread = store.get_thread(conversation.id)
+    assert thread is not None
+
+    if thread.message_store is None:
+        thread.message_store = ChatMessageStore()
+
+    # Simulate message with image and file
+    messages = [
+        ChatMessage(
+            role=Role.USER,
+            contents=[
+                {"type": "text", "text": "Check this image and PDF"},
+                {"type": "data", "uri": "data:image/png;base64,iVBORw0KGgo=", "media_type": "image/png"},
+                {"type": "data", "uri": "data:application/pdf;base64,JVBERi0=", "media_type": "application/pdf"},
+            ],
+        ),
+    ]
+
+    await thread.on_new_messages(messages)
+
+    # List items
+    items, has_more = await store.list_items(conversation.id)
+
+    assert len(items) == 1
+    assert items[0].type == "message"
+    assert items[0].role == "user"
+    assert len(items[0].content) == 3
+
+    # Check content types
+    assert items[0].content[0].type == "text"
+    text_content = cast(InputTextContent, items[0].content[0])
+    assert text_content.text == "Check this image and PDF"
+
+    assert items[0].content[1].type == "input_image"
+    image_content = cast(InputImageContent, items[0].content[1])
+    assert image_content.image_url == "data:image/png;base64,iVBORw0KGgo="
+    assert image_content.detail == "auto"
+
+    assert items[0].content[2].type == "input_file"
+    file_content = cast(InputFileContent, items[0].content[2])
+    assert file_content.file_url == "data:application/pdf;base64,JVBERi0="
@@ -14,9 +14,10 @@ from agent_framework_devui._discovery import EntityDiscovery
@pytest.fixture
 def test_entities_dir():
    """Use the samples directory which has proper entity structure."""
-    # Get the samples directory relative to the current test file
+    # Get the samples directory from the main python samples folder
    current_dir = Path(__file__).parent
-    samples_dir = current_dir.parent / "samples"
+    # Navigate to python/samples/getting_started/devui
+    samples_dir = current_dir.parent.parent.parent / "samples" / "getting_started" / "devui"
    return str(samples_dir.resolve())


@@ -66,6 +67,58 @@ async def test_empty_directory():
        assert len(entities) == 0


+async def test_discovery_accepts_agents_with_only_run():
+    """Test that discovery accepts agents with only run() method."""
+    import tempfile
+    from pathlib import Path
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_path = Path(temp_dir)
+
+        # Create agent with only run() method
+        agent_dir = temp_path / "non_streaming_agent"
+        agent_dir.mkdir()
+
+        init_file = agent_dir / "__init__.py"
+        init_file.write_text("""
+from agent_framework import AgentRunResponse, AgentThread, ChatMessage, Role, TextContent
+
+class NonStreamingAgent:
+    id = "non_streaming"
+    name = "Non-Streaming Agent"
+    description = "Agent without run_stream"
+
+    @property
+    def display_name(self):
+        return self.name
+
+    async def run(self, messages=None, *, thread=None, **kwargs):
+        return AgentRunResponse(
+            messages=[ChatMessage(
+                role=Role.ASSISTANT,
+                contents=[TextContent(text="response")]
+            )],
+            response_id="test"
+        )
+
+    def get_new_thread(self, **kwargs):
+        return AgentThread()
+
+agent = NonStreamingAgent()
+""")
+
+        discovery = EntityDiscovery(str(temp_path))
+        entities = await discovery.discover_entities()
+
+        # Should discover the non-streaming agent
+        agents = [e for e in entities if e.type == "agent"]
+        assert len(agents) == 1
+        # ID is auto-generated, just check it exists and starts with agent_
+        assert agents[0].id.startswith("agent_")
+        assert agents[0].name == "Non-Streaming Agent"
+        assert not agents[0].metadata.get("has_run_stream")
+
+
 if __name__ == "__main__":
    # Simple test runner
    async def run_tests():
@@ -12,7 +12,7 @@ import pytest
 from agent_framework_devui._discovery import EntityDiscovery
 from agent_framework_devui._executor import AgentFrameworkExecutor, EntityNotFoundError
 from agent_framework_devui._mapper import MessageMapper
-from agent_framework_devui.models._openai_custom import AgentFrameworkExtraBody, AgentFrameworkRequest
+from agent_framework_devui.models._openai_custom import AgentFrameworkRequest


 class _DummyStartExecutor:
@@ -38,8 +38,10 @@ class _DummyWorkflow:
@pytest.fixture
 def test_entities_dir():
    """Use the samples directory which has proper entity structure."""
+    # Get the samples directory from the main python samples folder
    current_dir = Path(__file__).parent
-    samples_dir = current_dir.parent / "samples"
+    # Navigate to python/samples/getting_started/devui
+    samples_dir = current_dir.parent.parent.parent / "samples" / "getting_started" / "devui"
    return str(samples_dir.resolve())


@@ -94,13 +96,17 @@ async def test_executor_sync_execution(executor):
    assert len(agents) > 0, "No agent entities found for testing"
    agent_id = agents[0].id

+    # Use simplified routing: model = entity_id
    request = AgentFrameworkRequest(
-        model="agent-framework", input="test data", stream=False, extra_body=AgentFrameworkExtraBody(entity_id=agent_id)
+        model=agent_id,  # Model IS the entity_id
+        input="test data",
+        stream=False,
    )

    response = await executor.execute_sync(request)

-    assert response.model == "agent-framework"
+    # With simplified routing, response.model reflects the actual agent_id
+    assert response.model == agent_id
    assert response.object == "response"
    assert len(response.output) > 0
    assert response.usage.total_tokens > 0
@@ -116,11 +122,11 @@ async def test_executor_streaming_execution(executor):
    assert len(agents) > 0, "No agent entities found for testing"
    agent_id = agents[0].id

+    # Use simplified routing: model = entity_id
    request = AgentFrameworkRequest(
-        model="agent-framework",
+        model=agent_id,  # Model IS the entity_id
        input="streaming test",
        stream=True,
-        extra_body=AgentFrameworkExtraBody(entity_id=agent_id),
    )

    event_count = 0
@@ -145,16 +151,16 @@ async def test_executor_invalid_entity_id(executor):


 async def test_executor_missing_entity_id(executor):
-    """Test execution without entity ID."""
+    """Test get_entity_id returns model field (simplified routing)."""
    request = AgentFrameworkRequest(
-        model="agent-framework",
+        model="my_agent",
        input="test",
        stream=False,
-        extra_body=None,  # Test case for missing entity_id
    )

+    # With simplified routing, model field IS the entity_id
    entity_id = request.get_entity_id()
-    assert entity_id is None
+    assert entity_id == "my_agent"


 def test_executor_get_start_executor_message_types_uses_handlers():
@@ -171,10 +177,11 @@ def test_executor_get_start_executor_message_types_uses_handlers():

 def test_executor_select_primary_input_prefers_string():
    """Select string input even when discovered after other handlers."""
-    executor = AgentFrameworkExecutor(EntityDiscovery(None), MessageMapper())
+    from agent_framework_devui._utils import select_primary_input_type
+
    placeholder_type = type("Placeholder", (), {})

-    chosen = executor._select_primary_input_type([placeholder_type, str])
+    chosen = select_primary_input_type([placeholder_type, str])

    assert chosen is str

@@ -201,6 +208,57 @@ def test_executor_parse_raw_falls_back_to_string():
    assert parsed == "hi there"


+async def test_executor_handles_non_streaming_agent():
+    """Test executor can handle agents with only run() method (no run_stream)."""
+    from agent_framework import AgentRunResponse, AgentThread, ChatMessage, Role, TextContent
+
+    class NonStreamingAgent:
+        """Agent with only run() method - does NOT satisfy full AgentProtocol."""
+
+        id = "non_streaming_test"
+        name = "Non-Streaming Test Agent"
+        description = "Test agent without run_stream()"
+
+        @property
+        def display_name(self):
+            return self.name
+
+        async def run(self, messages=None, *, thread=None, **kwargs):
+            return AgentRunResponse(
+                messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text=f"Processed: {messages}")])],
+                response_id="test_123",
+            )
+
+        def get_new_thread(self, **kwargs):
+            return AgentThread()
+
+    # Create executor and register agent
+    discovery = EntityDiscovery(None)
+    mapper = MessageMapper()
+    executor = AgentFrameworkExecutor(discovery, mapper)
+
+    agent = NonStreamingAgent()
+    entity_info = await discovery.create_entity_info_from_object(agent, source="test")
+    discovery.register_entity(entity_info.id, entity_info, agent)
+
+    # Execute non-streaming agent (use simplified routing)
+    request = AgentFrameworkRequest(
+        model=entity_info.id,  # Model IS the entity_id
+        input="hello",
+        stream=True,  # DevUI always streams
+    )
+
+    events = []
+    async for event in executor.execute_streaming(request):
+        events.append(event)
+
+    # Should get events even though agent doesn't stream
+    assert len(events) > 0
+    text_events = [e for e in events if hasattr(e, "type") and e.type == "response.output_text.delta"]
+    assert len(text_events) > 0
+    assert "Processed: hello" in text_events[0].delta
+
+
 if __name__ == "__main__":
    # Simple test runner
    async def run_tests():
@@ -227,12 +285,11 @@ class StreamingAgent:
            entities = await executor.discover_entities()

            if entities:
-                # Test sync execution
+                # Test sync execution (use simplified routing)
                request = AgentFrameworkRequest(
-                    model="agent-framework",
+                    model=entities[0].id,  # Model IS the entity_id
                    input="test input",
                    stream=False,
-                    extra_body=AgentFrameworkExtraBody(entity_id=entities[0].id),
                )

                await executor.execute_sync(request)
@@ -16,7 +16,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent.parent / "main"))
 from agent_framework._types import AgentRunResponseUpdate, ErrorContent, FunctionCallContent, Role, TextContent

 from agent_framework_devui._mapper import MessageMapper
-from agent_framework_devui.models._openai_custom import AgentFrameworkExtraBody, AgentFrameworkRequest
+from agent_framework_devui.models._openai_custom import AgentFrameworkRequest


 def create_test_content(content_type: str, **kwargs: Any) -> Any:
@@ -48,11 +48,11 @@ def mapper() -> MessageMapper:

@pytest.fixture
 def test_request() -> AgentFrameworkRequest:
+    # Use simplified routing: model = entity_id
    return AgentFrameworkRequest(
-        model="agent-framework",
+        model="test_agent",  # Model IS the entity_id
        input="Test input",
        stream=True,
-        extra_body=AgentFrameworkExtraBody(entity_id="test_agent"),
    )


@@ -97,11 +97,14 @@ async def test_function_call_mapping(mapper: MessageMapper, test_request: AgentF

    events = await mapper.convert_event(update, test_request)

-    assert len(events) >= 1
-    assert all(event.type == "response.function_call_arguments.delta" for event in events)
+    # Should generate: response.output_item.added + response.function_call_arguments.delta
+    assert len(events) >= 2
+    assert events[0].type == "response.output_item.added"
+    assert events[1].type == "response.function_call_arguments.delta"

-    # Check JSON is chunked
-    full_json = "".join(event.delta for event in events)
+    # Check JSON is in delta event
+    delta_events = [e for e in events if e.type == "response.function_call_arguments.delta"]
+    full_json = "".join(event.delta for event in delta_events)
    assert "TestCity" in full_json


@@ -158,12 +161,35 @@ async def test_unknown_content_fallback(mapper: MessageMapper, test_request: Age
    assert "WeirdUnknownContent" in event.delta


+async def test_agent_run_response_mapping(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
+    """Test that mapper handles complete AgentRunResponse (non-streaming)."""
+    from agent_framework import AgentRunResponse, ChatMessage, Role, TextContent
+
+    # Create a complete response like agent.run() would return
+    message = ChatMessage(
+        role=Role.ASSISTANT,
+        contents=[TextContent(text="Complete response from run()")],
+    )
+    response = AgentRunResponse(messages=[message], response_id="test_resp_123")
+
+    # Mapper should convert it to streaming events
+    events = await mapper.convert_event(response, test_request)
+
+    assert len(events) > 0
+    # Should produce text delta events
+    text_events = [e for e in events if e.type == "response.output_text.delta"]
+    assert len(text_events) > 0
+    assert text_events[0].delta == "Complete response from run()"
+
+
 if __name__ == "__main__":
    # Simple test runner
    async def run_all_tests() -> None:
        mapper = MessageMapper()
        test_request = AgentFrameworkRequest(
-            model="agent-framework", input="Test", stream=True, extra_body=AgentFrameworkExtraBody(entity_id="test")
+            model="test",
+            input="Test",
+            stream=True,
        )

        tests = [
@@ -9,8 +9,8 @@ from pathlib import Path
 import pytest

 from agent_framework_devui import DevServer
-from agent_framework_devui._server import _extract_executor_message_types, _select_primary_input_type
-from agent_framework_devui.models._openai_custom import AgentFrameworkExtraBody, AgentFrameworkRequest
+from agent_framework_devui._utils import extract_executor_message_types, select_primary_input_type
+from agent_framework_devui.models._openai_custom import AgentFrameworkRequest


 class _StubExecutor:
@@ -26,8 +26,10 @@ class _StubExecutor:
@pytest.fixture
 def test_entities_dir():
    """Use the samples directory which has proper entity structure."""
+    # Get the samples directory from the main python samples folder
    current_dir = Path(__file__).parent
-    samples_dir = current_dir.parent / "samples"
+    # Navigate to python/samples/getting_started/devui
+    samples_dir = current_dir.parent.parent.parent / "samples" / "getting_started" / "devui"
    return str(samples_dir.resolve())


@@ -65,15 +67,15 @@ async def test_server_execution_sync(test_entities_dir):
    entities = await executor.discover_entities()
    agent_id = entities[0].id

+    # Use model as entity_id (new simplified routing)
    request = AgentFrameworkRequest(
-        model="agent-framework",
+        model=agent_id,  # model IS the entity_id now!
        input="San Francisco",
        stream=False,
-        extra_body=AgentFrameworkExtraBody(entity_id=agent_id),
    )

    response = await executor.execute_sync(request)
-    assert response.model == "agent-framework"
+    assert response.model == agent_id  # Should echo back the model (entity_id)
    assert len(response.output) > 0


@@ -85,8 +87,11 @@ async def test_server_execution_streaming(test_entities_dir):
    entities = await executor.discover_entities()
    agent_id = entities[0].id

+    # Use model as entity_id (new simplified routing)
    request = AgentFrameworkRequest(
-        model="agent-framework", input="New York", stream=True, extra_body=AgentFrameworkExtraBody(entity_id=agent_id)
+        model=agent_id,  # model IS the entity_id now!
+        input="New York",
+        stream=True,
    )

    event_count = 0
@@ -112,7 +117,7 @@ def test_extract_executor_message_types_prefers_input_types():
    """Input types property is used when available."""
    stub = _StubExecutor(input_types=[str, dict])

-    types = _extract_executor_message_types(stub)
+    types = extract_executor_message_types(stub)

    assert types == [str, dict]

@@ -121,7 +126,7 @@ def test_extract_executor_message_types_falls_back_to_handlers():
    """Handlers provide message metadata when input_types missing."""
    stub = _StubExecutor(handlers={str: object(), int: object()})

-    types = _extract_executor_message_types(stub)
+    types = extract_executor_message_types(stub)

    assert str in types
    assert int in types
@@ -129,9 +134,9 @@ def test_extract_executor_message_types_falls_back_to_handlers():

 def test_select_primary_input_type_prefers_string_and_dict():
    """Primary type selection prefers user-friendly primitives."""
-    string_first = _select_primary_input_type([dict[str, str], str])
-    dict_first = _select_primary_input_type([dict[str, str]])
-    fallback = _select_primary_input_type([int, float])
+    string_first = select_primary_input_type([dict[str, str], str])
+    dict_first = select_primary_input_type([dict[str, str]])
+    fallback = select_primary_input_type([int, float])

    assert string_first is str
    assert dict_first is dict
@@ -162,10 +167,9 @@ class WeatherAgent:

            if entities:
                request = AgentFrameworkRequest(
-                    model="agent-framework",
+                    model=entities[0].id,  # model IS the entity_id now!
                    input="test location",
                    stream=False,
-                    extra_body=AgentFrameworkExtraBody(entity_id=entities[0].id),
                )

                await executor.execute_sync(request)