mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: [BREAKING] Moved to a single get_response and run API (#3379)
* WIP * big update to new ResponseStream model * fixed tests and typing * fixed tests and typing * fixed tools typevar import * fix * mypy fix * mypy fixes and some cleanup * fix missing quoted names * and client * fix imports agui * fix anthropic override * fix agui * fix ag ui * fix import * fix anthropic types * fix mypy * refactoring * updated typing * fix 3.11 * fixes * redid layering of chat clients and agents * redid layering of chat clients and agents * Fix lint, type, and test issues after rebase - Add @overload decorators to AgentProtocol.run() for type compatibility - Add missing docstring params (middleware, function_invocation_configuration) - Fix TODO format (TD002) by adding author tags - Fix broken observability tests from upstream: - Replace non-existent use_instrumentation with direct instantiation - Replace non-existent use_agent_instrumentation with AgentTelemetryLayer mixin - Fix get_streaming_response to use get_response(stream=True) - Add AgentInitializationError import - Update streaming exception tests to match actual behavior * Fix AgentExecutionException import error in test_agents.py - Replace non-existent AgentExecutionException with AgentRunException * Fix test import and asyncio deprecation issues - Add 'tests' to pythonpath in ag-ui pyproject.toml for utils_test_ag_ui import - Replace deprecated asyncio.get_event_loop().run_until_complete with asyncio.run * Fix azure-ai test failures - Update _prepare_options patching to use correct class path - Fix test_to_azure_ai_agent_tools_web_search_missing_connection to clear env vars * Convert ag-ui utils_test_ag_ui.py to conftest.py - Move test utilities to conftest.py for proper pytest discovery - Update all test imports to use conftest instead of utils_test_ag_ui - Remove old utils_test_ag_ui.py file - Revert pythonpath change in pyproject.toml * fix: use relative imports for ag-ui test utilities * fix agui * Rename Bare*Client to Raw*Client and BaseChatClient - Renamed BareChatClient to BaseChatClient (abstract base class) - Renamed BareOpenAIChatClient to RawOpenAIChatClient - Renamed BareOpenAIResponsesClient to RawOpenAIResponsesClient - Renamed BareAzureAIClient to RawAzureAIClient - Added warning docstrings to Raw* classes about layer ordering - Updated README in samples/getting_started/agents/custom with layer docs - Added test for span ordering with function calling * Fix layer ordering: FunctionInvocationLayer before ChatTelemetryLayer This ensures each inner LLM call gets its own telemetry span, resulting in the correct span sequence: chat -> execute_tool -> chat Updated all production clients and test mocks to use correct ordering: - ChatMiddlewareLayer (first) - FunctionInvocationLayer (second) - ChatTelemetryLayer (third) - BaseChatClient/Raw...Client (fourth) * Remove run_stream usage * Fix conversation_id propagation * Python: Add BaseAgent implementation for Claude Agent SDK (#3509) * Added ClaudeAgent implementation * Updated streaming logic * Small updates * Small update * Fixes * Small fix * Naming improvements * Updated imports * Addressed comments * Updated package versions * Update Claude agent connector layering * fix test and plugin * Store function middleware in invocation layer * Fix telemetry streaming and ag-ui tests * Remove legacy ag-ui tests folder * updates * Remove terminate flag from FunctionInvocationContext, use MiddlewareTermination instead - Remove terminate attribute from FunctionInvocationContext - Add result attribute to MiddlewareTermination to carry function results - FunctionMiddlewarePipeline.execute() now lets MiddlewareTermination propagate - _auto_invoke_function captures context.result in exception before re-raising - _try_execute_function_calls catches MiddlewareTermination and sets should_terminate - Fix handoff middleware to append to chat_client.function_middleware directly - Update tests to use raise MiddlewareTermination instead of context.terminate - Add middleware flow documentation in samples/concepts/tools/README.md - Fix ag-ui to use FunctionMiddlewarePipeline instead of removed create_function_middleware_pipeline * fix: remove references to removed terminate flag in purview tests, add type ignore * fix: move _test_utils.py from package to test folder * fix: call get_final_response() to trigger context provider notification in streaming test * fix: correct broken links in tools README * docs: clarify default middleware behavior in summary table * fix: ensure inner stream result hooks are called when using map()/from_awaitable() * Fix mypy type errors * Address PR review comments on observability.py - Remove TODO comment about unconsumed streams, add explanatory note instead - Remove redundant _close_span cleanup hook (already called in _finalize_stream) - Clarify behavior: cleanup hooks run after stream iteration, if stream is not consumed the span remains open until garbage collected * Remove gen_ai.client.operation.duration from span attributes Duration is a metrics-only attribute per OpenTelemetry semantic conventions. It should be recorded to the histogram but not set as a span attribute. * Remove duration from _get_response_attributes, pass directly to _capture_response Duration is a metrics-only attribute. It's now passed directly to _capture_response instead of being included in the attributes dict that gets set on the span. * Remove redundant _close_span cleanup hook in AgentTelemetryLayer _finalize_stream already calls _close_span() in its finally block, so adding it as a separate cleanup hook is redundant. * Use weakref.finalize to close span when stream is garbage collected If a user creates a streaming response but never consumes it, the cleanup hooks won't run. Now we register a weak reference finalizer that will close the span when the stream object is garbage collected, ensuring spans don't leak in this scenario. * Fix _get_finalizers_from_stream to use _result_hooks attribute Renamed function to _get_result_hooks_from_stream and fixed it to look for the _result_hooks attribute which is the correct name in ResponseStream class. * Add missing asyncio import in test_request_info_mixin.py * Fix leftover merge conflict marker in image_generation sample * Update integration tests * Fix integration tests: increase max_iterations from 1 to 2 Tests with tool_choice options require at least 2 iterations: 1. First iteration to get function call and execute the tool 2. Second iteration to get the final text response With max_iterations=1, streaming tests would return early with only the function call/result but no final text content. * Fix duplicate function call error in conversation-based APIs When using conversation_id (for Responses/Assistants APIs), the server already has the function call message from the previous response. We should only send the new function result message, not all messages including the function call which would cause a duplicate ID error. Fix: When conversation_id is set, only send the last message (the tool result) instead of all response.messages. * Add regression test for conversation_id propagation between tool iterations Port test from PR #3664 with updates for new streaming API pattern. Tests that conversation_id is properly updated in options dict during function invocation loop iterations. * Fix tool_choice=required to return after tool execution When tool_choice is 'required', the user's intent is to force exactly one tool call. After the tool executes, return immediately with the function call and result - don't continue to call the model again. This fixes integration tests that were failing with empty text responses because with tool_choice=required, the model would keep returning function calls instead of text. Also adds regression tests for: - conversation_id propagation between tool iterations (from PR #3664) - tool_choice=required returns after tool execution * Document tool_choice behavior in tools README - Add table explaining tool_choice values (auto, none, required) - Explain why tool_choice=required returns immediately after tool execution - Add code example showing the difference between required and auto - Update flow diagram to show the early return path for tool_choice=required * Fix tool_choice=None behavior - don't default to 'auto' Remove the hardcoded default of 'auto' for tool_choice in ChatAgent init. When tool_choice is not specified (None), it will now not be sent to the API, allowing the API's default behavior to be used. Users who want tool_choice='auto' can still explicitly set it either in default_options or at runtime. Fixes #3585 * Fix tool_choice=none should not remove tools In OpenAI Assistants client, tools were not being sent when tool_choice='none'. This was incorrect - tool_choice='none' means the model won't call tools, but tools should still be available in the request (they may be used later in the conversation). Fixes #3585 * Add test for tool_choice=none preserving tools Adds a regression test to ensure that when tool_choice='none' is set but tools are provided, the tools are still sent to the API. This verifies the fix for #3585. * Fix tool_choice=none should not remove tools in all clients Apply the same fix to OpenAI Responses client and Azure AI client: - OpenAI Responses: Remove else block that popped tool_choice/parallel_tool_calls - Azure AI: Remove tool_choice != 'none' check when adding tools When tool_choice='none', the model won't call tools, but tools should still be sent to the API so they're available for future turns. Also update README to clarify tool_choice=required supports multiple tools. Fixes #3585 * Keep tool_choice even when tools is None Move tool_choice processing outside of the 'if tools' block in OpenAI Responses client so tool_choice is sent to the API even when no tools are provided. * Update test to match new parallel_tool_calls behavior Changed test_prepare_options_removes_parallel_tool_calls_when_no_tools to test_prepare_options_preserves_parallel_tool_calls_when_no_tools to reflect that parallel_tool_calls is now preserved even when no tools are present, consistent with the tool_choice behavior. * Fix ChatMessage API and Role enum usage after rebase - Update ChatMessage instantiation to use keyword args (role=, text=, contents=) - Fix Role enum comparisons to use .value for string comparison - Add created_at to AgentResponse in error handling - Fix AgentResponse.from_updates -> from_agent_run_response_updates - Fix DurableAgentStateMessage.from_chat_message to convert Role enum to string - Add Role import where needed * Fix additional ChatMessage API and method name changes - Fix ChatMessage usage in workflow files (use text= instead of contents= for strings) - Fix AgentResponse.from_updates -> from_agent_run_response_updates in workflow files - Fix test files for ChatMessage and Role enum usage * Fix remaining ChatMessage API usage in test files * Fix more ChatMessage and Role API changes in source and test files - Fix ChatMessage in _magentic.py replan method - Fix Role enum comparison in test assertions - Fix remaining test files with old ChatMessage syntax * Fix ChatMessage and Role API changes across packages - Add Role import where missing - Fix ChatMessage signature: positional args to keyword args (role=, text=, contents=) - Fix Role enum comparisons: .role.value instead of .role string - Fix FinishReason enum usage in ag-ui event converters - Rename AgentResponse.from_updates to from_agent_run_response_updates in ag-ui Fixes API compatibility after Types API Review improvements merge * Fix ChatMessage and Role API changes in github_copilot tests * Fix ChatMessage and Role API changes in redis and github_copilot packages - Fix redis provider: Role enum comparison using .value - Fix redis tests: ChatMessage signature and Role comparisons - Fix github_copilot tests: ChatMessage signature and Role comparisons - Update docstring examples in redis chat message store * Fix ChatMessage and Role API changes in devui package - Fix executor: ChatMessage signature change - Fix conversations: Role enum to string conversion in two places - Fix tests: ChatMessage signatures and Role comparisons * Fix ChatMessage and Role API changes in a2a and lab packages - Fix a2a tests: Role comparisons and ChatMessage signatures - Fix lab tau2 source: Role enum comparison in flip_messages, log_messages, sliding_window - Fix lab tau2 tests: ChatMessage signatures and Role comparisons * Remove duplicate test files from ag-ui/tests (tests are in ag_ui_tests) * Fix ChatMessage and Role API changes across packages After rebasing on upstream/main which merged PR #3647 (Types API Review improvements), fix all packages to use the new API: - ChatMessage: Use keyword args (role=, text=, contents=) instead of positional args - Role: Compare using .value attribute since it's now an enum Packages fixed: - ag-ui: Fixed Role value extraction bugs in _message_adapters.py - anthropic: Fixed ChatMessage and Role comparisons in tests - azure-ai: Fixed Role comparison in _client.py - azure-ai-search: Fixed ChatMessage and Role in source/tests - bedrock: Fixed ChatMessage signatures in tests - chatkit: Fixed ChatMessage and Role in source/tests - copilotstudio: Fixed ChatMessage and Role in tests - declarative: Fixed ChatMessage in _executors_agents.py - mem0: Fixed ChatMessage and Role in source/tests - purview: Fixed ChatMessage in source/tests * Fix mypy errors for ChatMessage and Role API changes - durabletask: Use str() fallback in role value extraction - core: Fix ChatMessage in _orchestrator_helpers.py to use keyword args - core: Add type ignore for _conversation_state.py contents deserialization - ag-ui: Fix type ignore comments (call-overload instead of arg-type) - azure-ai-search: Fix get_role_value type hint to accept Any - lab: Move get_role_value to module level with Any type hint * Improve CI test timeout configuration - Increase job timeout from 10 to 15 minutes - Reduce per-test timeout to 60s (was 900s/300s) - Add --timeout_method thread for better timeout handling - Add --timeout-verbose to see which tests are slow - Reduce retries from 3 to 2 and delay from 10s to 5s This ensures individual test timeouts are shorter than the job timeout, providing better visibility when tests hang. With 60s timeout and 2 retries, worst case per test is ~180s. * Fix ChatMessage API usage in docstrings and source - Fix ChatMessage positional args in docstrings: _serialization.py, _threads.py, _middleware.py - Fix ChatMessage in tau2 runner.py - Fix role comparison in _orchestrator_helpers.py to use .value - Fix role comparison in _group_chat.py docstring example - Fix role assertions in test_durable_entities.py to use .value * Revert tool_choice/parallel_tool_calls changes - must be removed when no tools OpenAI API requires tool_choice and parallel_tool_calls to only be present when tools are specified. Restored the logic that removes these options when there are no tools. - Restored check in _chat_client.py to remove tool_choice and parallel_tool_calls when no tools present - Restored same logic in _responses_client.py - Reverted test to expect the correct behavior * fixed issue in tests * fix: resolve merge conflict markers in ag-ui tests * fix: restructure ag-ui tests and fix Role/FinishReason to use string types * fix: streaming function invocation and middleware termination - Refactor streaming function invocation to use get_final_response() on inner streams - Fix MiddlewareTermination to accept result parameter for passing results - Fix _AutoHandoffMiddleware to use MiddlewareTermination instead of context.terminate - Fix AgentMiddlewareLayer.run() to properly forward function/chat middleware - Remove duplicate middleware registration in AgentMiddlewareLayer.__init__ - Fix exception handling in _auto_invoke_function to properly capture termination - Fix mypy errors in core package - Update tests to use stream=True parameter for unified run API * fix all tests command * Refactor integration tests to use pytest fixtures - Merge testutils.py into conftest.py for azurefunctions integration tests - Merge dt_testutils.py into conftest.py for durabletask integration tests - Convert all integration tests to use fixtures instead of direct imports (fixes ModuleNotFoundError with --import-mode=importlib) - Add sample_helper fixture for azurefunctions tests - Add agent_client_factory and orchestration_helper fixtures for durabletask - Integration tests now skip with descriptive messages when services unavailable - Restructure devui tests into tests/devui/ with proper conftest.py - Add test organization guidelines to CODING_STANDARD.md - Remove __init__.py from test directories per pytest best practices * Fix pytest_collection_modifyitems to only skip integration tests The hook was skipping all tests in the test session, not just integration tests. Now it only skips items in the integration_tests directory. * Fix mem0 tests failing on Python 3.13 Use patch.object on the imported module instead of @patch with string path to ensure the mock takes effect regardless of import timing. * fix mem0 * another attempt for mem0 * fix for mem0 * fix mem0 * Increase worker initialization wait time in durabletask tests Increase from 2 to 8 seconds to allow time for: - Python startup and module imports - Azure OpenAI client creation - Agent registration with DTS worker - Worker connection to DTS This helps prevent test failures in CI where the first tests may run before the worker is fully ready to process requests. * Fix streaming test to use ResponseStream with finalizer The _consume_stream method now expects a ResponseStream that can provide a final AgentResponse via get_final_response(). Update the test to use ResponseStream with AgentResponse.from_updates as the finalizer. * Fix MockToolCallingAgent to use new ResponseStream API and update samples * small updates to run_stream to run * fix sub workflow * temp fix for az func test --------- Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
d1205896a1
commit
3dc59c83b5
@@ -0,0 +1,227 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""
|
||||
Message Capture Script - Debug message flow
|
||||
- This script is intended to provide a reference for the types of events
|
||||
that are emitted by the server when agents and workflows are executed
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import http.client
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import uvicorn
|
||||
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 - updated path after samples were moved
|
||||
current_dir = Path(__file__).parent
|
||||
# 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(
|
||||
entities_dir=str(samples_dir.resolve()),
|
||||
host="127.0.0.1",
|
||||
port=8085, # Use different port
|
||||
ui_enabled=False,
|
||||
)
|
||||
|
||||
app = server.get_app()
|
||||
|
||||
server_config = uvicorn.Config(
|
||||
app=app,
|
||||
host="127.0.0.1",
|
||||
port=8085,
|
||||
# log_level="info", # More verbose to see tracing setup
|
||||
)
|
||||
server_instance = uvicorn.Server(server_config)
|
||||
|
||||
def run_server():
|
||||
asyncio.run(server_instance.serve())
|
||||
|
||||
server_thread = threading.Thread(target=run_server, daemon=True)
|
||||
server_thread.start()
|
||||
|
||||
# Wait for server to start
|
||||
time.sleep(5) # Increased wait time
|
||||
|
||||
# Verify server is running with retries
|
||||
max_retries = 10
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
conn = http.client.HTTPConnection("127.0.0.1", 8085, timeout=5)
|
||||
try:
|
||||
conn.request("GET", "/health")
|
||||
response = conn.getresponse()
|
||||
if response.status == 200:
|
||||
break
|
||||
finally:
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
if attempt < max_retries - 1:
|
||||
time.sleep(2)
|
||||
else:
|
||||
raise RuntimeError(f"Server failed to start after {max_retries} attempts: {e}") from e
|
||||
|
||||
return "http://127.0.0.1:8085", server_instance
|
||||
|
||||
|
||||
def capture_agent_stream_with_tracing(client: OpenAI, agent_id: str, scenario: str = "success") -> list[dict[str, Any]]:
|
||||
"""Capture agent streaming events."""
|
||||
|
||||
try:
|
||||
stream = client.responses.create(
|
||||
metadata={"entity_id": agent_id},
|
||||
input="Tell me about the weather in Tokyo. I want details.",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
# Just capture everything as-is
|
||||
if len(events) >= 200: # Increased limit
|
||||
break
|
||||
|
||||
return events
|
||||
|
||||
except Exception as e:
|
||||
# Return error information as events
|
||||
error_event = {
|
||||
"type": "error",
|
||||
"scenario": scenario,
|
||||
"error_message": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"timestamp": time.time(),
|
||||
}
|
||||
return [error_event]
|
||||
|
||||
|
||||
def capture_workflow_stream_with_tracing(
|
||||
client: OpenAI, workflow_id: str, scenario: str = "success"
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Capture workflow streaming events."""
|
||||
|
||||
try:
|
||||
stream = client.responses.create(
|
||||
metadata={"entity_id": workflow_id},
|
||||
input=(
|
||||
"Process this spam detection workflow with multiple emails: "
|
||||
"'Buy now!', 'Hello mom', 'URGENT: Click here!'"
|
||||
),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
# Just capture everything as-is
|
||||
if len(events) >= 200: # Increased limit
|
||||
break
|
||||
|
||||
return events
|
||||
|
||||
except Exception as e:
|
||||
# Return error information as events
|
||||
error_event = {
|
||||
"type": "error",
|
||||
"scenario": scenario,
|
||||
"error_message": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"timestamp": time.time(),
|
||||
"entity_type": "workflow",
|
||||
}
|
||||
return [error_event]
|
||||
|
||||
|
||||
def main():
|
||||
"""Main capture script - testing both success and failure scenarios."""
|
||||
|
||||
# Setup
|
||||
output_dir = Path(__file__).parent / "captured_messages"
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Start server
|
||||
base_url, server_instance = start_server()
|
||||
|
||||
try:
|
||||
# Create OpenAI client for success scenario
|
||||
client = OpenAI(base_url=f"{base_url}/v1", api_key="dummy-key")
|
||||
|
||||
# Discover entities
|
||||
conn = http.client.HTTPConnection("127.0.0.1", 8085, timeout=10)
|
||||
try:
|
||||
conn.request("GET", "/v1/entities")
|
||||
response = conn.getresponse()
|
||||
response_data = response.read().decode("utf-8")
|
||||
entities = json.loads(response_data)["entities"]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
all_results = {}
|
||||
|
||||
# Test each entity
|
||||
for entity in entities:
|
||||
entity_type = entity["type"]
|
||||
entity_id = entity["id"]
|
||||
|
||||
if entity_type == "agent":
|
||||
events = capture_agent_stream_with_tracing(client, entity_id, "success")
|
||||
elif entity_type == "workflow":
|
||||
events = capture_workflow_stream_with_tracing(client, entity_id, "success")
|
||||
else:
|
||||
continue
|
||||
|
||||
all_results[f"{entity_type}_{entity_id}"] = {"entity_info": entity, "events": events}
|
||||
# Save results
|
||||
file_path = output_dir / "entities_stream_events.json"
|
||||
with open(file_path, "w") as f:
|
||||
json.dump(
|
||||
{"timestamp": time.time(), "server_type": "DevServer", "entities_tested": all_results},
|
||||
f,
|
||||
indent=2,
|
||||
default=str,
|
||||
)
|
||||
|
||||
finally:
|
||||
# Cleanup server
|
||||
with contextlib.suppress(Exception):
|
||||
server_instance.should_exit = True
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,553 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Pytest configuration and fixtures for DevUI tests.
|
||||
|
||||
This module provides reusable test fixtures including:
|
||||
- Mock chat clients that don't require API keys
|
||||
- Real workflow event classes from agent_framework
|
||||
- Test agents and executors for workflow testing
|
||||
- Factory functions for test data
|
||||
"""
|
||||
|
||||
import sys
|
||||
from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any, Generic
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from agent_framework import (
|
||||
AgentResponse,
|
||||
AgentResponseUpdate,
|
||||
AgentThread,
|
||||
BaseAgent,
|
||||
BaseChatClient,
|
||||
ChatAgent,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
ResponseStream,
|
||||
)
|
||||
from agent_framework._clients import TOptions_co
|
||||
from agent_framework._workflows._agent_executor import AgentExecutorResponse
|
||||
from agent_framework._workflows._events import (
|
||||
ExecutorCompletedEvent,
|
||||
ExecutorFailedEvent,
|
||||
ExecutorInvokedEvent,
|
||||
WorkflowErrorDetails,
|
||||
)
|
||||
from agent_framework.orchestrations import ConcurrentBuilder, SequentialBuilder
|
||||
|
||||
from agent_framework_devui._discovery import EntityDiscovery
|
||||
from agent_framework_devui._executor import AgentFrameworkExecutor
|
||||
from agent_framework_devui._mapper import MessageMapper
|
||||
from agent_framework_devui.models._openai_custom import AgentFrameworkRequest
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # type: ignore[import] # pragma: no cover
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Mock Chat Clients (from core tests pattern)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class MockChatClient:
|
||||
"""Simple mock chat client that doesn't require API keys.
|
||||
|
||||
Configure responses by setting `responses` or `streaming_responses` lists.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.additional_properties: dict[str, Any] = {}
|
||||
self.call_count: int = 0
|
||||
self.responses: list[ChatResponse] = []
|
||||
self.streaming_responses: list[list[ChatResponseUpdate]] = []
|
||||
|
||||
async def get_response(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage],
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
self.call_count += 1
|
||||
if self.responses:
|
||||
return self.responses.pop(0)
|
||||
return ChatResponse(messages=ChatMessage("assistant", ["test response"]))
|
||||
|
||||
async def get_streaming_response(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage],
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
self.call_count += 1
|
||||
if self.streaming_responses:
|
||||
for update in self.streaming_responses.pop(0):
|
||||
yield update
|
||||
else:
|
||||
yield ChatResponseUpdate(contents=[Content.from_text(text="test streaming response")], role="assistant")
|
||||
|
||||
|
||||
class MockBaseChatClient(BaseChatClient[TOptions_co], Generic[TOptions_co]):
|
||||
"""Full BaseChatClient mock with middleware support.
|
||||
|
||||
Use this when testing features that require the full BaseChatClient interface.
|
||||
This goes through all the middleware, message normalization, etc. - only the
|
||||
actual LLM call is mocked.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs: Any):
|
||||
super().__init__(**kwargs)
|
||||
self.run_responses: list[ChatResponse] = []
|
||||
self.streaming_responses: list[list[ChatResponseUpdate]] = []
|
||||
self.call_count: int = 0
|
||||
self.received_messages: list[list[ChatMessage]] = []
|
||||
|
||||
@override
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[ChatMessage],
|
||||
stream: bool,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
if stream:
|
||||
return self._build_response_stream(self._stream_impl(messages))
|
||||
|
||||
async def _get() -> ChatResponse:
|
||||
self.call_count += 1
|
||||
self.received_messages.append(list(messages))
|
||||
if self.run_responses:
|
||||
return self.run_responses.pop(0)
|
||||
return ChatResponse(messages=ChatMessage("assistant", ["Mock response from ChatAgent"]))
|
||||
|
||||
return _get()
|
||||
|
||||
async def _stream_impl(self, messages: Sequence[ChatMessage]) -> AsyncIterable[ChatResponseUpdate]:
|
||||
self.call_count += 1
|
||||
self.received_messages.append(list(messages))
|
||||
if self.streaming_responses:
|
||||
for update in self.streaming_responses.pop(0):
|
||||
yield update
|
||||
else:
|
||||
# Simulate realistic streaming chunks
|
||||
yield ChatResponseUpdate(contents=[Content.from_text(text="Mock ")], role="assistant")
|
||||
yield ChatResponseUpdate(contents=[Content.from_text(text="streaming ")], role="assistant")
|
||||
yield ChatResponseUpdate(contents=[Content.from_text(text="response ")], role="assistant")
|
||||
yield ChatResponseUpdate(contents=[Content.from_text(text="from ChatAgent")], role="assistant")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Mock Agents (for workflow testing without API keys)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class MockAgent(BaseAgent):
|
||||
"""Mock agent that returns configurable responses without needing a chat client."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
response_text: str = "Mock agent response",
|
||||
streaming_chunks: list[str] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.response_text = response_text
|
||||
self.streaming_chunks = streaming_chunks or [response_text]
|
||||
self.call_count = 0
|
||||
|
||||
def run(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
stream: bool = False,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
|
||||
self.call_count += 1
|
||||
if stream:
|
||||
return self._run_stream(messages=messages, thread=thread, **kwargs)
|
||||
return self._run(messages=messages, thread=thread, **kwargs)
|
||||
|
||||
async def _run(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResponse:
|
||||
self.call_count += 1
|
||||
return AgentResponse(messages=[ChatMessage("assistant", [Content.from_text(text=self.response_text)])])
|
||||
|
||||
def _run_stream(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
|
||||
self.call_count += 1
|
||||
|
||||
async def _iter():
|
||||
for chunk in self.streaming_chunks:
|
||||
yield AgentResponseUpdate(contents=[Content.from_text(text=chunk)], role="assistant")
|
||||
|
||||
return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)
|
||||
|
||||
|
||||
class MockToolCallingAgent(BaseAgent):
|
||||
"""Mock agent that simulates tool calls and results in streaming mode."""
|
||||
|
||||
def __init__(self, **kwargs: Any):
|
||||
super().__init__(**kwargs)
|
||||
self.call_count = 0
|
||||
|
||||
def run(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
stream: bool = False,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
|
||||
self.call_count += 1
|
||||
if stream:
|
||||
return self._run_stream(messages=messages, thread=thread, **kwargs)
|
||||
return self._run(messages=messages, thread=thread, **kwargs)
|
||||
|
||||
async def _run(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResponse:
|
||||
return AgentResponse(messages=[ChatMessage("assistant", ["done"])])
|
||||
|
||||
def _run_stream(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
|
||||
async def _iter() -> AsyncIterable[AgentResponseUpdate]:
|
||||
# First: text
|
||||
yield AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Let me search for that...")],
|
||||
role="assistant",
|
||||
)
|
||||
# Second: tool call
|
||||
yield AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_123",
|
||||
name="search",
|
||||
arguments={"query": "weather"},
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
)
|
||||
# Third: tool result
|
||||
yield AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_result(
|
||||
call_id="call_123",
|
||||
result={"temperature": 72, "condition": "sunny"},
|
||||
)
|
||||
],
|
||||
role="tool",
|
||||
)
|
||||
# Fourth: final text
|
||||
yield AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="The weather is sunny, 72°F.")],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Helper Functions for Test Data Creation
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _create_agent_run_response(text: str = "Test response") -> AgentResponse:
|
||||
"""Create an AgentResponse with the given text."""
|
||||
return AgentResponse(messages=[ChatMessage("assistant", [Content.from_text(text=text)])])
|
||||
|
||||
|
||||
def _create_agent_executor_response(
|
||||
executor_id: str = "test_executor",
|
||||
response_text: str = "Executor response",
|
||||
) -> AgentExecutorResponse:
|
||||
"""Create an AgentExecutorResponse - the type that's nested in ExecutorCompletedEvent.data."""
|
||||
agent_response = _create_agent_run_response(response_text)
|
||||
return AgentExecutorResponse(
|
||||
executor_id=executor_id,
|
||||
agent_response=agent_response,
|
||||
full_conversation=[
|
||||
ChatMessage("user", [Content.from_text(text="User input")]),
|
||||
ChatMessage("assistant", [Content.from_text(text=response_text)]),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Public Factory Functions (for direct import in tests)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def create_agent_run_response(text: str = "Test response") -> AgentResponse:
|
||||
"""Create an AgentResponse with the given text."""
|
||||
return _create_agent_run_response(text)
|
||||
|
||||
|
||||
def create_executor_invoked_event(executor_id: str = "test_executor") -> ExecutorInvokedEvent:
|
||||
"""Create an ExecutorInvokedEvent."""
|
||||
return ExecutorInvokedEvent(executor_id=executor_id)
|
||||
|
||||
|
||||
def create_executor_completed_event(
|
||||
executor_id: str = "test_executor",
|
||||
with_agent_response: bool = True,
|
||||
) -> ExecutorCompletedEvent:
|
||||
"""Create an ExecutorCompletedEvent with realistic nested data.
|
||||
|
||||
This creates the exact data structure that caused the serialization bug:
|
||||
ExecutorCompletedEvent.data contains AgentExecutorResponse which contains
|
||||
AgentResponse and ChatMessage objects (SerializationMixin, not Pydantic).
|
||||
"""
|
||||
data = _create_agent_executor_response(executor_id) if with_agent_response else {"simple": "dict"}
|
||||
return ExecutorCompletedEvent(executor_id=executor_id, data=data)
|
||||
|
||||
|
||||
def create_executor_failed_event(
|
||||
executor_id: str = "test_executor",
|
||||
error_message: str = "Test error",
|
||||
) -> ExecutorFailedEvent:
|
||||
"""Create an ExecutorFailedEvent."""
|
||||
details = WorkflowErrorDetails(error_type="TestError", message=error_message)
|
||||
return ExecutorFailedEvent(executor_id=executor_id, details=details)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Pytest Fixtures
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mapper() -> MessageMapper:
|
||||
"""Create a fresh MessageMapper for each test."""
|
||||
return MessageMapper()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_request() -> AgentFrameworkRequest:
|
||||
"""Create a standard test request."""
|
||||
return AgentFrameworkRequest(
|
||||
metadata={"entity_id": "test_agent"},
|
||||
input="Test input",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_client() -> MockChatClient:
|
||||
"""Create a mock chat client."""
|
||||
return MockChatClient()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_base_chat_client() -> MockBaseChatClient:
|
||||
"""Create a mock BaseChatClient."""
|
||||
return MockBaseChatClient()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent() -> MockAgent:
|
||||
"""Create a mock agent."""
|
||||
return MockAgent(id="test_agent", name="TestAgent", response_text="Mock agent response")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_agent() -> MockToolCallingAgent:
|
||||
"""Create a mock agent that simulates tool calls."""
|
||||
return MockToolCallingAgent(id="tool_agent", name="ToolAgent")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent_run_response() -> AgentResponse:
|
||||
"""Create an AgentResponse with default text."""
|
||||
return _create_agent_run_response()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_completed_event() -> ExecutorCompletedEvent:
|
||||
"""Create an ExecutorCompletedEvent with realistic nested data.
|
||||
|
||||
This creates the exact data structure that caused the serialization bug:
|
||||
ExecutorCompletedEvent.data contains AgentExecutorResponse which contains
|
||||
AgentResponse and ChatMessage objects (SerializationMixin, not Pydantic).
|
||||
"""
|
||||
data = _create_agent_executor_response("test_executor")
|
||||
return ExecutorCompletedEvent(executor_id="test_executor", data=data)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_invoked_event() -> ExecutorInvokedEvent:
|
||||
"""Create an ExecutorInvokedEvent."""
|
||||
return ExecutorInvokedEvent(executor_id="test_executor")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_failed_event() -> ExecutorFailedEvent:
|
||||
"""Create an ExecutorFailedEvent."""
|
||||
details = WorkflowErrorDetails(error_type="TestError", message="Test error")
|
||||
return ExecutorFailedEvent(executor_id="test_executor", details=details)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_entities_dir() -> str:
|
||||
"""Use the samples directory which has proper entity structure."""
|
||||
current_dir = Path(__file__).parent
|
||||
# Navigate to python/samples/getting_started/devui
|
||||
samples_dir = current_dir.parent.parent.parent.parent / "samples" / "getting_started" / "devui"
|
||||
return str(samples_dir.resolve())
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Async Fixtures for Executor/Workflow Setup
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def executor_with_real_agent() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient]:
|
||||
"""Create an executor with a REAL ChatAgent using mock chat client.
|
||||
|
||||
This tests the full execution pipeline:
|
||||
- Real ChatAgent class
|
||||
- Real message handling and normalization
|
||||
- Real middleware pipeline
|
||||
- Only the LLM call is mocked
|
||||
|
||||
Returns tuple of (executor, entity_id, mock_client) so tests can access all components.
|
||||
"""
|
||||
mock_client = MockBaseChatClient()
|
||||
discovery = EntityDiscovery(None)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
# Create a REAL ChatAgent with mock client
|
||||
agent = ChatAgent(
|
||||
id="test_chat_agent",
|
||||
name="Test Chat Agent",
|
||||
description="A real ChatAgent for testing execution flow",
|
||||
chat_client=mock_client,
|
||||
system_message="You are a helpful test assistant.",
|
||||
)
|
||||
|
||||
# Register the real agent
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, source="test")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
return executor, entity_info.id, mock_client
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def sequential_workflow() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient, Any]:
|
||||
"""Create a realistic sequential workflow (Writer -> Reviewer).
|
||||
|
||||
This provides a reusable multi-agent workflow that:
|
||||
- Chains 2 ChatAgents sequentially
|
||||
- Writer generates content, Reviewer provides feedback
|
||||
- Pre-configures mock responses for both agents
|
||||
|
||||
Returns tuple of (executor, entity_id, mock_client, workflow) for test access.
|
||||
"""
|
||||
mock_client = MockBaseChatClient()
|
||||
mock_client.run_responses = [
|
||||
ChatResponse(messages=ChatMessage("assistant", ["Here's the draft content about the topic."])),
|
||||
ChatResponse(messages=ChatMessage("assistant", ["Review: Content is clear and well-structured."])),
|
||||
]
|
||||
|
||||
writer = ChatAgent(
|
||||
id="writer",
|
||||
name="Writer",
|
||||
description="Content writer agent",
|
||||
chat_client=mock_client,
|
||||
system_message="You are a content writer. Create clear, engaging content.",
|
||||
)
|
||||
reviewer = ChatAgent(
|
||||
id="reviewer",
|
||||
name="Reviewer",
|
||||
description="Content reviewer agent",
|
||||
chat_client=mock_client,
|
||||
system_message="You are a reviewer. Provide constructive feedback.",
|
||||
)
|
||||
|
||||
workflow = SequentialBuilder().participants([writer, reviewer]).build()
|
||||
|
||||
discovery = EntityDiscovery(None)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
entity_info = await discovery.create_entity_info_from_object(workflow, entity_type="workflow", source="test")
|
||||
discovery.register_entity(entity_info.id, entity_info, workflow)
|
||||
|
||||
return executor, entity_info.id, mock_client, workflow
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def concurrent_workflow() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient, Any]:
|
||||
"""Create a realistic concurrent workflow (Researcher | Analyst | Summarizer).
|
||||
|
||||
This provides a reusable fan-out/fan-in workflow that:
|
||||
- Runs 3 ChatAgents in parallel
|
||||
- Each agent processes the same input independently
|
||||
- Pre-configures mock responses for all agents
|
||||
|
||||
Returns tuple of (executor, entity_id, mock_client, workflow) for test access.
|
||||
"""
|
||||
mock_client = MockBaseChatClient()
|
||||
mock_client.run_responses = [
|
||||
ChatResponse(messages=ChatMessage("assistant", ["Research findings: Key data points identified."])),
|
||||
ChatResponse(messages=ChatMessage("assistant", ["Analysis: Trends indicate positive growth."])),
|
||||
ChatResponse(messages=ChatMessage("assistant", ["Summary: Overall outlook is favorable."])),
|
||||
]
|
||||
|
||||
researcher = ChatAgent(
|
||||
id="researcher",
|
||||
name="Researcher",
|
||||
description="Research agent",
|
||||
chat_client=mock_client,
|
||||
system_message="You are a researcher. Find key data and insights.",
|
||||
)
|
||||
analyst = ChatAgent(
|
||||
id="analyst",
|
||||
name="Analyst",
|
||||
description="Analysis agent",
|
||||
chat_client=mock_client,
|
||||
system_message="You are an analyst. Identify trends and patterns.",
|
||||
)
|
||||
summarizer = ChatAgent(
|
||||
id="summarizer",
|
||||
name="Summarizer",
|
||||
description="Summary agent",
|
||||
chat_client=mock_client,
|
||||
system_message="You are a summarizer. Provide concise summaries.",
|
||||
)
|
||||
|
||||
workflow = ConcurrentBuilder().participants([researcher, analyst, summarizer]).build()
|
||||
|
||||
discovery = EntityDiscovery(None)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
entity_info = await discovery.create_entity_info_from_object(workflow, entity_type="workflow", source="test")
|
||||
discovery.register_entity(entity_info.id, entity_info, workflow)
|
||||
|
||||
return executor, entity_info.id, mock_client, workflow
|
||||
@@ -0,0 +1,445 @@
|
||||
# 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,
|
||||
RequestInfoEvent,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
WorkflowStatusEvent,
|
||||
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."""
|
||||
|
||||
def __init__(self, id: str) -> None:
|
||||
super().__init__(id=id)
|
||||
self._data_value: str | None = None
|
||||
|
||||
@handler
|
||||
async def process(self, data: WorkflowTestData, ctx: WorkflowContext) -> None:
|
||||
"""Process data and request approval."""
|
||||
self._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."""
|
||||
value = self._data_value or ""
|
||||
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={}, 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={},
|
||||
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={},
|
||||
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={},
|
||||
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={},
|
||||
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.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={}, 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=True) 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={}, 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={},
|
||||
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.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=True) 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(WorkflowTestData(value="test"), stream=True):
|
||||
if isinstance(event, RequestInfoEvent):
|
||||
saw_request_event = True
|
||||
# Wait for IDLE_WITH_PENDING_REQUESTS status (comes after checkpoint creation)
|
||||
if isinstance(event, WorkflowStatusEvent) and "IDLE_WITH_PENDING_REQUESTS" in str(event.state):
|
||||
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
|
||||
@@ -0,0 +1,379 @@
|
||||
# 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 AgentResponse, ChatMessage, Content
|
||||
|
||||
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(self, messages=None, *, stream: bool = False, thread=None, **kwargs):
|
||||
"""Mock run method with streaming support."""
|
||||
if stream:
|
||||
|
||||
async def _stream():
|
||||
yield AgentResponse(
|
||||
messages=[ChatMessage(role="assistant", contents=[Content.from_text(text="Test response")])],
|
||||
)
|
||||
|
||||
return _stream()
|
||||
return AgentResponse(
|
||||
messages=[ChatMessage(role="assistant", contents=[Content.from_text(text="Test response")])],
|
||||
)
|
||||
|
||||
|
||||
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 AgentResponse, ChatMessage, Role, Content
|
||||
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(self, messages=None, *, stream: bool = False, thread=None, **kwargs):
|
||||
if stream:
|
||||
async def _stream():
|
||||
yield AgentResponse(
|
||||
messages=[ChatMessage(role="assistant", content=[Content.from_text(text="Test")])],
|
||||
inner_messages=[],
|
||||
)
|
||||
return _stream()
|
||||
return AgentResponse(
|
||||
messages=[ChatMessage(role="assistant", content=[Content.from_text(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"]
|
||||
@@ -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 = await 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 = await 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
|
||||
|
||||
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="user", contents=[{"type": "text", "text": "What's the weather in SF?"}]),
|
||||
ChatMessage(
|
||||
role="assistant",
|
||||
contents=[
|
||||
{
|
||||
"type": "function_call",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "San Francisco"}',
|
||||
"call_id": "call_test123",
|
||||
}
|
||||
],
|
||||
),
|
||||
ChatMessage(
|
||||
role="tool",
|
||||
contents=[
|
||||
{
|
||||
"type": "function_result",
|
||||
"call_id": "call_test123",
|
||||
"result": '{"temperature": 65, "condition": "sunny"}',
|
||||
}
|
||||
],
|
||||
),
|
||||
ChatMessage(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
|
||||
|
||||
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="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="
|
||||
@@ -0,0 +1,351 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Focused tests for entity discovery functionality."""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from agent_framework_devui._discovery import EntityDiscovery
|
||||
|
||||
# Note: test_entities_dir fixture is provided by conftest.py
|
||||
|
||||
|
||||
async def test_discover_agents(test_entities_dir):
|
||||
"""Test that agent discovery works and returns valid agent entities."""
|
||||
discovery = EntityDiscovery(test_entities_dir)
|
||||
entities = await discovery.discover_entities()
|
||||
|
||||
agents = [e for e in entities if e.type == "agent"]
|
||||
|
||||
# Test that we can discover agents (not specific count)
|
||||
assert len(agents) > 0, "Should discover at least one agent"
|
||||
|
||||
# Test agent structure/properties
|
||||
for agent in agents:
|
||||
assert agent.id, "Agent should have an ID"
|
||||
assert agent.name, "Agent should have a name"
|
||||
assert agent.type == "agent", "Should be identified as agent type"
|
||||
assert hasattr(agent, "description"), "Agent should have description attribute"
|
||||
|
||||
|
||||
async def test_discover_workflows(test_entities_dir):
|
||||
"""Test that workflow discovery works and returns valid workflow entities."""
|
||||
discovery = EntityDiscovery(test_entities_dir)
|
||||
entities = await discovery.discover_entities()
|
||||
|
||||
workflows = [e for e in entities if e.type == "workflow"]
|
||||
|
||||
# Test that we can discover workflows (not specific count)
|
||||
assert len(workflows) > 0, "Should discover at least one workflow"
|
||||
|
||||
# Test workflow structure/properties
|
||||
for workflow in workflows:
|
||||
assert workflow.id, "Workflow should have an ID"
|
||||
assert workflow.name, "Workflow should have a name"
|
||||
assert workflow.type == "workflow", "Should be identified as workflow type"
|
||||
assert hasattr(workflow, "description"), "Workflow should have description attribute"
|
||||
|
||||
|
||||
async def test_empty_directory():
|
||||
"""Test discovery with empty directory."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
discovery = EntityDiscovery(temp_dir)
|
||||
entities = await discovery.discover_entities()
|
||||
|
||||
assert len(entities) == 0
|
||||
|
||||
|
||||
async def test_discovery_accepts_agents_with_only_run():
|
||||
"""Test that discovery accepts agents with only run() method.
|
||||
|
||||
With lazy loading, entities with only __init__.py are discovered
|
||||
but marked as "unknown" type until loaded.
|
||||
"""
|
||||
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 AgentResponse, AgentThread, ChatMessage, Role, Content
|
||||
|
||||
class NonStreamingAgent:
|
||||
id = "non_streaming"
|
||||
name = "Non-Streaming Agent"
|
||||
description = "Agent with run() method"
|
||||
|
||||
async def run(self, messages=None, *, thread=None, **kwargs):
|
||||
return AgentResponse(
|
||||
messages=[ChatMessage(
|
||||
role="assistant",
|
||||
contents=[Content.from_text(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()
|
||||
|
||||
# With lazy loading, entity is discovered but type is "unknown"
|
||||
# (no agent.py or workflow.py to detect type from)
|
||||
assert len(entities) == 1
|
||||
entity = entities[0]
|
||||
assert entity.id == "non_streaming_agent"
|
||||
assert entity.type == "unknown" # Type not yet determined
|
||||
assert entity.tools == [] # Sparse metadata
|
||||
|
||||
# Trigger lazy loading to get full metadata
|
||||
agent_obj = await discovery.load_entity(entity.id)
|
||||
assert agent_obj is not None
|
||||
|
||||
# Now check enriched metadata after loading
|
||||
enriched = discovery.get_entity_info(entity.id)
|
||||
assert enriched.type == "agent" # Now correctly identified
|
||||
assert enriched.name == "Non-Streaming Agent"
|
||||
|
||||
|
||||
async def test_lazy_loading():
|
||||
"""Test that entities are loaded on-demand, not at discovery time."""
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create test workflow
|
||||
workflow_dir = temp_path / "test_workflow"
|
||||
workflow_dir.mkdir()
|
||||
(workflow_dir / "workflow.py").write_text("""
|
||||
from agent_framework import WorkflowBuilder, FunctionExecutor
|
||||
|
||||
# Create a simple workflow with a start executor
|
||||
def test_func(input: str) -> str:
|
||||
return f"Processed: {input}"
|
||||
|
||||
builder = WorkflowBuilder()
|
||||
executor = FunctionExecutor(id="test_executor", func=test_func)
|
||||
builder.set_start_executor(executor)
|
||||
workflow = builder.build()
|
||||
""")
|
||||
|
||||
discovery = EntityDiscovery(str(temp_path))
|
||||
|
||||
# Discovery should NOT import module
|
||||
entities = await discovery.discover_entities()
|
||||
assert len(entities) == 1
|
||||
assert entities[0].id == "test_workflow"
|
||||
assert entities[0].type == "workflow" # Type detected from filename
|
||||
assert entities[0].tools == [] # Sparse metadata (not loaded yet)
|
||||
|
||||
# Entity should NOT be in loaded_objects yet
|
||||
assert discovery.get_entity_object("test_workflow") is None
|
||||
|
||||
# Trigger lazy load
|
||||
workflow_obj = await discovery.load_entity("test_workflow")
|
||||
assert workflow_obj is not None
|
||||
|
||||
# Now in cache
|
||||
assert discovery.get_entity_object("test_workflow") is workflow_obj
|
||||
|
||||
# Second load is instant (from cache)
|
||||
workflow_obj2 = await discovery.load_entity("test_workflow")
|
||||
assert workflow_obj2 is workflow_obj # Same object
|
||||
|
||||
|
||||
async def test_type_detection():
|
||||
"""Test that entity types are detected from filenames."""
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create workflow with workflow.py
|
||||
workflow_dir = temp_path / "my_workflow"
|
||||
workflow_dir.mkdir()
|
||||
(workflow_dir / "workflow.py").write_text("""
|
||||
from agent_framework import WorkflowBuilder, FunctionExecutor
|
||||
|
||||
def test_func(input: str) -> str:
|
||||
return f"Processed: {input}"
|
||||
|
||||
builder = WorkflowBuilder()
|
||||
executor = FunctionExecutor(id="test_executor", func=test_func)
|
||||
builder.set_start_executor(executor)
|
||||
workflow = builder.build()
|
||||
""")
|
||||
|
||||
# Create agent with agent.py
|
||||
agent_dir = temp_path / "my_agent"
|
||||
agent_dir.mkdir()
|
||||
(agent_dir / "agent.py").write_text("""
|
||||
from agent_framework import AgentResponse, AgentThread, ChatMessage, Role, TextContent
|
||||
|
||||
class TestAgent:
|
||||
name = "Test Agent"
|
||||
|
||||
async def run(self, messages=None, *, thread=None, **kwargs):
|
||||
return AgentResponse(
|
||||
messages=[ChatMessage(role="assistant", contents=[Content.from_text(text="test")])],
|
||||
response_id="test"
|
||||
)
|
||||
|
||||
def get_new_thread(self, **kwargs):
|
||||
return AgentThread()
|
||||
|
||||
agent = TestAgent()
|
||||
""")
|
||||
|
||||
# Create ambiguous entity with __init__.py only
|
||||
unknown_dir = temp_path / "my_thing"
|
||||
unknown_dir.mkdir()
|
||||
(unknown_dir / "__init__.py").write_text("# thing")
|
||||
|
||||
discovery = EntityDiscovery(str(temp_path))
|
||||
entities = await discovery.discover_entities()
|
||||
|
||||
# Check types detected correctly
|
||||
by_id = {e.id: e for e in entities}
|
||||
|
||||
assert by_id["my_workflow"].type == "workflow"
|
||||
assert by_id["my_agent"].type == "agent"
|
||||
assert by_id["my_thing"].type == "unknown"
|
||||
|
||||
|
||||
async def test_hot_reload():
|
||||
"""Test that invalidate_entity() enables hot reload."""
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create workflow
|
||||
workflow_dir = temp_path / "test_workflow"
|
||||
workflow_dir.mkdir()
|
||||
workflow_file = workflow_dir / "workflow.py"
|
||||
workflow_file.write_text("""
|
||||
from agent_framework import WorkflowBuilder, FunctionExecutor
|
||||
|
||||
def test_func(input: str) -> str:
|
||||
return "v1"
|
||||
|
||||
builder = WorkflowBuilder()
|
||||
executor = FunctionExecutor(id="test_executor", func=test_func)
|
||||
builder.set_start_executor(executor)
|
||||
workflow = builder.build()
|
||||
""")
|
||||
|
||||
discovery = EntityDiscovery(str(temp_path))
|
||||
await discovery.discover_entities()
|
||||
|
||||
# Load entity
|
||||
workflow1 = await discovery.load_entity("test_workflow")
|
||||
assert workflow1 is not None
|
||||
|
||||
# Modify file to create a different workflow
|
||||
workflow_file.write_text("""
|
||||
from agent_framework import WorkflowBuilder, FunctionExecutor
|
||||
|
||||
def test_func(input: str) -> str:
|
||||
return "v2"
|
||||
|
||||
def test_func2(input: str) -> str:
|
||||
return "v2_extra"
|
||||
|
||||
builder = WorkflowBuilder()
|
||||
executor1 = FunctionExecutor(id="test_executor", func=test_func)
|
||||
executor2 = FunctionExecutor(id="test_executor2", func=test_func2)
|
||||
builder.set_start_executor(executor1)
|
||||
builder.add_edge(executor1, executor2)
|
||||
workflow = builder.build()
|
||||
""")
|
||||
|
||||
# Without invalidation, gets cached version
|
||||
workflow2 = await discovery.load_entity("test_workflow")
|
||||
assert workflow2 is workflow1 # Same object (cached)
|
||||
# Old workflow has 1 executor
|
||||
assert len(workflow2.get_executors_list()) == 1
|
||||
|
||||
# Invalidate cache
|
||||
discovery.invalidate_entity("test_workflow")
|
||||
|
||||
# Now reloads from disk
|
||||
workflow3 = await discovery.load_entity("test_workflow")
|
||||
assert workflow3 is not workflow1 # Different object
|
||||
# New workflow has 2 executors
|
||||
assert len(workflow3.get_executors_list()) == 2
|
||||
|
||||
|
||||
async def test_in_memory_entities_bypass_lazy_loading():
|
||||
"""Test that in-memory entities work as before (no lazy loading needed)."""
|
||||
from agent_framework import FunctionExecutor, WorkflowBuilder
|
||||
|
||||
# Create in-memory workflow
|
||||
def test_func(input: str) -> str:
|
||||
return f"Processed: {input}"
|
||||
|
||||
builder = WorkflowBuilder()
|
||||
executor = FunctionExecutor(id="test_executor", func=test_func)
|
||||
builder.set_start_executor(executor)
|
||||
workflow = builder.build()
|
||||
|
||||
discovery = EntityDiscovery()
|
||||
|
||||
# Register in-memory entity
|
||||
entity_info = await discovery.create_entity_info_from_object(workflow, entity_type="workflow", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, workflow)
|
||||
|
||||
# Should be immediately available (no lazy loading)
|
||||
loaded = discovery.get_entity_object(entity_info.id)
|
||||
assert loaded is workflow
|
||||
|
||||
# load_entity() should return immediately from cache
|
||||
loaded2 = await discovery.load_entity(entity_info.id)
|
||||
assert loaded2 is workflow # Same object (cache hit)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner
|
||||
async def run_tests():
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create test files
|
||||
agent_file = temp_path / "test_agent.py"
|
||||
agent_file.write_text("""
|
||||
class WeatherAgent:
|
||||
name = "Weather Agent"
|
||||
description = "Gets weather information"
|
||||
|
||||
def run(self, input_str, *, stream: bool = False, thread=None, **kwargs):
|
||||
return f"Weather in {input_str}"
|
||||
""")
|
||||
|
||||
workflow_file = temp_path / "test_workflow.py"
|
||||
workflow_file.write_text("""
|
||||
class DataWorkflow:
|
||||
name = "Data Processing Workflow"
|
||||
description = "Processes data"
|
||||
|
||||
def run(self, data):
|
||||
return f"Processed {data}"
|
||||
""")
|
||||
|
||||
discovery = EntityDiscovery(str(temp_path))
|
||||
await discovery.discover_entities()
|
||||
|
||||
asyncio.run(run_tests())
|
||||
@@ -0,0 +1,785 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Focused tests for execution flow functionality.
|
||||
|
||||
Tests include:
|
||||
- Entity discovery and info retrieval
|
||||
- Agent execution (sync and streaming) using real ChatAgent with mock LLM
|
||||
- Workflow execution using real WorkflowBuilder with FunctionExecutor
|
||||
- Edge cases like non-streaming agents
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentExecutor, ChatAgent, FunctionExecutor, WorkflowBuilder
|
||||
|
||||
# Import mock classes from conftest for direct use in some tests
|
||||
from conftest import MockBaseChatClient
|
||||
|
||||
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 AgentFrameworkRequest
|
||||
|
||||
# =============================================================================
|
||||
# Local Fixtures (module-specific)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def executor(test_entities_dir):
|
||||
"""Create configured executor."""
|
||||
discovery = EntityDiscovery(test_entities_dir)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
# Discover entities
|
||||
await executor.discover_entities()
|
||||
|
||||
return executor
|
||||
|
||||
|
||||
async def test_executor_entity_discovery(executor):
|
||||
"""Test executor entity discovery."""
|
||||
entities = await executor.discover_entities()
|
||||
|
||||
# Should find entities from samples directory
|
||||
assert len(entities) > 0, "Should discover at least one entity"
|
||||
|
||||
entity_types = [e.type for e in entities]
|
||||
assert "agent" in entity_types, "Should find at least one agent"
|
||||
assert "workflow" in entity_types, "Should find at least one workflow"
|
||||
|
||||
# Test entity structure
|
||||
for entity in entities:
|
||||
assert entity.id, "Entity should have an ID"
|
||||
assert entity.name, "Entity should have a name"
|
||||
# Entities with only an `__init__.py` file cannot have their type determined
|
||||
# until the module is imported during lazy loading. This is why 'unknown' type exists.
|
||||
assert entity.type in ["agent", "workflow", "unknown"], (
|
||||
"Entity should have valid type (unknown allowed during discovery phase)"
|
||||
)
|
||||
|
||||
|
||||
async def test_executor_get_entity_info(executor):
|
||||
"""Test getting entity info by ID."""
|
||||
entities = await executor.discover_entities()
|
||||
entity_id = entities[0].id
|
||||
|
||||
entity_info = executor.get_entity_info(entity_id)
|
||||
assert entity_info is not None
|
||||
assert entity_info.id == entity_id
|
||||
assert entity_info.type in ["agent", "workflow", "unknown"]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Agent Execution Tests (using real ChatAgent with mock LLM)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_agent_sync_execution(executor_with_real_agent):
|
||||
"""Test synchronous agent execution with REAL ChatAgent (mock LLM).
|
||||
|
||||
This tests the full execution pipeline without needing an API key:
|
||||
- Real ChatAgent class with middleware
|
||||
- Real message normalization
|
||||
- Mock chat client for LLM calls
|
||||
"""
|
||||
executor, entity_id, mock_client = executor_with_real_agent
|
||||
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": 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
|
||||
|
||||
# Verify mock client was called
|
||||
assert mock_client.call_count == 1
|
||||
|
||||
|
||||
async def test_agent_sync_execution_respects_model_field(executor_with_real_agent):
|
||||
"""Test synchronous execution respects the model field in the response."""
|
||||
executor, entity_id, mock_client = executor_with_real_agent
|
||||
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_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"
|
||||
assert response.object == "response"
|
||||
assert len(response.output) > 0
|
||||
|
||||
|
||||
async def test_chat_client_receives_correct_messages(executor_with_real_agent):
|
||||
"""Verify the mock chat client receives properly formatted messages.
|
||||
|
||||
This tests that the REAL ChatAgent properly:
|
||||
- Normalizes input messages
|
||||
- Formats messages for the chat client
|
||||
"""
|
||||
executor, entity_id, mock_client = executor_with_real_agent
|
||||
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_id},
|
||||
input="What is 2+2?",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
await executor.execute_sync(request)
|
||||
|
||||
# Verify chat client was called
|
||||
assert mock_client.call_count == 1
|
||||
|
||||
# Verify messages were received
|
||||
assert len(mock_client.received_messages) == 1
|
||||
messages = mock_client.received_messages[0]
|
||||
|
||||
# Should have at least one message
|
||||
assert len(messages) >= 1, f"Expected messages, got: {messages}"
|
||||
|
||||
# Verify the input text is present in the messages
|
||||
all_text = " ".join(m.text or "" for m in messages)
|
||||
assert "2+2" in all_text, f"Expected '2+2' in messages, got text: '{all_text}'"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Workflow Execution Tests (using real WorkflowBuilder with FunctionExecutor)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_workflow_streaming_execution():
|
||||
"""Test workflow streaming execution with REAL WorkflowBuilder and FunctionExecutor.
|
||||
|
||||
This tests the full workflow execution pipeline without needing an API key.
|
||||
Uses a simple function-based workflow that processes input.
|
||||
"""
|
||||
|
||||
# Create a simple workflow using real agent_framework classes
|
||||
def process_input(input_data: str) -> str:
|
||||
return f"Processed: {input_data}"
|
||||
|
||||
builder = WorkflowBuilder(name="Test Workflow", description="Test workflow for execution")
|
||||
start_executor = FunctionExecutor(id="process", func=process_input)
|
||||
builder.set_start_executor(start_executor)
|
||||
workflow = builder.build()
|
||||
|
||||
# Create executor and register workflow
|
||||
discovery = EntityDiscovery(None)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
entity_info = await discovery.create_entity_info_from_object(workflow, entity_type="workflow", source="test")
|
||||
discovery.register_entity(entity_info.id, entity_info, workflow)
|
||||
|
||||
# Execute workflow
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_info.id},
|
||||
input="hello workflow",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in executor.execute_streaming(request):
|
||||
events.append(event)
|
||||
|
||||
# Should get events from workflow execution
|
||||
assert len(events) > 0, "Should receive events from workflow"
|
||||
|
||||
# Check for workflow-specific events or completion
|
||||
event_types = [getattr(e, "type", None) for e in events]
|
||||
assert any(t is not None for t in event_types), f"Should have typed events, got: {event_types}"
|
||||
|
||||
|
||||
async def test_workflow_sync_execution():
|
||||
"""Test synchronous workflow execution."""
|
||||
|
||||
def echo(text: str) -> str:
|
||||
return f"Echo: {text}"
|
||||
|
||||
builder = WorkflowBuilder(name="Echo Workflow", description="Simple echo workflow")
|
||||
start_executor = FunctionExecutor(id="echo", func=echo)
|
||||
builder.set_start_executor(start_executor)
|
||||
workflow = builder.build()
|
||||
|
||||
# Create executor and register workflow
|
||||
discovery = EntityDiscovery(None)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
entity_info = await discovery.create_entity_info_from_object(workflow, entity_type="workflow", source="test")
|
||||
discovery.register_entity(entity_info.id, entity_info, workflow)
|
||||
|
||||
# Execute workflow synchronously
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_info.id},
|
||||
input="test input",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
response = await executor.execute_sync(request)
|
||||
|
||||
# Should get a valid response
|
||||
assert response.object == "response"
|
||||
assert len(response.output) > 0
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Full Pipeline Serialization Tests (Run + Map + JSON)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_full_pipeline_agent_events_are_json_serializable(executor_with_real_agent):
|
||||
"""CRITICAL TEST: Verify ALL events from agent execution can be JSON serialized.
|
||||
|
||||
This tests the exact code path that the server uses:
|
||||
1. Execute agent via executor.execute_streaming()
|
||||
2. Each event is converted by the mapper
|
||||
3. Server calls model_dump_json() on each event for SSE
|
||||
|
||||
If any event contains non-serializable objects (like AgentResponse),
|
||||
this test will fail - catching the bug before it hits production.
|
||||
"""
|
||||
executor, entity_id, mock_client = executor_with_real_agent
|
||||
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_id},
|
||||
input="Test message for serialization",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
events = []
|
||||
serialization_errors = []
|
||||
|
||||
async for event in executor.execute_streaming(request):
|
||||
events.append(event)
|
||||
|
||||
# This is EXACTLY what the server does before sending SSE
|
||||
try:
|
||||
if hasattr(event, "model_dump_json"):
|
||||
json_str = event.model_dump_json()
|
||||
assert json_str is not None
|
||||
assert len(json_str) > 0
|
||||
except Exception as e:
|
||||
serialization_errors.append(f"Event type={getattr(event, 'type', 'unknown')}: {e}")
|
||||
|
||||
# Should have received events
|
||||
assert len(events) > 0, "Should receive events from agent execution"
|
||||
|
||||
# NO serialization errors allowed
|
||||
assert len(serialization_errors) == 0, f"Found {len(serialization_errors)} serialization errors:\n" + "\n".join(
|
||||
serialization_errors
|
||||
)
|
||||
|
||||
|
||||
async def test_full_pipeline_workflow_events_are_json_serializable():
|
||||
"""CRITICAL TEST: Verify ALL events from workflow execution can be JSON serialized.
|
||||
|
||||
This is particularly important for workflows with AgentExecutor because:
|
||||
- AgentExecutor produces ExecutorCompletedEvent with AgentExecutorResponse
|
||||
- AgentExecutorResponse contains AgentResponse and ChatMessage objects
|
||||
- These are SerializationMixin objects, not Pydantic, which caused the original bug
|
||||
|
||||
This test ensures the ENTIRE streaming pipeline works end-to-end.
|
||||
"""
|
||||
# Create a workflow with AgentExecutor (the problematic case)
|
||||
mock_client = MockBaseChatClient()
|
||||
agent = ChatAgent(
|
||||
id="serialization_test_agent",
|
||||
name="Serialization Test Agent",
|
||||
description="Agent for testing serialization",
|
||||
chat_client=mock_client,
|
||||
system_message="You are a test assistant.",
|
||||
)
|
||||
|
||||
builder = WorkflowBuilder(name="Serialization Test Workflow", description="Test workflow")
|
||||
agent_executor = AgentExecutor(id="agent_node", agent=agent)
|
||||
builder.set_start_executor(agent_executor)
|
||||
workflow = builder.build()
|
||||
|
||||
# Create executor and register
|
||||
discovery = EntityDiscovery(None)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
entity_info = await discovery.create_entity_info_from_object(workflow, entity_type="workflow", source="test")
|
||||
discovery.register_entity(entity_info.id, entity_info, workflow)
|
||||
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_info.id},
|
||||
input="Test workflow serialization",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
events = []
|
||||
serialization_errors = []
|
||||
event_types_seen = []
|
||||
|
||||
async for event in executor.execute_streaming(request):
|
||||
events.append(event)
|
||||
event_type = getattr(event, "type", "unknown")
|
||||
event_types_seen.append(event_type)
|
||||
|
||||
# This is EXACTLY what the server does before sending SSE
|
||||
try:
|
||||
if hasattr(event, "model_dump_json"):
|
||||
json_str = event.model_dump_json()
|
||||
assert json_str is not None
|
||||
assert len(json_str) > 0
|
||||
except Exception as e:
|
||||
serialization_errors.append(f"Event type={event_type}: {e}")
|
||||
|
||||
# Should have received events
|
||||
assert len(events) > 0, "Should receive events from workflow execution"
|
||||
|
||||
# Verify we got workflow events (not just generic ones)
|
||||
assert any("output_item" in str(t) for t in event_types_seen), (
|
||||
f"Should see output_item events, got: {event_types_seen}"
|
||||
)
|
||||
|
||||
# NO serialization errors allowed - this is the critical assertion
|
||||
assert len(serialization_errors) == 0, (
|
||||
f"Found {len(serialization_errors)} serialization errors:\n"
|
||||
+ "\n".join(serialization_errors)
|
||||
+ f"\n\nEvent types seen: {event_types_seen}"
|
||||
)
|
||||
|
||||
# Also verify aggregate_to_response works (server calls this after streaming)
|
||||
final_response = await mapper.aggregate_to_response(events, request)
|
||||
assert final_response is not None
|
||||
|
||||
|
||||
async def test_get_entity_info_raises_for_invalid_id(executor):
|
||||
"""Test that get_entity_info raises EntityNotFoundError for invalid ID."""
|
||||
with pytest.raises(EntityNotFoundError):
|
||||
executor.get_entity_info("nonexistent_agent")
|
||||
|
||||
|
||||
async def test_request_extracts_entity_id_from_metadata(executor):
|
||||
"""Test that AgentFrameworkRequest extracts entity_id from metadata."""
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": "my_agent"},
|
||||
input="test",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# entity_id is extracted from metadata
|
||||
entity_id = request.get_entity_id()
|
||||
assert entity_id == "my_agent"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_get_start_executor_message_types(sequential_workflow):
|
||||
"""Test _get_start_executor_message_types with real workflow."""
|
||||
executor, _entity_id, _mock_client, workflow = sequential_workflow
|
||||
|
||||
start_exec, message_types = executor._get_start_executor_message_types(workflow)
|
||||
|
||||
assert start_exec is not None
|
||||
assert len(message_types) > 0
|
||||
# Real sequential workflows accept str input
|
||||
assert str in message_types
|
||||
|
||||
|
||||
def test_executor_select_primary_input_prefers_string():
|
||||
"""Select string input even when discovered after other handlers."""
|
||||
from agent_framework_devui._utils import select_primary_input_type
|
||||
|
||||
placeholder_type = type("Placeholder", (), {})
|
||||
|
||||
chosen = select_primary_input_type([placeholder_type, str])
|
||||
|
||||
assert chosen is str
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_parse_structured_extracts_input_for_string_workflow():
|
||||
"""Structured payloads extract 'input' field when workflow expects str."""
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
class StringInputExecutor(Executor):
|
||||
"""Executor that accepts string input directly."""
|
||||
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[Any, Any]) -> None:
|
||||
await ctx.yield_output(f"Got: {text}")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(name="String Workflow", description="Accepts string")
|
||||
.set_start_executor(StringInputExecutor(id="str_exec"))
|
||||
.build()
|
||||
)
|
||||
|
||||
executor = AgentFrameworkExecutor(EntityDiscovery(None), MessageMapper())
|
||||
|
||||
# When workflow expects str and receives {"input": "hello"}, extract "hello"
|
||||
parsed = executor._parse_structured_workflow_input(workflow, {"input": "hello"})
|
||||
assert parsed == "hello"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_parse_raw_string_for_string_workflow():
|
||||
"""Raw string inputs pass through for string-accepting workflows."""
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
class StringInputExecutor(Executor):
|
||||
"""Executor that accepts string input directly."""
|
||||
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[Any, Any]) -> None:
|
||||
await ctx.yield_output(f"Got: {text}")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(name="String Workflow", description="Accepts string")
|
||||
.set_start_executor(StringInputExecutor(id="str_exec"))
|
||||
.build()
|
||||
)
|
||||
|
||||
executor = AgentFrameworkExecutor(EntityDiscovery(None), MessageMapper())
|
||||
|
||||
# Raw string should pass through unchanged
|
||||
parsed = executor._parse_raw_workflow_input(workflow, "hi there")
|
||||
assert parsed == "hi there"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_parse_converts_to_chat_message_for_sequential_workflow(sequential_workflow):
|
||||
"""Sequential workflows convert string input to ChatMessage."""
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
executor, _entity_id, _mock_client, workflow = sequential_workflow
|
||||
|
||||
# Sequential workflows expect ChatMessage, so raw string becomes ChatMessage
|
||||
parsed = executor._parse_raw_workflow_input(workflow, "hello")
|
||||
|
||||
assert isinstance(parsed, ChatMessage)
|
||||
assert parsed.text == "hello"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_parse_stringified_json_workflow_input():
|
||||
"""Stringified JSON workflow input is parsed when workflow expects Pydantic model."""
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
from pydantic import BaseModel
|
||||
|
||||
class WorkflowInput(BaseModel):
|
||||
input: str
|
||||
metadata: dict | None = None
|
||||
|
||||
class PydanticInputExecutor(Executor):
|
||||
"""Executor that accepts a Pydantic model input."""
|
||||
|
||||
@handler
|
||||
async def process(self, data: WorkflowInput, ctx: WorkflowContext[Any, Any]) -> None:
|
||||
await ctx.yield_output(f"Got: {data.input}")
|
||||
|
||||
# Build workflow with Pydantic input type
|
||||
workflow = (
|
||||
WorkflowBuilder(name="Pydantic Workflow", description="Accepts Pydantic input")
|
||||
.set_start_executor(PydanticInputExecutor(id="pydantic_exec"))
|
||||
.build()
|
||||
)
|
||||
|
||||
executor = AgentFrameworkExecutor(EntityDiscovery(None), MessageMapper())
|
||||
|
||||
# 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"}
|
||||
|
||||
|
||||
def test_extract_workflow_hil_responses_handles_stringified_json():
|
||||
"""Test HIL response extraction handles both stringified and parsed JSON (regression test)."""
|
||||
from agent_framework_devui._discovery import EntityDiscovery
|
||||
from agent_framework_devui._executor import AgentFrameworkExecutor
|
||||
from agent_framework_devui._mapper import MessageMapper
|
||||
|
||||
executor = AgentFrameworkExecutor(EntityDiscovery(None), MessageMapper())
|
||||
|
||||
# Regression test: Frontend sends stringified JSON via streamWorkflowExecutionOpenAI
|
||||
stringified = '[{"type":"message","content":[{"type":"workflow_hil_response","responses":{"req_1":"spam"}}]}]'
|
||||
assert executor._extract_workflow_hil_responses(stringified) == {"req_1": "spam"}
|
||||
|
||||
# Ensure parsed format still works
|
||||
parsed = [{"type": "message", "content": [{"type": "workflow_hil_response", "responses": {"req_2": "ham"}}]}]
|
||||
assert executor._extract_workflow_hil_responses(parsed) == {"req_2": "ham"}
|
||||
|
||||
# Non-HIL inputs should return None
|
||||
assert executor._extract_workflow_hil_responses("plain text") is None
|
||||
assert executor._extract_workflow_hil_responses({"email": "test"}) is None
|
||||
|
||||
|
||||
async def test_executor_handles_streaming_agent():
|
||||
"""Test executor handles agents with run(stream=True) method."""
|
||||
from agent_framework import AgentResponse, AgentResponseUpdate, AgentThread, ChatMessage, Content
|
||||
|
||||
class StreamingAgent:
|
||||
"""Agent with run() method supporting stream parameter."""
|
||||
|
||||
id = "streaming_test"
|
||||
name = "Streaming Test Agent"
|
||||
description = "Test agent with run(stream=True)"
|
||||
|
||||
def run(self, messages=None, *, stream=False, thread=None, **kwargs):
|
||||
if stream:
|
||||
# Return an async generator for streaming
|
||||
return self._stream_impl(messages)
|
||||
# Return awaitable for non-streaming
|
||||
return self._run_impl(messages)
|
||||
|
||||
async def _run_impl(self, messages):
|
||||
return AgentResponse(
|
||||
messages=[ChatMessage(role="assistant", contents=[Content.from_text(text=f"Processed: {messages}")])],
|
||||
response_id="test_123",
|
||||
)
|
||||
|
||||
async def _stream_impl(self, messages):
|
||||
yield AgentResponseUpdate(
|
||||
contents=[Content.from_text(text=f"Processed: {messages}")],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
def get_new_thread(self, **kwargs):
|
||||
return AgentThread()
|
||||
|
||||
# Create executor and register agent
|
||||
discovery = EntityDiscovery(None)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
agent = StreamingAgent()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, source="test")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Execute streaming agent (use metadata.entity_id for routing)
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_info.id},
|
||||
input="hello",
|
||||
stream=True, # DevUI always streams
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in executor.execute_streaming(request):
|
||||
events.append(event)
|
||||
|
||||
# Should get events from streaming agent
|
||||
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
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Full Pipeline Tests for SequentialBuilder
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_pipeline_sequential_workflow(sequential_workflow):
|
||||
"""Test SequentialBuilder workflow full pipeline with JSON serialization.
|
||||
|
||||
Uses the shared sequential_workflow fixture (Writer → Reviewer) from conftest.
|
||||
Tests that all events can be JSON serialized for SSE streaming.
|
||||
"""
|
||||
executor, entity_id, mock_client, _workflow = sequential_workflow
|
||||
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_id},
|
||||
input="Write about testing best practices",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
events = []
|
||||
serialization_errors = []
|
||||
|
||||
async for event in executor.execute_streaming(request):
|
||||
events.append(event)
|
||||
event_type = getattr(event, "type", "unknown")
|
||||
|
||||
# Verify JSON serialization (exactly what server does for SSE)
|
||||
try:
|
||||
if hasattr(event, "model_dump_json"):
|
||||
json_str = event.model_dump_json()
|
||||
assert json_str is not None
|
||||
except Exception as e:
|
||||
serialization_errors.append(f"Event type={event_type}: {e}")
|
||||
|
||||
assert len(events) > 0, "Should receive events from sequential workflow"
|
||||
assert len(serialization_errors) == 0, f"Serialization errors: {serialization_errors}"
|
||||
assert mock_client.call_count >= 2, f"Expected both agents called, got {mock_client.call_count}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_pipeline_concurrent_workflow(concurrent_workflow):
|
||||
"""Test ConcurrentBuilder workflow full pipeline with JSON serialization.
|
||||
|
||||
Uses the shared concurrent_workflow fixture (Researcher | Analyst | Summarizer) from conftest.
|
||||
Tests fan-out/fan-in pattern with parallel agent execution.
|
||||
"""
|
||||
executor, entity_id, mock_client, _workflow = concurrent_workflow
|
||||
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_id},
|
||||
input="Analyze market trends for Q4",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
events = []
|
||||
serialization_errors = []
|
||||
|
||||
async for event in executor.execute_streaming(request):
|
||||
events.append(event)
|
||||
event_type = getattr(event, "type", "unknown")
|
||||
|
||||
# Verify JSON serialization
|
||||
try:
|
||||
if hasattr(event, "model_dump_json"):
|
||||
json_str = event.model_dump_json()
|
||||
assert json_str is not None
|
||||
except Exception as e:
|
||||
serialization_errors.append(f"Event type={event_type}: {e}")
|
||||
|
||||
assert len(events) > 0, "Should receive events from concurrent workflow"
|
||||
assert len(serialization_errors) == 0, f"Serialization errors: {serialization_errors}"
|
||||
assert mock_client.call_count >= 3, f"Expected all 3 agents called, got {mock_client.call_count}"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Full Pipeline Test for Workflow with Output Events
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_pipeline_workflow_output_event_serialization():
|
||||
"""Test that WorkflowOutputEvent from ctx.yield_output() serializes correctly.
|
||||
|
||||
This tests the pattern where executors yield output via ctx.yield_output(),
|
||||
which emits WorkflowOutputEvent that DevUI must serialize for SSE.
|
||||
"""
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
class OutputtingExecutor(Executor):
|
||||
"""Executor that yields multiple outputs."""
|
||||
|
||||
@handler
|
||||
async def process(self, input_text: str, ctx: WorkflowContext[Any, Any]) -> None:
|
||||
await ctx.yield_output(f"First output: {input_text}")
|
||||
await ctx.yield_output("Second output: processed")
|
||||
await ctx.yield_output({"final": "result", "data": [1, 2, 3]})
|
||||
|
||||
# Build workflow
|
||||
workflow = (
|
||||
WorkflowBuilder(name="Output Workflow", description="Tests yield_output")
|
||||
.set_start_executor(OutputtingExecutor(id="outputter"))
|
||||
.build()
|
||||
)
|
||||
|
||||
# Create DevUI executor and register workflow
|
||||
discovery = EntityDiscovery(None)
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
entity_info = await discovery.create_entity_info_from_object(workflow, entity_type="workflow", source="test")
|
||||
discovery.register_entity(entity_info.id, entity_info, workflow)
|
||||
|
||||
# Execute with streaming
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entity_info.id},
|
||||
input="Test output events",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
events = []
|
||||
output_events = []
|
||||
serialization_errors = []
|
||||
|
||||
async for event in executor.execute_streaming(request):
|
||||
events.append(event)
|
||||
event_type = getattr(event, "type", "")
|
||||
|
||||
# Track output item events
|
||||
if "output_item" in event_type:
|
||||
output_events.append(event)
|
||||
|
||||
try:
|
||||
if hasattr(event, "model_dump_json"):
|
||||
event.model_dump_json()
|
||||
except Exception as e:
|
||||
serialization_errors.append(f"Event type={event_type}: {e}")
|
||||
|
||||
assert len(events) > 0, "Should receive events"
|
||||
assert len(serialization_errors) == 0, f"Serialization errors: {serialization_errors}"
|
||||
|
||||
# Should have received output events for the yield_output calls
|
||||
assert len(output_events) >= 3, f"Expected 3+ output events for yield_output calls, got {len(output_events)}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner
|
||||
async def run_tests():
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create test agent
|
||||
agent_file = temp_path / "streaming_agent.py"
|
||||
agent_file.write_text("""
|
||||
class StreamingAgent:
|
||||
name = "Streaming Test Agent"
|
||||
description = "Test agent for streaming"
|
||||
|
||||
async def run(self, input_str, *, stream: bool = False, thread=None, **kwargs):
|
||||
if stream:
|
||||
async def _stream():
|
||||
for i, word in enumerate(f"Processing {input_str}".split()):
|
||||
yield f"word_{i}: {word} "
|
||||
return _stream()
|
||||
return f"Processing {input_str}"
|
||||
""")
|
||||
|
||||
discovery = EntityDiscovery(str(temp_path))
|
||||
mapper = MessageMapper()
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
# Test discovery
|
||||
entities = await executor.discover_entities()
|
||||
|
||||
if entities:
|
||||
# Test sync execution (use metadata.entity_id for routing)
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entities[0].id},
|
||||
input="test input",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
await executor.execute_sync(request)
|
||||
|
||||
# Test streaming execution
|
||||
request.stream = True
|
||||
event_count = 0
|
||||
async for _event in executor.execute_streaming(request):
|
||||
event_count += 1
|
||||
if event_count > 5: # Limit for testing
|
||||
break
|
||||
|
||||
asyncio.run(run_tests())
|
||||
@@ -0,0 +1,755 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for message mapping functionality.
|
||||
|
||||
This module tests the MessageMapper which converts Agent Framework events
|
||||
to OpenAI-compatible streaming events. Tests use REAL classes from
|
||||
agent_framework, not mocks, to ensure proper serialization.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
# Import Agent Framework types
|
||||
from agent_framework._types import (
|
||||
AgentResponseUpdate,
|
||||
Content,
|
||||
)
|
||||
|
||||
# Import real workflow event classes - NOT mocks!
|
||||
from agent_framework._workflows._events import (
|
||||
ExecutorCompletedEvent,
|
||||
WorkflowStartedEvent,
|
||||
WorkflowStatusEvent,
|
||||
)
|
||||
|
||||
# Import factory functions from conftest for parameterized test data creation
|
||||
from conftest import (
|
||||
create_agent_run_response,
|
||||
create_executor_completed_event,
|
||||
create_executor_failed_event,
|
||||
create_executor_invoked_event,
|
||||
)
|
||||
|
||||
from agent_framework_devui._mapper import MessageMapper
|
||||
from agent_framework_devui.models._openai_custom import (
|
||||
AgentCompletedEvent,
|
||||
AgentFailedEvent,
|
||||
AgentFrameworkRequest,
|
||||
AgentStartedEvent,
|
||||
)
|
||||
|
||||
# Note: mapper and test_request fixtures are provided by conftest.py
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Test Helpers
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def create_test_content(content_type: str, **kwargs: Any) -> Any:
|
||||
"""Create test content objects."""
|
||||
if content_type == "text":
|
||||
return Content.from_text(text=kwargs.get("text", "Hello, world!"))
|
||||
if content_type == "function_call":
|
||||
return Content.from_function_call(
|
||||
call_id=kwargs.get("call_id", "test_call_id"),
|
||||
name=kwargs.get("name", "test_func"),
|
||||
arguments=kwargs.get("arguments", {"param": "value"}),
|
||||
)
|
||||
if content_type == "error":
|
||||
return Content.from_error(
|
||||
message=kwargs.get("message", "Test error"), error_code=kwargs.get("code", "test_error")
|
||||
)
|
||||
raise ValueError(f"Unknown content type: {content_type}")
|
||||
|
||||
|
||||
def create_test_agent_update(contents: list[Any]) -> AgentResponseUpdate:
|
||||
"""Create test AgentResponseUpdate."""
|
||||
return AgentResponseUpdate(contents=contents, role="assistant", message_id="test_msg", response_id="test_resp")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Basic Content Mapping Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_critical_isinstance_bug_detection(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""CRITICAL: Test that would have caught the isinstance vs hasattr bug."""
|
||||
content = create_test_content("text", text="Bug detection test")
|
||||
update = create_test_agent_update([content])
|
||||
|
||||
# Key assertions that would have caught the bug
|
||||
assert hasattr(update, "contents") # Real attribute
|
||||
assert not hasattr(update, "response") # Fake attribute should not exist
|
||||
|
||||
# Test isinstance works with real types
|
||||
assert isinstance(update, AgentResponseUpdate)
|
||||
|
||||
# Test mapper conversion - should NOT produce "Unknown event"
|
||||
events = await mapper.convert_event(update, test_request)
|
||||
|
||||
assert len(events) > 0
|
||||
assert all(hasattr(event, "type") for event in events)
|
||||
assert all(event.type != "unknown" for event in events)
|
||||
|
||||
|
||||
async def test_text_content_mapping(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test TextContent mapping with proper OpenAI event hierarchy."""
|
||||
content = create_test_content("text", text="Hello, clean test!")
|
||||
update = create_test_agent_update([content])
|
||||
|
||||
events = await mapper.convert_event(update, test_request)
|
||||
|
||||
# With proper OpenAI hierarchy, we expect 3 events:
|
||||
# 1. response.output_item.added (message)
|
||||
# 2. response.content_part.added (text part)
|
||||
# 3. response.output_text.delta (actual text)
|
||||
assert len(events) == 3
|
||||
|
||||
# Check message output item
|
||||
assert events[0].type == "response.output_item.added"
|
||||
assert events[0].item.type == "message"
|
||||
assert events[0].item.role == "assistant"
|
||||
|
||||
# Check content part
|
||||
assert events[1].type == "response.content_part.added"
|
||||
assert events[1].part.type == "output_text"
|
||||
|
||||
# Check text delta
|
||||
assert events[2].type == "response.output_text.delta"
|
||||
assert events[2].delta == "Hello, clean test!"
|
||||
|
||||
|
||||
async def test_function_call_mapping(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test FunctionCallContent mapping."""
|
||||
content = create_test_content("function_call", name="test_func", arguments={"location": "TestCity"})
|
||||
update = create_test_agent_update([content])
|
||||
|
||||
events = await mapper.convert_event(update, test_request)
|
||||
|
||||
# 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 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
|
||||
|
||||
|
||||
async def test_function_result_content_with_string_result(
|
||||
mapper: MessageMapper, test_request: AgentFrameworkRequest
|
||||
) -> None:
|
||||
"""Test FunctionResultContent with plain string result (regular tools)."""
|
||||
content = Content.from_function_result(
|
||||
call_id="test_call_123",
|
||||
result="Hello, World!",
|
||||
)
|
||||
update = create_test_agent_update([content])
|
||||
|
||||
events = await mapper.convert_event(update, test_request)
|
||||
|
||||
assert len(events) >= 1
|
||||
result_events = [e for e in events if e.type == "response.function_result.complete"]
|
||||
assert len(result_events) == 1
|
||||
assert result_events[0].output == "Hello, World!"
|
||||
assert result_events[0].call_id == "test_call_123"
|
||||
assert result_events[0].status == "completed"
|
||||
|
||||
|
||||
async def test_function_result_content_with_nested_content_objects(
|
||||
mapper: MessageMapper, test_request: AgentFrameworkRequest
|
||||
) -> None:
|
||||
"""Test FunctionResultContent with nested Content objects (MCP tools case)."""
|
||||
content = Content.from_function_result(
|
||||
call_id="mcp_call_456",
|
||||
result=[Content.from_text(text="Hello from MCP!")],
|
||||
)
|
||||
update = create_test_agent_update([content])
|
||||
|
||||
events = await mapper.convert_event(update, test_request)
|
||||
|
||||
assert len(events) >= 1
|
||||
result_events = [e for e in events if e.type == "response.function_result.complete"]
|
||||
assert len(result_events) == 1
|
||||
assert "Hello from MCP!" in result_events[0].output
|
||||
assert result_events[0].call_id == "mcp_call_456"
|
||||
|
||||
|
||||
async def test_error_content_mapping(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test ErrorContent mapping."""
|
||||
content = create_test_content("error", message="Test error", code="test_code")
|
||||
update = create_test_agent_update([content])
|
||||
|
||||
events = await mapper.convert_event(update, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "error"
|
||||
assert events[0].message == "Test error"
|
||||
assert events[0].code == "test_code"
|
||||
|
||||
|
||||
async def test_mixed_content_types(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test multiple content types together."""
|
||||
contents = [
|
||||
create_test_content("text", text="Starting..."),
|
||||
create_test_content("function_call", name="process", arguments={"data": "test"}),
|
||||
create_test_content("text", text="Done!"),
|
||||
]
|
||||
update = create_test_agent_update(contents)
|
||||
|
||||
events = await mapper.convert_event(update, test_request)
|
||||
|
||||
assert len(events) >= 3
|
||||
event_types = {event.type for event in events}
|
||||
assert "response.output_text.delta" in event_types
|
||||
assert "response.function_call_arguments.delta" in event_types
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Agent Lifecycle Event Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_agent_lifecycle_events(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test that agent lifecycle events are properly converted to OpenAI format."""
|
||||
# Test AgentStartedEvent
|
||||
start_event = AgentStartedEvent()
|
||||
events = await mapper.convert_event(start_event, test_request)
|
||||
|
||||
assert len(events) == 2 # 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"
|
||||
assert events[0].response.status == "in_progress"
|
||||
|
||||
# Test AgentCompletedEvent
|
||||
complete_event = AgentCompletedEvent()
|
||||
events = await mapper.convert_event(complete_event, test_request)
|
||||
# AgentCompletedEvent no longer emits response.completed to avoid duplicates
|
||||
assert len(events) == 0
|
||||
|
||||
# Test AgentFailedEvent
|
||||
error = Exception("Test error")
|
||||
failed_event = AgentFailedEvent(error=error)
|
||||
events = await mapper.convert_event(failed_event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.failed"
|
||||
assert events[0].response.status == "failed"
|
||||
assert events[0].response.error.message == "Test error"
|
||||
|
||||
|
||||
async def test_agent_run_response_mapping(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test that mapper handles complete AgentResponse (non-streaming)."""
|
||||
response = create_agent_run_response("Complete response from run()")
|
||||
|
||||
events = await mapper.convert_event(response, test_request)
|
||||
|
||||
assert len(events) > 0
|
||||
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()"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Workflow Executor Event Tests (using REAL classes, not mocks!)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_executor_invoked_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test ExecutorInvokedEvent using the REAL class from agent_framework."""
|
||||
# Use real class, not mock!
|
||||
event = create_executor_invoked_event(executor_id="exec_123")
|
||||
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.output_item.added"
|
||||
# Access as dict since item might be ExecutorActionItem
|
||||
item = events[0].item if isinstance(events[0].item, dict) else events[0].item.model_dump()
|
||||
assert item["type"] == "executor_action"
|
||||
assert item["executor_id"] == "exec_123"
|
||||
assert item["status"] == "in_progress"
|
||||
|
||||
|
||||
async def test_executor_completed_event_simple_data(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test ExecutorCompletedEvent with simple dict data."""
|
||||
# Create event with simple data
|
||||
event = ExecutorCompletedEvent(executor_id="exec_123", data={"simple": "result"})
|
||||
|
||||
# First need to invoke the executor to set up context
|
||||
invoke_event = create_executor_invoked_event(executor_id="exec_123")
|
||||
await mapper.convert_event(invoke_event, test_request)
|
||||
|
||||
# Now complete it
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.output_item.done"
|
||||
item = events[0].item if isinstance(events[0].item, dict) else events[0].item.model_dump()
|
||||
assert item["type"] == "executor_action"
|
||||
assert item["executor_id"] == "exec_123"
|
||||
assert item["status"] == "completed"
|
||||
# Result should be serialized
|
||||
assert item["result"] == {"simple": "result"}
|
||||
|
||||
|
||||
async def test_executor_completed_event_with_agent_response(
|
||||
mapper: MessageMapper, test_request: AgentFrameworkRequest
|
||||
) -> None:
|
||||
"""Test ExecutorCompletedEvent with nested AgentExecutorResponse.
|
||||
|
||||
This is a REGRESSION TEST for the serialization bug where
|
||||
ExecutorCompletedEvent.data contained AgentExecutorResponse with nested
|
||||
AgentResponse and ChatMessage objects (SerializationMixin) that
|
||||
Pydantic couldn't serialize.
|
||||
"""
|
||||
# Create event with realistic nested data - the exact structure that caused the bug
|
||||
event = create_executor_completed_event(executor_id="exec_agent", with_agent_response=True)
|
||||
|
||||
# Verify the data has the problematic structure
|
||||
assert hasattr(event.data, "agent_response")
|
||||
assert hasattr(event.data, "full_conversation")
|
||||
|
||||
# First invoke the executor
|
||||
invoke_event = create_executor_invoked_event(executor_id="exec_agent")
|
||||
await mapper.convert_event(invoke_event, test_request)
|
||||
|
||||
# Now complete - this should NOT raise serialization errors
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.output_item.done"
|
||||
|
||||
# Get the item (might be Pydantic model or dict)
|
||||
item = events[0].item if isinstance(events[0].item, dict) else events[0].item.model_dump()
|
||||
assert item["type"] == "executor_action"
|
||||
assert item["executor_id"] == "exec_agent"
|
||||
assert item["status"] == "completed"
|
||||
|
||||
# The result should be serialized (converted to dict)
|
||||
result = item["result"]
|
||||
assert result is not None
|
||||
# Should be a dict or list, not the original object
|
||||
assert isinstance(result, (dict, list))
|
||||
|
||||
|
||||
async def test_executor_completed_event_serialization_to_json(
|
||||
mapper: MessageMapper, test_request: AgentFrameworkRequest
|
||||
) -> None:
|
||||
"""REGRESSION TEST: Verify the full JSON serialization works.
|
||||
|
||||
This tests the exact failure mode from the bug: calling model_dump_json()
|
||||
on the event containing nested SerializationMixin objects.
|
||||
"""
|
||||
# Create the problematic event
|
||||
event = create_executor_completed_event(executor_id="exec_json_test", with_agent_response=True)
|
||||
|
||||
# Invoke first
|
||||
invoke_event = create_executor_invoked_event(executor_id="exec_json_test")
|
||||
await mapper.convert_event(invoke_event, test_request)
|
||||
|
||||
# Complete
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
done_event = events[0]
|
||||
|
||||
# This is the critical test - model_dump_json() should NOT raise
|
||||
# "Unable to serialize unknown type: <class 'agent_framework._types.AgentResponse'>"
|
||||
try:
|
||||
json_str = done_event.model_dump_json()
|
||||
assert json_str is not None
|
||||
assert len(json_str) > 0
|
||||
# Verify it's valid JSON by checking it contains expected fields
|
||||
assert "executor_action" in json_str
|
||||
assert "exec_json_test" in json_str
|
||||
assert "completed" in json_str
|
||||
except Exception as e:
|
||||
pytest.fail(f"model_dump_json() raised an exception: {e}")
|
||||
|
||||
|
||||
async def test_executor_failed_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test ExecutorFailedEvent using the REAL class."""
|
||||
# First invoke the executor
|
||||
invoke_event = create_executor_invoked_event(executor_id="exec_fail")
|
||||
await mapper.convert_event(invoke_event, test_request)
|
||||
|
||||
# Now fail it
|
||||
event = create_executor_failed_event(executor_id="exec_fail", error_message="Executor failed")
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.output_item.done"
|
||||
item = events[0].item if isinstance(events[0].item, dict) else events[0].item.model_dump()
|
||||
assert item["type"] == "executor_action"
|
||||
assert item["executor_id"] == "exec_fail"
|
||||
assert item["status"] == "failed"
|
||||
assert "Executor failed" in str(item["error"])
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Workflow Lifecycle Event Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_workflow_started_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowStartedEvent using the REAL class."""
|
||||
|
||||
event = WorkflowStartedEvent(data=None)
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# WorkflowStartedEvent should emit response.created and response.in_progress
|
||||
assert len(events) == 2
|
||||
assert events[0].type == "response.created"
|
||||
assert events[1].type == "response.in_progress"
|
||||
|
||||
|
||||
async def test_workflow_status_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowStatusEvent using the REAL class."""
|
||||
from agent_framework._workflows._events import WorkflowRunState
|
||||
|
||||
event = WorkflowStatusEvent(state=WorkflowRunState.IN_PROGRESS)
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# Should emit some status-related event
|
||||
assert len(events) >= 0 # May emit events or may be filtered
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Magentic Event Tests - Testing WorkflowOutputEvent with additional_properties
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_magentic_agent_run_update_event_with_agent_delta_metadata(
|
||||
mapper: MessageMapper, test_request: AgentFrameworkRequest
|
||||
) -> None:
|
||||
"""Test that WorkflowOutputEvent with magentic_event_type='agent_delta' is handled correctly.
|
||||
|
||||
This tests the ACTUAL event format Magentic emits - not a fake MagenticAgentDeltaEvent class.
|
||||
Magentic uses WorkflowOutputEvent wrapping AgentResponseUpdate with additional_properties.
|
||||
"""
|
||||
from agent_framework._types import AgentResponseUpdate
|
||||
from agent_framework._workflows._events import WorkflowOutputEvent
|
||||
|
||||
# Create the REAL event format that Magentic emits
|
||||
update = AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Hello from agent")],
|
||||
role="assistant",
|
||||
author_name="Writer",
|
||||
additional_properties={
|
||||
"magentic_event_type": "agent_delta",
|
||||
"agent_id": "writer_agent",
|
||||
},
|
||||
)
|
||||
event = WorkflowOutputEvent(executor_id="magentic_executor", data=update)
|
||||
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# Should be treated as a regular WorkflowOutputEvent with text content
|
||||
# The mapper should emit text delta events
|
||||
assert len(events) >= 1
|
||||
text_events = [e for e in events if getattr(e, "type", "") == "response.output_text.delta"]
|
||||
assert len(text_events) >= 1
|
||||
assert text_events[0].delta == "Hello from agent"
|
||||
|
||||
|
||||
async def test_magentic_orchestrator_message_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test that WorkflowOutputEvent with magentic_event_type='orchestrator_message' is handled.
|
||||
|
||||
Magentic emits orchestrator planning/instruction messages using WorkflowOutputEvent
|
||||
wrapping AgentResponseUpdate with additional_properties.
|
||||
"""
|
||||
from agent_framework._types import AgentResponseUpdate
|
||||
from agent_framework._workflows._events import WorkflowOutputEvent
|
||||
|
||||
# Create orchestrator message event (REAL format from Magentic)
|
||||
update = AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Planning: First, the writer will create content...")],
|
||||
role="assistant",
|
||||
author_name="Orchestrator",
|
||||
additional_properties={
|
||||
"magentic_event_type": "orchestrator_message",
|
||||
"orchestrator_message_kind": "task_ledger",
|
||||
"orchestrator_id": "magentic_orchestrator",
|
||||
},
|
||||
)
|
||||
event = WorkflowOutputEvent(executor_id="magentic_orchestrator", data=update)
|
||||
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# Currently, mapper treats this as regular WorkflowOutputEvent (no special handling)
|
||||
# This test documents the current behavior
|
||||
assert len(events) >= 1
|
||||
text_events = [e for e in events if getattr(e, "type", "") == "response.output_text.delta"]
|
||||
assert len(text_events) >= 1
|
||||
assert "Planning:" in text_events[0].delta
|
||||
|
||||
|
||||
async def test_magentic_events_use_same_event_class_as_other_workflows(
|
||||
mapper: MessageMapper, test_request: AgentFrameworkRequest
|
||||
) -> None:
|
||||
"""Verify Magentic uses the same WorkflowOutputEvent class as other workflows.
|
||||
|
||||
This test documents that Magentic does NOT define separate event classes like
|
||||
MagenticAgentDeltaEvent - it reuses WorkflowOutputEvent with metadata in
|
||||
additional_properties. Any mapper code checking for 'MagenticAgentDeltaEvent'
|
||||
class names is dead code.
|
||||
"""
|
||||
from agent_framework._types import AgentResponseUpdate
|
||||
from agent_framework._workflows._events import WorkflowOutputEvent
|
||||
|
||||
# Create events the way different workflows do it
|
||||
# 1. Regular workflow (no additional_properties)
|
||||
regular_update = AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Regular workflow response")],
|
||||
role="assistant",
|
||||
)
|
||||
regular_event = WorkflowOutputEvent(executor_id="regular_executor", data=regular_update)
|
||||
|
||||
# 2. Magentic workflow (with additional_properties)
|
||||
magentic_update = AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Magentic workflow response")],
|
||||
role="assistant",
|
||||
additional_properties={"magentic_event_type": "agent_delta"},
|
||||
)
|
||||
magentic_event = WorkflowOutputEvent(executor_id="magentic_executor", data=magentic_update)
|
||||
|
||||
# Both should be the SAME class
|
||||
assert type(regular_event) is type(magentic_event)
|
||||
assert isinstance(regular_event, WorkflowOutputEvent)
|
||||
assert isinstance(magentic_event, WorkflowOutputEvent)
|
||||
|
||||
# Both should be handled by the same isinstance check in mapper
|
||||
regular_events = await mapper.convert_event(regular_event, test_request)
|
||||
magentic_events = await mapper.convert_event(magentic_event, test_request)
|
||||
|
||||
# Both produce text delta events
|
||||
regular_text = [e for e in regular_events if getattr(e, "type", "") == "response.output_text.delta"]
|
||||
magentic_text = [e for e in magentic_events if getattr(e, "type", "") == "response.output_text.delta"]
|
||||
|
||||
assert len(regular_text) >= 1
|
||||
assert len(magentic_text) >= 1
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Unknown Content Fallback Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_unknown_content_fallback(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test graceful handling of unknown content types."""
|
||||
|
||||
class MockUnknownContent:
|
||||
def __init__(self):
|
||||
self.__class__.__name__ = "WeirdUnknownContent"
|
||||
|
||||
context = mapper._get_or_create_context(test_request)
|
||||
unknown_content = MockUnknownContent()
|
||||
|
||||
event = await mapper._create_unknown_content_event(unknown_content, context)
|
||||
|
||||
assert event.type == "response.output_text.delta"
|
||||
assert "Unknown content type" in event.delta
|
||||
assert "WeirdUnknownContent" in event.delta
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# WorkflowOutputEvent Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_workflow_output_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowOutputEvent is converted to output_item.added."""
|
||||
from agent_framework._workflows._events import WorkflowOutputEvent
|
||||
|
||||
event = WorkflowOutputEvent(data="Final workflow output", executor_id="final_executor")
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# WorkflowOutputEvent should emit output_item.added
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.output_item.added"
|
||||
# Check item contains the output text
|
||||
item = events[0].item
|
||||
assert item.type == "message"
|
||||
assert any("Final workflow output" in str(c) for c in item.content)
|
||||
|
||||
|
||||
async def test_workflow_output_event_with_list_data(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowOutputEvent with list data (common for sequential/concurrent workflows)."""
|
||||
from agent_framework import ChatMessage
|
||||
from agent_framework._workflows._events import WorkflowOutputEvent
|
||||
|
||||
# Sequential/Concurrent workflows often output list[ChatMessage]
|
||||
messages = [
|
||||
ChatMessage(role="user", contents=[Content.from_text(text="Hello")]),
|
||||
ChatMessage(role="assistant", contents=[Content.from_text(text="World")]),
|
||||
]
|
||||
event = WorkflowOutputEvent(data=messages, executor_id="complete")
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.output_item.added"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# WorkflowFailedEvent Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_workflow_failed_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowFailedEvent is converted to response.failed."""
|
||||
from agent_framework._workflows._events import WorkflowErrorDetails, WorkflowFailedEvent
|
||||
|
||||
details = WorkflowErrorDetails(
|
||||
error_type="TestError",
|
||||
message="Workflow failed due to test error",
|
||||
executor_id="failing_executor",
|
||||
)
|
||||
event = WorkflowFailedEvent(details=details)
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# WorkflowFailedEvent should emit response.failed
|
||||
assert len(events) >= 1
|
||||
# Find the failed event
|
||||
failed_events = [e for e in events if getattr(e, "type", "") == "response.failed"]
|
||||
assert len(failed_events) == 1, f"Expected response.failed, got types: {[getattr(e, 'type', '') for e in events]}"
|
||||
# Check response contains error info
|
||||
response = failed_events[0].response
|
||||
assert response.status == "failed"
|
||||
assert response.error is not None
|
||||
# Verify error message is correctly extracted from details.message (not "Unknown error")
|
||||
assert "Workflow failed due to test error" in response.error.message
|
||||
assert "Unknown error" not in response.error.message
|
||||
|
||||
|
||||
async def test_workflow_failed_event_with_extra(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowFailedEvent includes extra context when available."""
|
||||
from agent_framework._workflows._events import WorkflowErrorDetails, WorkflowFailedEvent
|
||||
|
||||
details = WorkflowErrorDetails(
|
||||
error_type="ValidationError",
|
||||
message="Input validation failed",
|
||||
executor_id="validation_executor",
|
||||
extra={"field": "email", "reason": "invalid format"},
|
||||
)
|
||||
event = WorkflowFailedEvent(details=details)
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.failed"
|
||||
response = events[0].response
|
||||
# Verify both the message and extra context are included
|
||||
assert "Input validation failed" in response.error.message
|
||||
assert "extra:" in response.error.message
|
||||
assert "email" in response.error.message
|
||||
|
||||
|
||||
async def test_workflow_failed_event_with_traceback(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowFailedEvent includes traceback when available."""
|
||||
from agent_framework._workflows._events import WorkflowErrorDetails, WorkflowFailedEvent
|
||||
|
||||
details = WorkflowErrorDetails(
|
||||
error_type="ValueError",
|
||||
message="Invalid input provided",
|
||||
traceback="Traceback (most recent call last):\n File ...\nValueError: Invalid input",
|
||||
executor_id="validation_executor",
|
||||
)
|
||||
event = WorkflowFailedEvent(details=details)
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.failed"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# WorkflowWarningEvent and WorkflowErrorEvent Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_workflow_warning_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowWarningEvent is converted to trace event."""
|
||||
from agent_framework._workflows._events import WorkflowWarningEvent
|
||||
|
||||
event = WorkflowWarningEvent(data="This is a warning message")
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# WorkflowWarningEvent should emit a trace event
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.trace.completed"
|
||||
assert events[0].data["event_type"] == "WorkflowWarningEvent"
|
||||
|
||||
|
||||
async def test_workflow_error_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test WorkflowErrorEvent is converted to trace event."""
|
||||
from agent_framework._workflows._events import WorkflowErrorEvent
|
||||
|
||||
event = WorkflowErrorEvent(data=ValueError("Something went wrong"))
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# WorkflowErrorEvent should emit a trace event
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.trace.completed"
|
||||
assert events[0].data["event_type"] == "WorkflowErrorEvent"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# RequestInfoEvent Tests (Human-in-the-Loop)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_request_info_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test RequestInfoEvent is converted to HIL request event."""
|
||||
from agent_framework._workflows._events import RequestInfoEvent
|
||||
|
||||
event = RequestInfoEvent(
|
||||
request_id="req_123",
|
||||
source_executor_id="approval_executor",
|
||||
request_data={"action": "approve", "details": "Please approve this action"},
|
||||
response_type=str,
|
||||
)
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# RequestInfoEvent should emit response.request_info.requested
|
||||
assert len(events) >= 1
|
||||
# Check that request info is captured
|
||||
has_hil_event = any(getattr(e, "type", "") == "response.request_info.requested" for e in events)
|
||||
assert has_hil_event, f"Expected response.request_info.requested, got: {[getattr(e, 'type', '') for e in events]}"
|
||||
|
||||
# Verify the event contains the expected data
|
||||
hil_event = [e for e in events if getattr(e, "type", "") == "response.request_info.requested"][0]
|
||||
assert hil_event.request_id == "req_123"
|
||||
assert hil_event.source_executor_id == "approval_executor"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# SuperStep Event Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_superstep_started_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test SuperStepStartedEvent is handled gracefully."""
|
||||
from agent_framework._workflows._events import SuperStepStartedEvent
|
||||
|
||||
event = SuperStepStartedEvent(iteration=1)
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# SuperStepStartedEvent may not emit events (internal workflow signal)
|
||||
# Just ensure it doesn't crash
|
||||
assert isinstance(events, list)
|
||||
|
||||
|
||||
async def test_superstep_completed_event(mapper: MessageMapper, test_request: AgentFrameworkRequest) -> None:
|
||||
"""Test SuperStepCompletedEvent is handled gracefully."""
|
||||
from agent_framework._workflows._events import SuperStepCompletedEvent
|
||||
|
||||
event = SuperStepCompletedEvent(iteration=1)
|
||||
events = await mapper.convert_event(event, test_request)
|
||||
|
||||
# SuperStepCompletedEvent may not emit events (internal workflow signal)
|
||||
# Just ensure it doesn't crash
|
||||
assert isinstance(events, list)
|
||||
@@ -0,0 +1,152 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Test multimodal input handling for workflows.
|
||||
|
||||
This test verifies that workflows with AgentExecutor nodes correctly receive
|
||||
multimodal content (images, files) from the DevUI frontend.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from agent_framework_devui._discovery import EntityDiscovery
|
||||
from agent_framework_devui._executor import AgentFrameworkExecutor
|
||||
from agent_framework_devui._mapper import MessageMapper
|
||||
|
||||
# Create a small test image (1x1 red pixel PNG)
|
||||
TEST_IMAGE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
|
||||
TEST_IMAGE_DATA_URI = f"data:image/png;base64,{TEST_IMAGE_BASE64}"
|
||||
|
||||
|
||||
class TestMultimodalWorkflowInput:
|
||||
"""Test multimodal input handling for workflows."""
|
||||
|
||||
def test_is_openai_multimodal_format_detects_message_format(self):
|
||||
"""Test that _is_openai_multimodal_format correctly detects OpenAI format."""
|
||||
discovery = MagicMock(spec=EntityDiscovery)
|
||||
mapper = MagicMock(spec=MessageMapper)
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
# Valid OpenAI multimodal format
|
||||
valid_format = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "Describe this image"},
|
||||
{"type": "input_image", "image_url": TEST_IMAGE_DATA_URI},
|
||||
],
|
||||
}
|
||||
]
|
||||
assert executor._is_openai_multimodal_format(valid_format) is True
|
||||
|
||||
# Invalid formats
|
||||
assert executor._is_openai_multimodal_format({}) is False # dict, not list
|
||||
assert executor._is_openai_multimodal_format([]) is False # empty list
|
||||
assert executor._is_openai_multimodal_format("hello") is False # string
|
||||
assert executor._is_openai_multimodal_format([{"type": "other"}]) is False # wrong type
|
||||
assert executor._is_openai_multimodal_format([{"foo": "bar"}]) is False # no type field
|
||||
|
||||
def test_convert_openai_input_to_chat_message_with_image(self):
|
||||
"""Test that OpenAI format with image is converted to ChatMessage with DataContent."""
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
discovery = MagicMock(spec=EntityDiscovery)
|
||||
mapper = MagicMock(spec=MessageMapper)
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
# OpenAI format input with text and image (as sent by frontend)
|
||||
openai_input = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "Describe this image"},
|
||||
{"type": "input_image", "image_url": TEST_IMAGE_DATA_URI},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# Convert to ChatMessage
|
||||
result = executor._convert_input_to_chat_message(openai_input)
|
||||
|
||||
# Verify result is ChatMessage
|
||||
assert isinstance(result, ChatMessage), f"Expected ChatMessage, got {type(result)}"
|
||||
assert result.role == "user"
|
||||
|
||||
# Verify contents
|
||||
assert len(result.contents) == 2, f"Expected 2 contents, got {len(result.contents)}"
|
||||
|
||||
# First content should be text
|
||||
assert result.contents[0].type == "text"
|
||||
assert result.contents[0].text == "Describe this image"
|
||||
|
||||
# Second content should be image (DataContent)
|
||||
assert result.contents[1].type == "data"
|
||||
assert result.contents[1].media_type == "image/png"
|
||||
assert result.contents[1].uri == TEST_IMAGE_DATA_URI
|
||||
|
||||
async def test_parse_workflow_input_handles_json_string_with_multimodal(self):
|
||||
"""Test that _parse_workflow_input correctly handles JSON string with multimodal content."""
|
||||
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
discovery = MagicMock(spec=EntityDiscovery)
|
||||
mapper = MagicMock(spec=MessageMapper)
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
# This is what the frontend sends: JSON stringified OpenAI format
|
||||
openai_input = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "What is in this image?"},
|
||||
{"type": "input_image", "image_url": TEST_IMAGE_DATA_URI},
|
||||
],
|
||||
}
|
||||
]
|
||||
json_string_input = json.dumps(openai_input)
|
||||
|
||||
# Mock workflow
|
||||
mock_workflow = MagicMock()
|
||||
|
||||
# Parse the input
|
||||
result = await executor._parse_workflow_input(mock_workflow, json_string_input)
|
||||
|
||||
# Verify result is ChatMessage with multimodal content
|
||||
assert isinstance(result, ChatMessage), f"Expected ChatMessage, got {type(result)}"
|
||||
assert len(result.contents) == 2
|
||||
|
||||
# Verify text content
|
||||
assert result.contents[0].type == "text"
|
||||
assert result.contents[0].text == "What is in this image?"
|
||||
|
||||
# Verify image content
|
||||
assert result.contents[1].type == "data"
|
||||
assert result.contents[1].media_type == "image/png"
|
||||
|
||||
async def test_parse_workflow_input_still_handles_simple_dict(self):
|
||||
"""Test that simple dict input still works (backward compatibility)."""
|
||||
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
discovery = MagicMock(spec=EntityDiscovery)
|
||||
mapper = MagicMock(spec=MessageMapper)
|
||||
executor = AgentFrameworkExecutor(discovery, mapper)
|
||||
|
||||
# Simple dict input (old format)
|
||||
simple_input = {"text": "Hello world", "role": "user"}
|
||||
json_string_input = json.dumps(simple_input)
|
||||
|
||||
# Mock workflow with ChatMessage input type
|
||||
mock_workflow = MagicMock()
|
||||
mock_executor = MagicMock()
|
||||
mock_executor.input_types = [ChatMessage]
|
||||
mock_workflow.get_start_executor.return_value = mock_executor
|
||||
|
||||
# Parse the input
|
||||
result = await executor._parse_workflow_input(mock_workflow, json_string_input)
|
||||
|
||||
# Result should be ChatMessage (from _parse_structured_workflow_input)
|
||||
assert isinstance(result, ChatMessage), f"Expected ChatMessage, got {type(result)}"
|
||||
@@ -0,0 +1,266 @@
|
||||
# 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
|
||||
@@ -0,0 +1,231 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
"""Test schema generation for different input types."""
|
||||
|
||||
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
|
||||
|
||||
|
||||
@dataclass
|
||||
class InputData:
|
||||
text: str
|
||||
source: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class Address:
|
||||
street: str
|
||||
city: str
|
||||
zipcode: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class PersonData:
|
||||
name: str
|
||||
age: int
|
||||
address: Address
|
||||
|
||||
|
||||
def test_builtin_types_schema_generation():
|
||||
"""Test schema generation for built-in types."""
|
||||
# Test str schema
|
||||
str_schema = generate_input_schema(str)
|
||||
assert str_schema is not None
|
||||
assert isinstance(str_schema, dict)
|
||||
|
||||
# Test dict schema
|
||||
dict_schema = generate_input_schema(dict)
|
||||
assert dict_schema is not None
|
||||
assert isinstance(dict_schema, dict)
|
||||
|
||||
# Test int schema
|
||||
int_schema = generate_input_schema(int)
|
||||
assert int_schema is not None
|
||||
assert isinstance(int_schema, dict)
|
||||
|
||||
|
||||
def test_dataclass_schema_generation():
|
||||
"""Test schema generation for dataclass."""
|
||||
schema = generate_input_schema(InputData)
|
||||
|
||||
assert schema is not None
|
||||
assert isinstance(schema, dict)
|
||||
|
||||
# Basic schema structure checks
|
||||
if "properties" in schema:
|
||||
properties = schema["properties"]
|
||||
assert "text" in properties
|
||||
assert "source" in properties
|
||||
|
||||
|
||||
def test_chat_message_schema_generation():
|
||||
"""Test schema generation for ChatMessage (SerializationMixin)."""
|
||||
try:
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
schema = generate_input_schema(ChatMessage)
|
||||
assert schema is not None
|
||||
assert isinstance(schema, dict)
|
||||
|
||||
except ImportError:
|
||||
pytest.skip("ChatMessage not available - agent_framework not installed")
|
||||
|
||||
|
||||
def test_pydantic_model_schema_generation():
|
||||
"""Test schema generation for Pydantic models."""
|
||||
try:
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
class UserInput(BaseModel):
|
||||
name: str = Field(description="User's name")
|
||||
age: int = Field(description="User's age")
|
||||
email: str | None = Field(default=None, description="Optional email")
|
||||
|
||||
schema = generate_input_schema(UserInput)
|
||||
assert schema is not None
|
||||
assert isinstance(schema, dict)
|
||||
|
||||
# Check if properties exist
|
||||
if "properties" in schema:
|
||||
properties = schema["properties"]
|
||||
assert "name" in properties
|
||||
assert "age" in properties
|
||||
assert "email" in properties
|
||||
|
||||
except ImportError:
|
||||
pytest.skip("Pydantic not available")
|
||||
|
||||
|
||||
def test_nested_dataclass_schema_generation():
|
||||
"""Test schema generation for nested dataclass."""
|
||||
schema = generate_input_schema(PersonData)
|
||||
|
||||
assert schema is not None
|
||||
assert isinstance(schema, dict)
|
||||
|
||||
# Basic schema structure checks
|
||||
if "properties" in schema:
|
||||
properties = schema["properties"]
|
||||
assert "name" in properties
|
||||
assert "age" in properties
|
||||
assert "address" in properties
|
||||
|
||||
|
||||
def test_schema_generation_error_handling():
|
||||
"""Test schema generation with invalid inputs."""
|
||||
# Test with a non-type object - should handle gracefully
|
||||
try:
|
||||
# Use a non-type object that might cause issues
|
||||
schema = generate_input_schema("not_a_type") # type: ignore
|
||||
# If it doesn't raise an exception, the result should be valid
|
||||
if schema is not None:
|
||||
assert isinstance(schema, dict)
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
# It's acceptable for this to raise an error
|
||||
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"])
|
||||
@@ -0,0 +1,404 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Focused tests for server functionality."""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework_devui import DevServer
|
||||
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:
|
||||
"""Simple executor stub exposing handler metadata."""
|
||||
|
||||
def __init__(self, *, input_types=None, handlers=None):
|
||||
if input_types is not None:
|
||||
self.input_types = list(input_types)
|
||||
if handlers is not None:
|
||||
self._handlers = dict(handlers)
|
||||
|
||||
|
||||
# Note: test_entities_dir fixture is provided by conftest.py
|
||||
|
||||
|
||||
async def test_server_health_endpoint(test_entities_dir):
|
||||
"""Test /health endpoint."""
|
||||
server = DevServer(entities_dir=test_entities_dir)
|
||||
executor = await server._ensure_executor()
|
||||
|
||||
# Test entity count
|
||||
entities = await executor.discover_entities()
|
||||
assert len(entities) > 0
|
||||
# Framework name is now hardcoded since we simplified to single framework
|
||||
|
||||
|
||||
async def test_server_entities_endpoint(test_entities_dir):
|
||||
"""Test /v1/entities endpoint."""
|
||||
server = DevServer(entities_dir=test_entities_dir)
|
||||
executor = await server._ensure_executor()
|
||||
|
||||
entities = await executor.discover_entities()
|
||||
assert len(entities) >= 1
|
||||
# Should find at least one agent
|
||||
agent_entities = [e for e in entities if e.type == "agent"]
|
||||
assert len(agent_entities) >= 1, "Should discover at least one agent"
|
||||
# Verify agents have required properties
|
||||
for agent in agent_entities:
|
||||
assert agent.id, "Agent should have an ID"
|
||||
assert agent.name, "Agent should have a name"
|
||||
|
||||
|
||||
async def test_server_execution_sync(test_entities_dir):
|
||||
"""Test sync execution endpoint."""
|
||||
server = DevServer(entities_dir=test_entities_dir)
|
||||
executor = await server._ensure_executor()
|
||||
|
||||
entities = await executor.discover_entities()
|
||||
agent_id = entities[0].id
|
||||
|
||||
# Use metadata.entity_id for routing
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": agent_id},
|
||||
input="San Francisco",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
response = await executor.execute_sync(request)
|
||||
assert response.model == "devui" # Response model defaults to 'devui' when not specified
|
||||
assert len(response.output) > 0
|
||||
|
||||
|
||||
async def test_server_execution_streaming(test_entities_dir):
|
||||
"""Test streaming execution endpoint."""
|
||||
server = DevServer(entities_dir=test_entities_dir)
|
||||
executor = await server._ensure_executor()
|
||||
|
||||
entities = await executor.discover_entities()
|
||||
agent_id = entities[0].id
|
||||
|
||||
# Use metadata.entity_id for routing
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": agent_id},
|
||||
input="New York",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
event_count = 0
|
||||
async for _event in executor.execute_streaming(request):
|
||||
event_count += 1
|
||||
if event_count > 5: # Limit for testing
|
||||
break
|
||||
|
||||
assert event_count > 0
|
||||
|
||||
|
||||
def test_configuration():
|
||||
"""Test basic configuration."""
|
||||
server = DevServer(entities_dir="test", port=9000, host="localhost")
|
||||
assert server.port == 9000
|
||||
assert server.host == "localhost"
|
||||
assert server.entities_dir == "test"
|
||||
assert server.cors_origins == ["*"]
|
||||
assert server.ui_enabled
|
||||
|
||||
|
||||
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)
|
||||
|
||||
assert types == [str, dict]
|
||||
|
||||
|
||||
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)
|
||||
|
||||
assert str in types
|
||||
assert int in types
|
||||
|
||||
|
||||
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])
|
||||
|
||||
assert string_first is str
|
||||
assert dict_first is dict
|
||||
assert fallback is int
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credential_cleanup() -> None:
|
||||
"""Test that async credentials are properly closed during server cleanup."""
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
from agent_framework import ChatAgent
|
||||
|
||||
# Create mock credential with async close
|
||||
mock_credential = AsyncMock()
|
||||
mock_credential.close = AsyncMock()
|
||||
|
||||
# Create mock chat client with credential
|
||||
mock_client = Mock()
|
||||
mock_client.async_credential = mock_credential
|
||||
mock_client.model_id = "test-model"
|
||||
mock_client.function_invocation_configuration = None
|
||||
|
||||
# Create agent with mock client
|
||||
agent = ChatAgent(name="TestAgent", chat_client=mock_client, instructions="Test agent")
|
||||
|
||||
# Create DevUI server with agent
|
||||
server = DevServer()
|
||||
server._pending_entities = [agent]
|
||||
await server._ensure_executor()
|
||||
|
||||
# Run cleanup
|
||||
await server._cleanup_entities()
|
||||
|
||||
# Verify credential.close() was called
|
||||
assert mock_credential.close.called, "Async credential close should have been called"
|
||||
assert mock_credential.close.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credential_cleanup_error_handling() -> None:
|
||||
"""Test that credential cleanup errors are handled gracefully."""
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
from agent_framework import ChatAgent
|
||||
|
||||
# Create mock credential that raises error on close
|
||||
mock_credential = AsyncMock()
|
||||
mock_credential.close = AsyncMock(side_effect=Exception("Close failed"))
|
||||
|
||||
# Create mock chat client with credential
|
||||
mock_client = Mock()
|
||||
mock_client.async_credential = mock_credential
|
||||
mock_client.model_id = "test-model"
|
||||
mock_client.function_invocation_configuration = None
|
||||
|
||||
# Create agent with mock client
|
||||
agent = ChatAgent(name="TestAgent", chat_client=mock_client, instructions="Test agent")
|
||||
|
||||
# Create DevUI server with agent
|
||||
server = DevServer()
|
||||
server._pending_entities = [agent]
|
||||
await server._ensure_executor()
|
||||
|
||||
# Run cleanup - should not raise despite credential error
|
||||
await server._cleanup_entities()
|
||||
|
||||
# Verify close was attempted
|
||||
assert mock_credential.close.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_credential_attributes() -> None:
|
||||
"""Test that we check all common credential attribute names."""
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
from agent_framework import ChatAgent
|
||||
|
||||
# Create mock credentials
|
||||
mock_cred1 = Mock()
|
||||
mock_cred1.close = Mock()
|
||||
mock_cred2 = AsyncMock()
|
||||
mock_cred2.close = AsyncMock()
|
||||
|
||||
# Create mock chat client with multiple credential attributes
|
||||
mock_client = Mock()
|
||||
mock_client.credential = mock_cred1
|
||||
mock_client.async_credential = mock_cred2
|
||||
mock_client.model_id = "test-model"
|
||||
mock_client.function_invocation_configuration = None
|
||||
|
||||
# Create agent with mock client
|
||||
agent = ChatAgent(name="TestAgent", chat_client=mock_client, instructions="Test agent")
|
||||
|
||||
# Create DevUI server with agent
|
||||
server = DevServer()
|
||||
server._pending_entities = [agent]
|
||||
await server._ensure_executor()
|
||||
|
||||
# Run cleanup
|
||||
await server._cleanup_entities()
|
||||
|
||||
# Verify both credentials were closed
|
||||
assert mock_cred1.close.called, "Sync credential should be closed"
|
||||
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 accessible in both modes (UI needs this)
|
||||
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")
|
||||
# Should return 404 (entity doesn't exist) or 500 (other error), but NOT 403 (forbidden)
|
||||
# User mode needs entity info to display workflows/agents in the UI
|
||||
assert user_response.status_code in [200, 404, 500] # Not 403
|
||||
|
||||
# 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)
|
||||
# 200=success, 400=missing entity_id in metadata, 404=entity not found
|
||||
assert dev_response.status_code in [200, 400, 404]
|
||||
|
||||
user_response = user_client.post("/v1/responses", json=chat_payload)
|
||||
assert user_response.status_code in [200, 400, 404]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner
|
||||
async def run_tests():
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create test agent
|
||||
agent_file = temp_path / "weather_agent.py"
|
||||
agent_file.write_text("""
|
||||
class WeatherAgent:
|
||||
name = "Weather Agent"
|
||||
description = "Gets weather information"
|
||||
|
||||
def run(self, input_str, *, stream: bool = False, thread=None, **kwargs):
|
||||
return f"Weather in {input_str} is sunny"
|
||||
""")
|
||||
|
||||
server = DevServer(entities_dir=str(temp_path))
|
||||
executor = await server._ensure_executor()
|
||||
|
||||
entities = await executor.discover_entities()
|
||||
|
||||
if entities:
|
||||
request = AgentFrameworkRequest(
|
||||
metadata={"entity_id": entities[0].id},
|
||||
input="test location",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
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",
|
||||
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