diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 9f2ad10d85..8dc7c00655 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -1050,6 +1050,50 @@ class DataContent(BaseContent): def has_top_level_media_type(self, top_level_media_type: Literal["application", "audio", "image", "text"]) -> bool: return _has_top_level_media_type(self.media_type, top_level_media_type) + @staticmethod + def detect_image_format_from_base64(image_base64: str) -> str: + """Detect image format from base64 data by examining the binary header. + + Args: + image_base64: Base64 encoded image data + + Returns: + Image format as string (png, jpeg, webp, gif) with png as fallback + """ + try: + # Constants for image format detection + # ~75 bytes of binary data should be enough to detect most image formats + FORMAT_DETECTION_BASE64_CHARS = 100 + + # Decode a small portion to detect format + decoded_data = base64.b64decode(image_base64[:FORMAT_DETECTION_BASE64_CHARS]) + if decoded_data.startswith(b"\x89PNG"): + return "png" + if decoded_data.startswith(b"\xff\xd8\xff"): + return "jpeg" + if decoded_data.startswith(b"RIFF") and b"WEBP" in decoded_data[:12]: + return "webp" + if decoded_data.startswith(b"GIF87a") or decoded_data.startswith(b"GIF89a"): + return "gif" + return "png" # Default fallback + except Exception: + return "png" # Fallback if decoding fails + + @classmethod + def create_data_uri_from_base64(cls, image_base64: str) -> tuple[str, str]: + """Create a data URI and media type from base64 image data. + + Args: + image_base64: Base64 encoded image data + + Returns: + Tuple of (data_uri, media_type) + """ + format_type = cls.detect_image_format_from_base64(image_base64) + uri = f"data:image/{format_type};base64,{image_base64}" + media_type = f"image/{format_type}" + return uri, media_type + class UriContent(BaseContent): """Represents a URI content. diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 0c05abbb69..8fa85b7f84 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -2,11 +2,14 @@ import logging from dataclasses import dataclass -from typing import Any +from typing import Any, cast + +from agent_framework import FunctionApprovalRequestContent, FunctionApprovalResponseContent from .._agents import AgentProtocol, ChatAgent from .._threads import AgentThread from .._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage +from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value from ._conversation_state import encode_chat_messages from ._events import ( AgentRunEvent, @@ -14,6 +17,7 @@ from ._events import ( ) from ._executor import Executor, handler from ._message_utils import normalize_messages_input +from ._request_info_mixin import response_handler from ._workflow_context import WorkflowContext logger = logging.getLogger(__name__) @@ -83,6 +87,8 @@ class AgentExecutor(Executor): super().__init__(exec_id) self._agent = agent self._agent_thread = agent_thread or self._agent.get_new_thread() + self._pending_agent_requests: dict[str, FunctionApprovalRequestContent] = {} + self._pending_responses_to_agent: list[FunctionApprovalResponseContent] = [] self._output_response = output_response self._cache: list[ChatMessage] = [] @@ -93,50 +99,6 @@ class AgentExecutor(Executor): return [AgentRunResponse] return [] - async def _run_agent_and_emit(self, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse]) -> None: - """Execute the underlying agent, emit events, and enqueue response. - - Checks ctx.is_streaming() to determine whether to emit incremental AgentRunUpdateEvent - events (streaming mode) or a single AgentRunEvent (non-streaming mode). - """ - if ctx.is_streaming(): - # Streaming mode: emit incremental updates - updates: list[AgentRunResponseUpdate] = [] - async for update in self._agent.run_stream( - self._cache, - thread=self._agent_thread, - ): - updates.append(update) - await ctx.add_event(AgentRunUpdateEvent(self.id, update)) - - if isinstance(self._agent, ChatAgent): - response_format = self._agent.chat_options.response_format - response = AgentRunResponse.from_agent_run_response_updates( - updates, - output_format_type=response_format, - ) - else: - response = AgentRunResponse.from_agent_run_response_updates(updates) - else: - # Non-streaming mode: use run() and emit single event - response = await self._agent.run( - self._cache, - thread=self._agent_thread, - ) - await ctx.add_event(AgentRunEvent(self.id, response)) - - if self._output_response: - await ctx.yield_output(response) - - # Always construct a full conversation snapshot from inputs (cache) - # plus agent outputs (agent_run_response.messages). Do not mutate - # response.messages so AgentRunEvent remains faithful to the raw output. - full_conversation: list[ChatMessage] = list(self._cache) + list(response.messages) - - agent_response = AgentExecutorResponse(self.id, response, full_conversation=full_conversation) - await ctx.send_message(agent_response) - self._cache.clear() - @handler async def run( self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse] @@ -192,6 +154,31 @@ class AgentExecutor(Executor): self._cache = normalize_messages_input(messages) await self._run_agent_and_emit(ctx) + @response_handler + async def handle_user_input_response( + self, + original_request: FunctionApprovalRequestContent, + response: FunctionApprovalResponseContent, + ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse], + ) -> None: + """Handle user input responses for function approvals during agent execution. + + This will hold the executor's execution until all pending user input requests are resolved. + + Args: + original_request: The original function approval request sent by the agent. + response: The user's response to the function approval request. + ctx: The workflow context for emitting events and outputs. + """ + self._pending_responses_to_agent.append(response) + self._pending_agent_requests.pop(original_request.id, None) + + if not self._pending_agent_requests: + # All pending requests have been resolved; resume agent execution + self._cache = normalize_messages_input(ChatMessage(role="user", contents=self._pending_responses_to_agent)) + self._pending_responses_to_agent.clear() + await self._run_agent_and_emit(ctx) + async def snapshot_state(self) -> dict[str, Any]: """Capture current executor state for checkpointing. @@ -226,6 +213,8 @@ class AgentExecutor(Executor): return { "cache": encode_chat_messages(self._cache), "agent_thread": serialized_thread, + "pending_agent_requests": encode_checkpoint_value(self._pending_agent_requests), + "pending_responses_to_agent": encode_checkpoint_value(self._pending_responses_to_agent), } async def restore_state(self, state: dict[str, Any]) -> None: @@ -258,7 +247,109 @@ class AgentExecutor(Executor): else: self._agent_thread = self._agent.get_new_thread() + pending_requests_payload = state.get("pending_agent_requests") + if pending_requests_payload: + self._pending_agent_requests = decode_checkpoint_value(pending_requests_payload) + + pending_responses_payload = state.get("pending_responses_to_agent") + if pending_responses_payload: + self._pending_responses_to_agent = decode_checkpoint_value(pending_responses_payload) + def reset(self) -> None: """Reset the internal cache of the executor.""" logger.debug("AgentExecutor %s: Resetting cache", self.id) self._cache.clear() + + async def _run_agent_and_emit(self, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse]) -> None: + """Execute the underlying agent, emit events, and enqueue response. + + Checks ctx.is_streaming() to determine whether to emit incremental AgentRunUpdateEvent + events (streaming mode) or a single AgentRunEvent (non-streaming mode). + """ + if ctx.is_streaming(): + # Streaming mode: emit incremental updates + response = await self._run_agent_streaming(cast(WorkflowContext, ctx)) + else: + # Non-streaming mode: use run() and emit single event + response = await self._run_agent(cast(WorkflowContext, ctx)) + + if response is None: + # Agent did not complete (e.g., waiting for user input); do not emit response + logger.info("AgentExecutor %s: Agent did not complete, awaiting user input", self.id) + return + + if self._output_response: + await ctx.yield_output(response) + + # Always construct a full conversation snapshot from inputs (cache) + # plus agent outputs (agent_run_response.messages). Do not mutate + # response.messages so AgentRunEvent remains faithful to the raw output. + full_conversation: list[ChatMessage] = list(self._cache) + list(response.messages) + + agent_response = AgentExecutorResponse(self.id, response, full_conversation=full_conversation) + await ctx.send_message(agent_response) + self._cache.clear() + + async def _run_agent(self, ctx: WorkflowContext) -> AgentRunResponse | None: + """Execute the underlying agent in non-streaming mode. + + Args: + ctx: The workflow context for emitting events. + + Returns: + The complete AgentRunResponse, or None if waiting for user input. + """ + response = await self._agent.run( + self._cache, + thread=self._agent_thread, + ) + await ctx.add_event(AgentRunEvent(self.id, response)) + + # Handle any user input requests + if response.user_input_requests: + for user_input_request in response.user_input_requests: + self._pending_agent_requests[user_input_request.id] = user_input_request + await ctx.request_info(user_input_request, FunctionApprovalResponseContent) + return None + + return response + + async def _run_agent_streaming(self, ctx: WorkflowContext) -> AgentRunResponse | None: + """Execute the underlying agent in streaming mode and collect the full response. + + Args: + ctx: The workflow context for emitting events. + + Returns: + The complete AgentRunResponse, or None if waiting for user input. + """ + updates: list[AgentRunResponseUpdate] = [] + user_input_requests: list[FunctionApprovalRequestContent] = [] + async for update in self._agent.run_stream( + self._cache, + thread=self._agent_thread, + ): + updates.append(update) + await ctx.add_event(AgentRunUpdateEvent(self.id, update)) + + if update.user_input_requests: + user_input_requests.extend(update.user_input_requests) + + # Build the final AgentRunResponse from the collected updates + if isinstance(self._agent, ChatAgent): + response_format = self._agent.chat_options.response_format + response = AgentRunResponse.from_agent_run_response_updates( + updates, + output_format_type=response_format, + ) + else: + response = AgentRunResponse.from_agent_run_response_updates(updates) + + # Handle any user input requests after the streaming completes + if user_input_requests: + for user_input_request in user_input_requests: + self._pending_agent_requests[user_input_request.id] = user_input_request + await ctx.request_info(user_input_request, FunctionApprovalResponseContent) + return None + + return response diff --git a/python/packages/core/agent_framework/openai/_responses_client.py b/python/packages/core/agent_framework/openai/_responses_client.py index 279180e0ee..149fe4bfac 100644 --- a/python/packages/core/agent_framework/openai/_responses_client.py +++ b/python/packages/core/agent_framework/openai/_responses_client.py @@ -293,6 +293,14 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): # Map the parameter name and remove the old one mapped_tool[api_param] = mapped_tool.pop(user_param) + # Validate partial_images parameter for streaming image generation + # OpenAI API requires partial_images to be between 0-3 (inclusive) for image_generation tool + # Reference: https://platform.openai.com/docs/api-reference/responses/create#responses_create-tools-image_generation_tool-partial_images + if "partial_images" in mapped_tool: + partial_images = mapped_tool["partial_images"] + if not isinstance(partial_images, int) or partial_images < 0 or partial_images > 3: + raise ValueError("partial_images must be an integer between 0 and 3 (inclusive).") + response_tools.append(mapped_tool) else: response_tools.append(tool_dict) @@ -695,29 +703,8 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): uri = item.result media_type = None if not uri.startswith("data:"): - # Raw base64 string - convert to proper data URI format - # Detect format from base64 data - import base64 - - try: - # Decode a small portion to detect format - decoded_data = base64.b64decode(uri[:100]) # First ~75 bytes should be enough - if decoded_data.startswith(b"\x89PNG"): - format_type = "png" - elif decoded_data.startswith(b"\xff\xd8\xff"): - format_type = "jpeg" - elif decoded_data.startswith(b"RIFF") and b"WEBP" in decoded_data[:12]: - format_type = "webp" - elif decoded_data.startswith(b"GIF87a") or decoded_data.startswith(b"GIF89a"): - format_type = "gif" - else: - # Default to png if format cannot be detected - format_type = "png" - except Exception: - # Fallback to png if decoding fails - format_type = "png" - uri = f"data:image/{format_type};base64,{uri}" - media_type = f"image/{format_type}" + # Raw base64 string - convert to proper data URI format using helper + uri, media_type = DataContent.create_data_uri_from_base64(uri) else: # Parse media type from existing data URI try: @@ -933,6 +920,25 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): raw_representation=event, ) ) + case "response.image_generation_call.partial_image": + # Handle streaming partial image generation + image_base64 = event.partial_image_b64 + partial_index = event.partial_image_index + + # Use helper function to create data URI from base64 + uri, media_type = DataContent.create_data_uri_from_base64(image_base64) + + contents.append( + DataContent( + uri=uri, + media_type=media_type, + additional_properties={ + "partial_image_index": partial_index, + "is_partial_image": True, + }, + raw_representation=event, + ) + ) case _: logger.debug("Unparsed event of type: %s: %s", event.type, event) diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 909e72a0a0..38a3fe414e 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import base64 from collections.abc import AsyncIterable from typing import Any @@ -166,6 +167,57 @@ def test_data_content_empty(): DataContent(uri="") +def test_data_content_detect_image_format_from_base64(): + """Test the detect_image_format_from_base64 static method.""" + # Test each supported format + png_data = b"\x89PNG\r\n\x1a\n" + b"fake_data" + assert DataContent.detect_image_format_from_base64(base64.b64encode(png_data).decode()) == "png" + + jpeg_data = b"\xff\xd8\xff\xe0" + b"fake_data" + assert DataContent.detect_image_format_from_base64(base64.b64encode(jpeg_data).decode()) == "jpeg" + + webp_data = b"RIFF" + b"1234" + b"WEBP" + b"fake_data" + assert DataContent.detect_image_format_from_base64(base64.b64encode(webp_data).decode()) == "webp" + + gif_data = b"GIF89a" + b"fake_data" + assert DataContent.detect_image_format_from_base64(base64.b64encode(gif_data).decode()) == "gif" + + # Test fallback behavior + unknown_data = b"UNKNOWN_FORMAT" + assert DataContent.detect_image_format_from_base64(base64.b64encode(unknown_data).decode()) == "png" + + # Test error handling + assert DataContent.detect_image_format_from_base64("invalid_base64!") == "png" + assert DataContent.detect_image_format_from_base64("") == "png" + + +def test_data_content_create_data_uri_from_base64(): + """Test the create_data_uri_from_base64 class method.""" + # Test with PNG data + png_data = b"\x89PNG\r\n\x1a\n" + b"fake_data" + png_base64 = base64.b64encode(png_data).decode() + uri, media_type = DataContent.create_data_uri_from_base64(png_base64) + + assert uri == f"data:image/png;base64,{png_base64}" + assert media_type == "image/png" + + # Test with different format + jpeg_data = b"\xff\xd8\xff\xe0" + b"fake_data" + jpeg_base64 = base64.b64encode(jpeg_data).decode() + uri, media_type = DataContent.create_data_uri_from_base64(jpeg_base64) + + assert uri == f"data:image/jpeg;base64,{jpeg_base64}" + assert media_type == "image/jpeg" + + # Test fallback for unknown format + unknown_data = b"UNKNOWN_FORMAT" + unknown_base64 = base64.b64encode(unknown_data).decode() + uri, media_type = DataContent.create_data_uri_from_base64(unknown_base64) + + assert uri == f"data:image/png;base64,{unknown_base64}" + assert media_type == "image/png" + + # region UriContent diff --git a/python/packages/core/tests/openai/test_openai_responses_client.py b/python/packages/core/tests/openai/test_openai_responses_client.py index cb4f0dc0d3..5ff4bb3de3 100644 --- a/python/packages/core/tests/openai/test_openai_responses_client.py +++ b/python/packages/core/tests/openai/test_openai_responses_client.py @@ -36,6 +36,7 @@ from agent_framework import ( HostedMCPTool, HostedVectorStoreContent, HostedWebSearchTool, + MCPStreamableHTTPTool, Role, TextContent, TextReasoningContent, @@ -946,6 +947,505 @@ def test_streaming_response_basic_structure() -> None: assert response.raw_representation is mock_event +def test_service_response_exception_includes_original_error_details() -> None: + """Test that ServiceResponseException messages include original error details in the new format.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + messages = [ChatMessage(role="user", text="test message")] + + mock_response = MagicMock() + original_error_message = "Request rate limit exceeded" + mock_error = BadRequestError( + message=original_error_message, + response=mock_response, + body={"error": {"code": "rate_limit", "message": original_error_message}}, + ) + mock_error.code = "rate_limit" + + with ( + patch.object(client.client.responses, "parse", side_effect=mock_error), + pytest.raises(ServiceResponseException) as exc_info, + ): + asyncio.run(client.get_response(messages=messages, response_format=OutputStruct)) + + exception_message = str(exc_info.value) + assert "service failed to complete the prompt:" in exception_message + assert original_error_message in exception_message + + +def test_get_streaming_response_with_response_format() -> None: + """Test get_streaming_response with response_format.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + messages = [ChatMessage(role="user", text="Test streaming with format")] + + # It will fail due to invalid API key, but exercises the code path + with pytest.raises(ServiceResponseException): + + async def run_streaming(): + async for _ in client.get_streaming_response(messages=messages, response_format=OutputStruct): + pass + + asyncio.run(run_streaming()) + + +def test_openai_content_parser_image_content() -> None: + """Test _openai_content_parser with image content variations.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + + # Test image content with detail parameter and file_id + image_content_with_detail = UriContent( + uri="https://example.com/image.jpg", + media_type="image/jpeg", + additional_properties={"detail": "high", "file_id": "file_123"}, + ) + result = client._openai_content_parser(Role.USER, image_content_with_detail, {}) # type: ignore + assert result["type"] == "input_image" + assert result["image_url"] == "https://example.com/image.jpg" + assert result["detail"] == "high" + assert result["file_id"] == "file_123" + + # Test image content without additional properties (defaults) + image_content_basic = UriContent(uri="https://example.com/basic.png", media_type="image/png") + result = client._openai_content_parser(Role.USER, image_content_basic, {}) # type: ignore + assert result["type"] == "input_image" + assert result["detail"] == "auto" + assert result["file_id"] is None + + +def test_openai_content_parser_audio_content() -> None: + """Test _openai_content_parser with audio content variations.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + + # Test WAV audio content + wav_content = UriContent(uri="data:audio/wav;base64,abc123", media_type="audio/wav") + result = client._openai_content_parser(Role.USER, wav_content, {}) # type: ignore + assert result["type"] == "input_audio" + assert result["input_audio"]["data"] == "data:audio/wav;base64,abc123" + assert result["input_audio"]["format"] == "wav" + + # Test MP3 audio content + mp3_content = UriContent(uri="data:audio/mp3;base64,def456", media_type="audio/mp3") + result = client._openai_content_parser(Role.USER, mp3_content, {}) # type: ignore + assert result["type"] == "input_audio" + assert result["input_audio"]["format"] == "mp3" + + +def test_openai_content_parser_unsupported_content() -> None: + """Test _openai_content_parser with unsupported content types.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + + # Test unsupported audio format + unsupported_audio = UriContent(uri="data:audio/ogg;base64,ghi789", media_type="audio/ogg") + result = client._openai_content_parser(Role.USER, unsupported_audio, {}) # type: ignore + assert result == {} + + # Test non-media content + text_uri_content = UriContent(uri="https://example.com/document.txt", media_type="text/plain") + result = client._openai_content_parser(Role.USER, text_uri_content, {}) # type: ignore + assert result == {} + + +def test_create_streaming_response_content_code_interpreter() -> None: + """Test _create_streaming_response_content with code_interpreter_call.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + chat_options = ChatOptions() + function_call_ids: dict[int, tuple[str, str]] = {} + + mock_event_image = MagicMock() + mock_event_image.type = "response.output_item.added" + mock_item_image = MagicMock() + mock_item_image.type = "code_interpreter_call" + mock_image_output = MagicMock() + mock_image_output.type = "image" + mock_image_output.url = "https://example.com/plot.png" + mock_item_image.outputs = [mock_image_output] + mock_item_image.code = None + mock_event_image.item = mock_item_image + + result = client._create_streaming_response_content(mock_event_image, chat_options, function_call_ids) # type: ignore + assert len(result.contents) == 1 + assert isinstance(result.contents[0], UriContent) + assert result.contents[0].uri == "https://example.com/plot.png" + assert result.contents[0].media_type == "image" + + +def test_create_streaming_response_content_reasoning() -> None: + """Test _create_streaming_response_content with reasoning content.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + chat_options = ChatOptions() + function_call_ids: dict[int, tuple[str, str]] = {} + + mock_event_reasoning = MagicMock() + mock_event_reasoning.type = "response.output_item.added" + mock_item_reasoning = MagicMock() + mock_item_reasoning.type = "reasoning" + mock_reasoning_content = MagicMock() + mock_reasoning_content.text = "Analyzing the problem step by step..." + mock_item_reasoning.content = [mock_reasoning_content] + mock_item_reasoning.summary = ["Problem analysis summary"] + mock_event_reasoning.item = mock_item_reasoning + + result = client._create_streaming_response_content(mock_event_reasoning, chat_options, function_call_ids) # type: ignore + assert len(result.contents) == 1 + assert isinstance(result.contents[0], TextReasoningContent) + assert result.contents[0].text == "Analyzing the problem step by step..." + if result.contents[0].additional_properties: + assert result.contents[0].additional_properties["summary"] == "Problem analysis summary" + + +def test_openai_content_parser_text_reasoning_comprehensive() -> None: + """Test _openai_content_parser with TextReasoningContent all additional properties.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + + # Test TextReasoningContent with all additional properties + comprehensive_reasoning = TextReasoningContent( + text="Comprehensive reasoning summary", + additional_properties={ + "status": "in_progress", + "reasoning_text": "Step-by-step analysis", + "encrypted_content": "secure_data_456", + }, + ) + result = client._openai_content_parser(Role.ASSISTANT, comprehensive_reasoning, {}) # type: ignore + assert result["type"] == "reasoning" + assert result["summary"]["text"] == "Comprehensive reasoning summary" + assert result["status"] == "in_progress" + assert result["content"]["type"] == "reasoning_text" + assert result["content"]["text"] == "Step-by-step analysis" + assert result["encrypted_content"] == "secure_data_456" + + +def test_streaming_reasoning_text_delta_event() -> None: + """Test reasoning text delta event creates TextReasoningContent.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + chat_options = ChatOptions() + function_call_ids: dict[int, tuple[str, str]] = {} + + event = ResponseReasoningTextDeltaEvent( + type="response.reasoning_text.delta", + content_index=0, + item_id="reasoning_123", + output_index=0, + sequence_number=1, + delta="reasoning delta", + ) + + with patch.object(client, "_get_metadata_from_response", return_value={}) as mock_metadata: + response = client._create_streaming_response_content(event, chat_options, function_call_ids) # type: ignore + + assert len(response.contents) == 1 + assert isinstance(response.contents[0], TextReasoningContent) + assert response.contents[0].text == "reasoning delta" + assert response.contents[0].raw_representation == event + mock_metadata.assert_called_once_with(event) + + +def test_streaming_reasoning_text_done_event() -> None: + """Test reasoning text done event creates TextReasoningContent with complete text.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + chat_options = ChatOptions() + function_call_ids: dict[int, tuple[str, str]] = {} + + event = ResponseReasoningTextDoneEvent( + type="response.reasoning_text.done", + content_index=0, + item_id="reasoning_456", + output_index=0, + sequence_number=2, + text="complete reasoning", + ) + + with patch.object(client, "_get_metadata_from_response", return_value={"test": "data"}) as mock_metadata: + response = client._create_streaming_response_content(event, chat_options, function_call_ids) # type: ignore + + assert len(response.contents) == 1 + assert isinstance(response.contents[0], TextReasoningContent) + assert response.contents[0].text == "complete reasoning" + assert response.contents[0].raw_representation == event + mock_metadata.assert_called_once_with(event) + assert response.additional_properties == {"test": "data"} + + +def test_streaming_reasoning_summary_text_delta_event() -> None: + """Test reasoning summary text delta event creates TextReasoningContent.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + chat_options = ChatOptions() + function_call_ids: dict[int, tuple[str, str]] = {} + + event = ResponseReasoningSummaryTextDeltaEvent( + type="response.reasoning_summary_text.delta", + item_id="summary_789", + output_index=0, + sequence_number=3, + summary_index=0, + delta="summary delta", + ) + + with patch.object(client, "_get_metadata_from_response", return_value={}) as mock_metadata: + response = client._create_streaming_response_content(event, chat_options, function_call_ids) # type: ignore + + assert len(response.contents) == 1 + assert isinstance(response.contents[0], TextReasoningContent) + assert response.contents[0].text == "summary delta" + assert response.contents[0].raw_representation == event + mock_metadata.assert_called_once_with(event) + + +def test_streaming_reasoning_summary_text_done_event() -> None: + """Test reasoning summary text done event creates TextReasoningContent with complete text.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + chat_options = ChatOptions() + function_call_ids: dict[int, tuple[str, str]] = {} + + event = ResponseReasoningSummaryTextDoneEvent( + type="response.reasoning_summary_text.done", + item_id="summary_012", + output_index=0, + sequence_number=4, + summary_index=0, + text="complete summary", + ) + + with patch.object(client, "_get_metadata_from_response", return_value={"custom": "meta"}) as mock_metadata: + response = client._create_streaming_response_content(event, chat_options, function_call_ids) # type: ignore + + assert len(response.contents) == 1 + assert isinstance(response.contents[0], TextReasoningContent) + assert response.contents[0].text == "complete summary" + assert response.contents[0].raw_representation == event + mock_metadata.assert_called_once_with(event) + assert response.additional_properties == {"custom": "meta"} + + +def test_streaming_reasoning_events_preserve_metadata() -> None: + """Test that reasoning events preserve metadata like regular text events.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + chat_options = ChatOptions() + function_call_ids: dict[int, tuple[str, str]] = {} + + text_event = ResponseTextDeltaEvent( + type="response.output_text.delta", + content_index=0, + item_id="text_item", + output_index=0, + sequence_number=1, + logprobs=[], + delta="text", + ) + + reasoning_event = ResponseReasoningTextDeltaEvent( + type="response.reasoning_text.delta", + content_index=0, + item_id="reasoning_item", + output_index=0, + sequence_number=2, + delta="reasoning", + ) + + with patch.object(client, "_get_metadata_from_response", return_value={"test": "metadata"}): + text_response = client._create_streaming_response_content(text_event, chat_options, function_call_ids) # type: ignore + reasoning_response = client._create_streaming_response_content(reasoning_event, chat_options, function_call_ids) # type: ignore + + # Both should preserve metadata + assert text_response.additional_properties == {"test": "metadata"} + assert reasoning_response.additional_properties == {"test": "metadata"} + + # Content types should be different + assert isinstance(text_response.contents[0], TextContent) + assert isinstance(reasoning_response.contents[0], TextReasoningContent) + + +def test_create_response_content_image_generation_raw_base64(): + """Test image generation response parsing with raw base64 string.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + + # Create a mock response with raw base64 image data (PNG signature) + mock_response = MagicMock() + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.usage = None + mock_response.id = "test-response-id" + mock_response.model = "test-model" + mock_response.created_at = 1234567890 + + # Mock image generation output item with raw base64 (PNG format) + png_signature = b"\x89PNG\r\n\x1a\n" + mock_base64 = base64.b64encode(png_signature + b"fake_png_data_here").decode() + + mock_item = MagicMock() + mock_item.type = "image_generation_call" + mock_item.result = mock_base64 + + mock_response.output = [mock_item] + + with patch.object(client, "_get_metadata_from_response", return_value={}): + response = client._create_response_content(mock_response, chat_options=ChatOptions()) # type: ignore + + # Verify the response contains DataContent with proper URI and media_type + assert len(response.messages[0].contents) == 1 + content = response.messages[0].contents[0] + assert isinstance(content, DataContent) + assert content.uri.startswith("data:image/png;base64,") + assert content.media_type == "image/png" + + +def test_create_response_content_image_generation_existing_data_uri(): + """Test image generation response parsing with existing data URI.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + + # Create a mock response with existing data URI + mock_response = MagicMock() + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.usage = None + mock_response.id = "test-response-id" + mock_response.model = "test-model" + mock_response.created_at = 1234567890 + + # Mock image generation output item with existing data URI (valid WEBP header) + webp_signature = b"RIFF" + b"\x12\x00\x00\x00" + b"WEBP" + valid_webp_base64 = base64.b64encode(webp_signature + b"VP8 fake_data").decode() + mock_item = MagicMock() + mock_item.type = "image_generation_call" + mock_item.result = f"data:image/webp;base64,{valid_webp_base64}" + + mock_response.output = [mock_item] + + with patch.object(client, "_get_metadata_from_response", return_value={}): + response = client._create_response_content(mock_response, chat_options=ChatOptions()) # type: ignore + + # Verify the response contains DataContent with proper media_type parsed from URI + assert len(response.messages[0].contents) == 1 + content = response.messages[0].contents[0] + assert isinstance(content, DataContent) + assert content.uri == f"data:image/webp;base64,{valid_webp_base64}" + assert content.media_type == "image/webp" + + +def test_create_response_content_image_generation_format_detection(): + """Test different image format detection from base64 data.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + + # Test JPEG detection + jpeg_signature = b"\xff\xd8\xff" + mock_base64_jpeg = base64.b64encode(jpeg_signature + b"fake_jpeg_data").decode() + + mock_response_jpeg = MagicMock() + mock_response_jpeg.output_parsed = None + mock_response_jpeg.metadata = {} + mock_response_jpeg.usage = None + mock_response_jpeg.id = "test-id" + mock_response_jpeg.model = "test-model" + mock_response_jpeg.created_at = 1234567890 + + mock_item_jpeg = MagicMock() + mock_item_jpeg.type = "image_generation_call" + mock_item_jpeg.result = mock_base64_jpeg + mock_response_jpeg.output = [mock_item_jpeg] + + with patch.object(client, "_get_metadata_from_response", return_value={}): + response_jpeg = client._create_response_content(mock_response_jpeg, chat_options=ChatOptions()) # type: ignore + content_jpeg = response_jpeg.messages[0].contents[0] + assert isinstance(content_jpeg, DataContent) + assert content_jpeg.media_type == "image/jpeg" + assert "data:image/jpeg;base64," in content_jpeg.uri + + # Test WEBP detection + webp_signature = b"RIFF" + b"\x00\x00\x00\x00" + b"WEBP" + mock_base64_webp = base64.b64encode(webp_signature + b"fake_webp_data").decode() + + mock_response_webp = MagicMock() + mock_response_webp.output_parsed = None + mock_response_webp.metadata = {} + mock_response_webp.usage = None + mock_response_webp.id = "test-id" + mock_response_webp.model = "test-model" + mock_response_webp.created_at = 1234567890 + + mock_item_webp = MagicMock() + mock_item_webp.type = "image_generation_call" + mock_item_webp.result = mock_base64_webp + mock_response_webp.output = [mock_item_webp] + + with patch.object(client, "_get_metadata_from_response", return_value={}): + response_webp = client._create_response_content(mock_response_webp, chat_options=ChatOptions()) # type: ignore + content_webp = response_webp.messages[0].contents[0] + assert isinstance(content_webp, DataContent) + assert content_webp.media_type == "image/webp" + assert "data:image/webp;base64," in content_webp.uri + + +def test_create_response_content_image_generation_fallback(): + """Test image generation with invalid base64 falls back to PNG.""" + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + + # Create a mock response with invalid base64 + mock_response = MagicMock() + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.usage = None + mock_response.id = "test-response-id" + mock_response.model = "test-model" + mock_response.created_at = 1234567890 + + # Mock image generation output item with unrecognized format (should fall back to PNG) + unrecognized_data = b"UNKNOWN_FORMAT" + b"some_binary_data" + unrecognized_base64 = base64.b64encode(unrecognized_data).decode() + mock_item = MagicMock() + mock_item.type = "image_generation_call" + mock_item.result = unrecognized_base64 + + mock_response.output = [mock_item] + + with patch.object(client, "_get_metadata_from_response", return_value={}): + response = client._create_response_content(mock_response, chat_options=ChatOptions()) # type: ignore + + # Verify it falls back to PNG format for unrecognized binary data + assert len(response.messages[0].contents) == 1 + content = response.messages[0].contents[0] + assert isinstance(content, DataContent) + assert content.media_type == "image/png" + assert f"data:image/png;base64,{unrecognized_base64}" == content.uri + + +def test_prepare_options_store_parameter_handling() -> None: + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + messages = [ChatMessage(role="user", text="Test message")] + + test_conversation_id = "test-conversation-123" + chat_options = ChatOptions(store=True, conversation_id=test_conversation_id) + options = client._prepare_options(messages, chat_options) # type: ignore + assert options["store"] is True + assert options["previous_response_id"] == test_conversation_id + + chat_options = ChatOptions(store=False, conversation_id="") + options = client._prepare_options(messages, chat_options) # type: ignore + assert options["store"] is False + + chat_options = ChatOptions(store=None, conversation_id=None) + options = client._prepare_options(messages, chat_options) # type: ignore + assert options["store"] is False + assert "previous_response_id" not in options + + chat_options = ChatOptions() + options = client._prepare_options(messages, chat_options) # type: ignore + assert options["store"] is False + assert "previous_response_id" not in options + + +def test_openai_responses_client_with_callable_api_key() -> None: + """Test OpenAIResponsesClient initialization with callable API key.""" + + async def get_api_key() -> str: + return "test-api-key-123" + + client = OpenAIResponsesClient(model_id="gpt-4o", api_key=get_api_key) + + # Verify client was created successfully + assert client.model_id == "gpt-4o" + # OpenAI SDK now manages callable API keys internally + assert client.client is not None + + @pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_response() -> None: @@ -1615,500 +2115,28 @@ async def test_openai_responses_client_agent_hosted_mcp_tool() -> None: assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"]) -def test_service_response_exception_includes_original_error_details() -> None: - """Test that ServiceResponseException messages include original error details in the new format.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - messages = [ChatMessage(role="user", text="test message")] +@pytest.mark.flaky +@skip_if_openai_integration_tests_disabled +async def test_openai_responses_client_agent_local_mcp_tool() -> None: + """Integration test for MCPStreamableHTTPTool with OpenAI Response Agent using Microsoft Learn MCP.""" - mock_response = MagicMock() - original_error_message = "Request rate limit exceeded" - mock_error = BadRequestError( - message=original_error_message, - response=mock_response, - body={"error": {"code": "rate_limit", "message": original_error_message}}, - ) - mock_error.code = "rate_limit" - - with ( - patch.object(client.client.responses, "parse", side_effect=mock_error), - pytest.raises(ServiceResponseException) as exc_info, - ): - asyncio.run(client.get_response(messages=messages, response_format=OutputStruct)) - - exception_message = str(exc_info.value) - assert "service failed to complete the prompt:" in exception_message - assert original_error_message in exception_message - - -def test_get_streaming_response_with_response_format() -> None: - """Test get_streaming_response with response_format.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - messages = [ChatMessage(role="user", text="Test streaming with format")] - - # It will fail due to invalid API key, but exercises the code path - with pytest.raises(ServiceResponseException): - - async def run_streaming(): - async for _ in client.get_streaming_response(messages=messages, response_format=OutputStruct): - pass - - asyncio.run(run_streaming()) - - -def test_openai_content_parser_image_content() -> None: - """Test _openai_content_parser with image content variations.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - - # Test image content with detail parameter and file_id - image_content_with_detail = UriContent( - uri="https://example.com/image.jpg", - media_type="image/jpeg", - additional_properties={"detail": "high", "file_id": "file_123"}, - ) - result = client._openai_content_parser(Role.USER, image_content_with_detail, {}) # type: ignore - assert result["type"] == "input_image" - assert result["image_url"] == "https://example.com/image.jpg" - assert result["detail"] == "high" - assert result["file_id"] == "file_123" - - # Test image content without additional properties (defaults) - image_content_basic = UriContent(uri="https://example.com/basic.png", media_type="image/png") - result = client._openai_content_parser(Role.USER, image_content_basic, {}) # type: ignore - assert result["type"] == "input_image" - assert result["detail"] == "auto" - assert result["file_id"] is None - - -def test_openai_content_parser_audio_content() -> None: - """Test _openai_content_parser with audio content variations.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - - # Test WAV audio content - wav_content = UriContent(uri="data:audio/wav;base64,abc123", media_type="audio/wav") - result = client._openai_content_parser(Role.USER, wav_content, {}) # type: ignore - assert result["type"] == "input_audio" - assert result["input_audio"]["data"] == "data:audio/wav;base64,abc123" - assert result["input_audio"]["format"] == "wav" - - # Test MP3 audio content - mp3_content = UriContent(uri="data:audio/mp3;base64,def456", media_type="audio/mp3") - result = client._openai_content_parser(Role.USER, mp3_content, {}) # type: ignore - assert result["type"] == "input_audio" - assert result["input_audio"]["format"] == "mp3" - - -def test_openai_content_parser_unsupported_content() -> None: - """Test _openai_content_parser with unsupported content types.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - - # Test unsupported audio format - unsupported_audio = UriContent(uri="data:audio/ogg;base64,ghi789", media_type="audio/ogg") - result = client._openai_content_parser(Role.USER, unsupported_audio, {}) # type: ignore - assert result == {} - - # Test non-media content - text_uri_content = UriContent(uri="https://example.com/document.txt", media_type="text/plain") - result = client._openai_content_parser(Role.USER, text_uri_content, {}) # type: ignore - assert result == {} - - -def test_create_streaming_response_content_code_interpreter() -> None: - """Test _create_streaming_response_content with code_interpreter_call.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - chat_options = ChatOptions() - function_call_ids: dict[int, tuple[str, str]] = {} - - mock_event_image = MagicMock() - mock_event_image.type = "response.output_item.added" - mock_item_image = MagicMock() - mock_item_image.type = "code_interpreter_call" - mock_image_output = MagicMock() - mock_image_output.type = "image" - mock_image_output.url = "https://example.com/plot.png" - mock_item_image.outputs = [mock_image_output] - mock_item_image.code = None - mock_event_image.item = mock_item_image - - result = client._create_streaming_response_content(mock_event_image, chat_options, function_call_ids) # type: ignore - assert len(result.contents) == 1 - assert isinstance(result.contents[0], UriContent) - assert result.contents[0].uri == "https://example.com/plot.png" - assert result.contents[0].media_type == "image" - - -def test_create_streaming_response_content_reasoning() -> None: - """Test _create_streaming_response_content with reasoning content.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - chat_options = ChatOptions() - function_call_ids: dict[int, tuple[str, str]] = {} - - mock_event_reasoning = MagicMock() - mock_event_reasoning.type = "response.output_item.added" - mock_item_reasoning = MagicMock() - mock_item_reasoning.type = "reasoning" - mock_reasoning_content = MagicMock() - mock_reasoning_content.text = "Analyzing the problem step by step..." - mock_item_reasoning.content = [mock_reasoning_content] - mock_item_reasoning.summary = ["Problem analysis summary"] - mock_event_reasoning.item = mock_item_reasoning - - result = client._create_streaming_response_content(mock_event_reasoning, chat_options, function_call_ids) # type: ignore - assert len(result.contents) == 1 - assert isinstance(result.contents[0], TextReasoningContent) - assert result.contents[0].text == "Analyzing the problem step by step..." - if result.contents[0].additional_properties: - assert result.contents[0].additional_properties["summary"] == "Problem analysis summary" - - -def test_openai_content_parser_text_reasoning_comprehensive() -> None: - """Test _openai_content_parser with TextReasoningContent all additional properties.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - - # Test TextReasoningContent with all additional properties - comprehensive_reasoning = TextReasoningContent( - text="Comprehensive reasoning summary", - additional_properties={ - "status": "in_progress", - "reasoning_text": "Step-by-step analysis", - "encrypted_content": "secure_data_456", - }, - ) - result = client._openai_content_parser(Role.ASSISTANT, comprehensive_reasoning, {}) # type: ignore - assert result["type"] == "reasoning" - assert result["summary"]["text"] == "Comprehensive reasoning summary" - assert result["status"] == "in_progress" - assert result["content"]["type"] == "reasoning_text" - assert result["content"]["text"] == "Step-by-step analysis" - assert result["encrypted_content"] == "secure_data_456" - - -def test_streaming_reasoning_text_delta_event() -> None: - """Test reasoning text delta event creates TextReasoningContent.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - chat_options = ChatOptions() - function_call_ids: dict[int, tuple[str, str]] = {} - - event = ResponseReasoningTextDeltaEvent( - type="response.reasoning_text.delta", - content_index=0, - item_id="reasoning_123", - output_index=0, - sequence_number=1, - delta="reasoning delta", + mcp_tool = MCPStreamableHTTPTool( + name="Microsoft Learn MCP", + url="https://learn.microsoft.com/api/mcp", ) - with patch.object(client, "_get_metadata_from_response", return_value={}) as mock_metadata: - response = client._create_streaming_response_content(event, chat_options, function_call_ids) # type: ignore - - assert len(response.contents) == 1 - assert isinstance(response.contents[0], TextReasoningContent) - assert response.contents[0].text == "reasoning delta" - assert response.contents[0].raw_representation == event - mock_metadata.assert_called_once_with(event) - - -def test_streaming_reasoning_text_done_event() -> None: - """Test reasoning text done event creates TextReasoningContent with complete text.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - chat_options = ChatOptions() - function_call_ids: dict[int, tuple[str, str]] = {} - - event = ResponseReasoningTextDoneEvent( - type="response.reasoning_text.done", - content_index=0, - item_id="reasoning_456", - output_index=0, - sequence_number=2, - text="complete reasoning", - ) - - with patch.object(client, "_get_metadata_from_response", return_value={"test": "data"}) as mock_metadata: - response = client._create_streaming_response_content(event, chat_options, function_call_ids) # type: ignore - - assert len(response.contents) == 1 - assert isinstance(response.contents[0], TextReasoningContent) - assert response.contents[0].text == "complete reasoning" - assert response.contents[0].raw_representation == event - mock_metadata.assert_called_once_with(event) - assert response.additional_properties == {"test": "data"} - - -def test_streaming_reasoning_summary_text_delta_event() -> None: - """Test reasoning summary text delta event creates TextReasoningContent.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - chat_options = ChatOptions() - function_call_ids: dict[int, tuple[str, str]] = {} - - event = ResponseReasoningSummaryTextDeltaEvent( - type="response.reasoning_summary_text.delta", - item_id="summary_789", - output_index=0, - sequence_number=3, - summary_index=0, - delta="summary delta", - ) - - with patch.object(client, "_get_metadata_from_response", return_value={}) as mock_metadata: - response = client._create_streaming_response_content(event, chat_options, function_call_ids) # type: ignore - - assert len(response.contents) == 1 - assert isinstance(response.contents[0], TextReasoningContent) - assert response.contents[0].text == "summary delta" - assert response.contents[0].raw_representation == event - mock_metadata.assert_called_once_with(event) - - -def test_streaming_reasoning_summary_text_done_event() -> None: - """Test reasoning summary text done event creates TextReasoningContent with complete text.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - chat_options = ChatOptions() - function_call_ids: dict[int, tuple[str, str]] = {} - - event = ResponseReasoningSummaryTextDoneEvent( - type="response.reasoning_summary_text.done", - item_id="summary_012", - output_index=0, - sequence_number=4, - summary_index=0, - text="complete summary", - ) - - with patch.object(client, "_get_metadata_from_response", return_value={"custom": "meta"}) as mock_metadata: - response = client._create_streaming_response_content(event, chat_options, function_call_ids) # type: ignore - - assert len(response.contents) == 1 - assert isinstance(response.contents[0], TextReasoningContent) - assert response.contents[0].text == "complete summary" - assert response.contents[0].raw_representation == event - mock_metadata.assert_called_once_with(event) - assert response.additional_properties == {"custom": "meta"} - - -def test_streaming_reasoning_events_preserve_metadata() -> None: - """Test that reasoning events preserve metadata like regular text events.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - chat_options = ChatOptions() - function_call_ids: dict[int, tuple[str, str]] = {} - - text_event = ResponseTextDeltaEvent( - type="response.output_text.delta", - content_index=0, - item_id="text_item", - output_index=0, - sequence_number=1, - logprobs=[], - delta="text", - ) - - reasoning_event = ResponseReasoningTextDeltaEvent( - type="response.reasoning_text.delta", - content_index=0, - item_id="reasoning_item", - output_index=0, - sequence_number=2, - delta="reasoning", - ) - - with patch.object(client, "_get_metadata_from_response", return_value={"test": "metadata"}): - text_response = client._create_streaming_response_content(text_event, chat_options, function_call_ids) # type: ignore - reasoning_response = client._create_streaming_response_content(reasoning_event, chat_options, function_call_ids) # type: ignore - - # Both should preserve metadata - assert text_response.additional_properties == {"test": "metadata"} - assert reasoning_response.additional_properties == {"test": "metadata"} - - # Content types should be different - assert isinstance(text_response.contents[0], TextContent) - assert isinstance(reasoning_response.contents[0], TextReasoningContent) - - -def test_create_response_content_image_generation_raw_base64(): - """Test image generation response parsing with raw base64 string.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - - # Create a mock response with raw base64 image data (PNG signature) - mock_response = MagicMock() - mock_response.output_parsed = None - mock_response.metadata = {} - mock_response.usage = None - mock_response.id = "test-response-id" - mock_response.model = "test-model" - mock_response.created_at = 1234567890 - - # Mock image generation output item with raw base64 (PNG format) - png_signature = b"\x89PNG\r\n\x1a\n" - mock_base64 = base64.b64encode(png_signature + b"fake_png_data_here").decode() - - mock_item = MagicMock() - mock_item.type = "image_generation_call" - mock_item.result = mock_base64 - - mock_response.output = [mock_item] - - with patch.object(client, "_get_metadata_from_response", return_value={}): - response = client._create_response_content(mock_response, chat_options=ChatOptions()) # type: ignore - - # Verify the response contains DataContent with proper URI and media_type - assert len(response.messages[0].contents) == 1 - content = response.messages[0].contents[0] - assert isinstance(content, DataContent) - assert content.uri.startswith("data:image/png;base64,") - assert content.media_type == "image/png" - - -def test_create_response_content_image_generation_existing_data_uri(): - """Test image generation response parsing with existing data URI.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - - # Create a mock response with existing data URI - mock_response = MagicMock() - mock_response.output_parsed = None - mock_response.metadata = {} - mock_response.usage = None - mock_response.id = "test-response-id" - mock_response.model = "test-model" - mock_response.created_at = 1234567890 - - # Mock image generation output item with existing data URI (valid WEBP header) - webp_signature = b"RIFF" + b"\x12\x00\x00\x00" + b"WEBP" - valid_webp_base64 = base64.b64encode(webp_signature + b"VP8 fake_data").decode() - mock_item = MagicMock() - mock_item.type = "image_generation_call" - mock_item.result = f"data:image/webp;base64,{valid_webp_base64}" - - mock_response.output = [mock_item] - - with patch.object(client, "_get_metadata_from_response", return_value={}): - response = client._create_response_content(mock_response, chat_options=ChatOptions()) # type: ignore - - # Verify the response contains DataContent with proper media_type parsed from URI - assert len(response.messages[0].contents) == 1 - content = response.messages[0].contents[0] - assert isinstance(content, DataContent) - assert content.uri == f"data:image/webp;base64,{valid_webp_base64}" - assert content.media_type == "image/webp" - - -def test_create_response_content_image_generation_format_detection(): - """Test different image format detection from base64 data.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - - # Test JPEG detection - jpeg_signature = b"\xff\xd8\xff" - mock_base64_jpeg = base64.b64encode(jpeg_signature + b"fake_jpeg_data").decode() - - mock_response_jpeg = MagicMock() - mock_response_jpeg.output_parsed = None - mock_response_jpeg.metadata = {} - mock_response_jpeg.usage = None - mock_response_jpeg.id = "test-id" - mock_response_jpeg.model = "test-model" - mock_response_jpeg.created_at = 1234567890 - - mock_item_jpeg = MagicMock() - mock_item_jpeg.type = "image_generation_call" - mock_item_jpeg.result = mock_base64_jpeg - mock_response_jpeg.output = [mock_item_jpeg] - - with patch.object(client, "_get_metadata_from_response", return_value={}): - response_jpeg = client._create_response_content(mock_response_jpeg, chat_options=ChatOptions()) # type: ignore - content_jpeg = response_jpeg.messages[0].contents[0] - assert isinstance(content_jpeg, DataContent) - assert content_jpeg.media_type == "image/jpeg" - assert "data:image/jpeg;base64," in content_jpeg.uri - - # Test WEBP detection - webp_signature = b"RIFF" + b"\x00\x00\x00\x00" + b"WEBP" - mock_base64_webp = base64.b64encode(webp_signature + b"fake_webp_data").decode() - - mock_response_webp = MagicMock() - mock_response_webp.output_parsed = None - mock_response_webp.metadata = {} - mock_response_webp.usage = None - mock_response_webp.id = "test-id" - mock_response_webp.model = "test-model" - mock_response_webp.created_at = 1234567890 - - mock_item_webp = MagicMock() - mock_item_webp.type = "image_generation_call" - mock_item_webp.result = mock_base64_webp - mock_response_webp.output = [mock_item_webp] - - with patch.object(client, "_get_metadata_from_response", return_value={}): - response_webp = client._create_response_content(mock_response_webp, chat_options=ChatOptions()) # type: ignore - content_webp = response_webp.messages[0].contents[0] - assert isinstance(content_webp, DataContent) - assert content_webp.media_type == "image/webp" - assert "data:image/webp;base64," in content_webp.uri - - -def test_create_response_content_image_generation_fallback(): - """Test image generation with invalid base64 falls back to PNG.""" - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - - # Create a mock response with invalid base64 - mock_response = MagicMock() - mock_response.output_parsed = None - mock_response.metadata = {} - mock_response.usage = None - mock_response.id = "test-response-id" - mock_response.model = "test-model" - mock_response.created_at = 1234567890 - - # Mock image generation output item with unrecognized format (should fall back to PNG) - unrecognized_data = b"UNKNOWN_FORMAT" + b"some_binary_data" - unrecognized_base64 = base64.b64encode(unrecognized_data).decode() - mock_item = MagicMock() - mock_item.type = "image_generation_call" - mock_item.result = unrecognized_base64 - - mock_response.output = [mock_item] - - with patch.object(client, "_get_metadata_from_response", return_value={}): - response = client._create_response_content(mock_response, chat_options=ChatOptions()) # type: ignore - - # Verify it falls back to PNG format for unrecognized binary data - assert len(response.messages[0].contents) == 1 - content = response.messages[0].contents[0] - assert isinstance(content, DataContent) - assert content.media_type == "image/png" - assert f"data:image/png;base64,{unrecognized_base64}" == content.uri - - -def test_prepare_options_store_parameter_handling() -> None: - client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") - messages = [ChatMessage(role="user", text="Test message")] - - test_conversation_id = "test-conversation-123" - chat_options = ChatOptions(store=True, conversation_id=test_conversation_id) - options = client._prepare_options(messages, chat_options) # type: ignore - assert options["store"] is True - assert options["previous_response_id"] == test_conversation_id - - chat_options = ChatOptions(store=False, conversation_id="") - options = client._prepare_options(messages, chat_options) # type: ignore - assert options["store"] is False - - chat_options = ChatOptions(store=None, conversation_id=None) - options = client._prepare_options(messages, chat_options) # type: ignore - assert options["store"] is False - assert "previous_response_id" not in options - - chat_options = ChatOptions() - options = client._prepare_options(messages, chat_options) # type: ignore - assert options["store"] is False - assert "previous_response_id" not in options - - -def test_openai_responses_client_with_callable_api_key() -> None: - """Test OpenAIResponsesClient initialization with callable API key.""" - - async def get_api_key() -> str: - return "test-api-key-123" - - client = OpenAIResponsesClient(model_id="gpt-4o", api_key=get_api_key) - - # Verify client was created successfully - assert client.model_id == "gpt-4o" - # OpenAI SDK now manages callable API keys internally - assert client.client is not None + async with ChatAgent( + chat_client=OpenAIResponsesClient(), + instructions="You are a helpful assistant that can help with microsoft documentation questions.", + tools=[mcp_tool], + ) as agent: + response = await agent.run( + "How to create an Azure storage account using az cli?", + max_tokens=200, + ) + + assert isinstance(response, AgentRunResponse) + assert response.text is not None + assert len(response.text) > 0 + # Should contain Azure-related content since it's asking about Azure CLI + assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"]) diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py index 3bda2fcaad..77fd969f12 100644 --- a/python/packages/core/tests/workflow/test_agent_executor.py +++ b/python/packages/core/tests/workflow/test_agent_executor.py @@ -111,6 +111,10 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None: chat_store_state = thread_state["chat_message_store_state"] # type: ignore[index] assert "messages" in chat_store_state, "Message store state should include messages" + # Verify checkpoint contains pending requests from agents and responses to be sent + assert "pending_agent_requests" in executor_state + assert "pending_responses_to_agent" in executor_state + # Create a new agent and executor for restoration # This simulates starting from a fresh state and restoring from checkpoint restored_agent = _CountingAgent(id="test_agent", name="TestAgent") diff --git a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py index 8124f6253d..a7849120b0 100644 --- a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py +++ b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py @@ -5,19 +5,32 @@ from collections.abc import AsyncIterable from typing import Any +from typing_extensions import Never + from agent_framework import ( AgentExecutor, + AgentExecutorResponse, AgentRunResponse, AgentRunResponseUpdate, AgentRunUpdateEvent, AgentThread, BaseAgent, + ChatAgent, ChatMessage, + ChatResponse, + ChatResponseUpdate, + FunctionApprovalRequestContent, FunctionCallContent, FunctionResultContent, + RequestInfoEvent, Role, TextContent, WorkflowBuilder, + WorkflowContext, + WorkflowOutputEvent, + ai_function, + executor, + use_function_invocation, ) @@ -120,3 +133,235 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None: assert events[3].data is not None assert isinstance(events[3].data.contents[0], TextContent) assert "sunny" in events[3].data.contents[0].text + + +@ai_function(approval_mode="always_require") +def mock_tool_requiring_approval(query: str) -> str: + """Mock tool that requires approval before execution.""" + return f"Executed tool with query: {query}" + + +@use_function_invocation +class MockChatClient: + """Simple implementation of a chat client.""" + + def __init__(self, parallel_request: bool = False) -> None: + self.additional_properties: dict[str, Any] = {} + self._iteration: int = 0 + self._parallel_request: bool = parallel_request + + async def get_response( + self, + messages: str | ChatMessage | list[str] | list[ChatMessage], + **kwargs: Any, + ) -> ChatResponse: + if self._iteration == 0: + if self._parallel_request: + response = ChatResponse( + messages=ChatMessage( + role="assistant", + contents=[ + FunctionCallContent( + call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}' + ), + FunctionCallContent( + call_id="2", name="mock_tool_requiring_approval", arguments='{"query": "test"}' + ), + ], + ) + ) + else: + response = ChatResponse( + messages=ChatMessage( + role="assistant", + contents=[ + FunctionCallContent( + call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}' + ) + ], + ) + ) + else: + response = ChatResponse(messages=ChatMessage(role="assistant", text="Tool executed successfully.")) + + self._iteration += 1 + return response + + async def get_streaming_response( + self, + messages: str | ChatMessage | list[str] | list[ChatMessage], + **kwargs: Any, + ) -> AsyncIterable[ChatResponseUpdate]: + if self._iteration == 0: + if self._parallel_request: + yield ChatResponseUpdate( + contents=[ + FunctionCallContent( + call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}' + ), + FunctionCallContent( + call_id="2", name="mock_tool_requiring_approval", arguments='{"query": "test"}' + ), + ], + role="assistant", + ) + else: + yield ChatResponseUpdate( + contents=[ + FunctionCallContent( + call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}' + ) + ], + role="assistant", + ) + else: + yield ChatResponseUpdate(text=TextContent(text="Tool executed "), role="assistant") + yield ChatResponseUpdate(contents=[TextContent(text="successfully.")], role="assistant") + + self._iteration += 1 + + +@executor(id="test_executor") +async def test_executor(agent_executor_response: AgentExecutorResponse, ctx: WorkflowContext[Never, str]) -> None: + await ctx.yield_output(agent_executor_response.agent_run_response.text) + + +async def test_agent_executor_tool_call_with_approval() -> None: + """Test that AgentExecutor handles tool calls requiring approval.""" + # Arrange + agent = ChatAgent( + chat_client=MockChatClient(), + name="ApprovalAgent", + tools=[mock_tool_requiring_approval], + ) + + workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build() + + # Act + events = await workflow.run("Invoke tool requiring approval") + + # Assert + assert len(events.get_request_info_events()) == 1 + approval_request = events.get_request_info_events()[0] + assert isinstance(approval_request.data, FunctionApprovalRequestContent) + assert approval_request.data.function_call.name == "mock_tool_requiring_approval" + assert approval_request.data.function_call.arguments == '{"query": "test"}' + + # Act + events = await workflow.send_responses({approval_request.request_id: approval_request.data.create_response(True)}) + + # Assert + final_response = events.get_outputs() + assert len(final_response) == 1 + assert final_response[0] == "Tool executed successfully." + + +async def test_agent_executor_tool_call_with_approval_streaming() -> None: + """Test that AgentExecutor handles tool calls requiring approval in streaming mode.""" + # Arrange + agent = ChatAgent( + chat_client=MockChatClient(), + name="ApprovalAgent", + tools=[mock_tool_requiring_approval], + ) + + workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build() + + # Act + request_info_events: list[RequestInfoEvent] = [] + async for event in workflow.run_stream("Invoke tool requiring approval"): + if isinstance(event, RequestInfoEvent): + request_info_events.append(event) + + # Assert + assert len(request_info_events) == 1 + approval_request = request_info_events[0] + assert isinstance(approval_request.data, FunctionApprovalRequestContent) + assert approval_request.data.function_call.name == "mock_tool_requiring_approval" + assert approval_request.data.function_call.arguments == '{"query": "test"}' + + # Act + output: str | None = None + async for event in workflow.send_responses_streaming({ + approval_request.request_id: approval_request.data.create_response(True) + }): + if isinstance(event, WorkflowOutputEvent): + output = event.data + + # Assert + assert output is not None + assert output == "Tool executed successfully." + + +async def test_agent_executor_parallel_tool_call_with_approval() -> None: + """Test that AgentExecutor handles parallel tool calls requiring approval.""" + # Arrange + agent = ChatAgent( + chat_client=MockChatClient(parallel_request=True), + name="ApprovalAgent", + tools=[mock_tool_requiring_approval], + ) + + workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build() + + # Act + events = await workflow.run("Invoke tool requiring approval") + + # Assert + assert len(events.get_request_info_events()) == 2 + for approval_request in events.get_request_info_events(): + assert isinstance(approval_request.data, FunctionApprovalRequestContent) + assert approval_request.data.function_call.name == "mock_tool_requiring_approval" + assert approval_request.data.function_call.arguments == '{"query": "test"}' + + # Act + responses = { + approval_request.request_id: approval_request.data.create_response(True) # type: ignore + for approval_request in events.get_request_info_events() + } + events = await workflow.send_responses(responses) + + # Assert + final_response = events.get_outputs() + assert len(final_response) == 1 + assert final_response[0] == "Tool executed successfully." + + +async def test_agent_executor_parallel_tool_call_with_approval_streaming() -> None: + """Test that AgentExecutor handles parallel tool calls requiring approval in streaming mode.""" + # Arrange + agent = ChatAgent( + chat_client=MockChatClient(parallel_request=True), + name="ApprovalAgent", + tools=[mock_tool_requiring_approval], + ) + + workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build() + + # Act + request_info_events: list[RequestInfoEvent] = [] + async for event in workflow.run_stream("Invoke tool requiring approval"): + if isinstance(event, RequestInfoEvent): + request_info_events.append(event) + + # Assert + assert len(request_info_events) == 2 + for approval_request in request_info_events: + assert isinstance(approval_request.data, FunctionApprovalRequestContent) + assert approval_request.data.function_call.name == "mock_tool_requiring_approval" + assert approval_request.data.function_call.arguments == '{"query": "test"}' + + # Act + responses = { + approval_request.request_id: approval_request.data.create_response(True) # type: ignore + for approval_request in request_info_events + } + + output: str | None = None + async for event in workflow.send_responses_streaming(responses): + if isinstance(event, WorkflowOutputEvent): + output = event.data + + # Assert + assert output is not None + assert output == "Tool executed successfully." diff --git a/python/samples/README.md b/python/samples/README.md index bb127b1349..330d6f03fc 100644 --- a/python/samples/README.md +++ b/python/samples/README.md @@ -281,6 +281,7 @@ This directory contains samples demonstrating the capabilities of Microsoft Agen | File | Description | |------|-------------| | [`getting_started/workflows/human-in-the-loop/guessing_game_with_human_input.py`](./getting_started/workflows/human-in-the-loop/guessing_game_with_human_input.py) | Sample: Human in the loop guessing game | +| [`getting_started/workflows/human-in-the-loop/agents_with_approval_requests.py`](./getting_started/workflows/human-in-the-loop/agents_with_approval_requests.py) | Sample: Agents with Approval Requests in Workflows | ### Observability diff --git a/python/samples/getting_started/agents/openai/README.md b/python/samples/getting_started/agents/openai/README.md index ff4f46a84b..db71816558 100644 --- a/python/samples/getting_started/agents/openai/README.md +++ b/python/samples/getting_started/agents/openai/README.md @@ -23,6 +23,7 @@ This folder contains examples demonstrating different ways to create and use age | [`openai_responses_client_image_analysis.py`](openai_responses_client_image_analysis.py) | Demonstrates how to use vision capabilities with agents to analyze images. | | [`openai_responses_client_image_generation.py`](openai_responses_client_image_generation.py) | Demonstrates how to use image generation capabilities with OpenAI agents to create images based on text descriptions. Requires PIL (Pillow) for image display. | | [`openai_responses_client_reasoning.py`](openai_responses_client_reasoning.py) | Demonstrates how to use reasoning capabilities with OpenAI agents, showing how the agent can provide detailed reasoning for its responses. | +| [`openai_responses_client_streaming_image_generation.py`](openai_responses_client_streaming_image_generation.py) | Demonstrates streaming image generation with partial images for real-time image creation feedback and improved user experience. | | [`openai_responses_client_with_code_interpreter.py`](openai_responses_client_with_code_interpreter.py) | Shows how to use the HostedCodeInterpreterTool with OpenAI agents to write and execute Python code. Includes helper methods for accessing code interpreter data from response chunks. | | [`openai_responses_client_with_explicit_settings.py`](openai_responses_client_with_explicit_settings.py) | Shows how to initialize an agent with a specific responses client, configuring settings explicitly including API key and model ID. | | [`openai_responses_client_with_file_search.py`](openai_responses_client_with_file_search.py) | Demonstrates how to use file search capabilities with OpenAI agents, allowing the agent to search through uploaded files to answer questions. | diff --git a/python/samples/getting_started/agents/openai/openai_responses_client_streaming_image_generation.py b/python/samples/getting_started/agents/openai/openai_responses_client_streaming_image_generation.py new file mode 100644 index 0000000000..2d74429917 --- /dev/null +++ b/python/samples/getting_started/agents/openai/openai_responses_client_streaming_image_generation.py @@ -0,0 +1,96 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import base64 + +import anyio +from agent_framework import DataContent +from agent_framework.openai import OpenAIResponsesClient + +"""OpenAI Responses Client Streaming Image Generation Example + +Demonstrates streaming partial image generation using OpenAI's image generation tool. +Shows progressive image rendering with partial images for improved user experience. + +Note: The number of partial images received depends on generation speed: +- High quality/complex images: More partials (generation takes longer) +- Low quality/simple images: Fewer partials (generation completes quickly) +- You may receive fewer partial images than requested if generation is fast + +Important: The final partial image IS the complete, full-quality image. Each partial +represents a progressive refinement, with the last one being the finished result. +""" + + +async def save_image_from_data_uri(data_uri: str, filename: str) -> None: + """Save an image from a data URI to a file.""" + try: + if data_uri.startswith("data:image/"): + # Extract base64 data + base64_data = data_uri.split(",", 1)[1] + image_bytes = base64.b64decode(base64_data) + + # Save to file + await anyio.Path(filename).write_bytes(image_bytes) + print(f" Saved: {filename} ({len(image_bytes) / 1024:.1f} KB)") + except Exception as e: + print(f" Error saving {filename}: {e}") + + +async def main(): + """Demonstrate streaming image generation with partial images.""" + print("=== OpenAI Streaming Image Generation Example ===\n") + + # Create agent with streaming image generation enabled + agent = OpenAIResponsesClient().create_agent( + instructions="You are a helpful agent that can generate images.", + tools=[ + { + "type": "image_generation", + "size": "1024x1024", + "quality": "high", + "partial_images": 3, + } + ], + ) + + query = "Draw a beautiful sunset over a calm ocean with sailboats" + print(f" User: {query}") + print() + + # Track partial images + image_count = 0 + + # Create output directory + output_dir = anyio.Path("generated_images") + await output_dir.mkdir(exist_ok=True) + + print(" Streaming response:") + async for update in agent.run_stream(query): + for content in update.contents: + # Handle partial images + # The final partial image IS the complete, full-quality image. Each partial + # represents a progressive refinement, with the last one being the finished result. + if isinstance(content, DataContent) and content.additional_properties.get("is_partial_image"): + print(f" Image {image_count} received") + + # Extract file extension from media_type (e.g., "image/png" -> "png") + extension = "png" # Default fallback + if content.media_type and "/" in content.media_type: + extension = content.media_type.split("/")[-1] + + # Save images with correct extension + filename = output_dir / f"image{image_count}.{extension}" + await save_image_from_data_uri(content.uri, str(filename)) + + image_count += 1 + + # Summary + print("\n Summary:") + print(f" Images received: {image_count}") + print(" Output directory: generated_images") + print("\n Streaming image generation completed!") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/getting_started/workflows/README.md b/python/samples/getting_started/workflows/README.md index 503a8aacd7..ffbeedd751 100644 --- a/python/samples/getting_started/workflows/README.md +++ b/python/samples/getting_started/workflows/README.md @@ -78,6 +78,7 @@ Once comfortable with these, explore the rest of the samples below. |---|---|---| | Human-In-The-Loop (Guessing Game) | [human-in-the-loop/guessing_game_with_human_input.py](./human-in-the-loop/guessing_game_with_human_input.py) | Interactive request/response prompts with a human | | Azure Agents Tool Feedback Loop | [agents/azure_chat_agents_tool_calls_with_feedback.py](./agents/azure_chat_agents_tool_calls_with_feedback.py) | Two-agent workflow that streams tool calls and pauses for human guidance between passes | +| Agents with Approval Requests in Workflows | [human-in-the-loop/agents_with_approval_requests.py](./human-in-the-loop/agents_with_approval_requests.py) | Agents that create approval requests during workflow execution and wait for human approval to proceed | ### observability diff --git a/python/samples/getting_started/workflows/human-in-the-loop/agents_with_approval_requests.py b/python/samples/getting_started/workflows/human-in-the-loop/agents_with_approval_requests.py new file mode 100644 index 0000000000..a51088e886 --- /dev/null +++ b/python/samples/getting_started/workflows/human-in-the-loop/agents_with_approval_requests.py @@ -0,0 +1,340 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import json +from dataclasses import dataclass +from typing import Annotated, Never + +from agent_framework import ( + AgentExecutorResponse, + ChatMessage, + Executor, + FunctionApprovalRequestContent, + FunctionApprovalResponseContent, + WorkflowBuilder, + WorkflowContext, + ai_function, + executor, + handler, +) +from agent_framework.openai import OpenAIChatClient + +""" +Sample: Agents in a workflow with AI functions requiring approval + +This sample creates a workflow that automatically replies to incoming emails. +If historical email data is needed, it uses an AI function to read the data, +which requires human approval before execution. + +This sample works as follows: +1. An incoming email is received by the workflow. +2. The EmailPreprocessor executor preprocesses the email, adding special notes if the sender is important. +3. The preprocessed email is sent to the Email Writer agent, which generates a response. +4. If the agent needs to read historical email data, it calls the read_historical_email_data AI function, + which triggers an approval request. +5. The sample automatically approves the request for demonstration purposes. +6. Once approved, the AI function executes and returns the historical email data to the agent. +7. The agent uses the historical data to compose a comprehensive email response. +8. The response is sent to the conclude_workflow_executor, which yields the final response. + +Purpose: +Show how to integrate AI functions with approval requests into a workflow. + +Demonstrate: +- Creating AI functions that require approval before execution. +- Building a workflow that includes an agent and executors. +- Handling approval requests during workflow execution. + +Prerequisites: +- Azure AI Agent Service configured, along with the required environment variables. +- Authentication via azure-identity. Use AzureCliCredential and run az login before executing the sample. +- Basic familiarity with WorkflowBuilder, edges, events, RequestInfoEvent, and streaming runs. +""" + + +@ai_function +def get_current_date() -> str: + """Get the current date in YYYY-MM-DD format.""" + # For demonstration purposes, we return a fixed date. + return "2025-11-07" + + +@ai_function +def get_team_members_email_addresses() -> list[dict[str, str]]: + """Get the email addresses of team members.""" + # In a real implementation, this might query a database or directory service. + return [ + { + "name": "Alice", + "email": "alice@contoso.com", + "position": "Software Engineer", + "manager": "John Doe", + }, + { + "name": "Bob", + "email": "bob@contoso.com", + "position": "Product Manager", + "manager": "John Doe", + }, + { + "name": "Charlie", + "email": "charlie@contoso.com", + "position": "Senior Software Engineer", + "manager": "John Doe", + }, + { + "name": "Mike", + "email": "mike@contoso.com", + "position": "Principal Software Engineer Manager", + "manager": "VP of Engineering", + }, + ] + + +@ai_function +def get_my_information() -> dict[str, str]: + """Get my personal information.""" + return { + "name": "John Doe", + "email": "john@contoso.com", + "position": "Software Engineer Manager", + "manager": "Mike", + } + + +@ai_function(approval_mode="always_require") +async def read_historical_email_data( + email_address: Annotated[str, "The email address to read historical data from"], + start_date: Annotated[str, "The start date in YYYY-MM-DD format"], + end_date: Annotated[str, "The end date in YYYY-MM-DD format"], +) -> list[dict[str, str]]: + """Read historical email data for a given email address and date range.""" + historical_data = { + "alice@contoso.com": [ + { + "from": "alice@contoso.com", + "to": "john@contoso.com", + "date": "2025-11-05", + "subject": "Bug Bash Results", + "body": "We just completed the bug bash and found a few issues that need immediate attention.", + }, + { + "from": "alice@contoso.com", + "to": "john@contoso.com", + "date": "2025-11-03", + "subject": "Code Freeze", + "body": "We are entering code freeze starting tomorrow.", + }, + ], + "bob@contoso.com": [ + { + "from": "bob@contoso.com", + "to": "john@contoso.com", + "date": "2025-11-04", + "subject": "Team Outing", + "body": "Don't forget about the team outing this Friday!", + }, + { + "from": "bob@contoso.com", + "to": "john@contoso.com", + "date": "2025-11-02", + "subject": "Requirements Update", + "body": "The requirements for the new feature have been updated. Please review them.", + }, + ], + "charlie@contoso.com": [ + { + "from": "charlie@contoso.com", + "to": "john@contoso.com", + "date": "2025-11-05", + "subject": "Project Update", + "body": "The bug bash went well. A few critical bugs but should be fixed by the end of the week.", + }, + { + "from": "charlie@contoso.com", + "to": "john@contoso.com", + "date": "2025-11-06", + "subject": "Code Review", + "body": "Please review my latest code changes.", + }, + ], + } + + emails = historical_data.get(email_address, []) + return [email for email in emails if start_date <= email["date"] <= end_date] + + +@ai_function(approval_mode="always_require") +async def send_email( + to: Annotated[str, "The recipient email address"], + subject: Annotated[str, "The email subject"], + body: Annotated[str, "The email body"], +) -> str: + """Send an email.""" + await asyncio.sleep(1) # Simulate sending email + return "Email successfully sent." + + +@dataclass +class Email: + sender: str + subject: str + body: str + + +class EmailPreprocessor(Executor): + def __init__(self, special_email_addresses: set[str]) -> None: + super().__init__(id="email_preprocessor") + self.special_email_addresses = special_email_addresses + + @handler + async def preprocess(self, email: Email, ctx: WorkflowContext[str]) -> None: + """Preprocess the incoming email.""" + message = str(email) + if email.sender in self.special_email_addresses: + note = ( + "Pay special attention to this sender. This email is very important. " + "Gather relevant information from all previous emails within my team before responding." + ) + message = f"{note}\n\n{message}" + + await ctx.send_message(message) + + +@executor(id="conclude_workflow_executor") +async def conclude_workflow( + email_response: AgentExecutorResponse, + ctx: WorkflowContext[Never, str], +) -> None: + """Conclude the workflow by yielding the final email response.""" + await ctx.yield_output(email_response.agent_run_response.text) + + +async def main() -> None: + # Create the agent and executors + chat_client = OpenAIChatClient() + email_writer = chat_client.create_agent( + name="Email Writer", + instructions=("You are an excellent email assistant. You respond to incoming emails."), + # tools with `approval_mode="always_require"` will trigger approval requests + tools=[ + read_historical_email_data, + send_email, + get_current_date, + get_team_members_email_addresses, + get_my_information, + ], + ) + email_preprocessor = EmailPreprocessor(special_email_addresses={"mike@contoso.com"}) + + # Build the workflow + workflow = ( + WorkflowBuilder() + .set_start_executor(email_preprocessor) + .add_edge(email_preprocessor, email_writer) + .add_edge(email_writer, conclude_workflow) + .build() + ) + + # Simulate an incoming email + incoming_email = Email( + sender="mike@contoso.com", + subject="Important: Project Update", + body="Please provide your team's status update on the project since last week.", + ) + + responses: dict[str, FunctionApprovalResponseContent] = {} + output: list[ChatMessage] | None = None + while True: + if responses: + events = await workflow.send_responses(responses) + responses.clear() + else: + events = await workflow.run(incoming_email) + + request_info_events = events.get_request_info_events() + for request_info_event in request_info_events: + # We should only expect FunctionApprovalRequestContent in this sample + if not isinstance(request_info_event.data, FunctionApprovalRequestContent): + raise ValueError(f"Unexpected request info content type: {type(request_info_event.data)}") + + # Pretty print the function call details + arguments = json.dumps(request_info_event.data.function_call.parse_arguments(), indent=2) + print( + f"Received approval request for function: {request_info_event.data.function_call.name} " + f"with args:\n{arguments}" + ) + + # For demo purposes, we automatically approve the request + # The expected response type of the request is `FunctionApprovalResponseContent`, + # which can be created via `create_response` method on the request content + print("Performing automatic approval for demo purposes...") + responses[request_info_event.request_id] = request_info_event.data.create_response(approved=True) + + # Once we get an output event, we can conclude the workflow + # Outputs can only be produced by the conclude_workflow_executor in this sample + if outputs := events.get_outputs(): + # We expect only one output from the conclude_workflow_executor + output = outputs[0] + break + + if not output: + raise RuntimeError("Workflow did not produce any output event.") + + print("Final email response conversation:") + print(output) + + """ + Sample Output: + Received approval request for function: read_historical_email_data with args: + { + "email_address": "alice@contoso.com", + "start_date": "2025-10-31", + "end_date": "2025-11-07" + } + Performing automatic approval for demo purposes... + Received approval request for function: read_historical_email_data with args: + { + "email_address": "bob@contoso.com", + "start_date": "2025-10-31", + "end_date": "2025-11-07" + } + Performing automatic approval for demo purposes... + Received approval request for function: read_historical_email_data with args: + { + "email_address": "charlie@contoso.com", + "start_date": "2025-10-31", + "end_date": "2025-11-07" + } + Performing automatic approval for demo purposes... + Received approval request for function: send_email with args: + { + "to": "mike@contoso.com", + "subject": "Team's Status Update on the Project", + "body": " + Hi Mike, + + Here's the status update from our team: + - **Bug Bash and Code Freeze:** + - We recently completed a bug bash, during which several issues were identified. Alice and Charlie are working on fixing these critical bugs, and we anticipate resolving them by the end of this week. + - We have entered a code freeze as of November 4, 2025. + + - **Requirements Update:** + - Bob has updated the requirements for a new feature, and all team members are reviewing these changes to ensure alignment. + + - **Ongoing Reviews:** + - Charlie has submitted his latest code changes for review to ensure they meet our quality standards. + + Please let me know if you need more detailed information or have any questions. + + Best regards, + John" + } + Performing automatic approval for demo purposes... + Final email response conversation: + I've sent the status update to Mike with the relevant information from the team. Let me know if there's anything else you need + """ # noqa: E501 + + +if __name__ == "__main__": + asyncio.run(main())