From 4aa94596811acc9d39e90c9c8633621f35448cf1 Mon Sep 17 00:00:00 2001 From: chetantoshniwal Date: Sun, 24 May 2026 22:33:05 -0700 Subject: [PATCH] Fix allowed_tools filtering to correctly handle empty list The allowed_tools parameter in MCPTool.functions and MCP server config helpers used a falsy check (if not allowed_tools / if allowed_tools) which treats an empty list identically to None. This meant passing allowed_tools=[] would bypass filtering and return all tools, rather than returning no tools as expected. Changed to explicit None checks so that: - allowed_tools=None -> no filtering (return all tools) - allowed_tools=[] -> explicit empty allowlist (return no tools) - allowed_tools=[...] -> filter to listed names only Also normalized serialization to use list() for consistency with arbitrary Collection inputs. Co-authored-by: Azure SRE Agent --- .../agent_framework_anthropic/_chat_client.py | 251 +- python/packages/core/agent_framework/_mcp.py | 1075 ++++-- .../agent_framework_foundry/_chat_client.py | 1010 ++++++ .../agent_framework_openai/_chat_client.py | 3203 +++++++++++++++++ 4 files changed, 5134 insertions(+), 405 deletions(-) create mode 100644 python/packages/foundry/agent_framework_foundry/_chat_client.py create mode 100644 python/packages/openai/agent_framework_openai/_chat_client.py diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index a1915a69fb2..b0be6c249e8 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -8,7 +8,6 @@ from typing import Any, ClassVar, Final, Generic, Literal, TypedDict from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, Annotation, BaseChatClient, ChatAndFunctionMiddlewareTypes, @@ -28,10 +27,11 @@ tool, ) from agent_framework._settings import SecretString, load_settings +from agent_framework._telemetry import get_user_agent from agent_framework._tools import SHELL_TOOL_KIND_VALUE from agent_framework._types import _get_data_bytes_as_str # type: ignore from agent_framework.observability import ChatTelemetryLayer -from anthropic import AsyncAnthropic +from anthropic import AsyncAnthropic, AsyncAnthropicBedrock, AsyncAnthropicFoundry, AsyncAnthropicVertex from anthropic.types.beta import ( BetaContentBlock, BetaMessage, @@ -68,6 +68,7 @@ __all__ = [ "AnthropicChatOptions", "AnthropicClient", + "RawAnthropicClient", "ThinkingConfig", ] @@ -78,6 +79,7 @@ STRUCTURED_OUTPUTS_BETA_FLAG: Final[str] = "structured-outputs-2025-11-13" ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None) +AnthropicAsyncClient = AsyncAnthropic | AsyncAnthropicBedrock | AsyncAnthropicFoundry | AsyncAnthropicVertex # region Anthropic Chat Options TypedDict @@ -112,8 +114,6 @@ class AnthropicChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], a default of 1024 will be used. Keys: - model_id: The model to use for the request, - translates to ``model`` in Anthropic API. temperature: Sampling temperature between 0 and 1. top_p: Nucleus sampling parameter. max_tokens: Maximum number of tokens to generate (REQUIRED). @@ -168,12 +168,24 @@ class AnthropicChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], # Translation between framework options keys and Anthropic Messages API OPTION_TRANSLATIONS: dict[str, str] = { - "model_id": "model", "stop": "stop_sequences", "instructions": "system", } +def _apply_option_translations(options: dict[str, Any]) -> None: + """Translate framework option keys to Anthropic request keys in-place. + + When both the old and new key are present, the new key wins and the old key + is discarded to preserve explicit overrides. + """ + for old_key, new_key in OPTION_TRANSLATIONS.items(): + if old_key not in options or old_key == new_key: + continue + old_value = options.pop(old_key) + options.setdefault(new_key, old_value) + + # region Role and Finish Reason Maps @@ -203,21 +215,33 @@ class AnthropicSettings(TypedDict, total=False): Keys: api_key: The Anthropic API key. - chat_model_id: The Anthropic chat model ID. + chat_model: The Anthropic chat model. + base_url: Optional base URL for the Anthropic API endpoint. """ api_key: SecretString | None - chat_model_id: str | None + chat_model: str | None + base_url: str | None -class AnthropicClient( - ChatMiddlewareLayer[AnthropicOptionsT], - FunctionInvocationLayer[AnthropicOptionsT], - ChatTelemetryLayer[AnthropicOptionsT], +class RawAnthropicClient( BaseChatClient[AnthropicOptionsT], Generic[AnthropicOptionsT], ): - """Anthropic Chat client with middleware, telemetry, and function invocation support.""" + """Raw Anthropic chat client without middleware, telemetry, or function invocation support. + + Warning: + **This class should not normally be used directly.** It does not include middleware, + telemetry, or function invocation support that you most likely need. If you do use it, + you should consider which additional layers to apply. There is a defined ordering that + you should follow: + + 1. **FunctionInvocationLayer** - Owns the tool/function calling loop and routes function middleware + 2. **ChatMiddlewareLayer** - Applies chat middleware per model call and stays outside telemetry + 3. **ChatTelemetryLayer** - Must stay inside chat middleware for correct per-call telemetry + + Use ``AnthropicClient`` instead for a fully-featured client with all layers applied. + """ OTEL_PROVIDER_NAME: ClassVar[str] = "anthropic" # type: ignore[reportIncompatibleVariableOverride, misc] @@ -225,49 +249,55 @@ def __init__( self, *, api_key: str | None = None, - model_id: str | None = None, - anthropic_client: AsyncAnthropic | None = None, + model: str | None = None, + base_url: str | None = None, + anthropic_client: AnthropicAsyncClient | None = None, additional_beta_flags: list[str] | None = None, additional_properties: dict[str, Any] | None = None, - middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None, - function_invocation_configuration: FunctionInvocationConfiguration | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, ) -> None: - """Initialize an Anthropic Agent client. + """Initialize a raw Anthropic client. Keyword Args: api_key: The Anthropic API key to use for authentication. - model_id: The ID of the model to use. + model: The model to use. + base_url: Optional base URL for the Anthropic API endpoint. Useful for Foundry or + other compatible deployments. Falls back to ``ANTHROPIC_BASE_URL`` env variable. anthropic_client: An existing Anthropic client to use. If not provided, one will be created. This can be used to further configure the client before passing it in. For instance if you need to set a different base_url for testing or private deployments. additional_beta_flags: Additional beta flags to enable on the client. Default flags are: "mcp-client-2025-04-04", "code-execution-2025-08-25". additional_properties: Additional properties stored on the client instance. - middleware: Optional middleware to apply to the client. - function_invocation_configuration: Optional function invocation configuration override. env_file_path: Path to environment file for loading settings. env_file_encoding: Encoding of the environment file. Examples: .. code-block:: python - from agent_framework.anthropic import AnthropicClient + from agent_framework.anthropic import RawAnthropicClient from azure.identity.aio import DefaultAzureCredential # Using environment variables # Set ANTHROPIC_API_KEY=your_anthropic_api_key - # ANTHROPIC_CHAT_MODEL_ID=claude-sonnet-4-5-20250929 + # ANTHROPIC_CHAT_MODEL=claude-sonnet-4-5-20250929 # Or passing parameters directly - client = AnthropicClient( - model_id="claude-sonnet-4-5-20250929", + client = RawAnthropicClient( + model="claude-sonnet-4-5-20250929", + api_key="your_anthropic_api_key", + ) + + # Or with a custom base URL (e.g. for Foundry-compatible endpoints) + client = RawAnthropicClient( + model="claude-sonnet-4-5-20250929", api_key="your_anthropic_api_key", + base_url="https://custom-anthropic-endpoint.com", ) # Or loading from a .env file - client = AnthropicClient(env_file_path="path/to/.env") + client = RawAnthropicClient(env_file_path="path/to/.env") # Or passing in an existing client from anthropic import AsyncAnthropic @@ -275,8 +305,8 @@ def __init__( anthropic_client = AsyncAnthropic( api_key="your_anthropic_api_key", base_url="https://custom-anthropic-endpoint.com" ) - client = AnthropicClient( - model_id="claude-sonnet-4-5-20250929", + client = RawAnthropicClient( + model="claude-sonnet-4-5-20250929", anthropic_client=anthropic_client, ) @@ -289,7 +319,7 @@ class MyOptions(AnthropicChatOptions, total=False): my_custom_option: str - client: AnthropicClient[MyOptions] = AnthropicClient(model_id="claude-sonnet-4-5-20250929") + client: RawAnthropicClient[MyOptions] = RawAnthropicClient(model="claude-sonnet-4-5-20250929") response = await client.get_response("Hello", options={"my_custom_option": "value"}) """ @@ -297,13 +327,15 @@ class MyOptions(AnthropicChatOptions, total=False): AnthropicSettings, env_prefix="ANTHROPIC_", api_key=api_key, - chat_model_id=model_id, + chat_model=model, + base_url=base_url, env_file_path=env_file_path, env_file_encoding=env_file_encoding, ) api_key_secret = anthropic_settings.get("api_key") - model_id_setting = anthropic_settings.get("chat_model_id") + model_setting = anthropic_settings.get("chat_model") + base_url_setting = anthropic_settings.get("base_url") if anthropic_client is None: if api_key_secret is None: @@ -314,20 +346,19 @@ class MyOptions(AnthropicChatOptions, total=False): anthropic_client = AsyncAnthropic( api_key=api_key_secret.get_secret_value(), - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + base_url=base_url_setting, + default_headers={"User-Agent": get_user_agent()}, ) # Initialize parent super().__init__( additional_properties=additional_properties, - middleware=middleware, - function_invocation_configuration=function_invocation_configuration, ) # Initialize instance variables self.anthropic_client = anthropic_client self.additional_beta_flags = additional_beta_flags or [] - self.model_id = model_id_setting + self.model = model_setting # streaming requires tracking the last function call ID, name, and content type self._last_call_id_name: tuple[str, str] | None = None self._last_call_content_type: str | None = None @@ -481,8 +512,8 @@ def get_mcp_tool( "server_url": url, } - if allowed_tools: - result["allowed_tools"] = allowed_tools + if allowed_tools is not None: + result["allowed_tools"] = list(allowed_tools) if authorization_token: result["headers"] = {"authorization": authorization_token} @@ -508,7 +539,7 @@ def _inner_get_response( if stream: # Streaming mode async def _stream() -> AsyncIterable[ChatResponseUpdate]: - async for chunk in await self.anthropic_client.beta.messages.create(**run_options, stream=True): + async for chunk in await self.anthropic_client.beta.messages.create(**run_options, stream=True): # type: ignore[misc] parsed_chunk = self._process_stream_event(chunk) if parsed_chunk: yield parsed_chunk @@ -517,7 +548,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: # Non-streaming mode async def _get_response() -> ChatResponse: - message = await self.anthropic_client.beta.messages.create(**run_options, stream=False) + message = await self.anthropic_client.beta.messages.create(**run_options, stream=False) # type: ignore[misc] return self._process_message(message, options) return _get_response() @@ -556,16 +587,22 @@ def _prepare_options( # Stream mode is controlled explicitly at call sites. run_options.pop("stream", None) - # Translation between options keys and Anthropic Messages API - for old_key, new_key in OPTION_TRANSLATIONS.items(): - if old_key in run_options and old_key != new_key: - run_options[new_key] = run_options.pop(old_key) + _apply_option_translations(run_options) + + # Filter out framework kwargs that should not be passed to the Anthropic API. + # This includes underscore-prefixed internal objects (like _function_middleware_pipeline) + # and framework kwargs like 'thread' and 'middleware'. + filtered_kwargs = { + k: v for k, v in kwargs.items() if not k.startswith("_") and k not in {"thread", "middleware"} + } + _apply_option_translations(filtered_kwargs) + run_options.update(filtered_kwargs) - # model id + # model if not run_options.get("model"): - if not self.model_id: - raise ValueError("model_id must be a non-empty string") - run_options["model"] = self.model_id + if not self.model: + raise ValueError("model must be a non-empty string") + run_options["model"] = self.model # max_tokens - Anthropic requires this, default if not provided if not run_options.get("max_tokens"): @@ -582,7 +619,7 @@ def _prepare_options( run_options["betas"] = self._prepare_betas(options) # extra headers - run_options["extra_headers"] = {"User-Agent": AGENT_FRAMEWORK_USER_AGENT} + run_options["extra_headers"] = {"User-Agent": get_user_agent()} # Handle user option -> metadata.user_id (Anthropic uses metadata.user_id instead of user) if user := run_options.pop("user", None): @@ -602,13 +639,6 @@ def _prepare_options( # Add the structured outputs beta flag run_options["betas"].add(STRUCTURED_OUTPUTS_BETA_FLAG) - # Filter out framework kwargs that should not be passed to the Anthropic API. - # This includes underscore-prefixed internal objects (like _function_middleware_pipeline) - # and framework kwargs like 'thread' and 'middleware'. - filtered_kwargs = { - k: v for k, v in kwargs.items() if not k.startswith("_") and k not in {"thread", "middleware"} - } - run_options.update(filtered_kwargs) return run_options def _prepare_betas(self, options: Mapping[str, Any]) -> set[str]: @@ -857,6 +887,8 @@ def _prepare_tools_for_anthropic(self, options: Mapping[str, Any]) -> dict[str, tool_mode = validate_tool_mode(options.get("tool_choice")) if tool_mode is None: return result or None + if "allowed_tools" in tool_mode: + logger.warning("allowed_tools is not supported by Anthropic; the setting will be ignored") allow_multiple = options.get("allow_multiple_tool_calls") match tool_mode.get("mode"): case "auto": @@ -913,7 +945,7 @@ def _process_message(self, message: BetaMessage, options: Mapping[str, Any]) -> ) ], usage_details=self._parse_usage_from_anthropic(message.usage), - model_id=message.model, + model=message.model, finish_reason=FINISH_REASON_MAP.get(message.stop_reason) if message.stop_reason else None, response_format=options.get("response_format"), raw_representation=message, @@ -941,7 +973,7 @@ def _process_stream_event(self, event: BetaRawMessageStreamEvent) -> ChatRespons *self._parse_contents_from_anthropic(event.message.content), *usage_details, ], - model_id=event.message.model, + model=event.message.model, finish_reason=FINISH_REASON_MAP.get(event.message.stop_reason) if event.message.stop_reason else None, @@ -1266,8 +1298,8 @@ def _parse_contents_from_anthropic( ) ) case "input_json_delta": - # Skip argument deltas for MCP tools — execution is handled server-side. - if self._last_call_content_type == "mcp_tool_use": + # Skip argument deltas for MCP and server tools — execution is handled server-side. + if self._last_call_content_type in ("mcp_tool_use", "server_tool_use"): pass else: call_id = self._last_call_id_name[0] if self._last_call_id_name else "" @@ -1376,3 +1408,106 @@ def service_url(self) -> str: The service URL for the chat client, or None if not set. """ return str(self.anthropic_client.base_url) + + +class AnthropicClient( + FunctionInvocationLayer[AnthropicOptionsT], + ChatMiddlewareLayer[AnthropicOptionsT], + ChatTelemetryLayer[AnthropicOptionsT], + RawAnthropicClient[AnthropicOptionsT], + Generic[AnthropicOptionsT], +): + """Anthropic chat client with middleware, telemetry, and function invocation support.""" + + def __init__( + self, + *, + api_key: str | None = None, + model: str | None = None, + base_url: str | None = None, + anthropic_client: AnthropicAsyncClient | None = None, + additional_beta_flags: list[str] | None = None, + additional_properties: dict[str, Any] | None = None, + middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None, + function_invocation_configuration: FunctionInvocationConfiguration | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize an Anthropic client. + + Keyword Args: + api_key: The Anthropic API key to use for authentication. + model: The model to use. + base_url: Optional base URL for the Anthropic API endpoint. Useful for Foundry or + other compatible deployments. Falls back to ``ANTHROPIC_BASE_URL`` env variable. + anthropic_client: An existing Anthropic client to use. If not provided, one will be created. + This can be used to further configure the client before passing it in. + For instance if you need to set a different base_url for testing or private deployments. + additional_beta_flags: Additional beta flags to enable on the client. + Default flags are: "mcp-client-2025-04-04", "code-execution-2025-08-25". + additional_properties: Additional properties stored on the client instance. + middleware: Optional middleware to apply to the client. + function_invocation_configuration: Optional function invocation configuration override. + env_file_path: Path to environment file for loading settings. + env_file_encoding: Encoding of the environment file. + + Examples: + .. code-block:: python + + from agent_framework.anthropic import AnthropicClient + + # Using environment variables + # Set ANTHROPIC_API_KEY=your_anthropic_api_key + # ANTHROPIC_CHAT_MODEL=claude-sonnet-4-5-20250929 + + # Or passing parameters directly + client = AnthropicClient( + model="claude-sonnet-4-5-20250929", + api_key="your_anthropic_api_key", + ) + + # Or with a custom base URL (e.g. for Foundry-compatible endpoints) + client = AnthropicClient( + model="claude-sonnet-4-5-20250929", + api_key="your_anthropic_api_key", + base_url="https://custom-anthropic-endpoint.com", + ) + + # Or loading from a .env file + client = AnthropicClient(env_file_path="path/to/.env") + + # Or passing in an existing client + from anthropic import AsyncAnthropic + + anthropic_client = AsyncAnthropic( + api_key="your_anthropic_api_key", base_url="https://custom-anthropic-endpoint.com" + ) + client = AnthropicClient( + model="claude-sonnet-4-5-20250929", + anthropic_client=anthropic_client, + ) + + # Using custom ChatOptions with type safety: + from typing import TypedDict + from agent_framework.anthropic import AnthropicChatOptions + + + class MyOptions(AnthropicChatOptions, total=False): + my_custom_option: str + + + client: AnthropicClient[MyOptions] = AnthropicClient(model="claude-sonnet-4-5-20250929") + response = await client.get_response("Hello", options={"my_custom_option": "value"}) + """ + super().__init__( + api_key=api_key, + model=model, + base_url=base_url, + anthropic_client=anthropic_client, + additional_beta_flags=additional_beta_flags, + additional_properties=additional_properties, + middleware=middleware, + function_invocation_configuration=function_invocation_configuration, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 5901e34dd90..68b5e35e405 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -4,31 +4,23 @@ import asyncio import base64 +import contextvars import json import logging import re import sys from abc import abstractmethod -from collections.abc import Callable, Collection, Sequence +from collections.abc import Callable, Collection, Coroutine, Sequence from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ignore from datetime import timedelta from functools import partial -from typing import TYPE_CHECKING, Any, Literal, TypedDict - -import httpx -from anyio import ClosedResourceError -from mcp import types -from mcp.client.session import ClientSession -from mcp.client.stdio import StdioServerParameters, stdio_client -from mcp.client.streamable_http import streamable_http_client -from mcp.client.websocket import websocket_client -from mcp.shared.context import RequestContext -from mcp.shared.exceptions import McpError -from mcp.shared.session import RequestResponder +from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast + from opentelemetry import propagate from ._tools import FunctionTool from ._types import ( + ChatOptions, Content, Message, ) @@ -40,7 +32,17 @@ from typing_extensions import Self # pragma: no cover if TYPE_CHECKING: + from httpx import AsyncClient + from mcp import types + from mcp.client.session import ClientSession + from mcp.shared.context import RequestContext + from mcp.shared.session import RequestResponder + from ._clients import SupportsChatGetResponse + from ._middleware import FunctionInvocationContext + + +logger = logging.getLogger(__name__) class MCPSpecificApproval(TypedDict, total=False): @@ -57,13 +59,15 @@ class MCPSpecificApproval(TypedDict, total=False): never_require_approval: Collection[str] | None -logger = logging.getLogger(__name__) _MCP_REMOTE_NAME_KEY = "_mcp_remote_name" _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name" +_mcp_call_headers: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar("_mcp_call_headers") +MCP_DEFAULT_TIMEOUT = 30 +MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5 # region: Helpers -LOG_LEVEL_MAPPING: dict[types.LoggingLevel, int] = { +LOG_LEVEL_MAPPING: dict[str, int] = { "debug": logging.DEBUG, "info": logging.INFO, "notice": logging.INFO, @@ -75,269 +79,6 @@ class MCPSpecificApproval(TypedDict, total=False): } -def _parse_prompt_result_from_mcp( - mcp_type: types.GetPromptResult, -) -> str: - """Parse an MCP GetPromptResult directly into a string representation. - - Converts each message in the prompt result to its string form and combines them. - - Args: - mcp_type: The MCP GetPromptResult object to convert. - - Returns: - A string representation of the prompt result. - """ - parts: list[str] = [] - for message in mcp_type.messages: - content = message.content - if isinstance(content, types.TextContent): - parts.append(content.text) - elif isinstance(content, (types.ImageContent, types.AudioContent)): - parts.append( - json.dumps( - { - "type": "image" if isinstance(content, types.ImageContent) else "audio", - "data": content.data, - "mimeType": content.mimeType, - }, - default=str, - ) - ) - elif isinstance(content, types.EmbeddedResource): - match content.resource: - case types.TextResourceContents(): - parts.append(content.resource.text) - case types.BlobResourceContents(): - parts.append( - json.dumps( - { - "type": "blob", - "data": content.resource.blob, - "mimeType": content.resource.mimeType, - }, - default=str, - ) - ) - else: - parts.append(str(content)) - if not parts: - return "" - if len(parts) == 1: - return parts[0] - return json.dumps(parts, default=str) - - -def _parse_message_from_mcp( - mcp_type: types.PromptMessage | types.SamplingMessage, -) -> Message: - """Parse an MCP container type into an Agent Framework type.""" - return Message( - role=mcp_type.role, - contents=_parse_content_from_mcp(mcp_type.content), - raw_representation=mcp_type, - ) - - -def _parse_tool_result_from_mcp( - mcp_type: types.CallToolResult, -) -> list[Content]: - """Parse an MCP CallToolResult into a list of Content items. - - Converts each content item in the MCP result to its appropriate - Content form. Text items become ``Content(type="text")`` and media - items (images, audio) are preserved as rich Content. - - Args: - mcp_type: The MCP CallToolResult object to convert. - - Returns: - A list of Content items representing the tool result. - """ - result: list[Content] = [] - for item in mcp_type.content: - match item: - case types.TextContent(): - result.append(Content.from_text(item.text)) - case types.ImageContent() | types.AudioContent(): - decoded = base64.b64decode(item.data) - result.append( - Content.from_data( - data=decoded, - media_type=item.mimeType, - ) - ) - case types.ResourceLink(): - result.append( - Content.from_uri( - uri=str(item.uri), - media_type=item.mimeType, - ) - ) - case types.EmbeddedResource(): - match item.resource: - case types.TextResourceContents(): - result.append(Content.from_text(item.resource.text)) - case types.BlobResourceContents(): - blob = item.resource.blob - mime = item.resource.mimeType or "application/octet-stream" - if not blob.startswith("data:"): - blob = f"data:{mime};base64,{blob}" - result.append( - Content.from_uri( - uri=blob, - media_type=mime, - ) - ) - case _: - result.append(Content.from_text(str(item))) - - if not result: - result.append(Content.from_text("null")) - return result - - -def _parse_content_from_mcp( - mcp_type: types.ImageContent - | types.TextContent - | types.AudioContent - | types.EmbeddedResource - | types.ResourceLink - | types.ToolUseContent - | types.ToolResultContent - | Sequence[ - types.ImageContent - | types.TextContent - | types.AudioContent - | types.EmbeddedResource - | types.ResourceLink - | types.ToolUseContent - | types.ToolResultContent - ], -) -> list[Content]: - """Parse an MCP type into an Agent Framework type.""" - mcp_types = mcp_type if isinstance(mcp_type, Sequence) else [mcp_type] - return_types: list[Content] = [] - for mcp_type in mcp_types: - match mcp_type: - case types.TextContent(): - return_types.append(Content.from_text(text=mcp_type.text, raw_representation=mcp_type)) - case types.ImageContent() | types.AudioContent(): - # MCP protocol uses base64-encoded strings, convert to bytes - data_bytes = base64.b64decode(mcp_type.data) if isinstance(mcp_type.data, str) else mcp_type.data - return_types.append( - Content.from_data( - data=data_bytes, - media_type=mcp_type.mimeType, - raw_representation=mcp_type, - ) - ) - case types.ResourceLink(): - return_types.append( - Content.from_uri( - uri=str(mcp_type.uri), - media_type=mcp_type.mimeType or "application/json", - raw_representation=mcp_type, - ) - ) - case types.ToolUseContent(): - return_types.append( - Content.from_function_call( - call_id=mcp_type.id, - name=mcp_type.name, - arguments=mcp_type.input, - raw_representation=mcp_type, - ) - ) - case types.ToolResultContent(): - return_types.append( - Content.from_function_result( - call_id=mcp_type.toolUseId, - result=_parse_content_from_mcp(mcp_type.content) - if mcp_type.content - else mcp_type.structuredContent, - exception=str(Exception()) if mcp_type.isError else None, # type: ignore[arg-type] - raw_representation=mcp_type, - ) - ) - case types.EmbeddedResource(): - match mcp_type.resource: - case types.TextResourceContents(): - return_types.append( - Content.from_text( - text=mcp_type.resource.text, - raw_representation=mcp_type, - additional_properties=( - mcp_type.annotations.model_dump() if mcp_type.annotations else None - ), - ) - ) - case types.BlobResourceContents(): - return_types.append( - Content.from_uri( - uri=mcp_type.resource.blob, - media_type=mcp_type.resource.mimeType, - raw_representation=mcp_type, - additional_properties=( - mcp_type.annotations.model_dump() if mcp_type.annotations else None - ), - ) - ) - return return_types - - -def _prepare_content_for_mcp( - content: Content, -) -> types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink | None: - """Prepare an Agent Framework content type for MCP.""" - if content.type == "text": - return types.TextContent(type="text", text=content.text) # type: ignore[attr-defined] - if content.type == "data": - if content.media_type and content.media_type.startswith("image/"): # type: ignore[attr-defined] - return types.ImageContent(type="image", data=content.uri, mimeType=content.media_type) # type: ignore[attr-defined] - if content.media_type and content.media_type.startswith("audio/"): # type: ignore[attr-defined] - return types.AudioContent(type="audio", data=content.uri, mimeType=content.media_type) # type: ignore[attr-defined] - if content.media_type and content.media_type.startswith("application/"): # type: ignore[attr-defined] - return types.EmbeddedResource( - type="resource", - resource=types.BlobResourceContents( - blob=content.uri, # type: ignore[attr-defined] - mimeType=content.media_type, # type: ignore[attr-defined] - # uri's are not limited in MCP but they have to be set. - # the uri of data content, contains the data uri, which - # is not the uri meant here, UriContent would match this. - uri=( - content.additional_properties.get("uri", "af://binary") - if content.additional_properties - else "af://binary" - ), # type: ignore[reportArgumentType] - ), - ) - return None - if content.type == "uri": - return types.ResourceLink( - type="resource_link", - uri=content.uri, # type: ignore[reportArgumentType,attr-defined] - mimeType=content.media_type, # type: ignore[attr-defined] - name=(content.additional_properties.get("name", "Unknown") if content.additional_properties else "Unknown"), - ) - return None - - -def _prepare_message_for_mcp( - content: Message, -) -> list[types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink]: - """Prepare a Message for MCP format.""" - messages: list[ - types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink - ] = [] - for item in content.contents: - mcp_content = _prepare_content_for_mcp(item) - if mcp_content: - messages.append(mcp_content) - return messages - - def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> dict[str, Any]: """Get the input model from an MCP prompt. @@ -401,6 +142,38 @@ def _inject_otel_into_mcp_meta(meta: dict[str, Any] | None = None) -> dict[str, return meta +def streamable_http_client(*args: Any, **kwargs: Any) -> _AsyncGeneratorContextManager[Any, None]: + """Lazily import the MCP streamable HTTP transport.""" + try: + from mcp.client.streamable_http import streamable_http_client as _streamable_http_client + except ModuleNotFoundError as ex: + missing_name = ex.name or str(ex) + if missing_name == "mcp" or missing_name.startswith("mcp.") or "mcp" in missing_name: + raise ModuleNotFoundError("`MCPStreamableHTTPTool` requires `mcp`. Please install `mcp`.") from ex + raise ModuleNotFoundError( + f"`MCPStreamableHTTPTool` requires streamable HTTP transport support. " + f"The optional dependency `{missing_name}` is not installed. Please update your dependencies." + ) from ex + + return _streamable_http_client(*args, **kwargs) # type: ignore[return-value] + + +def _should_propagate_cancelled_error(ex: BaseException) -> bool: + """Return True if *ex* is a genuine task-cancellation that should propagate unchanged. + + On Python >= 3.11, ``task.cancelling() > 0`` distinguishes a real caller-driven + cancellation from a CancelledError raised internally by a library (e.g. via an + anyio cancel scope). On older Python versions the API is unavailable, so we + always return False and let callers wrap the error in ToolException instead. + """ + if not isinstance(ex, asyncio.CancelledError): + return False + if sys.version_info < (3, 11): + return False + task = asyncio.current_task() + return task is not None and task.cancelling() > 0 + + # region: MCP Plugin @@ -462,8 +235,8 @@ def __init__( ``Callable[[types.GetPromptResult], str]`` that overrides the default prompt result parsing. When ``None`` (the default), the built-in parser converts MCP prompt results to a string. If you need per-function result parsing, - access the ``.functions`` list after connecting and set ``result_parser`` on - individual ``FunctionTool`` instances. + access the ``.functions`` list after connecting and set ``result_parser`` on + individual ``FunctionTool`` instances. session: An existing MCP client session to use. request_timeout: Timeout in seconds for MCP requests. client: A chat client for sampling callbacks. @@ -482,25 +255,292 @@ def __init__( self._exit_stack = AsyncExitStack() self._lifecycle_lock = asyncio.Lock() self._lifecycle_request_lock = asyncio.Lock() - self._lifecycle_queue: asyncio.Queue[tuple[str, bool, asyncio.Future[None]]] | None = None + self._lifecycle_queue: asyncio.Queue[tuple[str, bool, bool, asyncio.Future[None]]] | None = None self._lifecycle_owner_task: asyncio.Task[None] | None = None self.session = session self.request_timeout = request_timeout self.client = client self._functions: list[FunctionTool] = [] + self._tool_call_meta_by_name: dict[str, dict[str, Any]] = {} self.is_connected: bool = False self._tools_loaded: bool = False self._prompts_loaded: bool = False + self._server_capabilities: types.ServerCapabilities | None = None + self._supports_tools: bool = True + self._supports_prompts: bool = True + self._supports_logging: bool | None = None + self._ping_available: bool = True + self._pending_reload_tasks: set[asyncio.Task[None]] = set() def __str__(self) -> str: return f"MCPTool(name={self.name}, description={self.description})" + def _parse_prompt_result_from_mcp( + self, + mcp_type: types.GetPromptResult, + ) -> str: + """Parse an MCP GetPromptResult directly into a string representation.""" + from mcp import types + + parts: list[str] = [] + for message in mcp_type.messages: + content = message.content + if isinstance(content, types.TextContent): + parts.append(content.text) + elif isinstance(content, (types.ImageContent, types.AudioContent)): + parts.append( + json.dumps( + { + "type": "image" if isinstance(content, types.ImageContent) else "audio", + "data": content.data, + "mimeType": content.mimeType, + }, + default=str, + ) + ) + elif isinstance(content, types.EmbeddedResource): + match content.resource: + case types.TextResourceContents(): + parts.append(content.resource.text) + case types.BlobResourceContents(): + parts.append( + json.dumps( + { + "type": "blob", + "data": content.resource.blob, + "mimeType": content.resource.mimeType, + }, + default=str, + ) + ) + else: + parts.append(str(content)) + if not parts: + return "" + if len(parts) == 1: + return parts[0] + return json.dumps(parts, default=str) + + def _parse_message_from_mcp( + self, + mcp_type: types.PromptMessage | types.SamplingMessage, + ) -> Message: + """Parse an MCP container type into an Agent Framework type.""" + return Message( + role=mcp_type.role, + contents=self._parse_content_from_mcp(mcp_type.content), + raw_representation=mcp_type, + ) + + def _parse_tool_result_from_mcp( + self, + mcp_type: types.CallToolResult, + ) -> list[Content]: + """Parse an MCP CallToolResult into a list of Content items.""" + from mcp import types + + result: list[Content] = [] + for item in mcp_type.content: + match item: + case types.TextContent(): + result.append(Content.from_text(item.text)) + case types.ImageContent() | types.AudioContent(): + decoded = base64.b64decode(item.data) + result.append( + Content.from_data( + data=decoded, + media_type=item.mimeType, + ) + ) + case types.ResourceLink(): + result.append( + Content.from_uri( + uri=str(item.uri), + media_type=item.mimeType, + ) + ) + case types.EmbeddedResource(): + match item.resource: + case types.TextResourceContents(): + result.append(Content.from_text(item.resource.text)) + case types.BlobResourceContents(): + blob = item.resource.blob + mime = item.resource.mimeType or "application/octet-stream" + if not blob.startswith("data:"): + blob = f"data:{mime};base64,{blob}" + result.append( + Content.from_uri( + uri=blob, + media_type=mime, + ) + ) + case _: + result.append(Content.from_text(str(item))) + + if not result: + result.append(Content.from_text("null")) + return result + + def _parse_content_from_mcp( + self, + mcp_type: types.ImageContent + | types.TextContent + | types.AudioContent + | types.EmbeddedResource + | types.ResourceLink + | types.ToolUseContent + | types.ToolResultContent + | Sequence[ + types.ImageContent + | types.TextContent + | types.AudioContent + | types.EmbeddedResource + | types.ResourceLink + | types.ToolUseContent + | types.ToolResultContent + ], + ) -> list[Content]: + """Parse an MCP type into an Agent Framework type.""" + from mcp import types + + mcp_content_types: Sequence[Any] = ( + cast(Sequence[Any], mcp_type) if isinstance(mcp_type, Sequence) else [mcp_type] + ) # type: ignore[redundant-cast] + return_types: list[Content] = [] + for mcp_type in mcp_content_types: + match mcp_type: + case types.TextContent(): + return_types.append(Content.from_text(text=mcp_type.text, raw_representation=mcp_type)) + case types.ImageContent() | types.AudioContent(): + data_bytes = base64.b64decode(mcp_type.data) if isinstance(mcp_type.data, str) else mcp_type.data + return_types.append( + Content.from_data( + data=data_bytes, + media_type=mcp_type.mimeType, + raw_representation=mcp_type, + ) + ) + case types.ResourceLink(): + return_types.append( + Content.from_uri( + uri=str(mcp_type.uri), + media_type=mcp_type.mimeType or "application/json", + raw_representation=mcp_type, + ) + ) + case types.ToolUseContent(): + return_types.append( + Content.from_function_call( + call_id=mcp_type.id, + name=mcp_type.name, + arguments=mcp_type.input, + raw_representation=mcp_type, + ) + ) + case types.ToolResultContent(): + return_types.append( + Content.from_function_result( + call_id=mcp_type.toolUseId, + result=self._parse_content_from_mcp(mcp_type.content) + if mcp_type.content + else mcp_type.structuredContent, + exception=str(Exception()) if mcp_type.isError else None, # type: ignore[arg-type] + raw_representation=mcp_type, + ) + ) + case types.EmbeddedResource(): + match mcp_type.resource: + case types.TextResourceContents(): + return_types.append( + Content.from_text( + text=mcp_type.resource.text, + raw_representation=mcp_type, + additional_properties=( + mcp_type.annotations.model_dump() if mcp_type.annotations else None + ), + ) + ) + case types.BlobResourceContents(): + return_types.append( + Content.from_uri( + uri=mcp_type.resource.blob, + media_type=mcp_type.resource.mimeType, + raw_representation=mcp_type, + additional_properties=( + mcp_type.annotations.model_dump() if mcp_type.annotations else None + ), + ) + ) + case _: + pass + return return_types + + def _prepare_content_for_mcp( + self, + content: Content, + ) -> ( + types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink | None + ): + """Prepare an Agent Framework content type for MCP.""" + from mcp import types + + if content.type == "text": + return types.TextContent(type="text", text=content.text) # type: ignore[attr-defined] + if content.type == "data": + if content.media_type and content.media_type.startswith("image/"): # type: ignore[attr-defined] + return types.ImageContent(type="image", data=content.uri, mimeType=content.media_type) # type: ignore[attr-defined] + if content.media_type and content.media_type.startswith("audio/"): # type: ignore[attr-defined] + return types.AudioContent(type="audio", data=content.uri, mimeType=content.media_type) # type: ignore[attr-defined] + if content.media_type and content.media_type.startswith("application/"): # type: ignore[attr-defined] + return types.EmbeddedResource( + type="resource", + resource=types.BlobResourceContents( + blob=content.uri, # type: ignore[attr-defined] + mimeType=content.media_type, # type: ignore[attr-defined] + uri=( + content.additional_properties.get("uri", "af://binary") + if content.additional_properties + else "af://binary" + ), # type: ignore[arg-type] + ), + ) + return None + if content.type == "uri": + resource_name = ( + content.additional_properties.get("name", "Unknown") if content.additional_properties else "Unknown" + ) + return types.ResourceLink( + type="resource_link", + uri=content.uri, # type: ignore[arg-type,attr-defined] + mimeType=content.media_type, # type: ignore[attr-defined] + name=resource_name, + ) + return None + + def _prepare_message_for_mcp( + self, + content: Message, + ) -> list[ + types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink + ]: + """Prepare a Message for MCP format.""" + messages: list[ + types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink + ] = [] + for item in content.contents: + mcp_content = self._prepare_content_for_mcp(item) + if mcp_content: + messages.append(mcp_content) + return messages + @property def functions(self) -> list[FunctionTool]: """Get the list of functions that are allowed.""" - if not self.allowed_tools: + if self.allowed_tools is None: return self._functions allowed_names = set(self.allowed_tools) + if not allowed_names: + return [] filtered_functions: list[FunctionTool] = [] for func in self._functions: additional_properties = func.additional_properties or {} @@ -533,11 +573,11 @@ async def _run_lifecycle_owner(self) -> None: stop_error: BaseException | None = None try: while True: - action, reset, future = await queue.get() + action, reset, load_configured, future = await queue.get() try: if action == "connect": - await self._connect_on_owner(reset=reset) + await self._connect_on_owner(reset=reset, load_configured=load_configured) elif action == "close": await self._close_on_owner() else: @@ -562,7 +602,7 @@ async def _run_lifecycle_owner(self) -> None: finally: while True: try: - _, _, future = queue.get_nowait() + _, _, _, future = queue.get_nowait() except asyncio.QueueEmpty: break if not future.done(): @@ -575,12 +615,18 @@ def _is_lifecycle_owner_task(self) -> bool: owner_task = self._lifecycle_owner_task return owner_task is not None and asyncio.current_task() is owner_task - async def _run_on_lifecycle_owner(self, action: str, *, reset: bool = False) -> None: + async def _run_on_lifecycle_owner( + self, + action: str, + *, + reset: bool = False, + load_configured: bool = True, + ) -> None: await self._ensure_lifecycle_owner() if self._is_lifecycle_owner_task(): if action == "connect": - await self._connect_on_owner(reset=reset) + await self._connect_on_owner(reset=reset, load_configured=load_configured) elif action == "close": await self._close_on_owner() else: @@ -592,7 +638,7 @@ async def _run_on_lifecycle_owner(self, action: str, *, reset: bool = False) -> raise RuntimeError("MCP lifecycle owner is not available.") future = asyncio.get_running_loop().create_future() - await queue.put((action, reset, future)) + await queue.put((action, reset, load_configured, future)) await future async def _safe_close_exit_stack(self) -> None: @@ -612,6 +658,43 @@ async def _safe_close_exit_stack(self) -> None: except asyncio.CancelledError: logger.warning("Could not cleanly close MCP exit stack because the lifecycle owner task was cancelled.") + async def _close_and_check_cancelled(self, ex: BaseException) -> bool: + """Close the exit stack and return True if *ex* is a genuine task cancellation. + + Callers should immediately re-raise when this returns True:: + + if await self._close_and_check_cancelled(ex): + raise + """ + await self._safe_close_exit_stack() + return _should_propagate_cancelled_error(ex) + + def _reset_session_state(self) -> None: + self._server_capabilities = None + self._supports_tools = True + self._supports_prompts = True + self._supports_logging = None + self._ping_available = True + + def _set_server_capabilities(self, capabilities: types.ServerCapabilities | None) -> None: + self._server_capabilities = capabilities + if capabilities is None: + self._supports_tools = False + self._supports_prompts = False + self._supports_logging = False + return + + self._supports_tools = getattr(capabilities, "tools", None) is not None + self._supports_prompts = getattr(capabilities, "prompts", None) is not None + self._supports_logging = getattr(capabilities, "logging", None) is not None + + async def _reconnect_without_loading(self) -> None: + if self._is_lifecycle_owner_task(): + await self._connect_on_owner(reset=True, load_configured=False) + return + + await self._run_on_lifecycle_owner("connect", reset=True, load_configured=False) + async def connect(self, *, reset: bool = False) -> None: if self._is_lifecycle_owner_task(): await self._connect_on_owner(reset=reset) @@ -620,7 +703,7 @@ async def connect(self, *, reset: bool = False) -> None: async with self._lifecycle_request_lock: await self._run_on_lifecycle_owner("connect", reset=reset) - async def _connect_on_owner(self, *, reset: bool = False) -> None: + async def _connect_on_owner(self, *, reset: bool = False, load_configured: bool = True) -> None: """Connect to the MCP server. Establishes a connection to the MCP server, initializes the session, @@ -628,6 +711,7 @@ async def _connect_on_owner(self, *, reset: bool = False) -> None: Keyword Args: reset: If True, forces a reconnection even if already connected. + load_configured: If True, loads tools and prompts according to the constructor flags. Raises: ToolException: If connection or session initialization fails. @@ -636,21 +720,46 @@ async def _connect_on_owner(self, *, reset: bool = False) -> None: await self._safe_close_exit_stack() self.session = None self.is_connected = False + self._reset_session_state() self._exit_stack = AsyncExitStack() if not self.session: try: transport = await self._exit_stack.enter_async_context(self.get_mcp_client()) - except Exception as ex: - await self._safe_close_exit_stack() + except (Exception, asyncio.CancelledError) as ex: + # On Python >= 3.11, re-raise genuine task cancellation (task.cancelling() > 0) + # instead of wrapping it in ToolException. On Python < 3.11, task.cancelling() + # is unavailable so MCP-internal CancelledErrors cannot be distinguished from + # caller-driven cancellation; they are wrapped as ToolException in that case. + if await self._close_and_check_cancelled(ex): + raise command = getattr(self, "command", None) if command: error_msg = f"Failed to start MCP server '{command}': {ex}" else: error_msg = f"Failed to connect to MCP server: {ex}" - raise ToolException(error_msg, inner_exception=ex) from ex + # CancelledError is a BaseException (not Exception) on Python >= 3.8, so + # inner_exception=None and ToolException.__init__ won't log exc_info. + if isinstance(ex, asyncio.CancelledError): + logger.debug(error_msg, exc_info=True) + raise ToolException(error_msg, inner_exception=ex if isinstance(ex, Exception) else None) from ex try: + try: + from mcp import types + from mcp.client.session import ClientSession as runtime_client_session + except ModuleNotFoundError as ex: + await self._safe_close_exit_stack() + raise ToolException( + "MCP support requires `mcp`. Please install `mcp`.", + inner_exception=ex, + ) from ex + + sampling_capabilities = None + if self.client is not None: + sampling_capabilities = types.SamplingCapability( + tools=types.SamplingToolsCapability(), + ) session = await self._exit_stack.enter_async_context( - ClientSession( + runtime_client_session( read_stream=transport[0], write_stream=transport[1], read_timeout_seconds=( @@ -659,18 +768,25 @@ async def _connect_on_owner(self, *, reset: bool = False) -> None: message_handler=self.message_handler, logging_callback=self.logging_callback, sampling_callback=self.sampling_callback, + sampling_capabilities=sampling_capabilities, ) ) - except Exception as ex: - await self._safe_close_exit_stack() + except (Exception, asyncio.CancelledError) as ex: + if await self._close_and_check_cancelled(ex): + raise + session_error_msg = f"Failed to create MCP session: {ex}" + if isinstance(ex, asyncio.CancelledError): + logger.debug(session_error_msg, exc_info=True) raise ToolException( - message="Failed to create MCP session. Please check your configuration.", - inner_exception=ex, + message=session_error_msg, + inner_exception=ex if isinstance(ex, Exception) else None, ) from ex try: - await session.initialize() - except Exception as ex: - await self._safe_close_exit_stack() + initialize_result = await session.initialize() + self._set_server_capabilities(getattr(initialize_result, "capabilities", None)) + except (Exception, asyncio.CancelledError) as ex: + if await self._close_and_check_cancelled(ex): + raise # Provide context about initialization failure command = getattr(self, "command", None) if command: @@ -679,25 +795,33 @@ async def _connect_on_owner(self, *, reset: bool = False) -> None: error_msg = f"MCP server '{full_command}' failed to initialize: {ex}" else: error_msg = f"MCP server failed to initialize: {ex}" - raise ToolException(error_msg, inner_exception=ex) from ex + if isinstance(ex, asyncio.CancelledError): + logger.debug(error_msg, exc_info=True) + raise ToolException(error_msg, inner_exception=ex if isinstance(ex, Exception) else None) from ex self.session = session - elif self.session._request_id == 0: # type: ignore[reportPrivateUsage] + elif self.session._request_id == 0: # type: ignore[attr-defined] # If the session is not initialized, we need to reinitialize it - await self.session.initialize() + initialize_result = await self.session.initialize() + self._set_server_capabilities(getattr(initialize_result, "capabilities", None)) + elif self._server_capabilities is None: + self._set_server_capabilities(getattr(self.session, "_server_capabilities", None)) logger.debug("Connected to MCP server: %s", self.session) self.is_connected = True - if self.load_tools_flag: - await self.load_tools() + if load_configured and self.load_tools_flag: + if self._supports_tools: + await self.load_tools() self._tools_loaded = True - if self.load_prompts_flag: - await self.load_prompts() + if load_configured and self.load_prompts_flag: + if self._supports_prompts: + await self.load_prompts() self._prompts_loaded = True - if logger.level != logging.NOTSET: + if logger.level != logging.NOTSET and self._supports_logging is not False: try: - await self.session.set_logging_level( - next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level) + level_name = cast( + Any, next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level) ) + await self.session.set_logging_level(level_name) except Exception as exc: logger.warning("Failed to set log level to %s", logger.level, exc_info=exc) @@ -723,6 +847,8 @@ async def sampling_callback( Returns: Either a CreateMessageResult with the generated message or ErrorData if generation fails. """ + from mcp import types + if not self.client: return types.ErrorData( code=types.INTERNAL_ERROR, @@ -731,15 +857,37 @@ async def sampling_callback( logger.debug("Sampling callback called with params: %s", params) messages: list[Message] = [] for msg in params.messages: - messages.append(_parse_message_from_mcp(msg)) + messages.append(self._parse_message_from_mcp(msg)) + + options: ChatOptions[None] = {} + if params.systemPrompt is not None: + options["instructions"] = params.systemPrompt + if params.tools is not None: + options["tools"] = [ + FunctionTool( + name=tool.name, + description=tool.description or "", + input_model=tool.inputSchema, + ) + for tool in params.tools + ] + if params.toolChoice is not None and params.toolChoice.mode is not None: + options["tool_choice"] = params.toolChoice.mode + + if params.temperature is not None: + options["temperature"] = params.temperature + options["max_tokens"] = params.maxTokens + if params.stopSequences is not None: + options["stop"] = params.stopSequences + try: - response = await self.client.get_response( + chat_client: Any = self.client + response: Any = await chat_client.get_response( messages, - temperature=params.temperature, - max_tokens=params.maxTokens, - stop=params.stopSequences, + options=options or None, ) except Exception as ex: + logger.debug("Sampling callback error: %s", ex, exc_info=True) return types.ErrorData( code=types.INTERNAL_ERROR, message=f"Failed to get chat message content: {ex}", @@ -749,7 +897,7 @@ async def sampling_callback( code=types.INTERNAL_ERROR, message="Failed to get chat message content.", ) - mcp_contents = _prepare_message_for_mcp(response.messages[0]) + mcp_contents = self._prepare_message_for_mcp(response.messages[0]) # grab the first content that is of type TextContent or ImageContent mcp_content = next( (content for content in mcp_contents if isinstance(content, (types.TextContent, types.ImageContent))), @@ -763,7 +911,7 @@ async def sampling_callback( return types.CreateMessageResult( role="assistant", content=mcp_content, - model=response.model_id or "unknown", + model=response.model or "unknown", ) async def logging_callback(self, params: types.LoggingMessageNotificationParams) -> None: @@ -798,18 +946,55 @@ async def message_handler( Args: message: The message from the MCP server (request responder, notification, or exception). """ + from mcp import types + if isinstance(message, Exception): logger.error("Error from MCP server: %s", message, exc_info=message) return if isinstance(message, types.ServerNotification): match message.root.method: case "notifications/tools/list_changed": - await self.load_tools() + self._schedule_reload(self.load_tools()) case "notifications/prompts/list_changed": - await self.load_prompts() + self._schedule_reload(self.load_prompts()) case _: logger.debug("Unhandled notification: %s", message.root.method) + def _schedule_reload(self, coro: Coroutine[Any, Any, None]) -> None: + """Schedule a reload coroutine as a background task. + + Reloads (load_tools / load_prompts) triggered by MCP server + notifications must NOT be awaited inside the message handler because + the handler runs on the MCP SDK's single-threaded receive loop. + Awaiting a session request (e.g. ``list_tools``) from within that loop + deadlocks: the receive loop cannot read the response while it is + blocked waiting for the handler to return. + + Instead we fire the reload as an independent ``asyncio.Task`` and keep + a strong reference in ``_pending_reload_tasks`` so it is not garbage- + collected before completion. Only one reload per kind (tools / prompts) + is kept in flight; a new notification cancels the previous pending task + for the same coroutine name to avoid unbounded growth. + """ + # Cancel-and-replace: only one reload per kind should be in flight. + reload_name = f"mcp-reload:{self.name}:{coro.__qualname__}" + for existing in list(self._pending_reload_tasks): + if existing.get_name() == reload_name and not existing.done(): + logger.debug("Cancelling in-flight reload %s; superseded by new notification", reload_name) + existing.cancel() + + async def _safe_reload() -> None: + try: + await coro + except asyncio.CancelledError: + raise + except Exception: + logger.warning("Background MCP reload failed", exc_info=True) + + task = asyncio.create_task(_safe_reload(), name=reload_name) + self._pending_reload_tasks.add(task) + task.add_done_callback(self._pending_reload_tasks.discard) + def _determine_approval_mode( self, *candidate_names: str, @@ -824,7 +1009,7 @@ def _determine_approval_mode( ): return "never_require" return None - return self.approval_mode # type: ignore[reportReturnType] + return self.approval_mode # type: ignore[return-value] async def load_prompts(self) -> None: """Load prompts from the MCP server. @@ -835,15 +1020,49 @@ async def load_prompts(self) -> None: Raises: ToolExecutionException: If the MCP server is not connected. """ + from anyio import ClosedResourceError + from mcp import types + + if not self._supports_prompts: + logger.debug("Skipping MCP prompt loading because the server did not advertise prompts support.") + return + # Track existing function names to prevent duplicates existing_names = {func.name for func in self._functions} params: types.PaginatedRequestParams | None = None while True: - # Ensure connection is still valid before each page request - await self._ensure_connected() + prompt_list: types.ListPromptsResult | None = None + for attempt in range(2): + try: + # Ensure connection is still valid before each page request + await self._ensure_connected() + if not self._supports_prompts: + logger.debug( + "Skipping MCP prompt loading because the server did not advertise prompts support." + ) + return + prompt_list = await self.session.list_prompts(params=params) # type: ignore[union-attr] + break + except ClosedResourceError as cl_ex: + if attempt == 0: + logger.info("MCP connection closed unexpectedly while loading prompts. Reconnecting...") + try: + await self._reconnect_without_loading() + except Exception as reconn_ex: + raise ToolExecutionException( + "Failed to reconnect to MCP server.", + inner_exception=reconn_ex, + ) from reconn_ex + continue + logger.error("MCP connection closed unexpectedly after reconnection: %s", cl_ex) + raise ToolExecutionException( + "Failed to load prompts - connection lost.", + inner_exception=cl_ex, + ) from cl_ex - prompt_list = await self.session.list_prompts(params=params) # type: ignore[union-attr] + if prompt_list is None: + raise ToolExecutionException("Failed to load prompts.") for prompt in prompt_list.prompts: normalized_name = _normalize_mcp_name(prompt.name) @@ -870,7 +1089,7 @@ async def load_prompts(self) -> None: existing_names.add(local_name) # Check if there are more pages - if not prompt_list or not prompt_list.nextCursor: + if not prompt_list.nextCursor: break params = types.PaginatedRequestParams(cursor=prompt_list.nextCursor) @@ -883,17 +1102,53 @@ async def load_tools(self) -> None: Raises: ToolExecutionException: If the MCP server is not connected. """ + from anyio import ClosedResourceError + from mcp import types + + if not self._supports_tools: + logger.debug("Skipping MCP tool loading because the server did not advertise tools support.") + return + # Track existing function names to prevent duplicates existing_names = {func.name for func in self._functions} + self._tool_call_meta_by_name.clear() params: types.PaginatedRequestParams | None = None while True: - # Ensure connection is still valid before each page request - await self._ensure_connected() + tool_list: types.ListToolsResult | None = None + for attempt in range(2): + try: + # Ensure connection is still valid before each page request + await self._ensure_connected() + if not self._supports_tools: + logger.debug("Skipping MCP tool loading because the server did not advertise tools support.") + return + tool_list = await self.session.list_tools(params=params) # type: ignore[union-attr] + break + except ClosedResourceError as cl_ex: + if attempt == 0: + logger.info("MCP connection closed unexpectedly while loading tools. Reconnecting...") + try: + await self._reconnect_without_loading() + except Exception as reconn_ex: + raise ToolExecutionException( + "Failed to reconnect to MCP server.", + inner_exception=reconn_ex, + ) from reconn_ex + continue + logger.error("MCP connection closed unexpectedly after reconnection: %s", cl_ex) + raise ToolExecutionException( + "Failed to load tools - connection lost.", + inner_exception=cl_ex, + ) from cl_ex - tool_list = await self.session.list_tools(params=params) # type: ignore[union-attr] + if tool_list is None: + raise ToolExecutionException("Failed to load tools.") for tool in tool_list.tools: + if tool.meta is not None: + self._tool_call_meta_by_name[tool.name] = dict(tool.meta) + normalized_name = _normalize_mcp_name(tool.name) local_name = _build_prefixed_mcp_name(normalized_name, self.tool_name_prefix) @@ -902,13 +1157,32 @@ async def load_tools(self) -> None: continue approval_mode = self._determine_approval_mode(local_name, normalized_name, tool.name) + # Normalize inputSchema: ensure "properties" exists for object schemas. + # Some MCP servers (e.g. zero-argument tools) omit "properties", + # which causes OpenAI API to reject the schema with a 400 error. + # Guard against non-conforming MCP servers that send inputSchema=None + # despite the MCP spec typing it as dict[str, Any]. + input_schema = dict(tool.inputSchema or {}) + if input_schema.get("type") == "object" and "properties" not in input_schema: + input_schema["properties"] = {} + + async def _call_tool_with_runtime_kwargs( + ctx: FunctionInvocationContext, + *, + _remote_tool_name: str = tool.name, + **kwargs: Any, + ) -> str | list[Content]: + call_kwargs = dict(ctx.kwargs) + call_kwargs.update(kwargs) + return await self.call_tool(_remote_tool_name, **call_kwargs) + # Create FunctionTools out of each tool func: FunctionTool = FunctionTool( - func=partial(self.call_tool, tool.name), + func=_call_tool_with_runtime_kwargs, name=local_name, description=tool.description or "", approval_mode=approval_mode, - input_model=tool.inputSchema, + input_model=input_schema, additional_properties={ _MCP_REMOTE_NAME_KEY: tool.name, _MCP_NORMALIZED_NAME_KEY: normalized_name, @@ -918,15 +1192,24 @@ async def load_tools(self) -> None: existing_names.add(local_name) # Check if there are more pages - if not tool_list or not tool_list.nextCursor: + if not tool_list.nextCursor: break params = types.PaginatedRequestParams(cursor=tool_list.nextCursor) async def _close_on_owner(self) -> None: + # Cancel any pending reload tasks before tearing down the session. + tasks = list(self._pending_reload_tasks) + for task in tasks: + task.cancel() + self._pending_reload_tasks.clear() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + await self._safe_close_exit_stack() self._exit_stack = AsyncExitStack() self.session = None self.is_connected = False + self._reset_session_state() async def close(self) -> None: """Disconnect from the MCP server. @@ -958,12 +1241,30 @@ async def _ensure_connected(self) -> None: Raises: ToolExecutionException: If reconnection fails. """ + from mcp.shared.exceptions import McpError + + if not self._ping_available: + return + try: await self.session.send_ping() # type: ignore[union-attr] + except McpError as mcp_exc: + if mcp_exc.error.code == -32601: + self._ping_available = False + logger.debug("Skipping future MCP pings because the server does not support ping.") + return + logger.info("MCP connection invalid or closed. Reconnecting...") + try: + await self._reconnect_without_loading() + except Exception as ex: + raise ToolExecutionException( + "Failed to establish MCP connection.", + inner_exception=ex, + ) from ex except Exception: logger.info("MCP connection invalid or closed. Reconnecting...") try: - await self.connect(reset=True) + await self._reconnect_without_loading() except Exception as ex: raise ToolExecutionException( "Failed to establish MCP connection.", @@ -988,6 +1289,9 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: ToolExecutionException: If the MCP server is not connected, tools are not loaded, or the tool call fails. """ + from anyio import ClosedResourceError + from mcp.shared.exceptions import McpError + if not self.load_tools_flag: raise ToolExecutionException( "Tools are not loaded for this server, please set load_tools=True in the constructor." @@ -1014,15 +1318,15 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: } } - # Inject OpenTelemetry trace context into MCP _meta for distributed tracing. - otel_meta = _inject_otel_into_mcp_meta() - - parser = self.parse_tool_results or _parse_tool_result_from_mcp + # Some MCP proxies require their tools/list metadata to be echoed on tools/call. + tool_meta = self._tool_call_meta_by_name.get(tool_name) + meta = _inject_otel_into_mcp_meta(dict(tool_meta) if tool_meta is not None else None) + parser = self.parse_tool_results or self._parse_tool_result_from_mcp # Try the operation, reconnecting once if the connection is closed for attempt in range(2): try: - result = await self.session.call_tool(tool_name, arguments=filtered_kwargs, meta=otel_meta) # type: ignore + result = await self.session.call_tool(tool_name, arguments=filtered_kwargs, meta=meta) # type: ignore if result.isError: parsed = parser(result) text = ( @@ -1054,7 +1358,8 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: inner_exception=cl_ex, ) from cl_ex except McpError as mcp_exc: - raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc + error_message = mcp_exc.error.message + raise ToolExecutionException(error_message, inner_exception=mcp_exc) from mcp_exc except Exception as ex: raise ToolExecutionException(f"Failed to call tool '{tool_name}'.", inner_exception=ex) from ex raise ToolExecutionException(f"Failed to call tool '{tool_name}' after retries.") @@ -1075,13 +1380,15 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: ToolExecutionException: If the MCP server is not connected, prompts are not loaded, or the prompt call fails. """ + from anyio import ClosedResourceError + from mcp.shared.exceptions import McpError + if not self.load_prompts_flag: raise ToolExecutionException( "Prompts are not loaded for this server, please set load_prompts=True in the constructor." ) - parser = self.parse_prompt_results or _parse_prompt_result_from_mcp - + parser = self.parse_prompt_results or self._parse_prompt_result_from_mcp # Try the operation, reconnecting once if the connection is closed for attempt in range(2): try: @@ -1107,7 +1414,8 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: inner_exception=cl_ex, ) from cl_ex except McpError as mcp_exc: - raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc + error_message = mcp_exc.error.message + raise ToolExecutionException(error_message, inner_exception=mcp_exc) from mcp_exc except Exception as ex: raise ToolExecutionException(f"Failed to call prompt '{prompt_name}'.", inner_exception=ex) from ex raise ToolExecutionException(f"Failed to get prompt '{prompt_name}' after retries.") @@ -1281,6 +1589,11 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: args["encoding"] = self.encoding if self._client_kwargs: args.update(self._client_kwargs) + try: + from mcp.client.stdio import StdioServerParameters, stdio_client + except ModuleNotFoundError as ex: + raise ModuleNotFoundError("`mcp` is required to use `MCPStdioTool`. Please install `mcp`.") from ex + return stdio_client(server=StdioServerParameters(**args)) @@ -1325,7 +1638,8 @@ def __init__( terminate_on_close: bool | None = None, client: SupportsChatGetResponse | None = None, additional_properties: dict[str, Any] | None = None, - http_client: httpx.AsyncClient | None = None, + http_client: AsyncClient | None = None, + header_provider: Callable[[dict[str, Any]], dict[str, str]] | None = None, **kwargs: Any, ) -> None: """Initialize the MCP streamable HTTP tool. @@ -1333,7 +1647,7 @@ def __init__( Note: The arguments are used to create a streamable HTTP client using the new ``mcp.client.streamable_http.streamable_http_client`` API. - If an httpx.AsyncClient is provided via ``http_client``, it will be used directly. + If an asyncClient is provided via ``http_client``, it will be used directly. Otherwise, the ``streamable_http_client`` API will create and manage a default client. Args: @@ -1369,10 +1683,15 @@ def __init__( additional_properties: Additional properties. terminate_on_close: Close the transport when the MCP client is terminated. client: The chat client to use for sampling. - http_client: Optional httpx.AsyncClient to use. If not provided, the + http_client: Optional asyncClient to use. If not provided, the ``streamable_http_client`` API will create and manage a default client. To configure headers, timeouts, or other HTTP client settings, create - and pass your own ``httpx.AsyncClient`` instance. + and pass your own ``asyncClient`` instance. + header_provider: Optional callable that receives the runtime keyword arguments + (from ``FunctionInvocationContext.kwargs``) and returns a ``dict[str, str]`` + of HTTP headers to inject into every outbound request to the MCP server. + Use this to forward per-request context (e.g. authentication tokens set in + agent middleware) without creating a separate ``httpx.AsyncClient``. kwargs: Additional keyword arguments (accepted for backward compatibility but not used). """ super().__init__( @@ -1392,7 +1711,8 @@ def __init__( ) self.url = url self.terminate_on_close = terminate_on_close - self._httpx_client: httpx.AsyncClient | None = http_client + self._httpx_client: AsyncClient | None = http_client + self._header_provider = header_provider def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: """Get an MCP streamable HTTP client. @@ -1400,13 +1720,59 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: Returns: An async context manager for the streamable HTTP client transport. """ - # Pass the http_client (which may be None) to streamable_http_client + from httpx import AsyncClient, Request, Timeout + + http_client = self._httpx_client + if self._header_provider is not None: + if http_client is None: + http_client = AsyncClient( + follow_redirects=True, + timeout=Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT), + ) + self._httpx_client = http_client + + if not hasattr(self, "_inject_headers_hook"): + + async def _inject_headers(request: Request) -> None: # noqa: RUF029 + headers = _mcp_call_headers.get({}) + for key, value in headers.items(): + request.headers[key] = value + + self._inject_headers_hook = _inject_headers # type: ignore[attr-defined] + http_client.event_hooks["request"].append(self._inject_headers_hook) # type: ignore[attr-defined] + return streamable_http_client( url=self.url, - http_client=self._httpx_client, + http_client=http_client, terminate_on_close=self.terminate_on_close if self.terminate_on_close is not None else True, ) + async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: + """Call a tool, injecting headers from the header_provider if configured. + + When a ``header_provider`` was supplied at construction time, the runtime + *kwargs* (originating from ``FunctionInvocationContext.kwargs``) are passed + to the provider. The returned headers are attached to every HTTP request + made during this tool call via a ``contextvars.ContextVar``. + + Args: + tool_name: The name of the tool to call. + + Keyword Args: + kwargs: Arguments to pass to the tool. + + Returns: + A list of Content items representing the tool output. + """ + if self._header_provider is not None: + headers = self._header_provider(kwargs) + token = _mcp_call_headers.set(headers) + try: + return await super().call_tool(tool_name, **kwargs) + finally: + _mcp_call_headers.reset(token) + return await super().call_tool(tool_name, **kwargs) + class MCPWebsocketTool(MCPTool): """MCP tool for connecting to WebSocket-based MCP servers. @@ -1514,6 +1880,21 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: Returns: An async context manager for the WebSocket client transport. """ + try: + from mcp.client.websocket import websocket_client + except ModuleNotFoundError as ex: + missing_name = ex.name or "mcp/websocket dependencies" + if missing_name == "mcp" or missing_name.startswith("mcp."): + reason = "The `mcp` package is not installed." + elif missing_name == "websockets" or missing_name.startswith("websockets."): + reason = "WebSocket transport support is not installed." + else: + reason = f"The optional dependency `{missing_name}` is not installed." + raise ModuleNotFoundError( + f"`MCPWebsocketTool` requires websocket transport support. {reason} " + "Please install `mcp[ws]` and update your dependencies." + ) from ex + args: dict[str, Any] = { "url": self.url, } diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py new file mode 100644 index 00000000000..6e524d62e6d --- /dev/null +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -0,0 +1,1010 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import logging +import sys +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal + +from agent_framework import ( + ChatMiddlewareLayer, + ChatResponseUpdate, + Content, + FunctionInvocationConfiguration, + FunctionInvocationLayer, + load_settings, +) +from agent_framework._compaction import CompactionStrategy, TokenizerProtocol +from agent_framework._feature_stage import ExperimentalFeature, experimental +from agent_framework._telemetry import get_user_agent +from agent_framework.observability import ChatTelemetryLayer +from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient +from azure.ai.projects.aio import AIProjectClient +from azure.ai.projects.models import ( + A2APreviewTool, + AISearchIndexResource, + AutoCodeInterpreterToolParam, + AzureAISearchTool, + AzureAISearchToolResource, + BingCustomSearchConfiguration, + BingCustomSearchPreviewTool, + BingCustomSearchToolParameters, + BingGroundingSearchConfiguration, + BingGroundingSearchToolParameters, + BingGroundingTool, + BrowserAutomationPreviewTool, + BrowserAutomationToolConnectionParameters, + BrowserAutomationToolParameters, + CodeInterpreterTool, + ComputerUsePreviewTool, + FabricDataAgentToolParameters, + ImageGenTool, + MemorySearchPreviewTool, + MicrosoftFabricPreviewTool, + SharepointGroundingToolParameters, + SharepointPreviewTool, + ToolProjectConnection, + WebSearchApproximateLocation, + WebSearchTool, + WebSearchToolFilters, +) +from azure.ai.projects.models import FileSearchTool as ProjectsFileSearchTool +from azure.ai.projects.models import MCPTool as FoundryMCPTool +from azure.core.credentials import TokenCredential +from azure.core.credentials_async import AsyncTokenCredential + +from agent_framework_foundry._oauth_helpers import try_parse_oauth_consent_event + +from ._tools import _sanitize_foundry_response_tool # pyright: ignore[reportPrivateUsage] + +if sys.version_info >= (3, 13): + from typing import TypeVar # type: ignore # pragma: no cover +else: + from typing_extensions import TypeVar # type: ignore # pragma: no cover +if sys.version_info >= (3, 12): + from typing import override # type: ignore # pragma: no cover +else: + from typing_extensions import override # type: ignore # pragma: no cover +if sys.version_info >= (3, 11): + from typing import TypedDict # type: ignore # pragma: no cover +else: + from typing_extensions import TypedDict # type: ignore # pragma: no cover + +if TYPE_CHECKING: + from agent_framework import ChatAndFunctionMiddlewareTypes, ToolTypes + +logger: logging.Logger = logging.getLogger("agent_framework.foundry") + +AzureTokenProvider = Callable[[], str | Awaitable[str]] +AzureCredentialTypes = TokenCredential | AsyncTokenCredential + + +class FoundrySettings(TypedDict, total=False): + """Settings for Microsoft FoundryChatClient resolved from args and environment. + + Keyword Args: + model: The model deployment name. + Can be set via environment variable FOUNDRY_MODEL. + project_endpoint: The Microsoft Foundry project endpoint URL. + Can be set via environment variable FOUNDRY_PROJECT_ENDPOINT. + """ + + model: str | None + project_endpoint: str | None + + +def resolve_file_ids(file_ids: Sequence[str | Content] | None) -> list[str] | None: + """Resolve file IDs from strings or hosted-file Content objects.""" + if not file_ids: + return None + + resolved: list[str] = [] + for item in file_ids: + if isinstance(item, str): + if not item: + raise ValueError("file_ids must not contain empty strings.") + resolved.append(item) + elif isinstance(item, Content): + if item.type != "hosted_file": + raise ValueError( + f"Unsupported Content type {item.type!r} for code interpreter file_ids. " + "Only Content.from_hosted_file() is supported." + ) + if item.file_id is None: + raise ValueError( + "Content.from_hosted_file() item is missing a file_id. " + "Ensure the Content object has a valid file_id before using it in file_ids." + ) + resolved.append(item.file_id) + + return resolved if resolved else None + + +FoundryChatOptionsT = TypeVar( + "FoundryChatOptionsT", + bound=TypedDict, # type: ignore[valid-type] + default="OpenAIChatOptions", + covariant=True, +) + +FoundryChatOptions = OpenAIChatOptions + + +class RawFoundryChatClient( # type: ignore[misc] + RawOpenAIChatClient[FoundryChatOptionsT], + Generic[FoundryChatOptionsT], +): + """Raw Microsoft Foundry chat client using the OpenAI Responses API via a Foundry project. + + This client creates an OpenAI-compatible client from a Foundry project + and delegates to ``RawOpenAIChatClient`` for request handling. + + Environment variables: + - ``FOUNDRY_PROJECT_ENDPOINT`` to provide the Foundry project endpoint. + - ``FOUNDRY_MODEL`` to provide the Foundry model deployment name. + + Warning: + **This class should not normally be used directly.** Use ``FoundryChatClient`` + for a fully-featured client with middleware, telemetry, and function invocation. + """ + + OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" # type: ignore[reportIncompatibleVariableOverride, misc] + SUPPORTS_RICH_FUNCTION_OUTPUT: ClassVar[bool] = False # type: ignore[reportIncompatibleVariableOverride, misc] + + def __init__( + self, + *, + project_endpoint: str | None = None, + project_client: AIProjectClient | None = None, + model: str | None = None, + credential: AzureCredentialTypes | AzureTokenProvider | None = None, + allow_preview: bool | None = None, + default_headers: Mapping[str, str] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + instruction_role: str | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + additional_properties: dict[str, Any] | None = None, + ) -> None: + """Initialize a raw Microsoft Foundry chat client. + + Keyword Args: + project_endpoint: The Foundry project endpoint URL. + Can also be set via environment variable FOUNDRY_PROJECT_ENDPOINT. + project_client: An existing AIProjectClient to use. If provided, + the OpenAI client will be obtained via ``project_client.get_openai_client()``. + model: The model deployment name. + Can also be set via environment variable FOUNDRY_MODEL. + credential: Azure credential or token provider for authentication. + Required when using ``project_endpoint`` without a ``project_client``. + allow_preview: Enables preview opt-in on internally-created AIProjectClient. + default_headers: Additional HTTP headers for requests made through the OpenAI client. + env_file_path: Path to .env file for settings. + env_file_encoding: Encoding for .env file. + instruction_role: The role to use for 'instruction' messages. + compaction_strategy: Optional per-client compaction override. + tokenizer: Optional tokenizer for compaction strategies. + additional_properties: Additional properties stored on the client instance. + """ + foundry_settings = load_settings( + FoundrySettings, + env_prefix="FOUNDRY_", + model=model, + project_endpoint=project_endpoint, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + + resolved_model = foundry_settings.get("model") + if not resolved_model: + raise ValueError("Model is required. Set via 'model' parameter or 'FOUNDRY_MODEL' environment variable.") + + project_endpoint = foundry_settings.get("project_endpoint") + + if project_endpoint is None and project_client is None: + raise ValueError( + "Either 'project_endpoint' or 'project_client' is required. " + "Set project_endpoint via parameter or 'FOUNDRY_PROJECT_ENDPOINT' environment variable." + ) + if not project_client: + if not project_endpoint: + raise ValueError( + "Azure AI project endpoint is required. Set via 'project_endpoint' parameter " + "or 'FOUNDRY_PROJECT_ENDPOINT' environment variable," + "or pass in a AIProjectClient." + ) + if not credential: + raise ValueError("Azure credential is required when using project_endpoint without a project_client.") + project_client_kwargs: dict[str, Any] = { + "endpoint": project_endpoint, + "credential": credential, # type: ignore[arg-type] + "user_agent": get_user_agent(), + } + if allow_preview is not None: + project_client_kwargs["allow_preview"] = allow_preview + project_client = AIProjectClient(**project_client_kwargs) + + openai_kwargs: dict[str, Any] = {} + if default_headers: + openai_kwargs["default_headers"] = default_headers + + super().__init__( + model=resolved_model, + async_client=project_client.get_openai_client(**openai_kwargs), + default_headers=default_headers, + instruction_role=instruction_role, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + additional_properties=additional_properties, + ) + self.project_client = project_client + + @override + def _check_model_presence(self, options: dict[str, Any]) -> None: + if not options.get("model"): + if not self.model: + raise ValueError("model must be a non-empty string") + options["model"] = self.model + + @override + def _prepare_tools_for_openai( + self, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, + ) -> list[Any]: + response_tools = super()._prepare_tools_for_openai(tools) + return [_sanitize_foundry_response_tool(tool_item) for tool_item in response_tools] + + @override + def _parse_chunk_from_openai( + self, + event: Any, + options: dict[str, Any], + function_call_ids: dict[int, tuple[str, str]], + seen_reasoning_delta_item_ids: set[str] | None = None, + ) -> ChatResponseUpdate: + """Parse streaming event, intercepting oauth_consent_request items.""" + update = try_parse_oauth_consent_event(event, self.model) + if update is not None: + return update + return super()._parse_chunk_from_openai(event, options, function_call_ids, seen_reasoning_delta_item_ids) + + async def configure_azure_monitor( + self, + enable_sensitive_data: bool = False, + **kwargs: Any, + ) -> None: + """Setup observability with Azure Monitor (Microsoft Foundry integration). + + This method configures Azure Monitor for telemetry collection using the + connection string from the Foundry project client. + + Args: + enable_sensitive_data: Enable sensitive data logging (prompts, responses). + Should only be enabled in development/test environments. Default is False. + **kwargs: Additional arguments passed to configure_azure_monitor(). + Common options include: + - enable_live_metrics (bool): Enable Azure Monitor Live Metrics + - credential (TokenCredential): Azure credential for Entra ID auth + - resource (Resource): Custom OpenTelemetry resource + + Raises: + ImportError: If azure-monitor-opentelemetry-exporter is not installed. + """ + from agent_framework.observability import ( + OBSERVABILITY_SETTINGS, + create_metric_views, + create_resource, + enable_instrumentation, + ) + from azure.core.exceptions import ResourceNotFoundError + + if OBSERVABILITY_SETTINGS.is_user_disabled: + logger.info( + "FoundryChatClient.configure_azure_monitor(): Skipping setup because instrumentation was " + "explicitly disabled via disable_instrumentation(). Call enable_instrumentation(force=True) " + "to re-enable, then re-invoke configure_azure_monitor()." + ) + return + + try: + conn_string = await self.project_client.telemetry.get_application_insights_connection_string() + except ResourceNotFoundError: + logger.warning( + "No Application Insights connection string found for the Foundry project. " + "Please ensure Application Insights is configured in your project, " + "or call configure_otel_providers() manually with custom exporters." + ) + return + + try: + from azure.monitor.opentelemetry import configure_azure_monitor # type: ignore[import] + except ImportError as exc: + raise ImportError( + "azure-monitor-opentelemetry is required for Azure Monitor integration. " + "Install it with: pip install azure-monitor-opentelemetry" + ) from exc + + if "resource" not in kwargs: + kwargs["resource"] = create_resource() + + configure_azure_monitor( + connection_string=conn_string, + views=create_metric_views(), + **kwargs, + ) + + enable_instrumentation(enable_sensitive_data=enable_sensitive_data) + + # region Tool factory methods (override OpenAI defaults with Foundry versions) + + @staticmethod + def get_code_interpreter_tool( # type: ignore[override] + *, + file_ids: list[str | Content] | None = None, + container: Literal["auto"] | dict[str, Any] = "auto", + **kwargs: Any, + ) -> CodeInterpreterTool: + """Create a code interpreter tool configuration for Foundry. + + Keyword Args: + file_ids: Optional list of file IDs or Content objects to make available. + container: Container configuration. Use "auto" for automatic management. + **kwargs: Additional arguments passed to the SDK CodeInterpreterTool constructor. + + Returns: + A CodeInterpreterTool ready to pass to an Agent. + """ + if file_ids is None and isinstance(container, dict): + file_ids = container.get("file_ids") + resolved = resolve_file_ids(file_ids) + tool_container = AutoCodeInterpreterToolParam(file_ids=resolved) + return CodeInterpreterTool(container=tool_container, **kwargs) + + @staticmethod + def get_file_search_tool( + *, + vector_store_ids: list[str], + max_num_results: int | None = None, + ranking_options: dict[str, Any] | None = None, + filters: dict[str, Any] | None = None, + **kwargs: Any, + ) -> ProjectsFileSearchTool: + """Create a file search tool configuration for Foundry. + + Keyword Args: + vector_store_ids: List of vector store IDs to search. + max_num_results: Maximum number of results to return (1-50). + ranking_options: Ranking options for search results. + filters: A filter to apply (ComparisonFilter or CompoundFilter). + **kwargs: Additional arguments passed to the SDK FileSearchTool constructor. + + Returns: + A FileSearchTool ready to pass to an Agent. + """ + if not vector_store_ids: + raise ValueError("File search tool requires 'vector_store_ids' to be specified.") + return ProjectsFileSearchTool( + vector_store_ids=vector_store_ids, + max_num_results=max_num_results, + ranking_options=ranking_options, # type: ignore[arg-type] + filters=filters, # type: ignore[arg-type] + **kwargs, + ) + + @staticmethod + def get_web_search_tool( # type: ignore[override] + *, + user_location: dict[str, str] | None = None, + search_context_size: Literal["low", "medium", "high"] | None = None, + allowed_domains: list[str] | None = None, + custom_search_configuration: dict[str, Any] | None = None, + **kwargs: Any, + ) -> WebSearchTool: + """Create a Web Search tool configuration for Microsoft Foundry. + + **Choosing a web grounding tool.** Foundry exposes three options that all reach + the public web via Bing. Pick the one that matches your scenario: + + * :py:meth:`get_web_search_tool` (this one, GA) — recommended starting point. + The Bing resource is managed by Microsoft, no extra Azure setup is required, + and only Azure OpenAI models are supported. Parameters are limited to + ``user_location`` and ``search_context_size``. + * :py:meth:`get_bing_grounding_tool` (preview) — use when you need finer Bing parameters (``count``, + ``freshness``, ``market``, ``set_lang``), want to ground non-OpenAI + Foundry models, or are migrating from Grounding with Bing Search on the + classic agents platform. You manage the Grounding with Bing Search + resource yourself (Contributor/Owner to create the resource, Foundry + Project Manager to wire the connection). + * :py:meth:`get_bing_custom_search_tool` (preview) — use when you need to + restrict grounding to a curated set of domains defined in a Bing Custom + Search instance. + + For all three, search data flows outside the Azure compliance boundary. See + https://learn.microsoft.com/azure/foundry/agents/how-to/tools/web-overview for + the full comparison. + + Keyword Args: + user_location: Location context with keys like ``"city"``, ``"country"``, + ``"region"``, ``"timezone"``. + search_context_size: Amount of context from search results + (``"low"``, ``"medium"``, ``"high"``). + allowed_domains: List of domains to restrict search results to. Wrapped + into ``WebSearchToolFilters`` and passed as the ``filters`` field on + the SDK ``WebSearchTool``. + custom_search_configuration: Custom Bing search configuration for + domain-restricted scenarios. + **kwargs: Additional arguments passed to the SDK ``WebSearchTool`` + constructor. + + Returns: + A ``WebSearchTool`` ready to pass to an Agent. + """ + ws_kwargs: dict[str, Any] = {**kwargs} + if search_context_size: + ws_kwargs["search_context_size"] = search_context_size + if allowed_domains: + ws_kwargs["filters"] = WebSearchToolFilters(allowed_domains=allowed_domains) + if custom_search_configuration: + ws_kwargs["custom_search_configuration"] = custom_search_configuration + if user_location: + ws_kwargs["user_location"] = WebSearchApproximateLocation( + city=user_location.get("city"), + country=user_location.get("country"), + region=user_location.get("region"), + timezone=user_location.get("timezone"), + ) + return WebSearchTool(**ws_kwargs) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_TOOLS) + def get_bing_grounding_tool( + *, + connection_id: str, + market: str | None = None, + set_lang: str | None = None, + count: int | None = None, + freshness: str | None = None, + **kwargs: Any, + ) -> BingGroundingTool: + """Create a Grounding with Bing Search tool configuration for Foundry. + + Use this factory when :py:meth:`get_web_search_tool` is too restrictive — for + example when you need ``count``/``freshness``/``market``/``set_lang`` + parameters, want to ground a non-OpenAI Foundry model, or are migrating an + agent that already uses Grounding with Bing Search on the classic agents + platform. You manage the Grounding with Bing Search Azure resource yourself + (Contributor or Owner to create the resource, Foundry Project Manager to + create the project connection). Search data flows outside the Azure + compliance boundary. + + For domain-restricted grounding to a curated allow-list, use + :py:meth:`get_bing_custom_search_tool` instead. For a zero-setup default that + works for most agents, see :py:meth:`get_web_search_tool`. The full + comparison lives at + https://learn.microsoft.com/azure/foundry/agents/how-to/tools/web-overview. + + Keyword Args: + connection_id: The Foundry project connection ID for the Grounding with + Bing Search resource. + market: Optional Bing market identifier (e.g. ``"en-US"``). + set_lang: Optional UI language code passed to the Bing API. + count: Optional number of search results to return. + freshness: Optional time-range filter for search results. See + https://learn.microsoft.com/bing/search-apis/bing-web-search/reference/query-parameters + for accepted values. + **kwargs: Additional arguments forwarded to the SDK + ``BingGroundingSearchConfiguration``. + + Returns: + A ``BingGroundingTool`` ready to pass to an Agent. + """ + config_kwargs: dict[str, Any] = { + **kwargs, + "project_connection_id": connection_id, + } + if market is not None: + config_kwargs["market"] = market + if set_lang is not None: + config_kwargs["set_lang"] = set_lang + if count is not None: + config_kwargs["count"] = count + if freshness is not None: + config_kwargs["freshness"] = freshness + return BingGroundingTool( + bing_grounding=BingGroundingSearchToolParameters( + search_configurations=[BingGroundingSearchConfiguration(**config_kwargs)], + ), + ) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_PREVIEW_TOOLS) + def get_bing_custom_search_tool( + *, + connection_id: str, + instance_name: str, + market: str | None = None, + set_lang: str | None = None, + count: int | None = None, + freshness: str | None = None, + **kwargs: Any, + ) -> BingCustomSearchPreviewTool: + """Create a Grounding with Bing Custom Search tool configuration for Foundry. + + Use this factory (preview) when you need to restrict grounding to a curated + list of domains. The allow/block list is defined ahead of time on a Bing + Custom Search resource (in the Bing portal) and referenced here by + ``instance_name``. Like the other Bing-backed tools, search data flows + outside the Azure compliance boundary, and you must create the Bing Custom + Search resource yourself. + + For unrestricted public-web grounding with no extra Azure setup, prefer + :py:meth:`get_web_search_tool`. For unrestricted grounding with finer Bing + parameters or non-OpenAI models, prefer :py:meth:`get_bing_grounding_tool`. + See + https://learn.microsoft.com/azure/foundry/agents/how-to/tools/web-overview + for the full comparison. + + Keyword Args: + connection_id: The Foundry project connection ID for the Grounding with + Bing Custom Search resource. + instance_name: The custom configuration instance name defined on the + Bing Custom Search resource. + market: Optional Bing market identifier (e.g. ``"en-US"``). + set_lang: Optional UI language code passed to the Bing API. + count: Optional number of search results to return. + freshness: Optional time-range filter for search results. + **kwargs: Additional arguments forwarded to the SDK + ``BingCustomSearchConfiguration``. + + Returns: + A ``BingCustomSearchPreviewTool`` ready to pass to an Agent. + """ + config_kwargs: dict[str, Any] = { + **kwargs, + "project_connection_id": connection_id, + "instance_name": instance_name, + } + if market is not None: + config_kwargs["market"] = market + if set_lang is not None: + config_kwargs["set_lang"] = set_lang + if count is not None: + config_kwargs["count"] = count + if freshness is not None: + config_kwargs["freshness"] = freshness + return BingCustomSearchPreviewTool( + bing_custom_search_preview=BingCustomSearchToolParameters( + search_configurations=[BingCustomSearchConfiguration(**config_kwargs)], + ), + ) + + @staticmethod + def get_image_generation_tool( # type: ignore[override] + *, + model: Literal["gpt-image-1"] | str | None = None, + size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] | None = None, + output_format: Literal["png", "webp", "jpeg"] | None = None, + quality: Literal["low", "medium", "high", "auto"] | None = None, + background: Literal["transparent", "opaque", "auto"] | None = None, + partial_images: int | None = None, + moderation: Literal["auto", "low"] | None = None, + output_compression: int | None = None, + **kwargs: Any, + ) -> ImageGenTool: + """Create an image generation tool configuration for Foundry. + + Keyword Args: + model: The model to use for image generation. + size: Output image size. + output_format: Output image format. + quality: Output image quality. + background: Background transparency setting. + partial_images: Number of partial images to return during generation. + moderation: Moderation level. + output_compression: Compression level. + **kwargs: Additional arguments passed to the SDK ImageGenTool constructor. + + Returns: + An ImageGenTool ready to pass to an Agent. + """ + return ImageGenTool( # type: ignore[misc] + model=model, # type: ignore[arg-type] + size=size, + output_format=output_format, + quality=quality, + background=background, + partial_images=partial_images, + moderation=moderation, + output_compression=output_compression, + **kwargs, + ) + + @staticmethod + def get_mcp_tool( + *, + name: str, + url: str | None = None, + description: str | None = None, + approval_mode: Literal["always_require", "never_require"] | dict[str, list[str]] | None = None, + allowed_tools: list[str] | None = None, + headers: dict[str, str] | None = None, + project_connection_id: str | None = None, + **kwargs: Any, + ) -> FoundryMCPTool: + """Create a hosted MCP tool configuration for Foundry. + + This configures an MCP server that runs remotely on Azure AI, not locally. + + Keyword Args: + name: A label/name for the MCP server. + url: The URL of the MCP server. Required if project_connection_id is not provided. + description: A description of what the MCP server provides. + approval_mode: Tool approval mode ("always_require", "never_require", or dict). + allowed_tools: List of allowed tool names from this MCP server. + headers: HTTP headers to include in requests to the MCP server. + project_connection_id: Foundry connection ID for managed MCP connections. + **kwargs: Additional arguments passed to the SDK MCPTool constructor. + + Returns: + An MCPTool configuration ready to pass to an Agent. + + Raises: + ValueError: If neither ``url`` nor ``project_connection_id`` is supplied + — one is required by the Foundry Responses API. + """ + if not url and not project_connection_id: + raise ValueError("MCP tool requires either 'url' or 'project_connection_id' to be specified.") + + mcp_kwargs: dict[str, Any] = {"server_label": name.replace(" ", "_"), **kwargs} + if url: + mcp_kwargs["server_url"] = url + mcp = FoundryMCPTool(**mcp_kwargs) + + if description: + mcp["server_description"] = description + if project_connection_id: + mcp["project_connection_id"] = project_connection_id + elif headers: + mcp["headers"] = headers + if allowed_tools is not None: + mcp["allowed_tools"] = list(allowed_tools) + if approval_mode: + if isinstance(approval_mode, str): + mcp["require_approval"] = "always" if approval_mode == "always_require" else "never" + else: + if always_require := approval_mode.get("always_require_approval"): + mcp["require_approval"] = {"always": {"tool_names": always_require}} + if never_require := approval_mode.get("never_require_approval"): + mcp["require_approval"] = {"never": {"tool_names": never_require}} + + return mcp + + # endregion + + # region Experimental Foundry tool factories (preview SDK types) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_TOOLS) + def get_azure_ai_search_tool( + *, + index_connection_id: str, + index_name: str, + query_type: str | None = None, + top_k: int | None = None, + filter: str | None = None, + index_asset_id: str | None = None, + **kwargs: Any, + ) -> AzureAISearchTool: + """Create an Azure AI Search tool configuration for Foundry. + + Keyword Args: + index_connection_id: The Foundry project connection ID for the Azure AI Search index. + index_name: The name of the index to search. + query_type: Optional query type (``"simple"``, ``"semantic"``, ``"vector"``, + ``"vector_simple_hybrid"``, or ``"vector_semantic_hybrid"``). + top_k: Optional number of documents to retrieve. + filter: Optional OData filter expression. + index_asset_id: Optional index asset id for the search resource. + **kwargs: Additional arguments forwarded to the SDK ``AISearchIndexResource``. + + Returns: + An ``AzureAISearchTool`` ready to pass to an Agent. + """ + index_kwargs: dict[str, Any] = { + **kwargs, + "project_connection_id": index_connection_id, + "index_name": index_name, + } + if query_type is not None: + index_kwargs["query_type"] = query_type + if top_k is not None: + index_kwargs["top_k"] = top_k + if filter is not None: + index_kwargs["filter"] = filter + if index_asset_id is not None: + index_kwargs["index_asset_id"] = index_asset_id + return AzureAISearchTool( + azure_ai_search=AzureAISearchToolResource(indexes=[AISearchIndexResource(**index_kwargs)]), + ) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_PREVIEW_TOOLS) + def get_sharepoint_tool( + *, + connection_id: str, + **kwargs: Any, + ) -> SharepointPreviewTool: + """Create a SharePoint grounding tool configuration for Foundry. + + Keyword Args: + connection_id: The Foundry project connection ID for the SharePoint resource. + **kwargs: Additional arguments forwarded to the SDK + ``SharepointGroundingToolParameters``. + + Returns: + A ``SharepointPreviewTool`` ready to pass to an Agent. + """ + return SharepointPreviewTool( + sharepoint_grounding_preview=SharepointGroundingToolParameters( + project_connections=[ToolProjectConnection(project_connection_id=connection_id)], + **kwargs, + ) + ) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_PREVIEW_TOOLS) + def get_fabric_tool( + *, + connection_id: str, + **kwargs: Any, + ) -> MicrosoftFabricPreviewTool: + """Create a Microsoft Fabric data agent tool configuration for Foundry. + + Keyword Args: + connection_id: The Foundry project connection ID for the Fabric data agent. + **kwargs: Additional arguments forwarded to the SDK + ``FabricDataAgentToolParameters``. + + Returns: + A ``MicrosoftFabricPreviewTool`` ready to pass to an Agent. + """ + return MicrosoftFabricPreviewTool( + fabric_dataagent_preview=FabricDataAgentToolParameters( + project_connections=[ToolProjectConnection(project_connection_id=connection_id)], + **kwargs, + ) + ) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_PREVIEW_TOOLS) + def get_memory_search_tool( + *, + memory_store_name: str, + scope: str, + search_options: Any | None = None, + update_delay: int | None = None, + **kwargs: Any, + ) -> MemorySearchPreviewTool: + """Create a Memory Search tool configuration for Foundry. + + Keyword Args: + memory_store_name: The name of the memory store to use. + scope: The namespace used to group and isolate memories (e.g. a user ID). + Use ``"{{$userId}}"`` to scope memories to the current signed-in user. + search_options: Optional ``MemorySearchOptions`` instance. + update_delay: Optional seconds to wait before updating memories after inactivity. + **kwargs: Additional arguments forwarded to the SDK ``MemorySearchPreviewTool``. + + Returns: + A ``MemorySearchPreviewTool`` ready to pass to an Agent. + """ + params: dict[str, Any] = { + **kwargs, + "memory_store_name": memory_store_name, + "scope": scope, + } + if search_options is not None: + params["search_options"] = search_options + if update_delay is not None: + params["update_delay"] = update_delay + return MemorySearchPreviewTool(**params) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_PREVIEW_TOOLS) + def get_computer_use_tool( + *, + environment: str, + display_width: int, + display_height: int, + **kwargs: Any, + ) -> ComputerUsePreviewTool: + """Create a Computer Use tool configuration for Foundry. + + Keyword Args: + environment: The computer environment to control. One of ``"windows"``, + ``"mac"``, ``"linux"``, ``"ubuntu"``, or ``"browser"``. + display_width: The width of the computer display. + display_height: The height of the computer display. + **kwargs: Additional arguments forwarded to the SDK ``ComputerUsePreviewTool``. + + Returns: + A ``ComputerUsePreviewTool`` ready to pass to an Agent. + """ + return ComputerUsePreviewTool( + environment=environment, + display_width=display_width, + display_height=display_height, + **kwargs, + ) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_PREVIEW_TOOLS) + def get_browser_automation_tool( + *, + connection_id: str, + **kwargs: Any, + ) -> BrowserAutomationPreviewTool: + """Create a Browser Automation tool configuration for Foundry. + + Keyword Args: + connection_id: The Foundry project connection ID for the Azure Playwright resource. + **kwargs: Additional arguments forwarded to the SDK + ``BrowserAutomationToolParameters``. + + Returns: + A ``BrowserAutomationPreviewTool`` ready to pass to an Agent. + """ + return BrowserAutomationPreviewTool( + browser_automation_preview=BrowserAutomationToolParameters( + connection=BrowserAutomationToolConnectionParameters(project_connection_id=connection_id), + **kwargs, + ) + ) + + @staticmethod + @experimental(feature_id=ExperimentalFeature.FOUNDRY_PREVIEW_TOOLS) + def get_a2a_tool( + *, + base_url: str | None = None, + agent_card_path: str | None = None, + project_connection_id: str | None = None, + **kwargs: Any, + ) -> A2APreviewTool: + """Create an Agent-to-Agent (A2A) tool configuration for Foundry. + + Keyword Args: + base_url: Base URL of the remote A2A agent. + agent_card_path: Path to the agent card relative to ``base_url``. + Defaults to ``"/.well-known/agent-card.json"`` server-side. + project_connection_id: Foundry connection ID for the A2A server. Stores + authentication and other connection details. + **kwargs: Additional arguments forwarded to the SDK ``A2APreviewTool``. + + Returns: + An ``A2APreviewTool`` ready to pass to an Agent. + """ + params: dict[str, Any] = dict(kwargs) + if base_url is not None: + params["base_url"] = base_url + if agent_card_path is not None: + params["agent_card_path"] = agent_card_path + if project_connection_id is not None: + params["project_connection_id"] = project_connection_id + return A2APreviewTool(**params) + + # endregion + + +class FoundryChatClient( # type: ignore[misc] + FunctionInvocationLayer[FoundryChatOptionsT], + ChatMiddlewareLayer[FoundryChatOptionsT], + ChatTelemetryLayer[FoundryChatOptionsT], + RawFoundryChatClient[FoundryChatOptionsT], + Generic[FoundryChatOptionsT], +): + """Microsoft Foundry chat client using the OpenAI Responses API. + + Creates an OpenAI-compatible client from a Foundry project + with middleware, telemetry, and function invocation support. + + Environment variables: + - ``FOUNDRY_PROJECT_ENDPOINT`` to provide the Foundry project endpoint. + - ``FOUNDRY_MODEL`` to provide the Foundry model deployment name. + + Keyword Args: + project_endpoint: The Foundry project endpoint URL. + Can also be set via environment variable ``FOUNDRY_PROJECT_ENDPOINT``. + project_client: An existing AIProjectClient to use. + model: The model deployment name. + Can also be set via environment variable ``FOUNDRY_MODEL``. + credential: Azure credential or token provider for authentication. + allow_preview: Enables preview opt-in on internally-created AIProjectClient. + env_file_path: Path to .env file for settings. + env_file_encoding: Encoding for .env file. + instruction_role: The role to use for 'instruction' messages. + middleware: Optional sequence of middleware. + function_invocation_configuration: Optional function invocation configuration. + + Examples: + .. code-block:: python + + from azure.identity import AzureCliCredential + from agent_framework_foundry import FoundryChatClient + + client = FoundryChatClient( + project_endpoint="https://your-project.services.ai.azure.com", + model="gpt-4o", + credential=AzureCliCredential(), + ) + + # Or using an existing AIProjectClient + from azure.ai.projects.aio import AIProjectClient + + project_client = AIProjectClient( + endpoint="https://your-project.services.ai.azure.com", + credential=AzureCliCredential(), + ) + client = FoundryChatClient( + project_client=project_client, + model="gpt-4o", + ) + """ + + OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" # type: ignore[reportIncompatibleVariableOverride, misc] + + def __init__( + self, + *, + project_endpoint: str | None = None, + project_client: AIProjectClient | None = None, + model: str | None = None, + credential: AzureCredentialTypes | AzureTokenProvider | None = None, + allow_preview: bool | None = None, + default_headers: Mapping[str, str] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + instruction_role: str | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + additional_properties: dict[str, Any] | None = None, + middleware: (Sequence[ChatAndFunctionMiddlewareTypes] | None) = None, + function_invocation_configuration: FunctionInvocationConfiguration | None = None, + ) -> None: + """Initialize a Foundry chat client. + + Keyword Args: + project_endpoint: The Foundry project endpoint URL. + Can also be set via environment variable ``FOUNDRY_PROJECT_ENDPOINT``. + project_client: An existing AIProjectClient to use. + model: The model deployment name. + Can also be set via environment variable ``FOUNDRY_MODEL``. + credential: Azure credential or token provider for authentication. + allow_preview: Enables preview opt-in on internally-created AIProjectClient. + default_headers: Additional HTTP headers for requests made through the OpenAI client. + env_file_path: Path to .env file for settings. + env_file_encoding: Encoding for .env file. + instruction_role: The role to use for 'instruction' messages. + compaction_strategy: Optional per-client compaction override. + tokenizer: Optional tokenizer for compaction strategies. + additional_properties: Additional properties stored on the client instance. + middleware: Optional sequence of middleware. + function_invocation_configuration: Optional function invocation configuration. + """ + super().__init__( + project_endpoint=project_endpoint, + project_client=project_client, + model=model, + credential=credential, + allow_preview=allow_preview, + default_headers=default_headers, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + instruction_role=instruction_role, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + additional_properties=additional_properties, + middleware=middleware, + function_invocation_configuration=function_invocation_configuration, + ) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py new file mode 100644 index 00000000000..988da69dcf6 --- /dev/null +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -0,0 +1,3203 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import json +import logging +import shlex +import sys +from collections.abc import ( + AsyncIterable, + Awaitable, + Callable, + Mapping, + MutableMapping, + Sequence, +) +from datetime import datetime, timezone +from itertools import chain +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Generic, + Literal, + NoReturn, + TypedDict, + cast, + overload, +) + +from agent_framework._clients import BaseChatClient +from agent_framework._compaction import CompactionStrategy, TokenizerProtocol +from agent_framework._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer +from agent_framework._settings import SecretString +from agent_framework._telemetry import USER_AGENT_KEY +from agent_framework._tools import ( + SHELL_TOOL_KIND_VALUE, + FunctionInvocationConfiguration, + FunctionInvocationLayer, + FunctionTool, + ToolTypes, + normalize_tools, + tool, +) +from agent_framework._types import ( + Annotation, + ChatOptions, + ChatResponse, + ChatResponseUpdate, + Content, + ContinuationToken, + Message, + ResponseStream, + Role, + TextSpanRegion, + UsageDetails, + detect_media_type_from_base64, + prepend_instructions_to_messages, + validate_tool_mode, +) +from agent_framework.exceptions import ( + ChatClientException, + ChatClientInvalidRequestException, +) +from agent_framework.observability import ChatTelemetryLayer +from openai import AsyncAzureOpenAI, AsyncOpenAI, BadRequestError +from openai.types.responses import FunctionShellTool +from openai.types.responses.file_search_tool_param import FileSearchToolParam +from openai.types.responses.function_tool_param import FunctionToolParam +from openai.types.responses.parsed_response import ( + ParsedResponse, +) +from openai.types.responses.response import Response as OpenAIResponse +from openai.types.responses.response_stream_event import ( + ResponseStreamEvent as OpenAIResponseStreamEvent, +) +from openai.types.responses.response_usage import ResponseUsage +from openai.types.responses.tool_param import ( + CodeInterpreter, + CodeInterpreterContainerCodeInterpreterToolAuto, + ImageGeneration, + Mcp, +) +from openai.types.responses.web_search_tool_param import WebSearchToolParam +from pydantic import BaseModel + +from ._exceptions import OpenAIContentFilterException +from ._shared import ( + AzureTokenProvider, + load_openai_service_settings, + maybe_append_azure_endpoint_guidance, +) + +if sys.version_info >= (3, 13): + from typing import TypeVar # type: ignore # pragma: no cover +else: + from typing_extensions import TypeVar # type: ignore # pragma: no cover +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 +if sys.version_info >= (3, 11): + from typing import TypedDict # type: ignore # pragma: no cover +else: + from typing_extensions import TypedDict # type: ignore # pragma: no cover + +if TYPE_CHECKING: + from azure.core.credentials import TokenCredential + from azure.core.credentials_async import AsyncTokenCredential + + AzureCredentialTypes = TokenCredential | AsyncTokenCredential + +logger = logging.getLogger("agent_framework.openai") + +DEFAULT_AZURE_OPENAI_RESPONSES_API_VERSION = "preview" + +OPENAI_SHELL_ENVIRONMENT_KEY = "openai.responses.shell.environment" +OPENAI_SHELL_OUTPUT_TYPE_KEY = "openai.responses.shell.output_type" +OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY = "openai.responses.local_shell.call_item_id" +OPENAI_LOCAL_SHELL_COMMAND_PARTS_KEY = "openai.local_shell_command_parts" +OPENAI_SHELL_OUTPUT_TYPE_SHELL_CALL = "shell_call_output" +OPENAI_SHELL_OUTPUT_TYPE_LOCAL_SHELL_CALL = "local_shell_call_output" + +# Internal marker emitted by `_prepare_content_for_openai` for an +# `mcp_server_tool_result` Content. The Responses API expects an `mcp_call` +# input item to carry both arguments and output as one item, so result +# Contents cannot be serialized standalone. `_prepare_messages_for_openai` +# coalesces these markers into the most recent matching `mcp_call` input +# item before returning, dropping any that are unmatched. +_AF_MCP_PENDING_OUTPUT_KEY = "__af_pending_mcp_result__" + + +class OpenAIContinuationToken(ContinuationToken): + """Continuation token for OpenAI Responses API background operations.""" + + response_id: str + """OpenAI Responses API response ID.""" + + +# region OpenAI Responses Options TypedDict + + +class ReasoningOptions(TypedDict, total=False): + """Configuration options for reasoning models (gpt-5, o-series). + + See: https://platform.openai.com/docs/guides/reasoning + """ + + effort: Literal["none", "low", "medium", "high", "xhigh"] + """The effort level for reasoning. Higher effort means more reasoning tokens.""" + + summary: Literal["auto", "concise", "detailed"] + """How to summarize reasoning in the response.""" + + +class StreamOptions(TypedDict, total=False): + """Options for streaming responses.""" + + include_usage: bool + """Whether to include usage statistics in stream events.""" + + +ResponseFormatT = TypeVar("ResponseFormatT", bound=BaseModel | None, default=None) + + +class OpenAIChatOptions(ChatOptions[ResponseFormatT], Generic[ResponseFormatT], total=False): + """OpenAI Responses API-specific chat options. + + Extends ChatOptions with options specific to OpenAI's Responses API. + These options provide fine-grained control over response generation, + reasoning, and API behavior. + + See: https://platform.openai.com/docs/api-reference/responses/create + """ + + # Responses API-specific parameters + + include: list[str] + """Additional output data to include in the response. + Supported values include: + - 'web_search_call.action.sources' + - 'code_interpreter_call.outputs' + - 'file_search_call.results' + - 'message.input_image.image_url' + - 'message.output_text.logprobs' + - 'reasoning.encrypted_content' + """ + + max_tool_calls: int + """Maximum number of total calls to built-in tools in a response.""" + + prompt: dict[str, Any] + """Reference to a prompt template and its variables. + Learn more: https://platform.openai.com/docs/guides/text#reusable-prompts""" + + prompt_cache_key: str + """Used by OpenAI to cache responses for similar requests. + Replaces the deprecated 'user' field for caching purposes.""" + + prompt_cache_retention: Literal["24h"] + """Retention policy for prompt cache. Set to '24h' for extended caching.""" + + reasoning: ReasoningOptions + """Configuration for reasoning models (gpt-5, o-series). + See: https://platform.openai.com/docs/guides/reasoning""" + + verbosity: Literal["low", "medium", "high"] + """Output verbosity for GPT-5 family models. Lower values yield shorter responses. + Translated to ``text.verbosity`` when sent to the Responses API. + See: https://developers.openai.com/cookbook/examples/gpt-5/gpt-5_new_params_and_tools#1-verbosity-parameter""" + + safety_identifier: str + """A stable identifier for detecting policy violations. + Recommend hashing username/email to avoid sending identifying info.""" + + service_tier: Literal["auto", "default", "flex", "priority"] + """Processing type for serving the request. + - 'auto': Use project settings + - 'default': Standard pricing/performance + - 'flex': Flexible processing + - 'priority': Priority processing""" + + stream_options: StreamOptions + """Options for streaming responses. Only set when stream=True.""" + + top_logprobs: int + """Number of most likely tokens (0-20) to return at each position.""" + + truncation: Literal["auto", "disabled"] + """Truncation strategy for model response. + - 'auto': Truncate from beginning if exceeds context + - 'disabled': Fail with 400 error if exceeds context""" + + background: bool + """Whether to run the model response in the background. + When True, the response returns immediately with a continuation token + that can be used to poll for the result. + See: https://platform.openai.com/docs/guides/background""" + + continuation_token: OpenAIContinuationToken + """Token for resuming or polling a long-running background operation. + Pass the ``continuation_token`` from a previous response to poll for + completion or resume a streaming response.""" + + +OpenAIChatOptionsT = TypeVar( + "OpenAIChatOptionsT", + bound=TypedDict, # type: ignore[valid-type] + default="OpenAIChatOptions", + covariant=True, +) + + +# endregion + + +# region Helpers + + +def _annotations_to_output_text(annotations: Sequence[Annotation] | None) -> list[dict[str, Any]]: + """Convert framework `Annotation` objects to Responses API `output_text` annotation dicts. + + Citations from `file_search`, `code_interpreter` file paths, and url citations all collapse + to `Annotation(type="citation", ...)` in the framework. The original API form is recovered + here so assistant messages roundtrip cleanly through history forwarding. + + Each Responses API annotation dict carries at most one `start_index`/`end_index` pair, so an + `Annotation` with multiple `annotated_regions` is fanned out into one entry per region. + Regions missing valid integer span bounds are skipped. + """ + if not annotations: + return [] + out: list[dict[str, Any]] = [] + for annotation in annotations: + if annotation.get("type") != "citation": + continue + props = annotation.get("additional_properties") or {} + regions = annotation.get("annotated_regions") or [] + file_id = annotation.get("file_id") + url = annotation.get("url") + title = annotation.get("title") + container_id = props.get("container_id") + + if container_id and file_id: + for region in regions: + start = region.get("start_index") + end = region.get("end_index") + if not (isinstance(start, int) and isinstance(end, int)): + continue + entry: dict[str, Any] = { + "type": "container_file_citation", + "container_id": container_id, + "file_id": file_id, + "start_index": start, + "end_index": end, + } + if url: + entry["filename"] = url + out.append(entry) + elif url and not file_id and regions: + for region in regions: + start = region.get("start_index") + end = region.get("end_index") + if not (isinstance(start, int) and isinstance(end, int)): + continue + out.append({ + "type": "url_citation", + "url": url, + "title": title or "", + "start_index": start, + "end_index": end, + }) + elif file_id and url: + entry = { + "type": "file_citation", + "file_id": file_id, + "filename": url, + } + if (idx := props.get("index")) is not None: + entry["index"] = idx + out.append(entry) + elif file_id: + entry = { + "type": "file_path", + "file_id": file_id, + } + if (idx := props.get("index")) is not None: + entry["index"] = idx + out.append(entry) + return out + + +# endregion + + +# region ResponsesClient + + +class RawOpenAIChatClient( # type: ignore[misc] + BaseChatClient[OpenAIChatOptionsT], + Generic[OpenAIChatOptionsT], +): + """Raw OpenAI Responses client without middleware, telemetry, or function invocation. + + Warning: + **This class should not normally be used directly.** It does not include middleware, + telemetry, or function invocation support that you most likely need. If you do use it, + you should consider which additional layers to apply. There is a defined ordering that + you should follow: + + 1. **FunctionInvocationLayer** - Owns the tool/function calling loop and routes function middleware + 2. **ChatMiddlewareLayer** - Applies chat middleware per model call and stays outside telemetry + 3. **ChatTelemetryLayer** - Must stay inside chat middleware for correct per-call telemetry + + Use ``OpenAIChatClient`` instead for a fully-featured client with all layers applied. + """ + + INJECTABLE: ClassVar[set[str]] = {"client"} + STORES_BY_DEFAULT: ClassVar[bool] = True # type: ignore[reportIncompatibleVariableOverride, misc] + SUPPORTS_RICH_FUNCTION_OUTPUT: ClassVar[bool] = True + + # Azure OpenAI Responses API may include this header in responses naming the actual model that + # served the request (e.g. ``gpt-5-nano-2025-08-07``), which can differ from the deployment alias + # that the request was addressed to and that ``response.model`` reports. When present, we use it + # as the value of ``ChatResponse.model`` / ``ChatResponseUpdate.model`` so telemetry and callers + # see the actually served model. (Chat Completions API already returns the snapshot in + # ``response.model``, so this header only matters for the Responses API.) + SERVED_MODEL_HEADER: ClassVar[str] = "x-ms-served-model" + + FILE_SEARCH_MAX_RESULTS: int = 50 + + @overload + def __init__( + self, + model: str | None = None, + *, + api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None, + org_id: str | None = None, + base_url: str | None = None, + default_headers: Mapping[str, str] | None = None, + async_client: AsyncOpenAI | None = None, + instruction_role: str | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + additional_properties: dict[str, Any] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize a raw OpenAI Chat client. + + Keyword Args: + model: Model identifier to use for the request. When not provided, the constructor + reads ``OPENAI_CHAT_MODEL`` and then ``OPENAI_MODEL``. + api_key: API key. When not provided explicitly, the constructor reads + ``OPENAI_API_KEY``. A callable API key is also supported. + org_id: OpenAI organization ID. When not provided explicitly, the constructor reads + ``OPENAI_ORG_ID``. + base_url: Base URL override. When not provided explicitly, the constructor reads + ``OPENAI_BASE_URL``. + default_headers: Additional HTTP headers. + async_client: Pre-configured OpenAI client. + instruction_role: Role for instruction messages (for example ``"system"``). + compaction_strategy: Optional per-client compaction override. + tokenizer: Optional tokenizer for compaction strategies. + additional_properties: Additional properties stored on the client instance. + env_file_path: Optional ``.env`` file that is checked before the process environment + for ``OPENAI_*`` values. + env_file_encoding: Encoding for the ``.env`` file. + """ + ... + + @overload + def __init__( + self, + model: str | None = None, + *, + azure_endpoint: str, + credential: AzureCredentialTypes | AzureTokenProvider | None = None, + api_version: str | None = None, + api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None, + base_url: str | None = None, + default_headers: Mapping[str, str] | None = None, + async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None, + instruction_role: str | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + additional_properties: dict[str, Any] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize a raw OpenAI Chat client. + + Keyword Args: + model: Model identifier to use for the request. When not provided, the constructor + reads ``AZURE_OPENAI_CHAT_MODEL`` and then + ``AZURE_OPENAI_MODEL``. + azure_endpoint: Azure resource endpoint. When not provided explicitly, the constructor + reads ``AZURE_OPENAI_ENDPOINT``. + credential: Azure credential or token provider for Entra auth. + api_version: Azure API version. When not provided explicitly, the constructor reads + ``AZURE_OPENAI_API_VERSION`` and then uses the Responses default. + api_key: API key. For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key + auth. A callable token provider is also accepted, + but ``credential`` is the preferred Azure auth surface. + base_url: Base URL override. When not provided explicitly, the constructor reads + ``AZURE_OPENAI_BASE_URL``. Use this instead of ``azure_endpoint`` when you want + to pass the full ``.../openai/v1`` base URL directly. + default_headers: Additional HTTP headers. + async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on + Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup. + instruction_role: Role for instruction messages (for example ``"system"``). + compaction_strategy: Optional per-client compaction override. + tokenizer: Optional tokenizer for compaction strategies. + additional_properties: Additional properties stored on the client instance. + env_file_path: Optional ``.env`` file that is checked before process environment + variables for ``AZURE_OPENAI_*`` values. + env_file_encoding: Encoding for the ``.env`` file. + """ + ... + + def __init__( + self, + model: str | None = None, + *, + api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None, + credential: AzureCredentialTypes | AzureTokenProvider | None = None, + org_id: str | None = None, + base_url: str | None = None, + azure_endpoint: str | None = None, + api_version: str | None = None, + default_headers: Mapping[str, str] | None = None, + async_client: AsyncOpenAI | None = None, + instruction_role: str | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + additional_properties: dict[str, Any] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize a raw OpenAI Chat client. + + Keyword Args: + model: Model identifier to use for the request. When not provided, the constructor + reads ``OPENAI_CHAT_MODEL`` and then ``OPENAI_MODEL`` for OpenAI, + or ``AZURE_OPENAI_CHAT_MODEL`` and then ``AZURE_OPENAI_MODEL`` for Azure. + api_key: API key override. For OpenAI this maps to ``OPENAI_API_KEY``. + For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key + auth. A callable token provider is also accepted for backwards compatibility, + but ``credential`` is the preferred Azure auth surface. + credential: Azure credential or token provider for Azure OpenAI auth. Passing this + is an explicit Azure signal, even when ``OPENAI_API_KEY`` is also configured. + Credential objects require the optional ``azure-identity`` package. + org_id: OpenAI organization ID. Used only for OpenAI and resolved from + ``OPENAI_ORG_ID`` when not provided. + base_url: Base URL override. For OpenAI this maps to ``OPENAI_BASE_URL``. + For Azure this may be used instead of ``azure_endpoint`` when you want + to pass the full ``.../openai/v1`` base URL directly. + azure_endpoint: Azure resource endpoint. When not provided explicitly, Azure + falls back to ``AZURE_OPENAI_ENDPOINT``. + api_version: Azure API version to use once Azure routing is selected. When + not provided explicitly, Azure routing falls back to + ``AZURE_OPENAI_API_VERSION`` and then the Responses default. + default_headers: Additional HTTP headers. + async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on + Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup. + instruction_role: Role for instruction messages (for example ``"system"``). + compaction_strategy: Optional per-client compaction override. + tokenizer: Optional tokenizer for compaction strategies. + additional_properties: Additional properties stored on the client instance. + env_file_path: Optional ``.env`` file that is checked before process environment + variables. The same file is used for both ``OPENAI_*`` and ``AZURE_OPENAI_*`` + lookups. + env_file_encoding: Encoding for the ``.env`` file. + + Notes: + Environment resolution and routing precedence are: + + 1. Explicit Azure inputs (``azure_endpoint`` or ``credential``) + 2. Explicit OpenAI API key or ``OPENAI_API_KEY`` + 3. Azure environment fallback + + OpenAI routing reads ``OPENAI_API_KEY``, ``OPENAI_CHAT_MODEL``, + ``OPENAI_MODEL``, ``OPENAI_ORG_ID``, and ``OPENAI_BASE_URL``. Azure routing + reads ``AZURE_OPENAI_ENDPOINT``, ``AZURE_OPENAI_BASE_URL``, + ``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_CHAT_MODEL``, + ``AZURE_OPENAI_MODEL``, and ``AZURE_OPENAI_API_VERSION``. + """ + settings, client, use_azure_client = load_openai_service_settings( + model=model, + api_key=api_key, + credential=credential, + org_id=org_id, + base_url=base_url, + endpoint=azure_endpoint, + api_version=api_version, + default_azure_api_version=DEFAULT_AZURE_OPENAI_RESPONSES_API_VERSION, + default_headers=default_headers, + client=async_client, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + openai_model_fields=("chat_model", "model"), + azure_model_fields=("chat_model", "model"), + responses_mode=True, + ) + + self.client = client + self.model: str = settings.get("model") or "" + + # Store configuration for serialization + self.org_id = settings.get("org_id") + self.base_url = settings.get("base_url") + self.azure_endpoint = settings.get("endpoint") + self.api_version = settings.get("api_version") + if default_headers: + self.default_headers: dict[str, Any] | None = { + k: v for k, v in default_headers.items() if k != USER_AGENT_KEY + } + else: + self.default_headers = None + self.instruction_role = instruction_role + if use_azure_client: + self.OTEL_PROVIDER_NAME = "azure.ai.openai" # type: ignore[misc] + + super().__init__( + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + additional_properties=additional_properties, + ) + + # region Inner Methods + + async def _prepare_request( + self, + messages: Sequence[Message], + options: Mapping[str, Any], + ) -> tuple[AsyncOpenAI, dict[str, Any], dict[str, Any]]: + """Validate options and prepare the request. + + Returns: + Tuple of (client, run_options, validated_options). + """ + client = self.client + validated_options = await self._validate_options(options) + run_options = await self._prepare_options(messages, validated_options) + return client, run_options, validated_options + + def _handle_request_error(self, ex: Exception) -> NoReturn: + """Convert exceptions to appropriate service exceptions. Always raises.""" + if isinstance(ex, BadRequestError) and ex.code == "content_filter": + raise OpenAIContentFilterException( + f"{type(self)} service encountered a content error: {ex}", + inner_exception=ex, + ) from ex + raise ChatClientException( + maybe_append_azure_endpoint_guidance( + f"{type(self)} service failed to complete the prompt: {ex}", + azure_endpoint=self.azure_endpoint, + ), + inner_exception=ex, + ) from ex + + @override + def _inner_get_response( + self, + *, + messages: Sequence[Message], + options: Mapping[str, Any], + stream: bool = False, + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + continuation_token: OpenAIContinuationToken | None = options.get("continuation_token") # type: ignore[assignment] + + if stream: + function_call_ids: dict[int, tuple[str, str]] = {} + seen_reasoning_delta_item_ids: set[str] = set() + validated_options: dict[str, Any] | None = None + # Captured once request options are validated/prepared so the streaming finalizer can + # still parse the aggregated response into structured output after the stream completes. + response_format: Any | None = None + + def _finalize_with_captured_format(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: + # ResponseStream only calls the finalizer after iterating or draining `_stream()`, + # so `response_format` has already been populated from the validated request state + # unless request setup failed before streaming began. + return self._finalize_response_updates(updates, response_format=response_format) + + async def _stream() -> AsyncIterable[ChatResponseUpdate]: + nonlocal response_format, validated_options + if continuation_token is not None: + # Resume a background streaming response by retrieving with stream=True + client = self.client + validated_options = await self._validate_options(options) + response_format = validated_options.get("response_format") + try: + raw_stream_response = await client.responses.with_raw_response.retrieve( + continuation_token["response_id"], + stream=True, + ) + served_model = self._extract_served_model(raw_stream_response.headers) + async with raw_stream_response.parse() as stream_response: + async for chunk in stream_response: + update = self._parse_chunk_from_openai( + chunk, + options=validated_options, + function_call_ids=function_call_ids, + seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, + ) + if served_model is not None: + update.model = served_model + yield update + except Exception as ex: + self._handle_request_error(ex) + else: + ( + client, + run_options, + validated_options, + ) = await self._prepare_request(messages, options) + response_format = validated_options.get("response_format") + try: + if "text_format" in run_options: + # The SDK's ``responses.stream(text_format=...)`` helper preserves + # client-side ``output_parsed`` partial parsing for structured outputs, + # but it does not expose the raw HTTP response (no ``x-ms-served-model`` + # access). We accept that trade-off: this single streaming path keeps + # the deployment alias as the reported model name. All other paths + # surface the served-model header. + async with client.responses.stream(**run_options) as response: + async for chunk in response: + yield self._parse_chunk_from_openai( + chunk, + options=validated_options, + function_call_ids=function_call_ids, + seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, + ) + else: + raw_create_response = await client.responses.with_raw_response.create( + stream=True, **run_options + ) + served_model = self._extract_served_model(raw_create_response.headers) + async with raw_create_response.parse() as stream_response: + async for chunk in stream_response: + update = self._parse_chunk_from_openai( + chunk, + options=validated_options, + function_call_ids=function_call_ids, + seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, + ) + if served_model is not None: + update.model = served_model + yield update + except Exception as ex: + self._handle_request_error(ex) + + return ResponseStream(_stream(), finalizer=_finalize_with_captured_format) + + # Non-streaming + async def _get_response() -> ChatResponse: + if continuation_token is not None: + # Poll a background response by retrieving without stream + client = self.client + validated_options = await self._validate_options(options) + try: + raw_response = await client.responses.with_raw_response.retrieve(continuation_token["response_id"]) + response = raw_response.parse() + except Exception as ex: + self._handle_request_error(ex) + chat_response = self._parse_response_from_openai(response, options=validated_options) + served_model = self._extract_served_model(raw_response.headers) + if served_model is not None: + chat_response.model = served_model + # Once the background response completes, drop the continuation_token from + # the caller's options dict. FunctionInvocationLayer reuses the same dict + # across tool-loop iterations, so leaving it in place makes the next iteration + # retrieve the same completed response again instead of POSTing tool results + # (issue #5394). Keep `background` so subsequent iterations still create + # background responses. + if chat_response.continuation_token is None and isinstance(options, dict): + options.pop("continuation_token", None) + return chat_response + client, run_options, validated_options = await self._prepare_request(messages, options) + try: + if "text_format" in run_options: + raw_response = await client.responses.with_raw_response.parse(stream=False, **run_options) # type: ignore + else: + raw_response = await client.responses.with_raw_response.create(stream=False, **run_options) # type: ignore + response = raw_response.parse() + except Exception as ex: + self._handle_request_error(ex) + chat_response = self._parse_response_from_openai(response, options=validated_options) + served_model = self._extract_served_model(raw_response.headers) + if served_model is not None: + chat_response.model = served_model + return chat_response + + return _get_response() + + @classmethod + def _extract_served_model(cls, headers: Any) -> str | None: + """Return the Azure OpenAI ``x-ms-served-model`` response header value when present. + + Azure OpenAI Responses API returns the deployment alias in ``response.model`` but the actual + snapshot served via the ``x-ms-served-model`` response header (e.g. ``gpt-5-nano-2025-08-07`` + vs deployment alias ``gpt-5-nano``). When present, the served snapshot is the source of truth + for observability and downstream callers. Empty/whitespace-only header values are rejected + here so every caller can simply check ``if served_model is not None``. + """ + if headers is None: + return None + served_model = headers.get(cls.SERVED_MODEL_HEADER) + if isinstance(served_model, str): + stripped = served_model.strip() + if stripped: + return stripped + return None + + def _prepare_response_and_text_format( + self, + *, + response_format: Any, + text_config: MutableMapping[str, Any] | None, + ) -> tuple[type[BaseModel] | None, dict[str, Any] | None]: + """Normalize response_format into Responses text configuration and parse target.""" + if text_config is not None and not isinstance(text_config, MutableMapping): + raise ChatClientInvalidRequestException("text must be a mapping when provided.") + text_config = cast(dict[str, Any], text_config) if isinstance(text_config, MutableMapping) else None + + if response_format is None: + return None, text_config + + if isinstance(response_format, type) and issubclass(response_format, BaseModel): + if text_config and "format" in text_config: + raise ChatClientInvalidRequestException("response_format cannot be combined with explicit text.format.") + return response_format, text_config + + if isinstance(response_format, Mapping): + format_config = self._convert_response_format(cast("Mapping[str, Any]", response_format)) + if text_config is None: + text_config = {} + elif "format" in text_config and text_config["format"] != format_config: + raise ChatClientInvalidRequestException("Conflicting response_format definitions detected.") + text_config["format"] = format_config + return None, text_config + + raise ChatClientInvalidRequestException("response_format must be a Pydantic model or mapping.") + + def _convert_response_format(self, response_format: Mapping[str, Any]) -> dict[str, Any]: + """Convert Chat style response_format into Responses text format config.""" + if "format" in response_format and isinstance(response_format["format"], Mapping): + return dict(cast("Mapping[str, Any]", response_format["format"])) + + format_type = response_format.get("type") + if format_type == "json_schema": + schema_section = response_format.get("json_schema", response_format) + if not isinstance(schema_section, Mapping): + raise ChatClientInvalidRequestException("json_schema response_format must be a mapping.") + schema_section_typed = cast("Mapping[str, Any]", schema_section) + schema: Any = schema_section_typed.get("schema") + if schema is None: + raise ChatClientInvalidRequestException("json_schema response_format requires a schema.") + name: str = str( + schema_section_typed.get("name") + or schema_section_typed.get("title") + or (cast("Mapping[str, Any]", schema).get("title") if isinstance(schema, Mapping) else None) + or "response" + ) + format_config: dict[str, Any] = { + "type": "json_schema", + "name": name, + "schema": schema, + } + if "strict" in schema_section: + format_config["strict"] = schema_section["strict"] + if "description" in schema_section and schema_section["description"] is not None: + format_config["description"] = schema_section["description"] + return format_config + + if format_type in {"json_object", "text"}: + return {"type": format_type} + + # Handle raw JSON schemas (e.g. {"type": "object", "properties": {...}}) + # by wrapping them in the expected json_schema envelope. + # Detect by checking for JSON Schema primitive types or known schema keywords. + json_schema_keywords = {"properties", "anyOf", "oneOf", "allOf", "$ref", "$defs"} + json_schema_primitive_types = {"object", "array", "string", "number", "integer", "boolean", "null"} + if format_type in json_schema_primitive_types or ( + format_type is None and any(k in response_format for k in json_schema_keywords) + ): + schema = dict(response_format) + if schema.get("type") == "object" and "additionalProperties" not in schema: + schema["additionalProperties"] = False + # Pop title from schema since OpenAI strict mode rejects unknown keys; + # use it as the schema name in the envelope instead. + name = str(schema.pop("title", None) or "response") + return { + "type": "json_schema", + "name": name, + "schema": schema, + "strict": True, + } + + raise ChatClientInvalidRequestException("Unsupported response_format provided for Responses client.") + + def _get_conversation_id( + self, response: OpenAIResponse | ParsedResponse[BaseModel], store: bool | None + ) -> str | None: + """Get the conversation ID from the response if store is True.""" + if store is False: + return None + # If conversation ID exists, it means that we operate with conversation + # so we use conversation ID as input and output. + if response.conversation and response.conversation.id: + return response.conversation.id + # If conversation ID doesn't exist, we operate with responses + # so we use response ID as input and output. + return response.id + + # region Prep methods + + def _prepare_tools_for_openai( + self, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, + ) -> list[Any]: + """Prepare tools for the OpenAI Responses API. + + Converts FunctionTool to Responses API format. Shell-enabled FunctionTools + with explicit shell environment metadata are mapped to OpenAI shell tools. + All other tools pass through unchanged. + + Args: + tools: A single tool or sequence of tools to prepare. + + Returns: + List of tool parameters ready for the OpenAI API. + """ + tools_list = normalize_tools(tools) + if not tools_list: + return [] + response_tools: list[Any] = [] + for tool_item in tools_list: + if isinstance(tool_item, FunctionTool) and tool_item.kind == SHELL_TOOL_KIND_VALUE: + shell_env = (tool_item.additional_properties or {}).get(OPENAI_SHELL_ENVIRONMENT_KEY) + response_tools.append( + FunctionShellTool( + type="shell", + environment=shell_env, # type: ignore[typeddict-item] + ) + ) + continue + if isinstance(tool_item, FunctionTool): + params = tool_item.parameters() + params["additionalProperties"] = False + response_tools.append( + FunctionToolParam( + name=tool_item.name, + parameters=params, + strict=False, + type="function", + description=tool_item.description, + ) + ) + else: + # Pass through all other tools (dicts, SDK types) unchanged + response_tools.append(tool_item) + return response_tools + + def _get_local_shell_tool_name( + self, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, + ) -> str | None: + """Return the name of the configured local shell tool function, if any.""" + for tool_item in normalize_tools(tools): + if not isinstance(tool_item, FunctionTool): + continue + if tool_item.kind != SHELL_TOOL_KIND_VALUE: + continue + shell_env = (tool_item.additional_properties or {}).get(OPENAI_SHELL_ENVIRONMENT_KEY) + if isinstance(shell_env, Mapping) and shell_env.get("type") == "local": # type: ignore[typeddict-item] + return tool_item.name + return None + + # region Hosted Tool Factory Methods + + @staticmethod + def get_code_interpreter_tool( + *, + file_ids: list[str] | None = None, + container: Literal["auto"] | CodeInterpreterContainerCodeInterpreterToolAuto = "auto", + ) -> Any: + """Create a code interpreter tool configuration for the Responses API. + + Keyword Args: + file_ids: List of file IDs to make available to the code interpreter. + container: Container configuration. Use "auto" for automatic container management, + or provide a TypedDict with custom container settings. + + Returns: + A CodeInterpreter tool parameter ready to pass to ChatAgent. + + Examples: + .. code-block:: python + + from agent_framework.openai import OpenAIChatClient + + # Basic code interpreter + tool = OpenAIChatClient.get_code_interpreter_tool() + + # With file access + tool = OpenAIChatClient.get_code_interpreter_tool(file_ids=["file-abc123"]) + + # Use with agent + agent = ChatAgent(client, tools=[tool]) + """ + container_config: CodeInterpreterContainerCodeInterpreterToolAuto = ( + container if isinstance(container, dict) else {"type": "auto"} + ) + + if file_ids: + container_config["file_ids"] = file_ids + + return CodeInterpreter(type="code_interpreter", container=container_config) + + @staticmethod + def get_web_search_tool( + *, + user_location: dict[str, str] | None = None, + search_context_size: Literal["low", "medium", "high"] | None = None, + filters: dict[str, Any] | None = None, + ) -> Any: + """Create a web search tool configuration for the Responses API. + + Keyword Args: + user_location: Location context for search results. Dict with keys like + "city", "country", "region", "timezone". + search_context_size: Amount of context to include from search results. + One of "low", "medium", or "high". + filters: Additional search filters. + + Returns: + A WebSearchToolParam dict ready to pass to ChatAgent. + + Examples: + .. code-block:: python + + from agent_framework.openai import OpenAIChatClient + + # Basic web search + tool = OpenAIChatClient.get_web_search_tool() + + # With location context + tool = OpenAIChatClient.get_web_search_tool( + user_location={"city": "Seattle", "country": "US"}, + search_context_size="medium", + ) + + agent = ChatAgent(client, tools=[tool]) + """ + web_search_tool = WebSearchToolParam(type="web_search") + + if user_location: + web_search_tool["user_location"] = { + "type": "approximate", + "city": user_location.get("city"), + "country": user_location.get("country"), + "region": user_location.get("region"), + "timezone": user_location.get("timezone"), + } + + if search_context_size: + web_search_tool["search_context_size"] = search_context_size + + if filters: + web_search_tool["filters"] = filters # type: ignore[typeddict-item] + + return web_search_tool + + @staticmethod + def get_image_generation_tool( + *, + size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] | None = None, + output_format: Literal["png", "jpeg", "webp"] | None = None, + model: Literal["gpt-image-1", "gpt-image-1-mini"] | str | None = None, + quality: Literal["low", "medium", "high", "auto"] | None = None, + partial_images: int | None = None, + background: Literal["transparent", "opaque", "auto"] | None = None, + moderation: Literal["auto", "low"] | None = None, + output_compression: int | None = None, + ) -> Any: + """Create an image generation tool configuration for the Responses API. + + Keyword Args: + size: Image dimensions. One of "1024x1024", "1024x1536", "1536x1024", or "auto". + output_format: Output image format. One of "png", "jpeg", or "webp". + model: Model to use for image generation. One of "gpt-image-1" or "gpt-image-1-mini". + quality: Image quality level. One of "low", "medium", "high", or "auto". + partial_images: Number of partial images to stream during generation. + background: Background type. One of "transparent", "opaque", or "auto". + moderation: Moderation level. One of "auto" or "low". + output_compression: Compression level for output (0-100). + + Returns: + An ImageGeneration tool parameter dict ready to pass to ChatAgent. + + Examples: + .. code-block:: python + + from agent_framework.openai import OpenAIChatClient + + # Basic image generation + tool = OpenAIChatClient.get_image_generation_tool() + + # High quality large image + tool = OpenAIChatClient.get_image_generation_tool( + size="1536x1024", + quality="high", + output_format="png", + ) + + agent = ChatAgent(client, tools=[tool]) + """ + tool: ImageGeneration = {"type": "image_generation"} + + if size: + tool["size"] = size + if output_format: + tool["output_format"] = output_format + if model: + tool["model"] = model # type: ignore + if quality: + tool["quality"] = quality + if partial_images is not None: + tool["partial_images"] = partial_images + if background: + tool["background"] = background + if moderation: + tool["moderation"] = moderation + if output_compression is not None: + tool["output_compression"] = output_compression + + return tool + + @staticmethod + def get_shell_tool( + *, + func: Callable[..., Any] | FunctionTool | None = None, + environment: Literal["auto"] | dict[str, Any] | None = "auto", + name: str | None = None, + description: str | None = None, + approval_mode: Literal["always_require", "never_require"] | None = None, + ) -> Any: + """Create a shell tool for the Responses API. + + - When ``func`` is ``None`` (default), returns an OpenAI hosted shell + tool declaration. + - When ``func`` is provided, returns a local FunctionTool that is + declared to OpenAI as a local shell tool and executed via the function + invocation layer. + + Keyword Args: + func: Optional local shell function or ``FunctionTool``. + environment: Container environment configuration. + Used only when ``func`` is ``None``. + Use ``"auto"`` (default) for managed containers, or provide a + dict with explicit hosted container settings. + name: Optional local tool name when ``func`` is provided. + description: Optional local tool description when ``func`` is provided. + approval_mode: Optional local tool approval mode. + + Returns: + A hosted shell declaration or a local shell FunctionTool. + + Examples: + .. code-block:: python + + from agent_framework.openai import OpenAIChatClient + + # Hosted shell (OpenAI container) + tool = OpenAIChatClient.get_shell_tool() + + # Hosted shell with custom environment + tool = OpenAIChatClient.get_shell_tool(environment={"type": "container_auto", "file_ids": ["file-abc"]}) + + # Local shell execution + tool = OpenAIChatClient.get_shell_tool( + func=my_shell_func, + ) + """ + if func is None: + env_config: dict[str, Any] = ( + dict(environment) if isinstance(environment, dict) else {"type": "container_auto"} + ) + if env_config.get("type") == "local": + raise ValueError("Local shell requires func. Provide func for local execution.") + return FunctionShellTool(type="shell", environment=env_config) # type: ignore[typeddict-item] + + if isinstance(environment, dict): + raise ValueError("When func is provided, environment config is not supported.") + local_env = {"type": "local"} + + base_tool: FunctionTool + if isinstance(func, FunctionTool): + base_tool = func + if name is not None: + base_tool.name = name + if description is not None: + base_tool.description = description + if approval_mode is not None: + base_tool.approval_mode = approval_mode + else: + base_tool = tool( + func=func, + name=name, + description=description, + approval_mode=approval_mode, + ) + + if base_tool.func is None: + raise ValueError("Shell tool requires an executable function.") + + additional_properties = dict(base_tool.additional_properties or {}) + additional_properties[OPENAI_SHELL_ENVIRONMENT_KEY] = local_env + base_tool.additional_properties = additional_properties + base_tool.kind = SHELL_TOOL_KIND_VALUE + return base_tool + + @staticmethod + def get_mcp_tool( + *, + name: str, + url: str, + description: str | None = None, + approval_mode: Literal["always_require", "never_require"] | dict[str, list[str]] | None = None, + allowed_tools: list[str] | None = None, + headers: dict[str, str] | None = None, + ) -> Any: + """Create a hosted MCP (Model Context Protocol) tool configuration for the Responses API. + + This configures an MCP server that will be called by OpenAI's service. + The tools from this MCP server are executed remotely by OpenAI, + not locally by your application. + + Note: + For local MCP execution where your application calls the MCP server + directly, use the MCP client tools instead of this method. + + Keyword Args: + name: A label/name for the MCP server. + url: The URL of the MCP server. + description: A description of what the MCP server provides. + approval_mode: Tool approval mode. Use "always_require" or "never_require" for all tools, + or provide a dict with "always_require_approval" and/or "never_require_approval" + keys mapping to lists of tool names. + allowed_tools: List of tool names that are allowed to be used from this MCP server. + headers: HTTP headers to include in requests to the MCP server. + + Returns: + An Mcp tool parameter dict ready to pass to ChatAgent. + + Examples: + .. code-block:: python + + from agent_framework.openai import OpenAIChatClient + + # Basic MCP tool + tool = OpenAIChatClient.get_mcp_tool( + name="my_mcp", + url="https://mcp.example.com", + ) + + # With approval settings + tool = OpenAIChatClient.get_mcp_tool( + name="github_mcp", + url="https://mcp.github.com", + description="GitHub MCP server", + approval_mode="always_require", + headers={"Authorization": "Bearer token"}, + ) + + # With specific tool approvals + tool = OpenAIChatClient.get_mcp_tool( + name="tools_mcp", + url="https://tools.example.com", + approval_mode={ + "always_require_approval": ["dangerous_tool"], + "never_require_approval": ["safe_tool"], + }, + ) + + agent = ChatAgent(client, tools=[tool]) + """ + mcp: Mcp = { + "type": "mcp", + "server_label": name.replace(" ", "_"), + "server_url": url, + } + + if description: + mcp["server_description"] = description + + if headers: + mcp["headers"] = headers + + if allowed_tools is not None: + mcp["allowed_tools"] = list(allowed_tools) + + if approval_mode: + if isinstance(approval_mode, str): + mcp["require_approval"] = "always" if approval_mode == "always_require" else "never" + else: + if always_require := approval_mode.get("always_require_approval"): + mcp["require_approval"] = {"always": {"tool_names": always_require}} + if never_require := approval_mode.get("never_require_approval"): + mcp["require_approval"] = {"never": {"tool_names": never_require}} + + return mcp + + @staticmethod + def get_file_search_tool( + *, + vector_store_ids: list[str], + max_num_results: int | None = None, + ) -> Any: + """Create a file search tool configuration for the Responses API. + + Keyword Args: + vector_store_ids: List of vector store IDs to search within. + max_num_results: Maximum number of results to return. Defaults to 50 if not specified. + + Returns: + A FileSearchToolParam dict ready to pass to ChatAgent. + + Examples: + .. code-block:: python + + from agent_framework.openai import OpenAIChatClient + + # Basic file search + tool = OpenAIChatClient.get_file_search_tool( + vector_store_ids=["vs_abc123"], + ) + + # With result limit + tool = OpenAIChatClient.get_file_search_tool( + vector_store_ids=["vs_abc123", "vs_def456"], + max_num_results=10, + ) + + agent = ChatAgent(client, tools=[tool]) + """ + tool = FileSearchToolParam( + type="file_search", + vector_store_ids=vector_store_ids, + ) + + if max_num_results is not None: + tool["max_num_results"] = max_num_results + + return tool + + # endregion + + async def _prepare_options( + self, + messages: Sequence[Message], + options: Mapping[str, Any], + ) -> dict[str, Any]: + """Take options dict and create the specific options for Responses API.""" + # Exclude keys that are not supported or handled separately + exclude_keys = { + "type", + "presence_penalty", # not supported + "frequency_penalty", # not supported + "logit_bias", # not supported + "seed", # not supported + "stop", # not supported + "instructions", # already added as system message + "response_format", # handled separately + "conversation_id", # handled separately + "tool_choice", # handled separately + "continuation_token", # handled separately in _inner_get_response + } + run_options: dict[str, Any] = {k: v for k, v in options.items() if k not in exclude_keys and v is not None} + + # messages + # Handle instructions by prepending to messages as system message + # Only prepend instructions for the first turn (when no conversation/response ID exists) + conversation_id = options.get("conversation_id") + if (instructions := options.get("instructions")) and not conversation_id: + # First turn: prepend instructions as system message + messages = prepend_instructions_to_messages(list(messages), instructions, role="system") + # Continuation turn: instructions already exist in conversation context, skip prepending + request_uses_service_side_storage = False + for key in ("conversation_id", "previous_response_id", "conversation"): + value = options.get(key) + if isinstance(value, str) and value: + request_uses_service_side_storage = True + break + request_input = self._prepare_messages_for_openai( + messages, + request_uses_service_side_storage=request_uses_service_side_storage, + ) + if not request_input: + raise ChatClientInvalidRequestException("Messages are required for chat completions") + conversation_id = options.get("conversation_id") + run_options["input"] = request_input + + # model id + self._check_model_presence(run_options) + + # translations between options and Responses API + translations = { + "allow_multiple_tool_calls": "parallel_tool_calls", + "conversation_id": "previous_response_id", + "max_tokens": "max_output_tokens", + } + for old_key, new_key in translations.items(): + if old_key in run_options and old_key != new_key: + run_options[new_key] = run_options.pop(old_key) + + # Handle different conversation ID formats + if conversation_id := options.get("conversation_id"): + if conversation_id.startswith("resp_"): + # For response IDs, set previous_response_id and remove conversation property + run_options["previous_response_id"] = conversation_id + elif conversation_id.startswith("conv_"): + # For conversation IDs, set conversation and remove previous_response_id property + run_options["conversation"] = conversation_id + else: + # If the format is unrecognized, default to previous_response_id + run_options["previous_response_id"] = conversation_id + + # tools + if tools := self._prepare_tools_for_openai(options.get("tools")): + run_options["tools"] = tools + # tool_choice: convert ToolMode to appropriate format + if tool_choice := options.get("tool_choice"): + tool_mode = validate_tool_mode(tool_choice) + if tool_mode is not None: + if (mode := tool_mode.get("mode")) == "required" and ( + func_name := tool_mode.get("required_function_name") + ) is not None: + run_options["tool_choice"] = { + "type": "function", + "name": func_name, + } + elif mode == "auto" and (allowed := tool_mode.get("allowed_tools")) is not None: + run_options["tool_choice"] = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": name} for name in allowed], + } + else: + run_options["tool_choice"] = mode + else: + run_options.pop("parallel_tool_calls", None) + run_options.pop("tool_choice", None) + + # response format and text config + response_format = options.get("response_format") + text_config = run_options.pop("text", None) + response_format, text_config = self._prepare_response_and_text_format( + response_format=response_format, text_config=text_config + ) + # The Responses API nests verbosity under ``text.verbosity``; surface it as a + # top-level option for parity with ``reasoning`` and translate here. + if (verbosity := run_options.pop("verbosity", None)) is not None: + text_config = dict(text_config) if text_config else {} + text_config["verbosity"] = verbosity + if text_config: + run_options["text"] = text_config + if response_format: + run_options["text_format"] = response_format + + return run_options + + def _check_model_presence(self, options: dict[str, Any]) -> None: + """Check if the 'model' param is present, and if not raise a Error. + + Subclasses can override this when they populate the model through a different option field. + """ + if not options.get("model"): + if not self.model: + raise ValueError("model must be a non-empty string") + options["model"] = self.model + + def _prepare_messages_for_openai( + self, + chat_messages: Sequence[Message], + *, + request_uses_service_side_storage: bool = True, + ) -> list[dict[str, Any]]: + """Prepare the chat messages for a request. + + Allowing customization of the key names for role/author, and optionally overriding the role. + + "tool" messages need to be formatted different than system/user/assistant messages: + They require a "tool_call_id" and (function) "name" key, and the "metadata" key should + be removed. The "encoding" key should also be removed. + + Override this method to customize the formatting of the chat history for a request. + + Args: + chat_messages: The chat history to prepare. + request_uses_service_side_storage: Whether this request continues a service-managed + response/conversation and can safely reference service-scoped response items. + + Returns: + The prepared chat messages for a request. + """ + list_of_list = [ + self._prepare_message_for_openai( + message, + request_uses_service_side_storage=request_uses_service_side_storage, + ) + for message in chat_messages + ] + # Flatten the list of lists into a single list + flat = list(chain.from_iterable(list_of_list)) + # Coalesce hosted-MCP result markers onto matching mcp_call input + # items (drop unmatched). See `_AF_MCP_PENDING_OUTPUT_KEY`. + return self._coalesce_pending_mcp_results(flat) + + def _prepare_message_for_openai( + self, + message: Message, + *, + request_uses_service_side_storage: bool = True, + ) -> list[dict[str, Any]]: + """Prepare a chat message for the OpenAI Responses API format.""" + all_messages: list[dict[str, Any]] = [] + args: dict[str, Any] = { + "type": "message", + "role": message.role, + } + additional_properties = message.additional_properties + replays_local_storage = "_attribution" in additional_properties + # Server-issued response item identities (function_call fc_*, reasoning rs_*, approval IDs, + # local-shell-call IDs) must not be re-sent inline when the request carries + # previous_response_id / conversation_id / conversation: the server already has them via + # the prior response and rejects duplicates with "Duplicate item found with id ...". + # function_result keeps its call_id and the server pairs it to the prior function_call via + # that key. See microsoft/agent-framework#3295. The strip is gated on the request-level + # flag, not a message-level one: HistoryProvider-attributed messages + # (replays_local_storage) still need stripping when the request also carries a continuation + # marker, since the server-stored items would otherwise duplicate the inline ones. Without + # storage, standalone reasoning items are invalid per the API ("reasoning was provided + # without its required following item"), so the reasoning branch always drops. + for content in message.contents: + match content.type: + case "text_reasoning": + continue + case "function_result": + if request_uses_service_side_storage: + props = content.additional_properties or {} + # Local-shell variant serializes as `local_shell_call` carrying a server-issued id; + # plain function_call_output pairs by call_id and is safe under storage. + if props.get( + OPENAI_SHELL_OUTPUT_TYPE_KEY + ) == OPENAI_SHELL_OUTPUT_TYPE_LOCAL_SHELL_CALL and props.get( + OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY + ): + continue + new_args: dict[str, Any] = {} + new_args.update( + self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) + ) + if new_args: + all_messages.append(new_args) + case "function_call": + if request_uses_service_side_storage: + continue + function_call = self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) + if function_call: + all_messages.append(function_call) + case "function_approval_response" | "function_approval_request": + if request_uses_service_side_storage: + continue + prepared = self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) + if prepared: + all_messages.append(prepared) + case "mcp_server_tool_call" | "mcp_server_tool_result": + # Hosted MCP call/result contents serialize as a single + # top-level mcp_call input item; the result side emits an + # internal marker that `_prepare_messages_for_openai` + # coalesces onto the matching call (or drops if unmatched). + # The mcp_call item carries the model-emitted call_id as its + # server-side `id`, so under continuation it would duplicate + # the prior response's items (#3295). Drop the call here; the + # orphan result is dropped by the coalesce step that follows. + if request_uses_service_side_storage: + continue + prepared_mcp = self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) + if prepared_mcp: + all_messages.append(prepared_mcp) + case _: + prepared_content = self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) + if prepared_content: + if "content" not in args: + args["content"] = [] + args["content"].append(prepared_content) # type: ignore[reportUnknownMemberType] + if "content" in args or "tool_calls" in args: + all_messages.append(args) + return all_messages + + def _prepare_content_for_openai( + self, + role: Role | str, + content: Content, + *, + replays_local_storage: bool = False, + ) -> dict[str, Any]: + """Prepare content for the OpenAI Responses API format.""" + role = Role(role) + match content.type: + case "text": + if role == "assistant": + # Assistant history is represented as output text items; Azure validation + # requires `annotations` to be present for this type. + return { + "type": "output_text", + "text": content.text, + "annotations": _annotations_to_output_text(getattr(content, "annotations", None)), + } + return { + "type": "input_text", + "text": content.text, + } + case "text_reasoning": + ret: dict[str, Any] = {"type": "reasoning", "summary": []} + if content.id: + ret["id"] = content.id + props: dict[str, Any] | None = getattr(content, "additional_properties", None) + if props: + if status := props.get("status"): + ret["status"] = status + if reasoning_text := props.get("reasoning_text"): + ret["content"] = [{"type": "reasoning_text", "text": reasoning_text}] + if encrypted_content := props.get("encrypted_content"): + ret["encrypted_content"] = encrypted_content + if content.text: + ret["summary"].append({"type": "summary_text", "text": content.text}) + return ret + case "data" | "uri": + if content.has_top_level_media_type("image"): + result: dict[str, Any] = { + "type": "input_image", + "image_url": content.uri, + "detail": content.additional_properties.get("detail", "auto") + if content.additional_properties + else "auto", + } + file_id = content.additional_properties.get("file_id") if content.additional_properties else None + if file_id is not None: + result["file_id"] = file_id + return result + if content.has_top_level_media_type("audio"): + if content.media_type and "wav" in content.media_type: + format = "wav" + elif content.media_type and "mp3" in content.media_type: + format = "mp3" + else: + logger.warning("Unsupported audio media type: %s", content.media_type) + return {} + return { + "type": "input_audio", + "input_audio": { + "data": content.uri, + "format": format, + }, + } + if content.has_top_level_media_type("application"): + filename = getattr(content, "filename", None) or ( + content.additional_properties.get("filename") + if hasattr(content, "additional_properties") and content.additional_properties + else None + ) + file_obj = { + "type": "input_file", + "file_data": content.uri, + } + if filename: + file_obj["filename"] = filename + return file_obj + return {} + case "function_call": + if not content.call_id: + logger.warning(f"FunctionCallContent missing call_id for function '{content.name}'") + return {} + fc_id = content.call_id + if not replays_local_storage and content.additional_properties: + live_fc_id = content.additional_properties.get("fc_id") + if isinstance(live_fc_id, str) and live_fc_id: + fc_id = live_fc_id + # OpenAI Responses API requires IDs to start with `fc_` + if not fc_id.startswith("fc_"): + fc_id = f"fc_{fc_id}" + + function_call_obj = { + "call_id": content.call_id, + "id": fc_id, + "type": "function_call", + "name": content.name, + "arguments": content.arguments, + } + if status := content.additional_properties.get("status"): + function_call_obj["status"] = status + return function_call_obj + case "function_result": + shell_output_type = ( + content.additional_properties.get(OPENAI_SHELL_OUTPUT_TYPE_KEY) + if content.additional_properties + else None + ) + if shell_output_type == OPENAI_SHELL_OUTPUT_TYPE_SHELL_CALL: + return { + "call_id": content.call_id, + "type": OPENAI_SHELL_OUTPUT_TYPE_SHELL_CALL, + "output": self._to_shell_call_output_payload(content), + } + local_shell_call_item_id = ( + content.additional_properties.get(OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY) + if content.additional_properties + else None + ) + if shell_output_type == OPENAI_SHELL_OUTPUT_TYPE_LOCAL_SHELL_CALL and local_shell_call_item_id: + return { + "id": local_shell_call_item_id, + "type": OPENAI_SHELL_OUTPUT_TYPE_LOCAL_SHELL_CALL, + "output": self._to_local_shell_output_payload(content), + } + # call_id for the result needs to be the same as the call_id for the function call + output: str | list[dict[str, Any]] = content.result or "" + if ( + self.SUPPORTS_RICH_FUNCTION_OUTPUT + and content.items + and any(item.type in ("data", "uri") for item in content.items) + ): + output_parts: list[dict[str, Any]] = [] + for item in content.items: + if item.type == "text": + output_parts.append({"type": "input_text", "text": item.text or ""}) + else: + part = self._prepare_content_for_openai("user", item) + if part: + output_parts.append(part) + if output_parts: + output = output_parts + return { + "call_id": content.call_id, + "type": "function_call_output", + "output": output, + } + case "function_approval_request": + return { + "type": "mcp_approval_request", + "id": content.id, # type: ignore[union-attr] + "arguments": content.function_call.arguments, # type: ignore[union-attr] + "name": content.function_call.name, # type: ignore[union-attr] + "server_label": content.function_call.additional_properties.get("server_label") # type: ignore[union-attr] + if content.function_call.additional_properties # type: ignore[union-attr] + else None, + } + case "function_approval_response": + return { + "type": "mcp_approval_response", + "approval_request_id": content.id, + "approve": content.approved, + } + case "mcp_server_tool_call": + if not content.call_id: + return {} + return { + "type": "mcp_call", + "id": content.call_id, + "server_label": content.server_name or "", + "name": content.tool_name or "", + "arguments": self._stringify_mcp_arguments(content.arguments), + } + case "mcp_server_tool_result": + if not content.call_id: + return {} + return { + _AF_MCP_PENDING_OUTPUT_KEY: True, + "call_id": content.call_id, + "output": self._stringify_mcp_output(content.output), + } + case "hosted_file": + # `input_file` is an input-only content type in the Responses API and is rejected + # inside an assistant message. Hosted-file content on an assistant message + # represents a citation produced by a hosted tool (e.g., file_search) and cannot be + # meaningfully replayed as input — drop it. The accompanying text annotations carry + # the citation context for round-tripping. + if role == "assistant": + return {} + return { + "type": "input_file", + "file_id": content.file_id, + } + case _: # should catch UsageDetails and ErrorContent and HostedVectorStoreContent + logger.debug("Unsupported content type passed (type: %s)", content.type) + return {} + + @staticmethod + def _to_local_shell_output_payload(content: Content) -> str: + """Convert function tool output to the local shell JSON payload format.""" + payload: dict[str, Any] + if isinstance(content.result, Mapping): + payload = dict(content.result) # type: ignore[assignment] + else: + payload = { + "stdout": "" if content.result is None else str(content.result), + } + if content.exception is not None and "stderr" not in payload: + payload["stderr"] = str(content.exception) + if "exit_code" not in payload: + payload["exit_code"] = 1 if content.exception else 0 + return json.dumps(payload, ensure_ascii=False) + + @staticmethod + def _to_shell_call_output_payload(content: Content) -> list[dict[str, Any]]: + """Convert function tool output to shell_call_output payload format.""" + payload: dict[str, Any] + if isinstance(content.result, Mapping): + payload = dict(content.result) # type: ignore[assignment] + else: + payload = { + "stdout": "" if content.result is None else str(content.result), + } + if content.exception is not None and "stderr" not in payload: + payload["stderr"] = str(content.exception) + + # Pass through native payload shape when tool already returns shell output entries. + direct_output = payload.get("output") + if isinstance(direct_output, list) and all(isinstance(item, Mapping) for item in direct_output): # type: ignore[reportUnknownMemberType] + return [dict(item) for item in direct_output] # type: ignore[reportUnknownMemberType] + + stdout = str(payload.get("stdout", "")) + stderr = str(payload.get("stderr", "")) + timed_out = bool(payload.get("timed_out", False)) + if timed_out: + outcome: dict[str, Any] = {"type": "timeout"} + else: + exit_code_raw = payload.get("exit_code") + try: + exit_code = int(exit_code_raw) if exit_code_raw is not None else (1 if content.exception else 0) + except (TypeError, ValueError): + exit_code = 1 if content.exception else 0 + outcome = {"type": "exit", "exit_code": exit_code} + return [ + { + "stdout": stdout, + "stderr": stderr, + "outcome": outcome, + } + ] + + @staticmethod + def _join_shell_commands(commands: Sequence[str]) -> str: + """Join shell commands into a single executable command string.""" + return "\n".join(command for command in commands if command).strip() + + @staticmethod + def _stringify_mcp_arguments(arguments: Any) -> str: + """Render hosted-MCP tool-call arguments as a JSON string for the Responses API.""" + if arguments is None: + return "" + if isinstance(arguments, str): + return arguments + try: + return json.dumps(arguments) + except (TypeError, ValueError): + return str(arguments) + + @staticmethod + def _stringify_mcp_output(output: Any) -> str: + """Render a hosted-MCP tool-call result into the string `mcp_call.output` field. + + Accepts a string, a list of text-bearing Content objects (the form + the chat client produces when parsing an `mcp_call` Responses item), + or any other value. List entries that are dicts with the canonical + MCP text-content shape (`{"text": "..."}`) are unwrapped to their + text. Anything else falls back to JSON encoding rather than Python + `repr`, so the wire payload stays parseable for downstream callers. + """ + if output is None: + return "" + if isinstance(output, str): + return output + if isinstance(output, Sequence) and not isinstance(output, (str, bytes, bytearray)): + # cast is for pyright (reportUnknownVariableType); mypy considers + # it redundant after the isinstance narrowing. + entries = cast(Sequence[Any], output) # type: ignore[redundant-cast] + parts: list[str] = [] + for entry in entries: + if isinstance(entry, str): + parts.append(entry) + continue + text = getattr(entry, "text", None) + if isinstance(text, str): + parts.append(text) + continue + if isinstance(entry, Mapping): + mapping_text = cast(Any, entry).get("text") + if isinstance(mapping_text, str): + parts.append(mapping_text) + continue + parts.append(json.dumps(entry, default=str)) + return "".join(parts) + return json.dumps(output, default=str) + + @staticmethod + def _coalesce_pending_mcp_results(items: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Merge pending hosted-MCP result markers onto matching mcp_call input items. + + See `_AF_MCP_PENDING_OUTPUT_KEY`. The Responses API expects a single + `mcp_call` input item carrying both `arguments` and `output`, so a + result Content cannot be its own input item. Any unmatched markers + are dropped (debug-logged); surfacing them as standalone items + would produce the orphan `function_call_output` / `mcp_call_output` + the API rejects. + """ + out: list[dict[str, Any]] = [] + for item in items: + if item.get(_AF_MCP_PENDING_OUTPUT_KEY): + target_call_id = item.get("call_id") + target = next( + ( + existing + for existing in reversed(out) + if existing.get("type") == "mcp_call" and existing.get("id") == target_call_id + ), + None, + ) + if target is not None: + if target.get("output") is None: + target["output"] = item.get("output") + else: + logger.debug( + "Dropping orphan mcp_server_tool_result for call_id=%s; " + "no matching mcp_call appeared in input.", + target_call_id, + ) + continue + out.append(item) + return out + + @staticmethod + def _serialize_provider_payload(value: Any) -> Any: + """Convert OpenAI SDK objects into JSON-serializable Python values.""" + if isinstance(value, BaseModel): + return value.model_dump(mode="json", exclude_none=True) + if isinstance(value, Mapping): + return {str(key): RawOpenAIChatClient._serialize_provider_payload(item) for key, item in value.items()} # type: ignore[reportUnknownVariableType] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return [RawOpenAIChatClient._serialize_provider_payload(item) for item in value] # type: ignore[reportUnknownVariableType] + return value + + @staticmethod + def _get_search_tool_name(item_type: str) -> str: + """Map OpenAI search output item types to unified content tool names.""" + return "web_search" if item_type == "web_search_call" else "file_search" + + def _parse_search_tool_call_content(self, item: Any) -> Content: + """Create unified search tool call content from an OpenAI search output item.""" + item_type = getattr(item, "type", "") + call_id = getattr(item, "id", None) or getattr(item, "call_id", None) or "" + if item_type == "web_search_call": + arguments = self._serialize_provider_payload(getattr(item, "action", None)) + else: + arguments = {"queries": list(getattr(item, "queries", []) or [])} + return Content.from_search_tool_call( + call_id=call_id, + tool_name=self._get_search_tool_name(item_type), + arguments=arguments, + status=getattr(item, "status", None), + raw_representation=item, + ) + + def _parse_search_tool_result_content(self, item: Any) -> Content: + """Create unified search tool result content from an OpenAI search output item.""" + item_type = getattr(item, "type", "") + call_id = getattr(item, "id", None) or getattr(item, "call_id", None) or "" + if item_type == "web_search_call": + result = {"action": self._serialize_provider_payload(getattr(item, "action", None))} + else: + result = {"results": self._serialize_provider_payload(getattr(item, "results", None))} + return Content.from_search_tool_result( + call_id=call_id, + tool_name=self._get_search_tool_name(item_type), + result=result, + status=getattr(item, "status", None), + raw_representation=item, + ) + + # region Parse methods + def _parse_response_from_openai( + self, + response: OpenAIResponse | ParsedResponse[BaseModel], + options: dict[str, Any], + ) -> ChatResponse: + """Parse an OpenAI Responses API response into a ChatResponse.""" + structured_response: BaseModel | None = response.output_parsed if isinstance(response, ParsedResponse) else None # type: ignore[reportUnknownMemberType] + + metadata: dict[str, Any] = response.metadata or {} + contents: list[Content] = [] + local_shell_tool_name = self._get_local_shell_tool_name(options.get("tools")) + for item in response.output: # type: ignore[reportUnknownMemberType] + match item.type: + # types: + # ParsedResponseOutputMessage[Unknown] | + # ParsedResponseFunctionToolCall | + # ResponseFileSearchToolCall | + # ResponseFunctionWebSearch | + # ResponseComputerToolCall | + # ResponseReasoningItem | + # MCPCall | + # MCPApprovalRequest | + # ImageGenerationCall | + # LocalShellCall | + # LocalShellCallAction | + # MCPListTools | + # ResponseCodeInterpreterToolCall | + # ResponseCustomToolCall | + # ParsedResponseOutputMessage[BaseModel] | + # ResponseOutputMessage | + # ResponseFunctionToolCall + case "message": # ResponseOutputMessage + for message_content in item.content: # type: ignore[reportMissingTypeArgument] + match message_content.type: + case "output_text": + text_content = Content.from_text( + text=message_content.text, + raw_representation=message_content, # type: ignore[reportUnknownArgumentType] + ) + metadata.update(self._get_metadata_from_response(message_content)) + if message_content.annotations: + text_content.annotations = [] + for annotation in message_content.annotations: + match annotation.type: + case "file_path": + text_content.annotations.append( # pyright: ignore[reportUnknownMemberType] + Annotation( + type="citation", + file_id=annotation.file_id, + additional_properties={ + "index": annotation.index, + }, + raw_representation=annotation, + ) + ) + case "file_citation": + text_content.annotations.append( # pyright: ignore[reportUnknownMemberType] + Annotation( + type="citation", + url=annotation.filename, + file_id=annotation.file_id, + raw_representation=annotation, + additional_properties={ + "index": annotation.index, + }, + ) + ) + case "url_citation": + text_content.annotations.append( # pyright: ignore[reportUnknownMemberType] + Annotation( + type="citation", + title=annotation.title, + url=annotation.url, + annotated_regions=[ + TextSpanRegion( + type="text_span", + start_index=annotation.start_index, + end_index=annotation.end_index, + ) + ], + raw_representation=annotation, + ) + ) + case "container_file_citation": + text_content.annotations.append( # pyright: ignore[reportUnknownMemberType] + Annotation( + type="citation", + file_id=annotation.file_id, + url=annotation.filename, + additional_properties={ + "container_id": annotation.container_id, + }, + annotated_regions=[ + TextSpanRegion( + type="text_span", + start_index=annotation.start_index, + end_index=annotation.end_index, + ) + ], + raw_representation=annotation, + ) + ) + case _: + logger.debug( + "Unparsed annotation type: %s", + annotation.type, + ) + contents.append(text_content) + case "refusal": + contents.append( + Content.from_text( + text=message_content.refusal, + raw_representation=message_content, + ) + ) + case "reasoning": # ResponseOutputReasoning + added_reasoning = False + if item_content := getattr(item, "content", None): + for index, reasoning_content in enumerate(item_content): + additional_properties: dict[str, Any] = {} + if hasattr(item, "summary") and item.summary and index < len(item.summary): + additional_properties["summary"] = item.summary[index] + contents.append( + Content.from_text_reasoning( + id=item.id, + text=reasoning_content.text, + raw_representation=reasoning_content, + additional_properties=additional_properties or None, + ) + ) + added_reasoning = True + if item_summary := getattr(item, "summary", None): + for summary in item_summary: + contents.append( + Content.from_text_reasoning( + id=item.id, + text=summary.text, + raw_representation=summary, # type: ignore[arg-type] + ) + ) + added_reasoning = True + if not added_reasoning: + # Reasoning item with no visible text (e.g. encrypted reasoning). + # Always emit an empty marker so co-occurrence detection can be done + additional_properties_empty: dict[str, Any] = {} + if encrypted := getattr(item, "encrypted_content", None): + additional_properties_empty["encrypted_content"] = encrypted + contents.append( + Content.from_text_reasoning( + id=item.id, + text="", + raw_representation=item, + additional_properties=additional_properties_empty or None, + ) + ) + case "code_interpreter_call": # ResponseOutputCodeInterpreterCall + call_id = getattr(item, "call_id", None) or getattr(item, "id", None) + outputs: list[Content] = [] + if item_outputs := getattr(item, "outputs", None): + for code_output in item_outputs: + if getattr(code_output, "type", None) == "logs": + outputs.append( + Content.from_text( + text=code_output.logs, + raw_representation=code_output, + ) + ) + elif getattr(code_output, "type", None) == "image": + outputs.append( + Content.from_uri( + uri=code_output.url, + raw_representation=code_output, + media_type="image", + ) + ) + if code := getattr(item, "code", None): + contents.append( + Content.from_code_interpreter_tool_call( + call_id=call_id, + inputs=[Content.from_text(text=code, raw_representation=item)], + raw_representation=item, + ) + ) + contents.append( + Content.from_code_interpreter_tool_result( + call_id=call_id, + outputs=outputs, + raw_representation=item, + ) + ) + case "function_call": # ResponseOutputFunctionCall + contents.append( + Content.from_function_call( + call_id=item.call_id, + name=item.name, + arguments=item.arguments, + additional_properties={"fc_id": item.id, "status": item.status}, + raw_representation=item, + ) + ) + case "web_search_call" | "file_search_call": + contents.append(self._parse_search_tool_call_content(item)) + contents.append(self._parse_search_tool_result_content(item)) + case "mcp_approval_request": # ResponseOutputMcpApprovalRequest + contents.append( + Content.from_function_approval_request( + id=item.id, + function_call=Content.from_function_call( + call_id=item.id, + name=item.name, + arguments=item.arguments, + additional_properties={"server_label": item.server_label}, + raw_representation=item, + ), + ) + ) + case "mcp_call": + call_id = getattr(item, "id", None) or getattr(item, "call_id", None) or "" + contents.append( + Content.from_mcp_server_tool_call( + call_id=call_id, + tool_name=item.name, + server_name=item.server_label, + arguments=item.arguments, + raw_representation=item, + ) + ) + if item.output is not None: + contents.append( + Content.from_mcp_server_tool_result( + call_id=call_id, + output=[Content.from_text(text=item.output)], + raw_representation=item, + ) + ) + case "image_generation_call": # ResponseOutputImageGenerationCall + image_output: Content | None = None + if item.result is not None: + # item.result contains raw base64 string + # so we call detect_media_type_from_base64 to get the media type and fallback to image/png + image_output = Content.from_uri( + uri=f"data:{detect_media_type_from_base64(data_str=item.result) or 'image/png'}" + f";base64,{item.result}", + raw_representation=item.result, + ) + image_id = item.id + contents.append( + Content.from_image_generation_tool_call( + image_id=image_id, + raw_representation=item, + ) + ) + contents.append( + Content.from_image_generation_tool_result( + image_id=image_id, + outputs=image_output, + raw_representation=item, + ) + ) + case "shell_call": # ResponseFunctionShellToolCall + shell_call_id = item.call_id if hasattr(item, "call_id") else "" + shell_commands: list[str] = [] + shell_timeout_ms: int | None = None + shell_max_output: int | None = None + if action := getattr(item, "action", None): + shell_commands = list(getattr(action, "commands", []) or []) + shell_timeout_ms = getattr(action, "timeout_ms", None) + shell_max_output = getattr(action, "max_output_length", None) + if local_shell_tool_name: + command_text = self._join_shell_commands(shell_commands) + contents.append( + Content.from_function_call( + call_id=shell_call_id, + name=local_shell_tool_name, + arguments=json.dumps({"command": command_text}), + additional_properties={ + OPENAI_SHELL_OUTPUT_TYPE_KEY: OPENAI_SHELL_OUTPUT_TYPE_SHELL_CALL, + OPENAI_LOCAL_SHELL_COMMAND_PARTS_KEY: shell_commands, + }, + raw_representation=item, + ) + ) + else: + contents.append( + Content.from_shell_tool_call( + call_id=shell_call_id, + commands=shell_commands, + timeout_ms=shell_timeout_ms, + max_output_length=shell_max_output, + status=getattr(item, "status", None), + raw_representation=item, + ) + ) + case "local_shell_call": + local_call_id = getattr(item, "call_id", None) or "" + local_command_parts = list(getattr(getattr(item, "action", None), "command", []) or []) + local_command = shlex.join(local_command_parts) if local_command_parts else "" + if local_shell_tool_name: + contents.append( + Content.from_function_call( + call_id=local_call_id, + name=local_shell_tool_name, + arguments=json.dumps({"command": local_command}), + additional_properties={ + OPENAI_SHELL_OUTPUT_TYPE_KEY: OPENAI_SHELL_OUTPUT_TYPE_LOCAL_SHELL_CALL, + OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY: getattr(item, "id", None), + OPENAI_LOCAL_SHELL_COMMAND_PARTS_KEY: local_command_parts, + }, + raw_representation=item, + ) + ) + else: + contents.append( + Content.from_shell_tool_call( + call_id=local_call_id, + commands=[local_command] if local_command else [], + timeout_ms=getattr(getattr(item, "action", None), "timeout_ms", None), + status=getattr(item, "status", None), + raw_representation=item, + ) + ) + case "shell_call_output": # ResponseFunctionShellToolCallOutput + shell_output_call_id = item.call_id if hasattr(item, "call_id") else "" + shell_outputs: list[Content] = [] + for shell_out in getattr(item, "output", []) or []: + s_exit_code: int | None = None + s_timed_out: bool | None = None + if outcome := getattr(shell_out, "outcome", None): + if getattr(outcome, "type", None) == "exit": + s_exit_code = getattr(outcome, "exit_code", None) + s_timed_out = False + elif getattr(outcome, "type", None) == "timeout": + s_timed_out = True + shell_outputs.append( + Content.from_shell_command_output( + stdout=getattr(shell_out, "stdout", None), + stderr=getattr(shell_out, "stderr", None), + exit_code=s_exit_code, + timed_out=s_timed_out, + raw_representation=shell_out, + ) + ) + contents.append( + Content.from_shell_tool_result( + call_id=shell_output_call_id, + outputs=shell_outputs, + max_output_length=getattr(item, "max_output_length", None), + raw_representation=item, + ) + ) + case _: + logger.debug("Unparsed output of type: %s: %s", item.type, item) + response_message = Message(role="assistant", contents=contents) + args: dict[str, Any] = { + "response_id": response.id, + "created_at": datetime.fromtimestamp(response.created_at, tz=timezone.utc).strftime( + "%Y-%m-%dT%H:%M:%S.%fZ" + ), + "messages": response_message, + "model": response.model, + "additional_properties": metadata, + "raw_representation": response, + } + + if conversation_id := self._get_conversation_id(response, options.get("store")): # pyright: ignore[reportUnknownArgumentType] + args["conversation_id"] = conversation_id + if response.usage and (usage_details := self._parse_usage_from_openai(response.usage)): + args["usage_details"] = usage_details + if structured_response: + args["value"] = structured_response + elif response_format := options.get("response_format"): + args["response_format"] = response_format + # Set continuation_token when background operation is still in progress + if response.status and response.status in ("in_progress", "queued"): + args["continuation_token"] = OpenAIContinuationToken(response_id=response.id) + return ChatResponse(**args) + + def _parse_chunk_from_openai( + self, + event: OpenAIResponseStreamEvent, + options: dict[str, Any], + function_call_ids: dict[int, tuple[str, str]], + seen_reasoning_delta_item_ids: set[str] | None = None, + ) -> ChatResponseUpdate: + """Parse an OpenAI Responses API streaming event into a ChatResponseUpdate.""" + metadata: dict[str, Any] = {} + contents: list[Content] = [] + local_shell_tool_name = self._get_local_shell_tool_name(options.get("tools")) + conversation_id: str | None = None + response_id: str | None = None + created_at: str | None = None + continuation_token: OpenAIContinuationToken | None = None + model = self.model + match event.type: + # types: + # ResponseAudioDeltaEvent, + # ResponseAudioDoneEvent, + # ResponseAudioTranscriptDeltaEvent, + # ResponseAudioTranscriptDoneEvent, + # ResponseCodeInterpreterCallCodeDeltaEvent, + # ResponseCodeInterpreterCallCodeDoneEvent, + # ResponseCodeInterpreterCallCompletedEvent, + # ResponseCodeInterpreterCallInProgressEvent, + # ResponseCodeInterpreterCallInterpretingEvent, + # ResponseCompletedEvent, + # ResponseContentPartAddedEvent, + # ResponseContentPartDoneEvent, + # ResponseCreatedEvent, + # ResponseErrorEvent, + # ResponseFileSearchCallCompletedEvent, + # ResponseFileSearchCallInProgressEvent, + # ResponseFileSearchCallSearchingEvent, + # ResponseFunctionCallArgumentsDeltaEvent, + # ResponseFunctionCallArgumentsDoneEvent, + # ResponseInProgressEvent, + # ResponseFailedEvent, + # ResponseIncompleteEvent, + # ResponseOutputItemAddedEvent, + # ResponseOutputItemDoneEvent, + # ResponseReasoningSummaryPartAddedEvent, + # ResponseReasoningSummaryPartDoneEvent, + # ResponseReasoningSummaryTextDeltaEvent, + # ResponseReasoningSummaryTextDoneEvent, + # ResponseReasoningTextDeltaEvent, + # ResponseReasoningTextDoneEvent, + # ResponseRefusalDeltaEvent, + # ResponseRefusalDoneEvent, + # ResponseTextDeltaEvent, + # ResponseTextDoneEvent, + # ResponseWebSearchCallCompletedEvent, + # ResponseWebSearchCallInProgressEvent, + # ResponseWebSearchCallSearchingEvent, + # ResponseImageGenCallCompletedEvent, + # ResponseImageGenCallGeneratingEvent, + # ResponseImageGenCallInProgressEvent, + # ResponseImageGenCallPartialImageEvent, + # ResponseMcpCallArgumentsDeltaEvent, + # ResponseMcpCallArgumentsDoneEvent, + # ResponseMcpCallCompletedEvent, + # ResponseMcpCallFailedEvent, + # ResponseMcpCallInProgressEvent, + # ResponseMcpListToolsCompletedEvent, + # ResponseMcpListToolsFailedEvent, + # ResponseMcpListToolsInProgressEvent, + # ResponseOutputTextAnnotationAddedEvent, + # ResponseQueuedEvent, + # ResponseCustomToolCallInputDeltaEvent, + # ResponseCustomToolCallInputDoneEvent, + case "response.content_part.added": + event_part = event.part + match event_part.type: + case "output_text": + contents.append(Content.from_text(text=event_part.text, raw_representation=event)) + metadata.update(self._get_metadata_from_response(event_part)) + case "refusal": + contents.append(Content.from_text(text=event_part.refusal, raw_representation=event)) + case _: + pass + case "response.output_text.delta": + contents.append(Content.from_text(text=event.delta, raw_representation=event)) + metadata.update(self._get_metadata_from_response(event)) + case "response.reasoning_text.delta": + if seen_reasoning_delta_item_ids is not None: + seen_reasoning_delta_item_ids.add(event.item_id) + contents.append( + Content.from_text_reasoning( + id=event.item_id, + text=event.delta, + raw_representation=event, + ) + ) + metadata.update(self._get_metadata_from_response(event)) + case "response.reasoning_text.done": + # Done event carries the full accumulated text. Emit it only as a + # fallback when no delta was already received for this item_id, to + # avoid duplicating content in downstream accumulators (e.g. ag-ui). + if seen_reasoning_delta_item_ids is None or event.item_id not in seen_reasoning_delta_item_ids: + contents.append( + Content.from_text_reasoning( + id=event.item_id, + text=event.text, + raw_representation=event, + ) + ) + metadata.update(self._get_metadata_from_response(event)) + case "response.reasoning_summary_text.delta": + if seen_reasoning_delta_item_ids is not None: + seen_reasoning_delta_item_ids.add(event.item_id) + contents.append( + Content.from_text_reasoning( + id=event.item_id, + text=event.delta, + raw_representation=event, + ) + ) + metadata.update(self._get_metadata_from_response(event)) + case "response.reasoning_summary_text.done": + # Done event carries the full accumulated text. Emit it only as a + # fallback when no delta was already received for this item_id, to + # avoid duplicating content in downstream accumulators (e.g. ag-ui). + if seen_reasoning_delta_item_ids is None or event.item_id not in seen_reasoning_delta_item_ids: + contents.append( + Content.from_text_reasoning( + id=event.item_id, + text=event.text, + raw_representation=event, + ) + ) + metadata.update(self._get_metadata_from_response(event)) + case "response.code_interpreter_call_code.delta": + call_id = getattr(event, "call_id", None) or getattr(event, "id", None) or event.item_id + ci_additional_properties = { + "output_index": event.output_index, + "sequence_number": event.sequence_number, + "item_id": event.item_id, + } + contents.append( + Content.from_code_interpreter_tool_call( + call_id=call_id, + inputs=[ + Content.from_text( + text=event.delta, + raw_representation=event, + additional_properties=ci_additional_properties, + ) + ], + raw_representation=event, + additional_properties=ci_additional_properties, + ) + ) + metadata.update(self._get_metadata_from_response(event)) + # NOTE: Unlike reasoning done events, code_interpreter done events always + # emit content because downstream consumers do not accumulate + # code_interpreter deltas the same way. + case "response.code_interpreter_call_code.done": + call_id = getattr(event, "call_id", None) or getattr(event, "id", None) or event.item_id + ci_additional_properties = { + "output_index": event.output_index, + "sequence_number": event.sequence_number, + "item_id": event.item_id, + } + contents.append( + Content.from_code_interpreter_tool_call( + call_id=call_id, + inputs=[ + Content.from_text( + text=event.code, + raw_representation=event, + additional_properties=ci_additional_properties, + ) + ], + raw_representation=event, + additional_properties=ci_additional_properties, + ) + ) + metadata.update(self._get_metadata_from_response(event)) + case "response.created": + response_id = event.response.id + conversation_id = self._get_conversation_id(event.response, options.get("store")) + if event.response.status and event.response.status in ( + "in_progress", + "queued", + ): + continuation_token = OpenAIContinuationToken(response_id=event.response.id) + case "response.in_progress": + response_id = event.response.id + conversation_id = self._get_conversation_id(event.response, options.get("store")) + continuation_token = OpenAIContinuationToken(response_id=event.response.id) + case "response.completed": + response_id = event.response.id + conversation_id = self._get_conversation_id(event.response, options.get("store")) + model = event.response.model + created_at = datetime.fromtimestamp(event.response.created_at, tz=timezone.utc).strftime( + "%Y-%m-%dT%H:%M:%S.%fZ" + ) + if event.response.usage: + usage = self._parse_usage_from_openai(event.response.usage) + if usage: + contents.append(Content.from_usage(usage_details=usage, raw_representation=event)) + case "response.output_item.added": + event_item = event.item + match event_item.type: + # types: + # ResponseOutputMessage, + # ResponseFileSearchToolCall, + # ResponseFunctionToolCall, + # ResponseFunctionWebSearch, + # ResponseComputerToolCall, + # ResponseReasoningItem, + # ImageGenerationCall, + # ResponseCodeInterpreterToolCall, + # LocalShellCall, + # McpCall, + # McpListTools, + # McpApprovalRequest, + # ResponseCustomToolCall, + case "function_call": + function_call_ids[event.output_index] = ( + event_item.call_id, + event_item.name, + ) + case "mcp_approval_request": + contents.append( + Content.from_function_approval_request( + id=event_item.id, + function_call=Content.from_function_call( + call_id=event_item.id, + name=event_item.name, + arguments=event_item.arguments, + additional_properties={"server_label": event_item.server_label}, + raw_representation=event_item, + ), + ) + ) + case "mcp_call": + call_id = getattr(event_item, "id", None) or getattr(event_item, "call_id", None) or "" + contents.append( + Content.from_mcp_server_tool_call( + call_id=call_id, + tool_name=getattr(event_item, "name", "") or "", + server_name=getattr(event_item, "server_label", None), + arguments=getattr(event_item, "arguments", None), + raw_representation=event_item, + ) + ) + # Result deferred to response.output_item.done + case "code_interpreter_call": # ResponseOutputCodeInterpreterCall + call_id = getattr(event_item, "call_id", None) or getattr(event_item, "id", None) + outputs: list[Content] = [] + if hasattr(event_item, "outputs") and event_item.outputs: + for code_output in event_item.outputs: + if getattr(code_output, "type", None) == "logs": + outputs.append( + Content.from_text( + text=cast(Any, code_output).logs, + raw_representation=code_output, + ) + ) + elif getattr(code_output, "type", None) == "image": + outputs.append( + Content.from_uri( + uri=cast(Any, code_output).url, + raw_representation=code_output, + media_type="image", + ) + ) + if hasattr(event_item, "code") and event_item.code: + contents.append( + Content.from_code_interpreter_tool_call( + call_id=call_id, + inputs=[ + Content.from_text( + text=event_item.code, + raw_representation=event_item, + ) + ], + raw_representation=event_item, + ) + ) + contents.append( + Content.from_code_interpreter_tool_result( + call_id=call_id, + outputs=outputs, + raw_representation=event_item, + ) + ) + case "shell_call": # ResponseFunctionShellToolCall + s_call_id = getattr(event_item, "call_id", None) or "" + s_commands: list[str] = [] + s_timeout_ms: int | None = None + s_max_output: int | None = None + if s_action := getattr(event_item, "action", None): + s_commands = list(getattr(s_action, "commands", []) or []) + s_timeout_ms = getattr(s_action, "timeout_ms", None) + s_max_output = getattr(s_action, "max_output_length", None) + if local_shell_tool_name: + command_text = self._join_shell_commands(s_commands) + contents.append( + Content.from_function_call( + call_id=s_call_id, + name=local_shell_tool_name, + arguments=json.dumps({"command": command_text}), + additional_properties={ + OPENAI_SHELL_OUTPUT_TYPE_KEY: OPENAI_SHELL_OUTPUT_TYPE_SHELL_CALL, + OPENAI_LOCAL_SHELL_COMMAND_PARTS_KEY: s_commands, + }, + raw_representation=event_item, + ) + ) + else: + contents.append( + Content.from_shell_tool_call( + call_id=s_call_id, + commands=s_commands, + timeout_ms=s_timeout_ms, + max_output_length=s_max_output, + status=getattr(event_item, "status", None), + raw_representation=event_item, + ) + ) + case "local_shell_call": + local_call_id = getattr(event_item, "call_id", None) or "" + local_command_parts = list(getattr(getattr(event_item, "action", None), "command", []) or []) + local_command = shlex.join(local_command_parts) if local_command_parts else "" + if local_shell_tool_name: + contents.append( + Content.from_function_call( + call_id=local_call_id, + name=local_shell_tool_name, + arguments=json.dumps({"command": local_command}), + additional_properties={ + OPENAI_SHELL_OUTPUT_TYPE_KEY: OPENAI_SHELL_OUTPUT_TYPE_LOCAL_SHELL_CALL, + OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY: getattr(event_item, "id", None), + OPENAI_LOCAL_SHELL_COMMAND_PARTS_KEY: local_command_parts, + }, + raw_representation=event_item, + ) + ) + else: + contents.append( + Content.from_shell_tool_call( + call_id=local_call_id, + commands=[local_command] if local_command else [], + timeout_ms=getattr( + getattr(event_item, "action", None), + "timeout_ms", + None, + ), + status=getattr(event_item, "status", None), + raw_representation=event_item, + ) + ) + case "shell_call_output": # ResponseFunctionShellToolCallOutput + s_out_call_id = getattr(event_item, "call_id", None) or "" + s_outputs: list[Content] = [] + for s_out in getattr(event_item, "output", []) or []: + s_exit_code: int | None = None + s_timed_out: bool | None = None + if s_outcome := getattr(s_out, "outcome", None): + if getattr(s_outcome, "type", None) == "exit": + s_exit_code = getattr(s_outcome, "exit_code", None) + s_timed_out = False + elif getattr(s_outcome, "type", None) == "timeout": + s_timed_out = True + s_outputs.append( + Content.from_shell_command_output( + stdout=getattr(s_out, "stdout", None), + stderr=getattr(s_out, "stderr", None), + exit_code=s_exit_code, + timed_out=s_timed_out, + raw_representation=s_out, + ) + ) + contents.append( + Content.from_shell_tool_result( + call_id=s_out_call_id, + outputs=s_outputs, + max_output_length=getattr(event_item, "max_output_length", None), + raw_representation=event_item, + ) + ) + case "reasoning": # ResponseOutputReasoning + reasoning_id = getattr(event_item, "id", None) + added_reasoning = False + if hasattr(event_item, "content") and event_item.content: + for index, reasoning_content in enumerate(event_item.content): + additional_properties: dict[str, Any] = {} + if ( + hasattr(event_item, "summary") + and event_item.summary + and index < len(event_item.summary) + ): + additional_properties["summary"] = event_item.summary[index] + contents.append( + Content.from_text_reasoning( + id=reasoning_id or None, + text=reasoning_content.text, + raw_representation=reasoning_content, + additional_properties=additional_properties or None, + ) + ) + added_reasoning = True + if not added_reasoning: + # Reasoning item with no visible text (e.g. encrypted reasoning). + # Always emit an empty marker so co-occurrence detection can occur. + additional_properties_empty: dict[str, Any] = {} + if encrypted := getattr(event_item, "encrypted_content", None): + additional_properties_empty["encrypted_content"] = encrypted + contents.append( + Content.from_text_reasoning( + id=reasoning_id or None, + text="", + raw_representation=event_item, + additional_properties=additional_properties_empty or None, + ) + ) + case "web_search_call" | "file_search_call": + contents.append(self._parse_search_tool_call_content(event_item)) + case _: + logger.debug("Unparsed event of type: %s: %s", event.type, event) + case ( + "response.web_search_call.in_progress" + | "response.web_search_call.searching" + | "response.web_search_call.completed" + | "response.file_search_call.in_progress" + | "response.file_search_call.searching" + | "response.file_search_call.completed" + ): + pass + case "response.function_call_arguments.delta": + call_id, name = function_call_ids.get(event.output_index, (None, None)) + if call_id and name: + contents.append( + Content.from_function_call( + call_id=call_id, + name=name, + arguments=event.delta, + additional_properties={ + "output_index": event.output_index, + "fc_id": event.item_id, + }, + 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 + image_output = Content.from_uri( + uri=f"data:{detect_media_type_from_base64(data_str=image_base64) or 'image/png'}" + f";base64,{image_base64}", + additional_properties={ + "partial_image_index": partial_index, + "is_partial_image": True, + }, + raw_representation=event, + ) + + image_id = getattr(event, "item_id", None) + contents.append( + Content.from_image_generation_tool_call( + image_id=image_id, + raw_representation=event, + ) + ) + contents.append( + Content.from_image_generation_tool_result( + image_id=image_id, + outputs=image_output, + raw_representation=event, + ) + ) + case "response.output_text.annotation.added": + # Handle streaming text annotations (file citations, file paths, etc.) + annotation: Any = event.annotation + + def _get_ann_value(key: str) -> Any: + """Extract value from annotation (dict or object).""" + if isinstance(annotation, dict): + return cast("dict[str, Any]", annotation).get(key) + return getattr(annotation, key, None) + + ann_type = _get_ann_value("type") + ann_file_id = _get_ann_value("file_id") + # Hosted-file citations attach as text annotations (matching the non-streaming path) + # so they don't roundtrip as standalone `input_file` items in assistant history. + if ann_type == "file_path": + if ann_file_id: + annotation_obj = Annotation( + type="citation", + file_id=str(ann_file_id), + additional_properties={ + "annotation_index": event.annotation_index, + "index": _get_ann_value("index"), + }, + raw_representation=annotation, + ) + contents.append( + Content.from_text(text="", annotations=[annotation_obj], raw_representation=event) + ) + elif ann_type == "file_citation": + if ann_file_id: + ann_filename = _get_ann_value("filename") + annotation_obj = Annotation( + type="citation", + file_id=str(ann_file_id), + url=ann_filename, + additional_properties={ + "annotation_index": event.annotation_index, + "index": _get_ann_value("index"), + }, + raw_representation=annotation, + ) + contents.append( + Content.from_text(text="", annotations=[annotation_obj], raw_representation=event) + ) + elif ann_type == "container_file_citation": + if ann_file_id: + ann_filename = _get_ann_value("filename") + ann_start = _get_ann_value("start_index") + ann_end = _get_ann_value("end_index") + annotation_obj = Annotation( + type="citation", + file_id=str(ann_file_id), + url=ann_filename, + additional_properties={ + "annotation_index": event.annotation_index, + "container_id": _get_ann_value("container_id"), + }, + raw_representation=annotation, + ) + if ann_start is not None and ann_end is not None: + annotation_obj["annotated_regions"] = [ + TextSpanRegion( + type="text_span", + start_index=ann_start, + end_index=ann_end, + ) + ] + contents.append( + Content.from_text(text="", annotations=[annotation_obj], raw_representation=event) + ) + elif ann_type == "url_citation": + ann_url = _get_ann_value("url") + if ann_url: + ann_start = _get_ann_value("start_index") + ann_end = _get_ann_value("end_index") + annotation_obj = Annotation( + type="citation", + title=_get_ann_value("title") or "", + url=str(ann_url), + additional_properties={"annotation_index": event.annotation_index}, + raw_representation=annotation, + ) + if ann_start is not None and ann_end is not None: + annotation_obj["annotated_regions"] = [ + TextSpanRegion( + type="text_span", + start_index=ann_start, + end_index=ann_end, + ) + ] + contents.append( + Content.from_text(text="", annotations=[annotation_obj], raw_representation=event) + ) + else: + logger.debug("Unparsed annotation type in streaming: %s", ann_type) + case "response.output_item.done": + done_item = event.item + if getattr(done_item, "type", None) == "mcp_call": + call_id = getattr(done_item, "id", None) or getattr(done_item, "call_id", None) or "" + output_text = getattr(done_item, "output", None) + parsed_output: list[Content] | None = ( + [Content.from_text(text=output_text)] if isinstance(output_text, str) else None + ) + contents.append( + Content.from_mcp_server_tool_result( + call_id=call_id, + output=parsed_output, + raw_representation=done_item, + ) + ) + elif getattr(done_item, "type", None) in ("web_search_call", "file_search_call"): + contents.append(self._parse_search_tool_result_content(done_item)) + case _: + logger.debug("Unparsed event of type: %s: %s", event.type, event) + + return ChatResponseUpdate( + contents=contents, + conversation_id=conversation_id, + response_id=response_id, + role="assistant", + model=model, + created_at=created_at, + continuation_token=continuation_token, + additional_properties=metadata, + raw_representation=event, + ) + + def _parse_usage_from_openai(self, usage: ResponseUsage) -> UsageDetails | None: + details = UsageDetails( + input_token_count=usage.input_tokens, + output_token_count=usage.output_tokens, + total_token_count=usage.total_tokens, + ) + if usage.input_tokens_details and usage.input_tokens_details.cached_tokens: + details["openai.cached_input_tokens"] = usage.input_tokens_details.cached_tokens # type: ignore[typeddict-unknown-key] + if usage.output_tokens_details and usage.output_tokens_details.reasoning_tokens: + details["openai.reasoning_tokens"] = usage.output_tokens_details.reasoning_tokens # type: ignore[typeddict-unknown-key] + return details + + def _get_metadata_from_response(self, output: Any) -> dict[str, Any]: + """Get metadata from a chat choice.""" + if logprobs := getattr(output, "logprobs", None): + return { + "logprobs": logprobs, + } + return {} + + +class OpenAIChatClient( # type: ignore[misc] + FunctionInvocationLayer[OpenAIChatOptionsT], + ChatMiddlewareLayer[OpenAIChatOptionsT], + ChatTelemetryLayer[OpenAIChatOptionsT], + RawOpenAIChatClient[OpenAIChatOptionsT], + Generic[OpenAIChatOptionsT], +): + """OpenAI Responses client class with middleware, telemetry, and function invocation support.""" + + OTEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc] + + @overload + def __init__( + self, + model: str | None = None, + *, + api_key: str | Callable[[], str | Awaitable[str]] | None = None, + org_id: str | None = None, + base_url: str | None = None, + default_headers: Mapping[str, str] | None = None, + async_client: AsyncOpenAI | None = None, + instruction_role: str | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None, + function_invocation_configuration: FunctionInvocationConfiguration | None = None, + additional_properties: dict[str, Any] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize an OpenAI Responses client. + + Keyword Args: + model: Model identifier to use for the request. When not provided, the constructor + reads ``OPENAI_CHAT_MODEL`` and then ``OPENAI_MODEL``. + api_key: API key. When not provided explicitly, the constructor reads + ``OPENAI_API_KEY``. A callable API key is also supported. + org_id: OpenAI organization ID. When not provided explicitly, the constructor reads + ``OPENAI_ORG_ID``. + base_url: Base URL override. When not provided explicitly, the constructor reads + ``OPENAI_BASE_URL``. + default_headers: Additional HTTP headers. + async_client: Pre-configured OpenAI client. + instruction_role: Role for instruction messages (for example ``"system"``). + compaction_strategy: Optional per-client compaction override. + tokenizer: Optional tokenizer for compaction strategies. + middleware: Optional middleware to apply to the client. + function_invocation_configuration: Optional function invocation configuration override. + additional_properties: Optional additional properties to include on all requests. + env_file_path: Optional ``.env`` file that is checked before the process environment + for ``OPENAI_*`` values. + env_file_encoding: Encoding for the ``.env`` file. + """ + ... + + @overload + def __init__( + self, + model: str | None = None, + *, + azure_endpoint: str | None = None, + credential: AzureCredentialTypes | AzureTokenProvider | None = None, + api_version: str | None = None, + api_key: str | Callable[[], str | Awaitable[str]] | None = None, + base_url: str | None = None, + default_headers: Mapping[str, str] | None = None, + async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None, + instruction_role: str | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None, + function_invocation_configuration: FunctionInvocationConfiguration | None = None, + additional_properties: dict[str, Any] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize an OpenAI Responses client. + + Keyword Args: + model: Model identifier to use for the request. When not provided, the constructor + reads ``AZURE_OPENAI_CHAT_MODEL`` and then + ``AZURE_OPENAI_MODEL``. + azure_endpoint: Azure resource endpoint. When not provided explicitly, the constructor + reads ``AZURE_OPENAI_ENDPOINT``. + credential: Azure credential or token provider for Entra auth. + api_version: Azure API version. When not provided explicitly, the constructor reads + ``AZURE_OPENAI_API_VERSION`` and then uses the Responses default. + api_key: API key. For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key + auth. A callable token provider is also accepted, but ``credential`` is the preferred + Azure auth surface. + base_url: Base URL override. When not provided explicitly, the constructor reads + ``AZURE_OPENAI_BASE_URL``. Use this instead of ``azure_endpoint`` when you want + to pass the full ``.../openai/v1`` base URL directly. + default_headers: Additional HTTP headers. + async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on + Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup. + instruction_role: Role for instruction messages (for example ``"system"``). + compaction_strategy: Optional per-client compaction override. + tokenizer: Optional tokenizer for compaction strategies. + middleware: Optional middleware to apply to the client. + function_invocation_configuration: Optional function invocation configuration override. + additional_properties: Optional additional properties to include on all requests. + env_file_path: Optional ``.env`` file that is checked before process environment + variables for ``AZURE_OPENAI_*`` values. + env_file_encoding: Encoding for the ``.env`` file. + """ + ... + + def __init__( + self, + model: str | None = None, + *, + api_key: str | Callable[[], str | Awaitable[str]] | None = None, + credential: AzureCredentialTypes | AzureTokenProvider | None = None, + org_id: str | None = None, + base_url: str | None = None, + azure_endpoint: str | None = None, + api_version: str | None = None, + default_headers: Mapping[str, str] | None = None, + async_client: AsyncOpenAI | None = None, + instruction_role: str | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None, + function_invocation_configuration: FunctionInvocationConfiguration | None = None, + additional_properties: dict[str, Any] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize an OpenAI Responses client. + + Keyword Args: + model: Model identifier to use for the request. When not provided, the constructor + reads ``OPENAI_CHAT_MODEL`` and then ``OPENAI_MODEL`` for OpenAI + routing, or ``AZURE_OPENAI_CHAT_MODEL`` and then + ``AZURE_OPENAI_MODEL`` for Azure routing. + api_key: API key override. For OpenAI routing this maps to ``OPENAI_API_KEY``. + For Azure routing this can be used instead of ``AZURE_OPENAI_API_KEY`` for key + auth. A callable token provider is also accepted for backwards compatibility, + but ``credential`` is the preferred Azure auth surface. + credential: Azure credential or token provider for Azure OpenAI auth. Passing this + is an explicit Azure signal, even when ``OPENAI_API_KEY`` is also configured. + Credential objects require the optional ``azure-identity`` package. + org_id: OpenAI organization ID. Used only for OpenAI routing and resolved from + ``OPENAI_ORG_ID`` when not provided. + base_url: Base URL override. For OpenAI routing this maps to ``OPENAI_BASE_URL``. + For Azure routing this may be used instead of ``azure_endpoint`` when you want + to pass the full ``.../openai/v1`` base URL directly. + azure_endpoint: Azure resource endpoint. When not provided explicitly, Azure routing + falls back to ``AZURE_OPENAI_ENDPOINT``. + api_version: Azure API version to use once Azure routing is selected. When + not provided explicitly, Azure routing falls back to + ``AZURE_OPENAI_API_VERSION`` and then the Responses default. + default_headers: Default HTTP headers that are merged into each request. + async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on + Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup. + instruction_role: Role to use for instruction messages (for example ``"system"``). + compaction_strategy: Optional per-client compaction override. + tokenizer: Optional tokenizer for compaction strategies. + middleware: Optional middleware to apply to the client. + function_invocation_configuration: Optional function invocation configuration override. + additional_properties: Additional properties stored on the client instance. + env_file_path: Optional ``.env`` file that is checked before process environment + variables. The same file is used for both ``OPENAI_*`` and ``AZURE_OPENAI_*`` + lookups. + env_file_encoding: Encoding for the ``.env`` file. + + Notes: + Environment resolution and routing precedence are: + + 1. Explicit Azure inputs (``azure_endpoint`` or ``credential``) + 2. Explicit OpenAI API key or ``OPENAI_API_KEY`` + 3. Azure environment fallback + + OpenAI routing reads ``OPENAI_API_KEY``, ``OPENAI_CHAT_MODEL``, + ``OPENAI_MODEL``, ``OPENAI_ORG_ID``, and ``OPENAI_BASE_URL``. Azure routing + reads ``AZURE_OPENAI_ENDPOINT``, ``AZURE_OPENAI_BASE_URL``, + ``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_CHAT_MODEL``, + ``AZURE_OPENAI_MODEL``, and ``AZURE_OPENAI_API_VERSION``. + + Examples: + .. code-block:: python + + from agent_framework.openai import OpenAIChatClient + + # Using environment variables + # Set OPENAI_API_KEY=sk-... + # Set OPENAI_MODEL=gpt-4o + client = OpenAIChatClient() + + # Or passing parameters directly + client = OpenAIChatClient(model="gpt-4o", api_key="sk-...") + + # Or loading from a .env file + client = OpenAIChatClient(env_file_path="path/to/.env") + + # Using custom ChatOptions with type safety: + from typing import TypedDict + from agent_framework.openai import OpenAIChatOptions + + + class MyOptions(OpenAIChatOptions, total=False): + my_custom_option: str + + + client: OpenAIChatClient[MyOptions] = OpenAIChatClient(model="gpt-4o") + response = await client.get_response("Hello", options={"my_custom_option": "value"}) + """ + super().__init__( + model=model, + api_key=api_key, + credential=credential, + org_id=org_id, + base_url=base_url, + azure_endpoint=azure_endpoint, + api_version=api_version, + default_headers=default_headers, + async_client=async_client, + instruction_role=instruction_role, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + middleware=middleware, + function_invocation_configuration=function_invocation_configuration, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + additional_properties=additional_properties, + ) + + +def _apply_openai_chat_client_docstrings() -> None: + """Align OpenAI Responses client docstrings with the raw implementation.""" + from agent_framework._clients import BaseChatClient + from agent_framework._docstrings import apply_layered_docstring + + apply_layered_docstring(RawOpenAIChatClient.get_response, BaseChatClient.get_response) + apply_layered_docstring( + OpenAIChatClient.get_response, + RawOpenAIChatClient.get_response, + extra_keyword_args={ + "middleware": """ + Optional per-call chat and function middleware. + This is merged with any middleware configured on the client for the current request. + """, + }, + ) + + +_apply_openai_chat_client_docstrings()