mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Dmytro Struk
@@ -89,7 +89,7 @@ def capture_agent_stream_with_tracing(client: OpenAI, agent_id: str, scenario: s
|
||||
|
||||
try:
|
||||
stream = client.responses.create(
|
||||
metadata={"entity_id": agent_id},
|
||||
model=agent_id, # DevUI uses model field as entity_id
|
||||
input="Tell me about the weather in Tokyo. I want details.",
|
||||
stream=True,
|
||||
)
|
||||
@@ -130,7 +130,7 @@ def capture_workflow_stream_with_tracing(
|
||||
|
||||
try:
|
||||
stream = client.responses.create(
|
||||
metadata={"entity_id": workflow_id},
|
||||
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!'"
|
||||
|
||||
@@ -1,443 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for checkpoint-as-conversation-items implementation."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from agent_framework import (
|
||||
Executor,
|
||||
InMemoryCheckpointStorage,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
handler,
|
||||
response_handler,
|
||||
)
|
||||
|
||||
from agent_framework_devui._conversations import (
|
||||
CheckpointConversationManager,
|
||||
InMemoryConversationStore,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkflowTestData:
|
||||
"""Simple test data."""
|
||||
|
||||
value: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkflowHILRequest:
|
||||
"""HIL request for testing."""
|
||||
|
||||
question: str
|
||||
|
||||
|
||||
class WorkflowTestExecutor(Executor):
|
||||
"""Test executor with HIL."""
|
||||
|
||||
@handler
|
||||
async def process(self, data: WorkflowTestData, ctx: WorkflowContext) -> None:
|
||||
"""Process data and request approval."""
|
||||
await ctx.set_executor_state({"data_value": data.value})
|
||||
|
||||
# Request HIL (checkpoint created here)
|
||||
await ctx.request_info(request_data=WorkflowHILRequest(question=f"Approve {data.value}?"), response_type=str)
|
||||
|
||||
@response_handler
|
||||
async def handle_response(
|
||||
self, original_request: WorkflowHILRequest, response: str, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
"""Handle HIL response."""
|
||||
state = await ctx.get_executor_state() or {}
|
||||
value = state.get("data_value", "")
|
||||
await ctx.send_message(f"{value}_approved" if response.lower() == "yes" else f"{value}_rejected")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def conversation_store():
|
||||
"""Create in-memory conversation store."""
|
||||
return InMemoryConversationStore()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def checkpoint_manager(conversation_store):
|
||||
"""Create checkpoint manager."""
|
||||
return CheckpointConversationManager(conversation_store)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_workflow():
|
||||
"""Create test workflow with checkpointing."""
|
||||
executor = WorkflowTestExecutor(id="test_executor")
|
||||
checkpoint_storage = InMemoryCheckpointStorage()
|
||||
|
||||
return (
|
||||
WorkflowBuilder(name="Test Workflow", description="Test checkpoint behavior")
|
||||
.set_start_executor(executor)
|
||||
.with_checkpointing(checkpoint_storage)
|
||||
.build()
|
||||
)
|
||||
|
||||
|
||||
class TestCheckpointConversationManager:
|
||||
"""Test CheckpointConversationManager functionality - CONVERSATION-SCOPED."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conversation_scoped_checkpoint_save(self, checkpoint_manager, test_workflow):
|
||||
"""Test checkpoint save in a specific conversation."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"conv_{entity_id}_test123"
|
||||
|
||||
# Create conversation first
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Create test checkpoint
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()), workflow_id=test_workflow.id, messages={}, shared_state={"test": "data"}
|
||||
)
|
||||
|
||||
# Get checkpoint storage for this conversation and save
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
checkpoint_id = await storage.save_checkpoint(checkpoint)
|
||||
|
||||
assert checkpoint_id == checkpoint.checkpoint_id
|
||||
|
||||
# Verify checkpoint stored in THIS conversation only
|
||||
checkpoints = await storage.list_checkpoints()
|
||||
assert len(checkpoints) == 1
|
||||
assert checkpoints[0].checkpoint_id == checkpoint.checkpoint_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conversation_isolation(self, checkpoint_manager, test_workflow):
|
||||
"""Test that conversations are isolated - checkpoints don't leak between conversations."""
|
||||
entity_id = "test_entity"
|
||||
conv_a = f"conv_{entity_id}_aaa"
|
||||
conv_b = f"conv_{entity_id}_bbb"
|
||||
|
||||
# Create two conversations
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conv_a
|
||||
)
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conv_b
|
||||
)
|
||||
|
||||
# Save checkpoint to conversation A
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint_a = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()),
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"conversation": "A"},
|
||||
)
|
||||
storage_a = checkpoint_manager.get_checkpoint_storage(conv_a)
|
||||
await storage_a.save_checkpoint(checkpoint_a)
|
||||
|
||||
# Verify conversation A has checkpoint
|
||||
checkpoints_a = await storage_a.list_checkpoints()
|
||||
assert len(checkpoints_a) == 1
|
||||
|
||||
# Verify conversation B has NO checkpoints (isolation)
|
||||
storage_b = checkpoint_manager.get_checkpoint_storage(conv_b)
|
||||
checkpoints_b = await storage_b.list_checkpoints()
|
||||
assert len(checkpoints_b) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_checkpoints_in_session(self, checkpoint_manager, test_workflow):
|
||||
"""Test listing checkpoints within a session."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_test456"
|
||||
|
||||
# Create session
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Save multiple checkpoints
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
checkpoint_ids = []
|
||||
for i in range(3):
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()),
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"iteration": i},
|
||||
)
|
||||
saved_id = await storage.save_checkpoint(checkpoint)
|
||||
checkpoint_ids.append(saved_id)
|
||||
|
||||
# List checkpoints using the storage
|
||||
checkpoints_list = await storage.list_checkpoints()
|
||||
assert len(checkpoints_list) == 3
|
||||
|
||||
# Verify all checkpoint IDs are present
|
||||
loaded_ids = [cp.checkpoint_id for cp in checkpoints_list]
|
||||
for saved_id in checkpoint_ids:
|
||||
assert saved_id in loaded_ids
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoints_appear_as_conversation_items(self, checkpoint_manager, test_workflow):
|
||||
"""Test that checkpoints appear as conversation items through the standard API."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_items_test"
|
||||
|
||||
# Create session
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Save multiple checkpoints
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
checkpoint_ids = []
|
||||
for i in range(2):
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=f"checkpoint_{i}",
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"iteration": i},
|
||||
)
|
||||
saved_id = await storage.save_checkpoint(checkpoint)
|
||||
checkpoint_ids.append(saved_id)
|
||||
|
||||
# List conversation items - should include checkpoints
|
||||
items, has_more = await checkpoint_manager.conversation_store.list_items(conversation_id)
|
||||
|
||||
# Filter for checkpoint items
|
||||
checkpoint_items = [item for item in items if (isinstance(item, dict) and item.get("type") == "checkpoint")]
|
||||
|
||||
# Verify we have the correct number of checkpoint items
|
||||
assert len(checkpoint_items) == 2, f"Expected 2 checkpoint items, got {len(checkpoint_items)}"
|
||||
|
||||
# Verify checkpoint items have correct structure
|
||||
for item in checkpoint_items:
|
||||
assert item.get("type") == "checkpoint"
|
||||
assert item.get("checkpoint_id") in checkpoint_ids
|
||||
assert item.get("workflow_id") == test_workflow.id
|
||||
assert "timestamp" in item
|
||||
assert item.get("id").startswith("checkpoint_") # ID format: checkpoint_{checkpoint_id}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_checkpoint_from_session(self, checkpoint_manager, test_workflow):
|
||||
"""Test loading checkpoint from a specific session."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_test789"
|
||||
|
||||
# Create session
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Create and save a checkpoint
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
original_checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()),
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"test_key": "test_value"},
|
||||
)
|
||||
|
||||
# Save to this session
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
await storage.save_checkpoint(original_checkpoint)
|
||||
|
||||
# Load checkpoint from this session
|
||||
loaded_checkpoint = await storage.load_checkpoint(original_checkpoint.checkpoint_id)
|
||||
|
||||
assert loaded_checkpoint is not None
|
||||
assert loaded_checkpoint.checkpoint_id == original_checkpoint.checkpoint_id
|
||||
assert loaded_checkpoint.workflow_id == original_checkpoint.workflow_id
|
||||
assert loaded_checkpoint.shared_state == {"test_key": "test_value"}
|
||||
|
||||
|
||||
class TestCheckpointStorage:
|
||||
"""Test InMemoryCheckpointStorage per conversation - SESSION-SCOPED."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_storage_protocol(self, checkpoint_manager, test_workflow):
|
||||
"""Test that adapter implements CheckpointStorage protocol."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_adapter_test"
|
||||
|
||||
# Create session
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Get storage adapter for this session
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
|
||||
# Create test checkpoint
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()), workflow_id=test_workflow.id, messages={}, shared_state={"test": "data"}
|
||||
)
|
||||
|
||||
# Test save_checkpoint
|
||||
checkpoint_id = await storage.save_checkpoint(checkpoint)
|
||||
assert checkpoint_id == checkpoint.checkpoint_id
|
||||
|
||||
# Test load_checkpoint
|
||||
loaded = await storage.load_checkpoint(checkpoint_id)
|
||||
assert loaded is not None
|
||||
assert loaded.checkpoint_id == checkpoint_id
|
||||
|
||||
# Test list_checkpoint_ids
|
||||
ids = await storage.list_checkpoint_ids(workflow_id=test_workflow.id)
|
||||
assert checkpoint_id in ids
|
||||
|
||||
# Test list_checkpoints
|
||||
checkpoints_list = await storage.list_checkpoints(workflow_id=test_workflow.id)
|
||||
assert len(checkpoints_list) >= 1
|
||||
assert any(cp.checkpoint_id == checkpoint_id for cp in checkpoints_list)
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests for checkpoint workflow execution."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manual_checkpoint_save_via_injected_storage(self, checkpoint_manager, test_workflow):
|
||||
"""Test manual checkpoint save via build-time storage injection."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_integration_test1"
|
||||
|
||||
# Create session conversation
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Get checkpoint storage for this session
|
||||
checkpoint_storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
|
||||
# Set build-time storage (equivalent to .with_checkpointing() at build time)
|
||||
# Note: In production, DevUI uses runtime injection via run_stream() parameter
|
||||
if hasattr(test_workflow, "_runner") and hasattr(test_workflow._runner, "context"):
|
||||
test_workflow._runner.context._checkpoint_storage = checkpoint_storage
|
||||
|
||||
# Create and save a checkpoint via injected storage
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()), workflow_id=test_workflow.id, messages={}, shared_state={"injected": True}
|
||||
)
|
||||
await checkpoint_storage.save_checkpoint(checkpoint)
|
||||
|
||||
# Verify checkpoint is accessible via storage (in this session)
|
||||
storage_checkpoints = await checkpoint_storage.list_checkpoints()
|
||||
assert len(storage_checkpoints) > 0
|
||||
assert storage_checkpoints[0].checkpoint_id == checkpoint.checkpoint_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_roundtrip_via_storage(self, checkpoint_manager, test_workflow):
|
||||
"""Test checkpoint save/load roundtrip via storage adapter."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_integration_test2"
|
||||
|
||||
# Create session conversation
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Set build-time storage for testing
|
||||
checkpoint_storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
test_workflow._runner.context._checkpoint_storage = checkpoint_storage
|
||||
|
||||
# Create checkpoint
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()),
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"ready_to_resume": True},
|
||||
)
|
||||
checkpoint_id = await checkpoint_storage.save_checkpoint(checkpoint)
|
||||
|
||||
# Verify checkpoint can be loaded for resume
|
||||
loaded = await checkpoint_storage.load_checkpoint(checkpoint_id)
|
||||
assert loaded is not None
|
||||
assert loaded.checkpoint_id == checkpoint_id
|
||||
assert loaded.shared_state == {"ready_to_resume": True}
|
||||
|
||||
# Verify checkpoint is accessible via storage (for UI to list checkpoints)
|
||||
checkpoints = await checkpoint_storage.list_checkpoints()
|
||||
assert len(checkpoints) > 0
|
||||
assert checkpoints[0].checkpoint_id == checkpoint_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_auto_saves_checkpoints_to_injected_storage(self, checkpoint_manager, test_workflow):
|
||||
"""Test that workflows automatically save checkpoints to our conversation-backed storage.
|
||||
|
||||
This is the critical end-to-end test that verifies the entire checkpoint flow:
|
||||
1. Storage is set as build-time storage (simulates .with_checkpointing())
|
||||
2. Workflow runs and pauses at HIL point (IDLE_WITH_PENDING_REQUESTS status)
|
||||
3. Framework automatically saves checkpoint to our storage
|
||||
4. Checkpoint is accessible via manager for UI to list/resume
|
||||
|
||||
Note: In production, DevUI passes checkpoint_storage to run_stream() as runtime parameter.
|
||||
This test uses build-time injection to verify framework's checkpoint auto-save behavior.
|
||||
"""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_integration_test3"
|
||||
|
||||
# Create session conversation
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Set build-time storage to test automatic checkpoint saves
|
||||
checkpoint_storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
test_workflow._runner.context._checkpoint_storage = checkpoint_storage
|
||||
|
||||
# Verify no checkpoints initially
|
||||
checkpoints_before = await checkpoint_storage.list_checkpoints()
|
||||
assert len(checkpoints_before) == 0
|
||||
|
||||
# Run workflow until it reaches IDLE_WITH_PENDING_REQUESTS (after checkpoint is created)
|
||||
saw_request_event = False
|
||||
async for event in test_workflow.run_stream(WorkflowTestData(value="test")):
|
||||
if hasattr(event, "__class__"):
|
||||
if event.__class__.__name__ == "RequestInfoEvent":
|
||||
saw_request_event = True
|
||||
# Wait for IDLE_WITH_PENDING_REQUESTS status (comes after checkpoint creation)
|
||||
is_status_event = event.__class__.__name__ == "WorkflowStatusEvent"
|
||||
has_pending_status = hasattr(event, "status") and "IDLE_WITH_PENDING_REQUESTS" in str(event.status)
|
||||
if is_status_event and has_pending_status:
|
||||
break
|
||||
|
||||
assert saw_request_event, "Test workflow should have emitted RequestInfoEvent"
|
||||
|
||||
# Verify checkpoint was AUTOMATICALLY saved to our storage by the framework
|
||||
checkpoints_after = await checkpoint_storage.list_checkpoints()
|
||||
assert len(checkpoints_after) > 0, "Workflow should have auto-saved checkpoint at HIL pause"
|
||||
|
||||
# Verify checkpoint has correct workflow_id
|
||||
checkpoint = checkpoints_after[0]
|
||||
assert checkpoint.workflow_id == test_workflow.id
|
||||
@@ -1,365 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for cleanup hook registration and execution."""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentRunResponse, ChatMessage, Role, TextContent
|
||||
|
||||
from agent_framework_devui import register_cleanup
|
||||
from agent_framework_devui._discovery import EntityDiscovery
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cleanup_registry():
|
||||
"""Clear the cleanup registry before each test."""
|
||||
import agent_framework_devui
|
||||
|
||||
agent_framework_devui._cleanup_registry.clear()
|
||||
yield
|
||||
agent_framework_devui._cleanup_registry.clear()
|
||||
|
||||
|
||||
class MockAgent:
|
||||
"""Mock agent for testing."""
|
||||
|
||||
def __init__(self, name: str = "TestAgent"):
|
||||
self.id = f"test-{name.lower()}"
|
||||
self.name = name
|
||||
self.description = "Test agent for cleanup hooks"
|
||||
self.cleanup_called = False
|
||||
self.async_cleanup_called = False
|
||||
|
||||
async def run_stream(self, messages=None, *, thread=None, **kwargs):
|
||||
"""Mock streaming run method."""
|
||||
yield AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, content=[TextContent(text="Test response")])],
|
||||
inner_messages=[],
|
||||
)
|
||||
|
||||
|
||||
class MockCredential:
|
||||
"""Mock credential object for testing cleanup."""
|
||||
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
"""Mock async close method."""
|
||||
self.closed = True
|
||||
|
||||
|
||||
class MockSyncResource:
|
||||
"""Mock synchronous resource for testing cleanup."""
|
||||
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
def close(self):
|
||||
"""Mock sync close method."""
|
||||
self.closed = True
|
||||
|
||||
|
||||
# Test 1: Register single cleanup hook
|
||||
async def test_register_cleanup_single_hook():
|
||||
"""Test registering a single cleanup hook for an entity."""
|
||||
agent = MockAgent("SingleHook")
|
||||
credential = MockCredential()
|
||||
|
||||
# Register cleanup
|
||||
register_cleanup(agent, credential.close)
|
||||
|
||||
# Verify credential not closed yet
|
||||
assert not credential.closed
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Get cleanup hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 1
|
||||
|
||||
# Execute hook
|
||||
await hooks[0]()
|
||||
assert credential.closed
|
||||
|
||||
|
||||
# Test 2: Register multiple cleanup hooks
|
||||
async def test_register_cleanup_multiple_hooks():
|
||||
"""Test registering multiple cleanup hooks for a single entity."""
|
||||
agent = MockAgent("MultipleHooks")
|
||||
credential1 = MockCredential()
|
||||
credential2 = MockCredential()
|
||||
sync_resource = MockSyncResource()
|
||||
|
||||
# Register multiple hooks at once
|
||||
register_cleanup(agent, credential1.close, credential2.close, sync_resource.close)
|
||||
|
||||
# Verify nothing closed yet
|
||||
assert not credential1.closed
|
||||
assert not credential2.closed
|
||||
assert not sync_resource.closed
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Get and execute hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 3
|
||||
|
||||
# Execute all hooks
|
||||
for hook in hooks:
|
||||
if asyncio.iscoroutinefunction(hook):
|
||||
await hook()
|
||||
else:
|
||||
hook()
|
||||
|
||||
assert credential1.closed
|
||||
assert credential2.closed
|
||||
assert sync_resource.closed
|
||||
|
||||
|
||||
# Test 3: Register cleanup hooks incrementally
|
||||
async def test_register_cleanup_incremental():
|
||||
"""Test registering cleanup hooks in multiple calls."""
|
||||
agent = MockAgent("IncrementalHooks")
|
||||
credential1 = MockCredential()
|
||||
credential2 = MockCredential()
|
||||
|
||||
# Register hooks incrementally
|
||||
register_cleanup(agent, credential1.close)
|
||||
register_cleanup(agent, credential2.close)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Should have both hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 2
|
||||
|
||||
# Execute all hooks
|
||||
for hook in hooks:
|
||||
await hook()
|
||||
|
||||
assert credential1.closed
|
||||
assert credential2.closed
|
||||
|
||||
|
||||
# Test 4: Test with no cleanup hooks
|
||||
async def test_no_cleanup_hooks():
|
||||
"""Test entity without any cleanup hooks registered."""
|
||||
agent = MockAgent("NoHooks")
|
||||
|
||||
# Don't register any cleanup hooks
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Should return empty list
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 0
|
||||
|
||||
|
||||
# Test 5: Test cleanup with async and sync hooks mixed
|
||||
async def test_mixed_async_sync_hooks():
|
||||
"""Test that both async and sync cleanup hooks work together."""
|
||||
agent = MockAgent("MixedHooks")
|
||||
async_resource = MockCredential()
|
||||
sync_resource = MockSyncResource()
|
||||
|
||||
# Register both types
|
||||
register_cleanup(agent, async_resource.close, sync_resource.close)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Get and execute hooks with proper async/sync handling
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 2
|
||||
|
||||
import inspect
|
||||
|
||||
for hook in hooks:
|
||||
if inspect.iscoroutinefunction(hook):
|
||||
await hook()
|
||||
else:
|
||||
hook()
|
||||
|
||||
assert async_resource.closed
|
||||
assert sync_resource.closed
|
||||
|
||||
|
||||
# Test 6: Test error handling in cleanup hooks
|
||||
async def test_cleanup_hook_error_handling():
|
||||
"""Test that errors in cleanup hooks don't break execution."""
|
||||
agent = MockAgent("ErrorHooks")
|
||||
credential = MockCredential()
|
||||
|
||||
def failing_hook():
|
||||
raise RuntimeError("Intentional error for testing")
|
||||
|
||||
# Register failing hook and valid hook
|
||||
register_cleanup(agent, failing_hook, credential.close)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Get hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 2
|
||||
|
||||
# Execute hooks with error handling (like _server.py does)
|
||||
import inspect
|
||||
|
||||
for hook in hooks:
|
||||
try:
|
||||
if inspect.iscoroutinefunction(hook):
|
||||
await hook()
|
||||
else:
|
||||
hook()
|
||||
except Exception:
|
||||
pass # Ignore errors like the server does
|
||||
|
||||
# Second hook should still execute despite first one failing
|
||||
await credential.close()
|
||||
assert credential.closed
|
||||
|
||||
|
||||
# Test 7: Test ValueError when no hooks provided
|
||||
def test_register_cleanup_no_hooks_error():
|
||||
"""Test that register_cleanup raises ValueError when no hooks provided."""
|
||||
agent = MockAgent("NoHooksError")
|
||||
|
||||
with pytest.raises(ValueError, match="At least one cleanup hook required"):
|
||||
register_cleanup(agent)
|
||||
|
||||
|
||||
# Test 8: Test file-based discovery with cleanup hooks
|
||||
async def test_cleanup_with_file_based_discovery():
|
||||
"""Test that cleanup hooks work with file-based entity discovery."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create agent directory
|
||||
agent_dir = temp_path / "test_agent"
|
||||
agent_dir.mkdir()
|
||||
|
||||
# Write agent module with cleanup registration
|
||||
agent_file = agent_dir / "__init__.py"
|
||||
agent_file.write_text("""
|
||||
from agent_framework import AgentRunResponse, ChatMessage, Role, TextContent
|
||||
from agent_framework_devui import register_cleanup
|
||||
|
||||
class MockCredential:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
# Create credential and agent
|
||||
credential = MockCredential()
|
||||
|
||||
class TestAgent:
|
||||
id = "test-agent"
|
||||
name = "Test Agent"
|
||||
description = "Test agent with cleanup"
|
||||
|
||||
async def run_stream(self, messages=None, *, thread=None, **kwargs):
|
||||
yield AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, content=[TextContent(text="Test")])],
|
||||
inner_messages=[],
|
||||
)
|
||||
|
||||
agent = TestAgent()
|
||||
|
||||
# Register cleanup at module level
|
||||
register_cleanup(agent, credential.close)
|
||||
""")
|
||||
|
||||
# Discover entities
|
||||
discovery = EntityDiscovery(str(temp_path))
|
||||
await discovery.discover_entities()
|
||||
|
||||
# Load the entity (triggers module import)
|
||||
await discovery.load_entity("test_agent")
|
||||
|
||||
# Verify cleanup hooks were registered
|
||||
hooks = discovery.get_cleanup_hooks("test_agent")
|
||||
assert len(hooks) == 1
|
||||
|
||||
|
||||
# Test 9: Test cleanup execution order
|
||||
async def test_cleanup_execution_order():
|
||||
"""Test that cleanup hooks execute in registration order."""
|
||||
agent = MockAgent("OrderTest")
|
||||
execution_order = []
|
||||
|
||||
def hook1():
|
||||
execution_order.append(1)
|
||||
|
||||
def hook2():
|
||||
execution_order.append(2)
|
||||
|
||||
def hook3():
|
||||
execution_order.append(3)
|
||||
|
||||
# Register in specific order
|
||||
register_cleanup(agent, hook1, hook2, hook3)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Execute hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
for hook in hooks:
|
||||
hook()
|
||||
|
||||
# Verify execution order
|
||||
assert execution_order == [1, 2, 3]
|
||||
|
||||
|
||||
# Test 10: Test custom cleanup logic
|
||||
async def test_custom_cleanup_logic():
|
||||
"""Test registering custom cleanup function with complex logic."""
|
||||
agent = MockAgent("CustomCleanup")
|
||||
cleanup_executed = False
|
||||
resources_closed = []
|
||||
|
||||
async def custom_cleanup():
|
||||
nonlocal cleanup_executed
|
||||
cleanup_executed = True
|
||||
resources_closed.append("credential")
|
||||
resources_closed.append("session")
|
||||
resources_closed.append("cache")
|
||||
|
||||
register_cleanup(agent, custom_cleanup)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Execute hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 1
|
||||
|
||||
await hooks[0]()
|
||||
|
||||
assert cleanup_executed
|
||||
assert resources_closed == ["credential", "session", "cache"]
|
||||
@@ -100,43 +100,17 @@ async def test_executor_sync_execution(executor):
|
||||
assert len(agents) > 0, "No agent entities found for testing"
|
||||
agent_id = agents[0].id
|
||||
|
||||
# Use metadata.entity_id for routing
|
||||
# Use simplified routing: model = entity_id
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": agent_id},
|
||||
model=agent_id, # Model IS the entity_id
|
||||
input="test data",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
response = await executor.execute_sync(request)
|
||||
|
||||
# Response model should be 'devui' when not specified
|
||||
assert response.model == "devui"
|
||||
assert response.object == "response"
|
||||
assert len(response.output) > 0
|
||||
assert response.usage.total_tokens > 0
|
||||
|
||||
|
||||
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="requires OpenAI API key")
|
||||
async def test_executor_sync_execution_with_model(executor):
|
||||
"""Test synchronous execution with model field specified."""
|
||||
entities = await executor.discover_entities()
|
||||
# Find an agent entity to test with
|
||||
agents = [e for e in entities if e.type == "agent"]
|
||||
assert len(agents) > 0, "No agent entities found for testing"
|
||||
agent_id = agents[0].id
|
||||
|
||||
# Use metadata.entity_id for routing AND specify a model
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": agent_id},
|
||||
model="custom-model-name",
|
||||
input="test data",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
response = await executor.execute_sync(request)
|
||||
|
||||
# Response model should reflect the specified model
|
||||
assert response.model == "custom-model-name"
|
||||
# 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
|
||||
@@ -152,9 +126,9 @@ async def test_executor_streaming_execution(executor):
|
||||
assert len(agents) > 0, "No agent entities found for testing"
|
||||
agent_id = agents[0].id
|
||||
|
||||
# Use metadata.entity_id for routing
|
||||
# Use simplified routing: model = entity_id
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": agent_id},
|
||||
model=agent_id, # Model IS the entity_id
|
||||
input="streaming test",
|
||||
stream=True,
|
||||
)
|
||||
@@ -181,14 +155,14 @@ async def test_executor_invalid_entity_id(executor):
|
||||
|
||||
|
||||
async def test_executor_missing_entity_id(executor):
|
||||
"""Test get_entity_id returns metadata.entity_id."""
|
||||
"""Test get_entity_id returns model field (simplified routing)."""
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": "my_agent"},
|
||||
model="my_agent",
|
||||
input="test",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# entity_id is extracted from metadata
|
||||
# With simplified routing, model field IS the entity_id
|
||||
entity_id = request.get_entity_id()
|
||||
assert entity_id == "my_agent"
|
||||
|
||||
@@ -238,29 +212,6 @@ def test_executor_parse_raw_falls_back_to_string():
|
||||
assert parsed == "hi there"
|
||||
|
||||
|
||||
def test_executor_parse_stringified_json_workflow_input():
|
||||
"""Stringified JSON workflow input (from frontend JSON.stringify) is correctly parsed."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class WorkflowInput(BaseModel):
|
||||
input: str
|
||||
metadata: dict | None = None
|
||||
|
||||
executor = AgentFrameworkExecutor(EntityDiscovery(None), MessageMapper())
|
||||
start_executor = _DummyStartExecutor(handlers={WorkflowInput: lambda *_: None})
|
||||
workflow = _DummyWorkflow(start_executor)
|
||||
|
||||
# Simulate frontend sending JSON.stringify({"input": "testing!", "metadata": {"key": "value"}})
|
||||
stringified_json = '{"input": "testing!", "metadata": {"key": "value"}}'
|
||||
|
||||
parsed = executor._parse_raw_workflow_input(workflow, stringified_json)
|
||||
|
||||
# Should parse into WorkflowInput object
|
||||
assert isinstance(parsed, WorkflowInput)
|
||||
assert parsed.input == "testing!"
|
||||
assert parsed.metadata == {"key": "value"}
|
||||
|
||||
|
||||
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
|
||||
@@ -294,9 +245,9 @@ async def test_executor_handles_non_streaming_agent():
|
||||
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 metadata.entity_id for routing)
|
||||
# Execute non-streaming agent (use simplified routing)
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_info.id},
|
||||
model=entity_info.id, # Model IS the entity_id
|
||||
input="hello",
|
||||
stream=True, # DevUI always streams
|
||||
)
|
||||
@@ -338,9 +289,9 @@ class StreamingAgent:
|
||||
entities = await executor.discover_entities()
|
||||
|
||||
if entities:
|
||||
# Test sync execution (use metadata.entity_id for routing)
|
||||
# Test sync execution (use simplified routing)
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entities[0].id},
|
||||
model=entities[0].id, # Model IS the entity_id
|
||||
input="test input",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
@@ -55,9 +55,9 @@ def mapper() -> MessageMapper:
|
||||
|
||||
@pytest.fixture
|
||||
def test_request() -> AgentFrameworkRequest:
|
||||
# Use metadata.entity_id for routing
|
||||
# Use simplified routing: model = entity_id
|
||||
return AgentFrameworkRequest(
|
||||
metadata={"entity_id": "test_agent"},
|
||||
model="test_agent", # Model IS the entity_id
|
||||
input="Test input",
|
||||
stream=True,
|
||||
)
|
||||
@@ -292,7 +292,7 @@ async def test_agent_lifecycle_events(mapper: MessageMapper, test_request: Agent
|
||||
assert len(events) == 2 # Should emit response.created and response.in_progress
|
||||
assert events[0].type == "response.created"
|
||||
assert events[1].type == "response.in_progress"
|
||||
assert events[0].response.model == "devui" # Should use 'devui' when model not specified in request
|
||||
assert events[0].response.model == "test_agent" # Should use model from request
|
||||
assert events[0].response.status == "in_progress"
|
||||
|
||||
# Test AgentCompletedEvent
|
||||
@@ -415,62 +415,12 @@ async def test_executor_action_events(mapper: MessageMapper, test_request: Agent
|
||||
assert "Executor failed" in str(events[0].item["error"]["message"])
|
||||
|
||||
|
||||
async def test_magentic_agent_delta_creates_message_container(
|
||||
mapper: MessageMapper, test_request: AgentFrameworkRequest
|
||||
) -> None:
|
||||
"""Test that MagenticAgentDeltaEvent creates message containers (Option A implementation)."""
|
||||
|
||||
# Create mock MagenticAgentDeltaEvent that mimics the real class
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from agent_framework import WorkflowEvent
|
||||
|
||||
@dataclass
|
||||
class MagenticAgentDeltaEvent(WorkflowEvent): # Inherit from WorkflowEvent
|
||||
agent_id: str
|
||||
text: str | None = None
|
||||
|
||||
except ImportError:
|
||||
# Fallback if WorkflowEvent is not available
|
||||
@dataclass
|
||||
class MagenticAgentDeltaEvent: # Use the expected name directly
|
||||
agent_id: str
|
||||
text: str | None = None
|
||||
|
||||
# First delta should create message container
|
||||
first_delta = MagenticAgentDeltaEvent(agent_id="test_agent", text="Hello ")
|
||||
events = await mapper.convert_event(first_delta, test_request)
|
||||
|
||||
# Should emit 3 events: message container, content part, and text delta
|
||||
assert len(events) == 3
|
||||
assert events[0].type == "response.output_item.added"
|
||||
assert events[0].item.type == "message" # Message, not executor_action!
|
||||
assert events[0].item.metadata["agent_id"] == "test_agent"
|
||||
assert events[0].item.metadata["source"] == "magentic"
|
||||
message_id = events[0].item.id
|
||||
|
||||
# Check text delta references the message ID
|
||||
assert events[2].type == "response.output_text.delta"
|
||||
assert events[2].item_id == message_id
|
||||
assert events[2].delta == "Hello "
|
||||
|
||||
# Second delta should NOT create new container
|
||||
second_delta = MagenticAgentDeltaEvent(agent_id="test_agent", text="world!")
|
||||
events = await mapper.convert_event(second_delta, test_request)
|
||||
|
||||
# Only text delta, no new container
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.output_text.delta"
|
||||
assert events[0].item_id == message_id
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner
|
||||
async def run_all_tests() -> None:
|
||||
mapper = MessageMapper()
|
||||
test_request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": "test"},
|
||||
model="test",
|
||||
input="Test",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
@@ -1,266 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Integration tests using the official OpenAI SDK to call DevUI."""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import http.client
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import pytest
|
||||
import uvicorn
|
||||
from openai import OpenAI
|
||||
|
||||
from agent_framework_devui import DevServer
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def devui_server() -> Generator[str, None, None]:
|
||||
"""Start a DevUI server for testing.
|
||||
|
||||
Yields:
|
||||
Base URL of the running server.
|
||||
"""
|
||||
# Get samples directory
|
||||
current_dir = Path(__file__).parent
|
||||
samples_dir = current_dir.parent.parent.parent / "samples" / "getting_started" / "devui"
|
||||
|
||||
if not samples_dir.exists():
|
||||
pytest.skip(f"Samples directory not found: {samples_dir}")
|
||||
|
||||
# Create and start server with port 0 to get a random available port
|
||||
server = DevServer(
|
||||
entities_dir=str(samples_dir.resolve()),
|
||||
host="127.0.0.1",
|
||||
port=0, # Use 0 to let OS assign a random available port
|
||||
ui_enabled=False,
|
||||
)
|
||||
|
||||
app = server.get_app()
|
||||
|
||||
server_config = uvicorn.Config(
|
||||
app=app,
|
||||
host="127.0.0.1",
|
||||
port=0, # Use 0 to let OS assign a random available port
|
||||
log_level="error",
|
||||
ws="none", # Disable websockets to avoid deprecation warnings
|
||||
)
|
||||
server_instance = uvicorn.Server(server_config)
|
||||
|
||||
def run_server() -> None:
|
||||
asyncio.run(server_instance.serve())
|
||||
|
||||
server_thread = threading.Thread(target=run_server, daemon=True)
|
||||
server_thread.start()
|
||||
|
||||
# Wait for server to start and get the actual port
|
||||
max_retries = 20
|
||||
actual_port = None
|
||||
for _ in range(max_retries):
|
||||
time.sleep(0.5)
|
||||
# Get the actual port from the server instance
|
||||
if hasattr(server_instance, "servers") and server_instance.servers:
|
||||
for srv in server_instance.servers:
|
||||
for socket in srv.sockets:
|
||||
actual_port = socket.getsockname()[1]
|
||||
break
|
||||
if actual_port:
|
||||
break
|
||||
|
||||
if actual_port:
|
||||
# Verify server is responding
|
||||
try:
|
||||
conn = http.client.HTTPConnection("127.0.0.1", actual_port, timeout=5)
|
||||
try:
|
||||
conn.request("GET", "/health")
|
||||
response = conn.getresponse()
|
||||
if response.status == 200:
|
||||
break
|
||||
finally:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not actual_port:
|
||||
pytest.skip("Server failed to start - could not determine port")
|
||||
|
||||
yield f"http://127.0.0.1:{actual_port}"
|
||||
|
||||
# Cleanup
|
||||
with contextlib.suppress(Exception):
|
||||
server_instance.should_exit = True
|
||||
|
||||
|
||||
def test_openai_sdk_responses_create_with_entity_id(devui_server: str) -> None:
|
||||
"""Test using OpenAI SDK with entity_id in metadata (no model parameter)."""
|
||||
base_url = devui_server
|
||||
client = OpenAI(base_url=f"{base_url}/v1", api_key="not-needed")
|
||||
|
||||
# Get available entities - extract host and port from base_url
|
||||
parsed = urlparse(base_url)
|
||||
conn = http.client.HTTPConnection(parsed.hostname, parsed.port, timeout=10)
|
||||
try:
|
||||
conn.request("GET", "/v1/entities")
|
||||
response = conn.getresponse()
|
||||
entities = json.loads(response.read().decode("utf-8"))["entities"]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
assert len(entities) > 0, "No entities discovered"
|
||||
|
||||
# Find an agent entity
|
||||
agent = next((e for e in entities if e["type"] == "agent"), None)
|
||||
if not agent:
|
||||
pytest.skip("No agent entities found")
|
||||
|
||||
agent_id = agent["id"]
|
||||
|
||||
# Test non-streaming request with entity_id in metadata
|
||||
response = client.responses.create(
|
||||
metadata={"entity_id": agent_id},
|
||||
input="What is 2+2?",
|
||||
)
|
||||
|
||||
assert response.object == "response"
|
||||
assert len(response.output) > 0
|
||||
assert response.output[0].content is not None
|
||||
|
||||
|
||||
def test_openai_sdk_responses_create_streaming(devui_server: str) -> None:
|
||||
"""Test using OpenAI SDK with streaming enabled."""
|
||||
base_url = devui_server
|
||||
client = OpenAI(base_url=f"{base_url}/v1", api_key="not-needed")
|
||||
|
||||
# Get available entities - extract host and port from base_url
|
||||
parsed = urlparse(base_url)
|
||||
conn = http.client.HTTPConnection(parsed.hostname, parsed.port, timeout=10)
|
||||
try:
|
||||
conn.request("GET", "/v1/entities")
|
||||
response = conn.getresponse()
|
||||
entities = json.loads(response.read().decode("utf-8"))["entities"]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
assert len(entities) > 0, "No entities discovered"
|
||||
|
||||
# Find an agent entity
|
||||
agent = next((e for e in entities if e["type"] == "agent"), None)
|
||||
if not agent:
|
||||
pytest.skip("No agent entities found")
|
||||
|
||||
agent_id = agent["id"]
|
||||
|
||||
# Test streaming request
|
||||
stream = client.responses.create(
|
||||
metadata={"entity_id": agent_id},
|
||||
input="Count to 3",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
events = []
|
||||
for event in stream:
|
||||
events.append(event)
|
||||
if len(events) >= 100: # Limit for safety
|
||||
break
|
||||
|
||||
assert len(events) > 0, "No events received from stream"
|
||||
|
||||
# Check that we got various event types
|
||||
event_types = {event.type for event in events}
|
||||
# Should have at least response.completed or some content events
|
||||
assert len(event_types) > 0
|
||||
|
||||
|
||||
def test_openai_sdk_with_conversations(devui_server: str) -> None:
|
||||
"""Test using OpenAI SDK with conversation continuity."""
|
||||
base_url = devui_server
|
||||
client = OpenAI(base_url=f"{base_url}/v1", api_key="not-needed")
|
||||
|
||||
# Get available entities - extract host and port from base_url
|
||||
parsed = urlparse(base_url)
|
||||
conn = http.client.HTTPConnection(parsed.hostname, parsed.port, timeout=10)
|
||||
try:
|
||||
conn.request("GET", "/v1/entities")
|
||||
response = conn.getresponse()
|
||||
entities = json.loads(response.read().decode("utf-8"))["entities"]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
assert len(entities) > 0, "No entities discovered"
|
||||
|
||||
# Find an agent entity
|
||||
agent = next((e for e in entities if e["type"] == "agent"), None)
|
||||
if not agent:
|
||||
pytest.skip("No agent entities found")
|
||||
|
||||
agent_id = agent["id"]
|
||||
|
||||
# Create a conversation
|
||||
conversation = client.conversations.create(metadata={"agent_id": agent_id})
|
||||
|
||||
assert conversation.id is not None
|
||||
|
||||
# First turn
|
||||
response1 = client.responses.create(
|
||||
metadata={"entity_id": agent_id},
|
||||
input="My name is Alice",
|
||||
conversation=conversation.id,
|
||||
)
|
||||
|
||||
assert response1.object == "response"
|
||||
assert len(response1.output) > 0
|
||||
|
||||
# Second turn - test conversation continuity
|
||||
response2 = client.responses.create(
|
||||
metadata={"entity_id": agent_id},
|
||||
input="What is my name?",
|
||||
conversation=conversation.id,
|
||||
)
|
||||
|
||||
assert response2.object == "response"
|
||||
assert len(response2.output) > 0
|
||||
# The agent should remember the name from the previous turn
|
||||
# Note: This may not work with all agents, so we just verify we got a response
|
||||
assert response2.output[0].content is not None
|
||||
|
||||
|
||||
def test_openai_sdk_with_model_and_entity_id(devui_server: str) -> None:
|
||||
"""Test that both model and entity_id can be specified together."""
|
||||
base_url = devui_server
|
||||
client = OpenAI(base_url=f"{base_url}/v1", api_key="not-needed")
|
||||
|
||||
# Get available entities - extract host and port from base_url
|
||||
parsed = urlparse(base_url)
|
||||
conn = http.client.HTTPConnection(parsed.hostname, parsed.port, timeout=10)
|
||||
try:
|
||||
conn.request("GET", "/v1/entities")
|
||||
response = conn.getresponse()
|
||||
entities = json.loads(response.read().decode("utf-8"))["entities"]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
assert len(entities) > 0, "No entities discovered"
|
||||
|
||||
# Find an agent entity
|
||||
agent = next((e for e in entities if e["type"] == "agent"), None)
|
||||
if not agent:
|
||||
pytest.skip("No agent entities found")
|
||||
|
||||
agent_id = agent["id"]
|
||||
|
||||
# Test with both model and entity_id - entity_id should be used for routing
|
||||
response = client.responses.create(
|
||||
metadata={"entity_id": agent_id},
|
||||
model="custom-model-name",
|
||||
input="Hello",
|
||||
)
|
||||
|
||||
assert response.object == "response"
|
||||
# The response model should reflect what was specified
|
||||
assert response.model == "custom-model-name"
|
||||
assert len(response.output) > 0
|
||||
@@ -4,14 +4,13 @@
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
import pytest
|
||||
|
||||
# Add parent package to path
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from agent_framework_devui._utils import extract_response_type_from_executor, generate_input_schema
|
||||
from agent_framework_devui._utils import generate_input_schema
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -133,99 +132,6 @@ def test_schema_generation_error_handling():
|
||||
pass
|
||||
|
||||
|
||||
def test_extract_response_type_from_executor():
|
||||
"""Test extraction of response type from @response_handler methods."""
|
||||
try:
|
||||
from agent_framework import Executor, WorkflowContext, handler, response_handler
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# Define test request and response types
|
||||
@dataclass
|
||||
class TestApprovalRequest:
|
||||
"""Test request for approval."""
|
||||
|
||||
prompt: str
|
||||
context: str
|
||||
|
||||
class TestDecision(BaseModel):
|
||||
"""Test decision response."""
|
||||
|
||||
decision: Literal["approve", "reject"] = Field(description="User's decision")
|
||||
reason: str = Field(description="Reason for decision", default="")
|
||||
|
||||
# Create test executor with @response_handler
|
||||
class TestExecutor(Executor):
|
||||
"""Test executor with response handler."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
|
||||
"""Regular handler to satisfy executor requirements."""
|
||||
# Request info that will be handled by response_handler
|
||||
request = TestApprovalRequest(prompt="Test", context="Test context")
|
||||
await ctx.request_info(request, TestDecision)
|
||||
|
||||
@response_handler
|
||||
async def handle_approval(
|
||||
self, original_request: TestApprovalRequest, response: TestDecision, ctx: WorkflowContext
|
||||
) -> None:
|
||||
"""Handle approval response."""
|
||||
pass
|
||||
|
||||
# Test extraction
|
||||
executor = TestExecutor()
|
||||
extracted_type = extract_response_type_from_executor(executor, TestApprovalRequest)
|
||||
|
||||
# Verify correct type was extracted
|
||||
assert extracted_type is not None, "Should extract response type from @response_handler"
|
||||
assert extracted_type == TestDecision, f"Expected TestDecision, got {extracted_type}"
|
||||
|
||||
# Test full schema generation pipeline
|
||||
schema = generate_input_schema(extracted_type)
|
||||
assert schema is not None
|
||||
assert isinstance(schema, dict)
|
||||
assert "properties" in schema
|
||||
assert "decision" in schema["properties"]
|
||||
assert "enum" in schema["properties"]["decision"]
|
||||
assert schema["properties"]["decision"]["enum"] == ["approve", "reject"]
|
||||
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Required dependencies not available: {e}")
|
||||
|
||||
|
||||
def test_extract_response_type_no_match():
|
||||
"""Test that extraction returns None when no matching handler exists."""
|
||||
try:
|
||||
from agent_framework import Executor, WorkflowContext, handler
|
||||
|
||||
@dataclass
|
||||
class UnmatchedRequest:
|
||||
"""Request type with no handler."""
|
||||
|
||||
data: str
|
||||
|
||||
class MinimalExecutor(Executor):
|
||||
"""Executor with a handler but no matching response_handler."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(id="minimal_executor")
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
|
||||
"""Regular handler."""
|
||||
pass
|
||||
|
||||
executor = MinimalExecutor()
|
||||
extracted_type = extract_response_type_from_executor(executor, UnmatchedRequest)
|
||||
|
||||
assert extracted_type is None, "Should return None when no matching handler exists"
|
||||
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Required dependencies not available: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner for manual execution
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
@@ -67,15 +67,15 @@ async def test_server_execution_sync(test_entities_dir):
|
||||
entities = await executor.discover_entities()
|
||||
agent_id = entities[0].id
|
||||
|
||||
# Use metadata.entity_id for routing
|
||||
# Use model as entity_id (new simplified routing)
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": agent_id},
|
||||
model=agent_id, # model IS the entity_id now!
|
||||
input="San Francisco",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
response = await executor.execute_sync(request)
|
||||
assert response.model == "devui" # Response model defaults to 'devui' when not specified
|
||||
assert response.model == agent_id # Should echo back the model (entity_id)
|
||||
assert len(response.output) > 0
|
||||
|
||||
|
||||
@@ -87,9 +87,9 @@ async def test_server_execution_streaming(test_entities_dir):
|
||||
entities = await executor.discover_entities()
|
||||
agent_id = entities[0].id
|
||||
|
||||
# Use metadata.entity_id for routing
|
||||
# Use model as entity_id (new simplified routing)
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": agent_id},
|
||||
model=agent_id, # model IS the entity_id now!
|
||||
input="New York",
|
||||
stream=True,
|
||||
)
|
||||
@@ -241,100 +241,6 @@ async def test_multiple_credential_attributes() -> None:
|
||||
assert mock_cred2.close.called, "Async credential should be closed"
|
||||
|
||||
|
||||
def test_ui_mode_configuration():
|
||||
"""Test UI mode configuration."""
|
||||
dev_server = DevServer(mode="developer")
|
||||
assert dev_server.mode == "developer"
|
||||
|
||||
user_server = DevServer(mode="user")
|
||||
assert user_server.mode == "user"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_restrictions_in_user_mode():
|
||||
"""Test that developer APIs are restricted in user mode."""
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
# Create servers with different modes
|
||||
dev_server = DevServer(mode="developer")
|
||||
user_server = DevServer(mode="user")
|
||||
|
||||
dev_app = dev_server.create_app()
|
||||
user_app = user_server.create_app()
|
||||
|
||||
dev_client = TestClient(dev_app)
|
||||
user_client = TestClient(user_app)
|
||||
|
||||
# Test 1: Health endpoint should work in both modes
|
||||
assert dev_client.get("/health").status_code == 200
|
||||
assert user_client.get("/health").status_code == 200
|
||||
|
||||
# Test 2: Meta endpoint should reflect correct mode
|
||||
dev_meta = dev_client.get("/meta").json()
|
||||
assert dev_meta["ui_mode"] == "developer"
|
||||
|
||||
user_meta = user_client.get("/meta").json()
|
||||
assert user_meta["ui_mode"] == "user"
|
||||
|
||||
# Test 3: Entity listing should work in both modes
|
||||
assert dev_client.get("/v1/entities").status_code == 200
|
||||
assert user_client.get("/v1/entities").status_code == 200
|
||||
|
||||
# Test 4: Entity info should be restricted in user mode
|
||||
dev_response = dev_client.get("/v1/entities/test_agent/info")
|
||||
assert dev_response.status_code in [200, 404, 500] # Not 403
|
||||
|
||||
user_response = user_client.get("/v1/entities/test_agent/info")
|
||||
assert user_response.status_code == 403
|
||||
error_data = user_response.json()
|
||||
# FastAPI wraps HTTPException detail in 'detail' field
|
||||
error = error_data.get("detail", {}).get("error") or error_data.get("error")
|
||||
assert error is not None
|
||||
assert "developer mode" in error["message"].lower()
|
||||
assert error["code"] == "developer_mode_required"
|
||||
|
||||
# Test 5: Hot reload should be restricted in user mode
|
||||
dev_response = dev_client.post("/v1/entities/test_agent/reload")
|
||||
assert dev_response.status_code in [200, 404, 500] # Not 403
|
||||
|
||||
user_response = user_client.post("/v1/entities/test_agent/reload")
|
||||
assert user_response.status_code == 403
|
||||
error_data = user_response.json()
|
||||
error = error_data.get("detail", {}).get("error") or error_data.get("error")
|
||||
assert "developer mode" in error["message"].lower()
|
||||
|
||||
# Test 6: Deployment endpoints should be restricted in user mode
|
||||
# List deployments (simplest test - no payload needed)
|
||||
user_response = user_client.get("/v1/deployments")
|
||||
assert user_response.status_code == 403
|
||||
error_data = user_response.json()
|
||||
error = error_data.get("detail", {}).get("error") or error_data.get("error")
|
||||
assert "developer mode" in error["message"].lower()
|
||||
|
||||
# Get deployment
|
||||
user_response = user_client.get("/v1/deployments/test-id")
|
||||
assert user_response.status_code == 403
|
||||
|
||||
# Delete deployment
|
||||
user_response = user_client.delete("/v1/deployments/test-id")
|
||||
assert user_response.status_code == 403
|
||||
|
||||
# Test 7: Conversation endpoints should work in both modes
|
||||
dev_response = dev_client.post("/v1/conversations", json={})
|
||||
assert dev_response.status_code == 200
|
||||
|
||||
user_response = user_client.post("/v1/conversations", json={})
|
||||
assert user_response.status_code == 200
|
||||
|
||||
# Test 8: Chat endpoint should work in both modes
|
||||
chat_payload = {"model": "test_agent", "input": "Hello"}
|
||||
dev_response = dev_client.post("/v1/responses", json=chat_payload)
|
||||
assert dev_response.status_code in [200, 404] # 404 if agent doesn't exist
|
||||
|
||||
user_response = user_client.post("/v1/responses", json=chat_payload)
|
||||
assert user_response.status_code in [200, 404]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner
|
||||
async def run_tests():
|
||||
@@ -359,7 +265,7 @@ class WeatherAgent:
|
||||
|
||||
if entities:
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entities[0].id},
|
||||
model=entities[0].id, # model IS the entity_id now!
|
||||
input="test location",
|
||||
stream=False,
|
||||
)
|
||||
@@ -367,44 +273,3 @@ class WeatherAgent:
|
||||
await executor.execute_sync(request)
|
||||
|
||||
asyncio.run(run_tests())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_api_endpoints(test_entities_dir):
|
||||
"""Test checkpoint list and delete API endpoints."""
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
server = DevServer(entities_dir=test_entities_dir)
|
||||
executor = await server._ensure_executor()
|
||||
|
||||
# Create a conversation
|
||||
conversation = executor.conversation_store.create_conversation(metadata={"name": "Test Session"})
|
||||
conv_id = conversation.id
|
||||
|
||||
# Get checkpoint storage and add a checkpoint
|
||||
storage = executor.checkpoint_manager.get_checkpoint_storage(conv_id)
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id="test_checkpoint_1",
|
||||
workflow_id="test_workflow",
|
||||
shared_state={"key": "value"},
|
||||
iteration_count=1,
|
||||
)
|
||||
await storage.save_checkpoint(checkpoint)
|
||||
|
||||
# Test list checkpoints endpoint
|
||||
checkpoints = await storage.list_checkpoints()
|
||||
assert len(checkpoints) == 1
|
||||
assert checkpoints[0].checkpoint_id == "test_checkpoint_1"
|
||||
assert checkpoints[0].workflow_id == "test_workflow"
|
||||
|
||||
# Test delete checkpoint endpoint
|
||||
deleted = await storage.delete_checkpoint("test_checkpoint_1")
|
||||
assert deleted is True
|
||||
|
||||
# Verify checkpoint was deleted
|
||||
remaining = await storage.list_checkpoints()
|
||||
assert len(remaining) == 0
|
||||
|
||||
# Test delete non-existent checkpoint
|
||||
deleted = await storage.delete_checkpoint("nonexistent")
|
||||
assert deleted is False
|
||||
|
||||
Reference in New Issue
Block a user