agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: cleanup and refactoring of chat clients (#2937)

Browse Source 
* refactoring and unifying naming schemes of internal methods of chat clients

* set tool_choice to auto

* fix for mypy

* added note on naming and fix #2951

* fix responses

* fixes in azure ai agents client
This commit is contained in:
Eduard van Valkenburg
2025-12-18 12:02:23 +00:00
committed by GitHub
parent a71f768331
commit e5c11d38d6
26 changed files with 1128 additions and 1068 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
+190 -141
View File
@@ -278,22 +278,13 @@ class AzureAIAgentClient(BaseChatClient):
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
# Extract necessary state from messages and options
run_options, required_action_results = await self._create_run_options(messages, chat_options, **kwargs)
# Get the thread ID
thread_id: str | None = (
chat_options.conversation_id
if chat_options.conversation_id is not None
else run_options.get("conversation_id", self.thread_id)
)
# Determine which agent to use and create if needed
# prepare
run_options, required_action_results = await self._prepare_options(messages, chat_options, **kwargs)
agent_id = await self._get_agent_id_or_create(run_options)
# Process and yield each update from the stream
# execute and process
async for update in self._process_stream(
*(await self._create_agent_stream(thread_id, agent_id, run_options, required_action_results))
*(await self._create_agent_stream(agent_id, run_options, required_action_results))
):
yield update
@@ -342,7 +333,6 @@ class AzureAIAgentClient(BaseChatClient):
async def _create_agent_stream(
self,
thread_id: str | None,
agent_id: str,
run_options: dict[str, Any],
required_action_results: list[FunctionResultContent | FunctionApprovalResponseContent] | None,
@@ -352,14 +342,14 @@ class AzureAIAgentClient(BaseChatClient):
Returns:
tuple: (stream, final_thread_id)
"""
thread_id = run_options.pop("thread_id", None)
# Get any active run for this thread
thread_run = await self._get_active_thread_run(thread_id)
stream: AsyncAgentRunStream[AsyncAgentEventHandler[Any]] | AsyncAgentEventHandler[Any]
handler: AsyncAgentEventHandler[Any] = AsyncAgentEventHandler()
tool_run_id, tool_outputs, tool_approvals = self._convert_required_action_to_tool_output(
required_action_results
)
tool_run_id, tool_outputs, tool_approvals = self._prepare_tool_outputs_for_azure_ai(required_action_results)
if (
thread_run is not None
@@ -421,19 +411,11 @@ class AzureAIAgentClient(BaseChatClient):
# No thread ID was provided, so create a new thread.
thread = await self.agents_client.threads.create(
tool_resources=run_options.get("tool_resources"), metadata=run_options.get("metadata")
tool_resources=run_options.get("tool_resources"),
metadata=run_options.get("metadata"),
messages=run_options.get("additional_messages"),
)
thread_id = thread.id
# workaround for: https://github.com/Azure/azure-sdk-for-python/issues/42805
# this occurs when otel is enabled
# once fixed, in the function above, readd:
# `messages=run_options.pop("additional_messages")`
for msg in run_options.pop("additional_messages", []):
await self.agents_client.messages.create(
thread_id=thread_id, role=msg.role, content=msg.content, metadata=msg.metadata
)
# and remove until here.
return thread_id
return thread.id
def _extract_url_citations(
self, message_delta_chunk: MessageDeltaChunk, azure_search_tool_calls: list[dict[str, Any]]
@@ -611,7 +593,7 @@ class AzureAIAgentClient(BaseChatClient):
"submit_tool_outputs",
"submit_tool_approval",
]:
function_call_contents = self._create_function_call_contents(
function_call_contents = self._parse_function_calls_from_azure_ai(
event_data, response_id
)
if function_call_contents:
@@ -753,8 +735,8 @@ class AzureAIAgentClient(BaseChatClient):
except Exception as ex:
logger.debug(f"Failed to capture Azure AI Search tool call: {ex}")
def _create_function_call_contents(self, event_data: ThreadRun, response_id: str | None) -> list[Contents]:
"""Create function call contents from a tool action event."""
def _parse_function_calls_from_azure_ai(self, event_data: ThreadRun, response_id: str | None) -> list[Contents]:
"""Parse function call contents from an Azure AI tool action event."""
if isinstance(event_data, ThreadRun) and event_data.required_action is not None:
if isinstance(event_data.required_action, SubmitToolOutputsAction):
return [
@@ -815,117 +797,197 @@ class AzureAIAgentClient(BaseChatClient):
chat_options.tool_choice = chat_tool_mode
async def _create_run_options(
async def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions | None,
chat_options: ChatOptions,
**kwargs: Any,
) -> tuple[dict[str, Any], list[FunctionResultContent | FunctionApprovalResponseContent] | None]:
run_options: dict[str, Any] = {**kwargs}
agent_definition = await self._load_agent_definition_if_needed()
if chat_options is not None:
run_options["max_completion_tokens"] = chat_options.max_tokens
if chat_options.model_id is not None:
run_options["model"] = chat_options.model_id
else:
run_options["model"] = self.model_id
run_options["top_p"] = chat_options.top_p
run_options["temperature"] = chat_options.temperature
run_options["parallel_tool_calls"] = chat_options.allow_multiple_tool_calls
# Use to_dict with exclusions for properties handled separately
run_options: dict[str, Any] = chat_options.to_dict(
exclude={
"type",
"instructions", # handled via messages
"tools", # handled separately
"tool_choice", # handled separately
"response_format", # handled separately
"additional_properties", # handled separately
"frequency_penalty", # not supported
"presence_penalty", # not supported
"user", # not supported
"stop", # not supported
"logit_bias", # not supported
"seed", # not supported
"store", # not supported
}
)
tool_definitions: list[ToolDefinition | dict[str, Any]] = []
# Translation between ChatOptions and Azure AI Agents API
translations = {
"model_id": "model",
"allow_multiple_tool_calls": "parallel_tool_calls",
"max_tokens": "max_completion_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)
# Add tools from existing agent
if agent_definition is not None:
# Don't include function tools, since they will be passed through chat_options.tools
agent_tools = [tool for tool in agent_definition.tools if not isinstance(tool, FunctionToolDefinition)]
if agent_tools:
tool_definitions.extend(agent_tools)
if agent_definition.tool_resources:
run_options["tool_resources"] = agent_definition.tool_resources
# model id fallback
if not run_options.get("model"):
run_options["model"] = self.model_id
if chat_options.tool_choice is not None:
if chat_options.tool_choice != "none" and chat_options.tools:
# Add run tools
tool_definitions.extend(await self._prep_tools(chat_options.tools, run_options))
# tools and tool_choice
if tool_definitions := await self._prepare_tool_definitions_and_resources(
chat_options, agent_definition, run_options
):
run_options["tools"] = tool_definitions
# Handle MCP tool resources for approval mode
mcp_tools = [tool for tool in chat_options.tools if isinstance(tool, HostedMCPTool)]
if mcp_tools:
mcp_resources = []
for mcp_tool in mcp_tools:
server_label = mcp_tool.name.replace(" ", "_")
mcp_resource: dict[str, Any] = {"server_label": server_label}
if tool_choice := self._prepare_tool_choice_mode(chat_options):
run_options["tool_choice"] = tool_choice
# Add headers if they exist
if mcp_tool.headers:
mcp_resource["headers"] = mcp_tool.headers
if mcp_tool.approval_mode is not None:
match mcp_tool.approval_mode:
case str():
# Map agent framework approval modes to Azure AI approval modes
approval_mode = (
"always" if mcp_tool.approval_mode == "always_require" else "never"
)
mcp_resource["require_approval"] = approval_mode
case _:
if "always_require_approval" in mcp_tool.approval_mode:
mcp_resource["require_approval"] = {
"always": mcp_tool.approval_mode["always_require_approval"]
}
elif "never_require_approval" in mcp_tool.approval_mode:
mcp_resource["require_approval"] = {
"never": mcp_tool.approval_mode["never_require_approval"]
}
mcp_resources.append(mcp_resource)
# Add MCP resources to tool_resources
if "tool_resources" not in run_options:
run_options["tool_resources"] = {}
run_options["tool_resources"]["mcp"] = mcp_resources
if chat_options.tool_choice == "none":
run_options["tool_choice"] = AgentsToolChoiceOptionMode.NONE
elif chat_options.tool_choice == "auto":
run_options["tool_choice"] = AgentsToolChoiceOptionMode.AUTO
elif (
isinstance(chat_options.tool_choice, ToolMode)
and chat_options.tool_choice == "required"
and chat_options.tool_choice.required_function_name is not None
):
run_options["tool_choice"] = AgentsNamedToolChoice(
type=AgentsNamedToolChoiceType.FUNCTION,
function=FunctionName(name=chat_options.tool_choice.required_function_name),
)
if tool_definitions:
run_options["tools"] = tool_definitions
if chat_options.response_format is not None:
run_options["response_format"] = ResponseFormatJsonSchemaType(
json_schema=ResponseFormatJsonSchema(
name=chat_options.response_format.__name__,
schema=chat_options.response_format.model_json_schema(),
)
# response format
if chat_options.response_format is not None:
run_options["response_format"] = ResponseFormatJsonSchemaType(
json_schema=ResponseFormatJsonSchema(
name=chat_options.response_format.__name__,
schema=chat_options.response_format.model_json_schema(),
)
)
# messages
additional_messages, instructions, required_action_results = self._prepare_messages(messages)
if additional_messages:
run_options["additional_messages"] = additional_messages
# Add instruction from existing agent at the beginning
if (
agent_definition is not None
and agent_definition.instructions
and agent_definition.instructions not in instructions
):
instructions.insert(0, agent_definition.instructions)
if instructions:
run_options["instructions"] = "\n".join(instructions)
# thread_id resolution (conversation_id takes precedence, then kwargs, then instance default)
run_options["thread_id"] = chat_options.conversation_id or kwargs.get("conversation_id") or self.thread_id
return run_options, required_action_results
def _prepare_tool_choice_mode(
self, chat_options: ChatOptions
) -> AgentsToolChoiceOptionMode | AgentsNamedToolChoice | None:
"""Prepare the tool choice mode for Azure AI Agents API."""
if chat_options.tool_choice is None:
return None
if chat_options.tool_choice == "none":
return AgentsToolChoiceOptionMode.NONE
if chat_options.tool_choice == "auto":
return AgentsToolChoiceOptionMode.AUTO
if (
isinstance(chat_options.tool_choice, ToolMode)
and chat_options.tool_choice == "required"
and chat_options.tool_choice.required_function_name is not None
):
return AgentsNamedToolChoice(
type=AgentsNamedToolChoiceType.FUNCTION,
function=FunctionName(name=chat_options.tool_choice.required_function_name),
)
return None
async def _prepare_tool_definitions_and_resources(
self,
chat_options: ChatOptions,
agent_definition: Agent | None,
run_options: dict[str, Any],
) -> list[ToolDefinition | dict[str, Any]]:
"""Prepare tool definitions and resources for the run options."""
tool_definitions: list[ToolDefinition | dict[str, Any]] = []
# Add tools from existing agent (exclude function tools - passed via chat_options.tools)
if agent_definition is not None:
agent_tools = [tool for tool in agent_definition.tools if not isinstance(tool, FunctionToolDefinition)]
if agent_tools:
tool_definitions.extend(agent_tools)
if agent_definition.tool_resources:
run_options["tool_resources"] = agent_definition.tool_resources
# Add run tools if tool_choice allows
if chat_options.tool_choice is not None and chat_options.tool_choice != "none" and chat_options.tools:
tool_definitions.extend(await self._prepare_tools_for_azure_ai(chat_options.tools, run_options))
# Handle MCP tool resources
mcp_resources = self._prepare_mcp_resources(chat_options.tools)
if mcp_resources:
if "tool_resources" not in run_options:
run_options["tool_resources"] = {}
run_options["tool_resources"]["mcp"] = mcp_resources
return tool_definitions
def _prepare_mcp_resources(
self, tools: Sequence["ToolProtocol | MutableMapping[str, Any]"]
) -> list[dict[str, Any]]:
"""Prepare MCP tool resources for approval mode configuration."""
mcp_tools = [tool for tool in tools if isinstance(tool, HostedMCPTool)]
if not mcp_tools:
return []
mcp_resources: list[dict[str, Any]] = []
for mcp_tool in mcp_tools:
server_label = mcp_tool.name.replace(" ", "_")
mcp_resource: dict[str, Any] = {"server_label": server_label}
if mcp_tool.headers:
mcp_resource["headers"] = mcp_tool.headers
if mcp_tool.approval_mode is not None:
match mcp_tool.approval_mode:
case str():
# Map agent framework approval modes to Azure AI approval modes
approval_mode = "always" if mcp_tool.approval_mode == "always_require" else "never"
mcp_resource["require_approval"] = approval_mode
case _:
if "always_require_approval" in mcp_tool.approval_mode:
mcp_resource["require_approval"] = {
"always": mcp_tool.approval_mode["always_require_approval"]
}
elif "never_require_approval" in mcp_tool.approval_mode:
mcp_resource["require_approval"] = {
"never": mcp_tool.approval_mode["never_require_approval"]
}
mcp_resources.append(mcp_resource)
return mcp_resources
def _prepare_messages(
self, messages: MutableSequence[ChatMessage]
) -> tuple[
list[ThreadMessageOptions] | None,
list[str],
list[FunctionResultContent | FunctionApprovalResponseContent] | None,
]:
"""Prepare messages for Azure AI Agents API.
System/developer messages are turned into instructions, since there is no such message roles in Azure AI.
All other messages are added 1:1, treating assistant messages as agent messages
and everything else as user messages.
Returns:
Tuple of (additional_messages, instructions, required_action_results)
"""
instructions: list[str] = []
required_action_results: list[FunctionResultContent | FunctionApprovalResponseContent] | None = None
additional_messages: list[ThreadMessageOptions] | None = None
# System/developer messages are turned into instructions, since there is no such message roles in Azure AI.
# All other messages are added 1:1, treating assistant messages as agent messages
# and everything else as user messages.
for chat_message in messages:
if chat_message.role.value in ["system", "developer"]:
for text_content in [content for content in chat_message.contents if isinstance(content, TextContent)]:
instructions.append(text_content.text)
continue
message_contents: list[MessageInputContentBlock] = []
@@ -942,7 +1004,7 @@ class AzureAIAgentClient(BaseChatClient):
elif isinstance(content.raw_representation, MessageInputContentBlock):
message_contents.append(content.raw_representation)
if len(message_contents) > 0:
if message_contents:
if additional_messages is None:
additional_messages = []
additional_messages.append(
@@ -952,26 +1014,12 @@ class AzureAIAgentClient(BaseChatClient):
)
)
if additional_messages is not None:
run_options["additional_messages"] = additional_messages
return additional_messages, instructions, required_action_results
# Add instruction from existing agent at the beginning
if (
agent_definition is not None
and agent_definition.instructions
and agent_definition.instructions not in instructions
):
instructions.insert(0, agent_definition.instructions)
if len(instructions) > 0:
run_options["instructions"] = "".join(instructions)
return run_options, required_action_results
async def _prep_tools(
async def _prepare_tools_for_azure_ai(
self, tools: Sequence["ToolProtocol | MutableMapping[str, Any]"], run_options: dict[str, Any] | None = None
) -> list[ToolDefinition | dict[str, Any]]:
"""Prepare tool definitions for the run options."""
"""Prepare tool definitions for the Azure AI Agents API."""
tool_definitions: list[ToolDefinition | dict[str, Any]] = []
for tool in tools:
match tool:
@@ -1044,10 +1092,11 @@ class AzureAIAgentClient(BaseChatClient):
raise ServiceInitializationError(f"Unsupported tool type: {type(tool)}")
return tool_definitions
def _convert_required_action_to_tool_output(
def _prepare_tool_outputs_for_azure_ai(
self,
required_action_results: list[FunctionResultContent | FunctionApprovalResponseContent] | None,
) -> tuple[str | None, list[ToolOutput] | None, list[ToolApproval] | None]:
"""Prepare function results and approvals for submission to the Azure AI API."""
run_id: str | None = None
tool_outputs: list[ToolOutput] | None = None
tool_approvals: list[ToolApproval] | None = None
python/packages/azure-ai/agent_framework_azure_ai/_client.py
+40 -58
View File
@@ -28,10 +28,6 @@ from azure.ai.projects.models import (
)
from azure.core.credentials_async import AsyncTokenCredential
from azure.core.exceptions import ResourceNotFoundError
from openai.types.responses.parsed_response import (
ParsedResponse,
)
from openai.types.responses.response import Response as OpenAIResponse
from pydantic import BaseModel, ValidationError
from ._shared import AzureAISettings
@@ -41,6 +37,11 @@ if sys.version_info >= (3, 11):
else:
from typing_extensions import Self # 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
logger = get_logger("agent_framework.azure")
@@ -368,7 +369,38 @@ class AzureAIClient(OpenAIBaseResponsesClient):
if self._should_close_client:
await self.project_client.close()
def _prepare_input(self, messages: MutableSequence[ChatMessage]) -> tuple[list[ChatMessage], str | None]:
@override
async def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> dict[str, Any]:
"""Take ChatOptions and create the specific options for Azure AI."""
prepared_messages, instructions = self._prepare_messages_for_azure_ai(messages)
run_options = await super()._prepare_options(prepared_messages, chat_options, **kwargs)
if not self._is_application_endpoint:
# Application-scoped response APIs do not support "agent" property.
agent_reference = await self._get_agent_reference_or_create(run_options, instructions)
run_options["extra_body"] = {"agent": agent_reference}
# Remove properties that are not supported on request level
# but were configured on agent level
exclude = ["model", "tools", "response_format", "temperature", "top_p"]
for property in exclude:
run_options.pop(property, None)
return run_options
@override
def _get_current_conversation_id(self, chat_options: ChatOptions, **kwargs: Any) -> str | None:
"""Get the current conversation ID from chat options or kwargs."""
return chat_options.conversation_id or kwargs.get("conversation_id") or self.conversation_id
def _prepare_messages_for_azure_ai(
self, messages: MutableSequence[ChatMessage]
) -> tuple[list[ChatMessage], str | None]:
"""Prepare input from messages and convert system/developer messages to instructions."""
result: list[ChatMessage] = []
instructions_list: list[str] = []
@@ -387,44 +419,7 @@ class AzureAIClient(OpenAIBaseResponsesClient):
return result, instructions
async def prepare_options(
self,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> dict[str, Any]:
"""Take ChatOptions and create the specific options for Azure AI."""
prepared_messages, instructions = self._prepare_input(messages)
run_options = await super().prepare_options(prepared_messages, chat_options, **kwargs)
if not self._is_application_endpoint:
# Application-scoped response APIs do not support "agent" property.
agent_reference = await self._get_agent_reference_or_create(run_options, instructions)
run_options["extra_body"] = {"agent": agent_reference}
conversation_id = chat_options.conversation_id or self.conversation_id
# Handle different conversation ID formats
if conversation_id:
if conversation_id.startswith("resp_"):
# For response IDs, set previous_response_id and remove conversation property
run_options.pop("conversation", None)
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.pop("previous_response_id", None)
run_options["conversation"] = conversation_id
# Remove properties that are not supported on request level
# but were configured on agent level
exclude = ["model", "tools", "response_format", "temperature", "top_p"]
for property in exclude:
run_options.pop(property, None)
return run_options
async def initialize_client(self) -> None:
async def _initialize_client(self) -> None:
"""Initialize OpenAI client."""
self.client = self.project_client.get_openai_client() # type: ignore
@@ -442,7 +437,8 @@ class AzureAIClient(OpenAIBaseResponsesClient):
if description and not self.agent_description:
self.agent_description = description
def get_mcp_tool(self, tool: HostedMCPTool) -> Any:
@staticmethod
def _prepare_mcp_tool(tool: HostedMCPTool) -> MCPTool: # type: ignore[override]
"""Get MCP tool from HostedMCPTool."""
mcp = MCPTool(server_label=tool.name.replace(" ", "_"), server_url=str(tool.url))
@@ -460,17 +456,3 @@ class AzureAIClient(OpenAIBaseResponsesClient):
mcp["require_approval"] = {"never": {"tool_names": list(never_require_approvals)}}
return mcp
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
python/packages/azure-ai/tests/test_azure_ai_agent_client.py
+95 -89
View File
@@ -367,33 +367,33 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_missing_model(
await chat_client._get_agent_id_or_create() # type: ignore
async def test_azure_ai_chat_client_create_run_options_basic(mock_agents_client: MagicMock) -> None:
"""Test _create_run_options with basic ChatOptions."""
async def test_azure_ai_chat_client_prepare_options_basic(mock_agents_client: MagicMock) -> None:
"""Test _prepare_options with basic ChatOptions."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
messages = [ChatMessage(role=Role.USER, text="Hello")]
chat_options = ChatOptions(max_tokens=100, temperature=0.7)
run_options, tool_results = await chat_client._create_run_options(messages, chat_options) # type: ignore
run_options, tool_results = await chat_client._prepare_options(messages, chat_options) # type: ignore
assert run_options is not None
assert tool_results is None
async def test_azure_ai_chat_client_create_run_options_no_chat_options(mock_agents_client: MagicMock) -> None:
"""Test _create_run_options with no ChatOptions."""
async def test_azure_ai_chat_client_prepare_options_no_chat_options(mock_agents_client: MagicMock) -> None:
"""Test _prepare_options with default ChatOptions."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
messages = [ChatMessage(role=Role.USER, text="Hello")]
run_options, tool_results = await chat_client._create_run_options(messages, None) # type: ignore
run_options, tool_results = await chat_client._prepare_options(messages, ChatOptions()) # type: ignore
assert run_options is not None
assert tool_results is None
async def test_azure_ai_chat_client_create_run_options_with_image_content(mock_agents_client: MagicMock) -> None:
"""Test _create_run_options with image content."""
async def test_azure_ai_chat_client_prepare_options_with_image_content(mock_agents_client: MagicMock) -> None:
"""Test _prepare_options with image content."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
@@ -403,7 +403,7 @@ async def test_azure_ai_chat_client_create_run_options_with_image_content(mock_a
image_content = UriContent(uri="https://example.com/image.jpg", media_type="image/jpeg")
messages = [ChatMessage(role=Role.USER, contents=[image_content])]
run_options, _ = await chat_client._create_run_options(messages, None) # type: ignore
run_options, _ = await chat_client._prepare_options(messages, ChatOptions()) # type: ignore
assert "additional_messages" in run_options
assert len(run_options["additional_messages"]) == 1
@@ -412,11 +412,11 @@ async def test_azure_ai_chat_client_create_run_options_with_image_content(mock_a
assert len(message.content) == 1
def test_azure_ai_chat_client_convert_function_results_to_tool_output_none(mock_agents_client: MagicMock) -> None:
"""Test _convert_required_action_to_tool_output with None input."""
def test_azure_ai_chat_client_prepare_tool_outputs_for_azure_ai_none(mock_agents_client: MagicMock) -> None:
"""Test _prepare_tool_outputs_for_azure_ai with None input."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
run_id, tool_outputs, tool_approvals = chat_client._convert_required_action_to_tool_output(None) # type: ignore
run_id, tool_outputs, tool_approvals = chat_client._prepare_tool_outputs_for_azure_ai(None) # type: ignore
assert run_id is None
assert tool_outputs is None
@@ -484,8 +484,8 @@ def test_azure_ai_chat_client_update_agent_name_and_description_with_none_input(
assert chat_client.agent_description is None
async def test_azure_ai_chat_client_create_run_options_with_messages(mock_agents_client: MagicMock) -> None:
"""Test _create_run_options with different message types."""
async def test_azure_ai_chat_client_prepare_options_with_messages(mock_agents_client: MagicMock) -> None:
"""Test _prepare_options with different message types."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
# Test with system message (becomes instruction)
@@ -494,7 +494,7 @@ async def test_azure_ai_chat_client_create_run_options_with_messages(mock_agents
ChatMessage(role=Role.USER, text="Hello"),
]
run_options, _ = await chat_client._create_run_options(messages, None) # type: ignore
run_options, _ = await chat_client._prepare_options(messages, ChatOptions()) # type: ignore
assert "instructions" in run_options
assert "You are a helpful assistant" in run_options["instructions"]
@@ -565,8 +565,8 @@ async def test_azure_ai_chat_client_prepare_thread_cancels_active_run(mock_agent
mock_agents_client.runs.cancel.assert_called_once_with("test-thread", "run_123")
def test_azure_ai_chat_client_create_function_call_contents_basic(mock_agents_client: MagicMock) -> None:
"""Test _create_function_call_contents with basic function call."""
def test_azure_ai_chat_client_parse_function_calls_from_azure_ai_basic(mock_agents_client: MagicMock) -> None:
"""Test _parse_function_calls_from_azure_ai with basic function call."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
mock_tool_call = MagicMock(spec=RequiredFunctionToolCall)
@@ -580,7 +580,7 @@ def test_azure_ai_chat_client_create_function_call_contents_basic(mock_agents_cl
mock_event_data = MagicMock(spec=ThreadRun)
mock_event_data.required_action = mock_submit_action
result = chat_client._create_function_call_contents(mock_event_data, "response_123") # type: ignore
result = chat_client._parse_function_calls_from_azure_ai(mock_event_data, "response_123") # type: ignore
assert len(result) == 1
assert isinstance(result[0], FunctionCallContent)
@@ -588,22 +588,24 @@ def test_azure_ai_chat_client_create_function_call_contents_basic(mock_agents_cl
assert result[0].call_id == '["response_123", "call_123"]'
def test_azure_ai_chat_client_create_function_call_contents_no_submit_action(mock_agents_client: MagicMock) -> None:
"""Test _create_function_call_contents when required_action is not SubmitToolOutputsAction."""
def test_azure_ai_chat_client_parse_function_calls_from_azure_ai_no_submit_action(
mock_agents_client: MagicMock,
) -> None:
"""Test _parse_function_calls_from_azure_ai when required_action is not SubmitToolOutputsAction."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
mock_event_data = MagicMock(spec=ThreadRun)
mock_event_data.required_action = MagicMock()
result = chat_client._create_function_call_contents(mock_event_data, "response_123") # type: ignore
result = chat_client._parse_function_calls_from_azure_ai(mock_event_data, "response_123") # type: ignore
assert result == []
def test_azure_ai_chat_client_create_function_call_contents_non_function_tool_call(
def test_azure_ai_chat_client_parse_function_calls_from_azure_ai_non_function_tool_call(
mock_agents_client: MagicMock,
) -> None:
"""Test _create_function_call_contents with non-function tool call."""
"""Test _parse_function_calls_from_azure_ai with non-function tool call."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
mock_tool_call = MagicMock()
@@ -614,37 +616,37 @@ def test_azure_ai_chat_client_create_function_call_contents_non_function_tool_ca
mock_event_data = MagicMock(spec=ThreadRun)
mock_event_data.required_action = mock_submit_action
result = chat_client._create_function_call_contents(mock_event_data, "response_123") # type: ignore
result = chat_client._parse_function_calls_from_azure_ai(mock_event_data, "response_123") # type: ignore
assert result == []
async def test_azure_ai_chat_client_create_run_options_with_none_tool_choice(
async def test_azure_ai_chat_client_prepare_options_with_none_tool_choice(
mock_agents_client: MagicMock,
) -> None:
"""Test _create_run_options with tool_choice set to 'none'."""
"""Test _prepare_options with tool_choice set to 'none'."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
chat_options = ChatOptions()
chat_options.tool_choice = "none"
run_options, _ = await chat_client._create_run_options([], chat_options) # type: ignore
run_options, _ = await chat_client._prepare_options([], chat_options) # type: ignore
from azure.ai.agents.models import AgentsToolChoiceOptionMode
assert run_options["tool_choice"] == AgentsToolChoiceOptionMode.NONE
async def test_azure_ai_chat_client_create_run_options_with_auto_tool_choice(
async def test_azure_ai_chat_client_prepare_options_with_auto_tool_choice(
mock_agents_client: MagicMock,
) -> None:
"""Test _create_run_options with tool_choice set to 'auto'."""
"""Test _prepare_options with tool_choice set to 'auto'."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
chat_options = ChatOptions()
chat_options.tool_choice = "auto"
run_options, _ = await chat_client._create_run_options([], chat_options) # type: ignore
run_options, _ = await chat_client._prepare_options([], chat_options) # type: ignore
from azure.ai.agents.models import AgentsToolChoiceOptionMode
@@ -669,10 +671,10 @@ async def test_azure_ai_chat_client_prepare_tool_choice_none_string(
assert chat_options.tool_choice == ToolMode.NONE.mode
async def test_azure_ai_chat_client_create_run_options_tool_choice_required_specific_function(
async def test_azure_ai_chat_client_prepare_options_tool_choice_required_specific_function(
mock_agents_client: MagicMock,
) -> None:
"""Test _create_run_options with ToolMode.REQUIRED specifying a specific function name."""
"""Test _prepare_options with ToolMode.REQUIRED specifying a specific function name."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
required_tool_mode = ToolMode.REQUIRED("specific_function_name")
@@ -682,7 +684,7 @@ async def test_azure_ai_chat_client_create_run_options_tool_choice_required_spec
chat_options = ChatOptions(tools=[dict_tool], tool_choice=required_tool_mode)
messages = [ChatMessage(role=Role.USER, text="Hello")]
run_options, _ = await chat_client._create_run_options(messages, chat_options) # type: ignore
run_options, _ = await chat_client._prepare_options(messages, chat_options) # type: ignore
# Verify tool_choice is set to the specific named function
assert "tool_choice" in run_options
@@ -692,10 +694,10 @@ async def test_azure_ai_chat_client_create_run_options_tool_choice_required_spec
assert tool_choice.function.name == "specific_function_name" # type: ignore
async def test_azure_ai_chat_client_create_run_options_with_response_format(
async def test_azure_ai_chat_client_prepare_options_with_response_format(
mock_agents_client: MagicMock,
) -> None:
"""Test _create_run_options with response_format configured."""
"""Test _prepare_options with response_format configured."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
class TestResponseModel(BaseModel):
@@ -704,7 +706,7 @@ async def test_azure_ai_chat_client_create_run_options_with_response_format(
chat_options = ChatOptions()
chat_options.response_format = TestResponseModel
run_options, _ = await chat_client._create_run_options([], chat_options) # type: ignore
run_options, _ = await chat_client._prepare_options([], chat_options) # type: ignore
assert "response_format" in run_options
response_format = run_options["response_format"]
@@ -720,8 +722,8 @@ def test_azure_ai_chat_client_service_url_method(mock_agents_client: MagicMock)
assert url == "https://test-endpoint.com/"
async def test_azure_ai_chat_client_prep_tools_ai_function(mock_agents_client: MagicMock) -> None:
"""Test _prep_tools with AIFunction tool."""
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_ai_function(mock_agents_client: MagicMock) -> None:
"""Test _prepare_tools_for_azure_ai with AIFunction tool."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
@@ -729,28 +731,28 @@ async def test_azure_ai_chat_client_prep_tools_ai_function(mock_agents_client: M
mock_ai_function = MagicMock(spec=AIFunction)
mock_ai_function.to_json_schema_spec.return_value = {"type": "function", "function": {"name": "test_function"}}
result = await chat_client._prep_tools([mock_ai_function]) # type: ignore
result = await chat_client._prepare_tools_for_azure_ai([mock_ai_function]) # type: ignore
assert len(result) == 1
assert result[0] == {"type": "function", "function": {"name": "test_function"}}
mock_ai_function.to_json_schema_spec.assert_called_once()
async def test_azure_ai_chat_client_prep_tools_code_interpreter(mock_agents_client: MagicMock) -> None:
"""Test _prep_tools with HostedCodeInterpreterTool."""
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_code_interpreter(mock_agents_client: MagicMock) -> None:
"""Test _prepare_tools_for_azure_ai with HostedCodeInterpreterTool."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
code_interpreter_tool = HostedCodeInterpreterTool()
result = await chat_client._prep_tools([code_interpreter_tool]) # type: ignore
result = await chat_client._prepare_tools_for_azure_ai([code_interpreter_tool]) # type: ignore
assert len(result) == 1
assert isinstance(result[0], CodeInterpreterToolDefinition)
async def test_azure_ai_chat_client_prep_tools_mcp_tool(mock_agents_client: MagicMock) -> None:
"""Test _prep_tools with HostedMCPTool."""
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_mcp_tool(mock_agents_client: MagicMock) -> None:
"""Test _prepare_tools_for_azure_ai with HostedMCPTool."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
@@ -762,7 +764,7 @@ async def test_azure_ai_chat_client_prep_tools_mcp_tool(mock_agents_client: Magi
mock_mcp_tool.definitions = [{"type": "mcp", "name": "test_mcp"}]
mock_mcp_tool_class.return_value = mock_mcp_tool
result = await chat_client._prep_tools([mcp_tool]) # type: ignore
result = await chat_client._prepare_tools_for_azure_ai([mcp_tool]) # type: ignore
assert len(result) == 1
assert result[0] == {"type": "mcp", "name": "test_mcp"}
@@ -774,8 +776,8 @@ async def test_azure_ai_chat_client_prep_tools_mcp_tool(mock_agents_client: Magi
assert set(call_args["allowed_tools"]) == {"tool1", "tool2"}
async def test_azure_ai_chat_client_create_run_options_mcp_never_require(mock_agents_client: MagicMock) -> None:
"""Test _create_run_options with HostedMCPTool having never_require approval mode."""
async def test_azure_ai_chat_client_prepare_options_mcp_never_require(mock_agents_client: MagicMock) -> None:
"""Test _prepare_options with HostedMCPTool having never_require approval mode."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
mcp_tool = HostedMCPTool(name="Test MCP Tool", url="https://example.com/mcp", approval_mode="never_require")
@@ -784,12 +786,12 @@ async def test_azure_ai_chat_client_create_run_options_mcp_never_require(mock_ag
chat_options = ChatOptions(tools=[mcp_tool], tool_choice="auto")
with patch("agent_framework_azure_ai._chat_client.McpTool") as mock_mcp_tool_class:
# Mock _prep_tools to avoid actual tool preparation
# Mock _prepare_tools_for_azure_ai to avoid actual tool preparation
mock_mcp_tool_instance = MagicMock()
mock_mcp_tool_instance.definitions = [{"type": "mcp", "name": "test_mcp"}]
mock_mcp_tool_class.return_value = mock_mcp_tool_instance
run_options, _ = await chat_client._create_run_options(messages, chat_options) # type: ignore
run_options, _ = await chat_client._prepare_options(messages, chat_options) # type: ignore
# Verify tool_resources is created with correct MCP approval structure
assert "tool_resources" in run_options, (
@@ -803,8 +805,8 @@ async def test_azure_ai_chat_client_create_run_options_mcp_never_require(mock_ag
assert mcp_resource["require_approval"] == "never"
async def test_azure_ai_chat_client_create_run_options_mcp_with_headers(mock_agents_client: MagicMock) -> None:
"""Test _create_run_options with HostedMCPTool having headers."""
async def test_azure_ai_chat_client_prepare_options_mcp_with_headers(mock_agents_client: MagicMock) -> None:
"""Test _prepare_options with HostedMCPTool having headers."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
# Test with headers
@@ -817,12 +819,12 @@ async def test_azure_ai_chat_client_create_run_options_mcp_with_headers(mock_age
chat_options = ChatOptions(tools=[mcp_tool], tool_choice="auto")
with patch("agent_framework_azure_ai._chat_client.McpTool") as mock_mcp_tool_class:
# Mock _prep_tools to avoid actual tool preparation
# Mock _prepare_tools_for_azure_ai to avoid actual tool preparation
mock_mcp_tool_instance = MagicMock()
mock_mcp_tool_instance.definitions = [{"type": "mcp", "name": "test_mcp"}]
mock_mcp_tool_class.return_value = mock_mcp_tool_instance
run_options, _ = await chat_client._create_run_options(messages, chat_options) # type: ignore
run_options, _ = await chat_client._prepare_options(messages, chat_options) # type: ignore
# Verify tool_resources is created with headers
assert "tool_resources" in run_options
@@ -835,8 +837,10 @@ async def test_azure_ai_chat_client_create_run_options_mcp_with_headers(mock_age
assert mcp_resource["headers"] == headers
async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding(mock_agents_client: MagicMock) -> None:
"""Test _prep_tools with HostedWebSearchTool using Bing Grounding."""
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_web_search_bing_grounding(
mock_agents_client: MagicMock,
) -> None:
"""Test _prepare_tools_for_azure_ai with HostedWebSearchTool using Bing Grounding."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
@@ -856,7 +860,7 @@ async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding(mock_ag
mock_bing_tool.definitions = [{"type": "bing_grounding"}]
mock_bing_grounding.return_value = mock_bing_tool
result = await chat_client._prep_tools([web_search_tool]) # type: ignore
result = await chat_client._prepare_tools_for_azure_ai([web_search_tool]) # type: ignore
assert len(result) == 1
assert result[0] == {"type": "bing_grounding"}
@@ -868,10 +872,10 @@ async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding(mock_ag
assert "connection_id" in call_args
async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding_with_connection_id(
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_web_search_bing_grounding_with_connection_id(
mock_agents_client: MagicMock,
) -> None:
"""Test _prep_tools with HostedWebSearchTool using Bing Grounding with connection_id (no HTTP call)."""
"""Test _prepare_tools_... with HostedWebSearchTool using Bing Grounding with connection_id (no HTTP call)."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
@@ -888,15 +892,17 @@ async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding_with_co
mock_bing_tool.definitions = [{"type": "bing_grounding"}]
mock_bing_grounding.return_value = mock_bing_tool
result = await chat_client._prep_tools([web_search_tool]) # type: ignore
result = await chat_client._prepare_tools_for_azure_ai([web_search_tool]) # type: ignore
assert len(result) == 1
assert result[0] == {"type": "bing_grounding"}
mock_bing_grounding.assert_called_once_with(connection_id="direct-connection-id", count=3)
async def test_azure_ai_chat_client_prep_tools_web_search_custom_bing(mock_agents_client: MagicMock) -> None:
"""Test _prep_tools with HostedWebSearchTool using Custom Bing Search."""
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_web_search_custom_bing(
mock_agents_client: MagicMock,
) -> None:
"""Test _prepare_tools_for_azure_ai with HostedWebSearchTool using Custom Bing Search."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
@@ -914,16 +920,16 @@ async def test_azure_ai_chat_client_prep_tools_web_search_custom_bing(mock_agent
mock_custom_tool.definitions = [{"type": "bing_custom_search"}]
mock_custom_bing.return_value = mock_custom_tool
result = await chat_client._prep_tools([web_search_tool]) # type: ignore
result = await chat_client._prepare_tools_for_azure_ai([web_search_tool]) # type: ignore
assert len(result) == 1
assert result[0] == {"type": "bing_custom_search"}
async def test_azure_ai_chat_client_prep_tools_file_search_with_vector_stores(
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_file_search_with_vector_stores(
mock_agents_client: MagicMock,
) -> None:
"""Test _prep_tools with HostedFileSearchTool using vector stores."""
"""Test _prepare_tools_for_azure_ai with HostedFileSearchTool using vector stores."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
@@ -938,7 +944,7 @@ async def test_azure_ai_chat_client_prep_tools_file_search_with_vector_stores(
mock_file_search.return_value = mock_file_tool
run_options = {}
result = await chat_client._prep_tools([file_search_tool], run_options) # type: ignore
result = await chat_client._prepare_tools_for_azure_ai([file_search_tool], run_options) # type: ignore
assert len(result) == 1
assert result[0] == {"type": "file_search"}
@@ -973,7 +979,7 @@ async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals(
with patch("azure.ai.agents.models.AsyncAgentEventHandler", return_value=mock_handler):
stream, final_thread_id = await chat_client._create_agent_stream( # type: ignore
"test-thread", "test-agent", {}, [approval_response]
"test-agent", {"thread_id": "test-thread"}, [approval_response]
)
# Verify the approvals path was taken
@@ -987,26 +993,26 @@ async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals(
assert call_args["tool_approvals"][0].approve is True
async def test_azure_ai_chat_client_prep_tools_dict_tool(mock_agents_client: MagicMock) -> None:
"""Test _prep_tools with dictionary tool definition."""
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_dict_tool(mock_agents_client: MagicMock) -> None:
"""Test _prepare_tools_for_azure_ai with dictionary tool definition."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
dict_tool = {"type": "custom_tool", "config": {"param": "value"}}
result = await chat_client._prep_tools([dict_tool]) # type: ignore
result = await chat_client._prepare_tools_for_azure_ai([dict_tool]) # type: ignore
assert len(result) == 1
assert result[0] == dict_tool
async def test_azure_ai_chat_client_prep_tools_unsupported_tool(mock_agents_client: MagicMock) -> None:
"""Test _prep_tools with unsupported tool type."""
async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_unsupported_tool(mock_agents_client: MagicMock) -> None:
"""Test _prepare_tools_for_azure_ai with unsupported tool type."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
unsupported_tool = "not_a_tool"
with pytest.raises(ServiceInitializationError, match="Unsupported tool type: <class 'str'>"):
await chat_client._prep_tools([unsupported_tool]) # type: ignore
await chat_client._prepare_tools_for_azure_ai([unsupported_tool]) # type: ignore
async def test_azure_ai_chat_client_get_active_thread_run_with_active_run(mock_agents_client: MagicMock) -> None:
@@ -1072,16 +1078,16 @@ async def test_azure_ai_chat_client_service_url(mock_agents_client: MagicMock) -
assert result == "https://test-endpoint.com/"
async def test_azure_ai_chat_client_convert_required_action_to_tool_output_function_result(
async def test_azure_ai_chat_client_prepare_tool_outputs_for_azure_ai_function_result(
mock_agents_client: MagicMock,
) -> None:
"""Test _convert_required_action_to_tool_output with FunctionResultContent."""
"""Test _prepare_tool_outputs_for_azure_ai with FunctionResultContent."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
# Test with simple result
function_result = FunctionResultContent(call_id='["run_123", "call_456"]', result="Simple result")
run_id, tool_outputs, tool_approvals = chat_client._convert_required_action_to_tool_output([function_result]) # type: ignore
run_id, tool_outputs, tool_approvals = chat_client._prepare_tool_outputs_for_azure_ai([function_result]) # type: ignore
assert run_id == "run_123"
assert tool_approvals is None
@@ -1092,7 +1098,7 @@ async def test_azure_ai_chat_client_convert_required_action_to_tool_output_funct
async def test_azure_ai_chat_client_convert_required_action_invalid_call_id(mock_agents_client: MagicMock) -> None:
"""Test _convert_required_action_to_tool_output with invalid call_id format."""
"""Test _prepare_tool_outputs_for_azure_ai with invalid call_id format."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
@@ -1100,19 +1106,19 @@ async def test_azure_ai_chat_client_convert_required_action_invalid_call_id(mock
function_result = FunctionResultContent(call_id="invalid_json", result="result")
with pytest.raises(json.JSONDecodeError):
chat_client._convert_required_action_to_tool_output([function_result]) # type: ignore
chat_client._prepare_tool_outputs_for_azure_ai([function_result]) # type: ignore
async def test_azure_ai_chat_client_convert_required_action_invalid_structure(
mock_agents_client: MagicMock,
) -> None:
"""Test _convert_required_action_to_tool_output with invalid call_id structure."""
"""Test _prepare_tool_outputs_for_azure_ai with invalid call_id structure."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
# Valid JSON but invalid structure (missing second element)
function_result = FunctionResultContent(call_id='["run_123"]', result="result")
run_id, tool_outputs, tool_approvals = chat_client._convert_required_action_to_tool_output([function_result]) # type: ignore
run_id, tool_outputs, tool_approvals = chat_client._prepare_tool_outputs_for_azure_ai([function_result]) # type: ignore
# Should return None values when structure is invalid
assert run_id is None
@@ -1123,7 +1129,7 @@ async def test_azure_ai_chat_client_convert_required_action_invalid_structure(
async def test_azure_ai_chat_client_convert_required_action_serde_model_results(
mock_agents_client: MagicMock,
) -> None:
"""Test _convert_required_action_to_tool_output with BaseModel results."""
"""Test _prepare_tool_outputs_for_azure_ai with BaseModel results."""
class MockResult(SerializationMixin):
def __init__(self, name: str, value: int):
@@ -1136,7 +1142,7 @@ async def test_azure_ai_chat_client_convert_required_action_serde_model_results(
mock_result = MockResult(name="test", value=42)
function_result = FunctionResultContent(call_id='["run_123", "call_456"]', result=mock_result)
run_id, tool_outputs, tool_approvals = chat_client._convert_required_action_to_tool_output([function_result]) # type: ignore
run_id, tool_outputs, tool_approvals = chat_client._prepare_tool_outputs_for_azure_ai([function_result]) # type: ignore
assert run_id == "run_123"
assert tool_approvals is None
@@ -1151,7 +1157,7 @@ async def test_azure_ai_chat_client_convert_required_action_serde_model_results(
async def test_azure_ai_chat_client_convert_required_action_multiple_results(
mock_agents_client: MagicMock,
) -> None:
"""Test _convert_required_action_to_tool_output with multiple results."""
"""Test _prepare_tool_outputs_for_azure_ai with multiple results."""
class MockResult(SerializationMixin):
def __init__(self, data: str):
@@ -1164,7 +1170,7 @@ async def test_azure_ai_chat_client_convert_required_action_multiple_results(
results_list = [mock_basemodel, {"key": "value"}, "string_result"]
function_result = FunctionResultContent(call_id='["run_123", "call_456"]', result=results_list)
run_id, tool_outputs, tool_approvals = chat_client._convert_required_action_to_tool_output([function_result]) # type: ignore
run_id, tool_outputs, tool_approvals = chat_client._prepare_tool_outputs_for_azure_ai([function_result]) # type: ignore
assert run_id == "run_123"
assert tool_outputs is not None
@@ -1184,7 +1190,7 @@ async def test_azure_ai_chat_client_convert_required_action_multiple_results(
async def test_azure_ai_chat_client_convert_required_action_approval_response(
mock_agents_client: MagicMock,
) -> None:
"""Test _convert_required_action_to_tool_output with FunctionApprovalResponseContent."""
"""Test _prepare_tool_outputs_for_azure_ai with FunctionApprovalResponseContent."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
# Test with approval response - need to provide required fields
@@ -1194,7 +1200,7 @@ async def test_azure_ai_chat_client_convert_required_action_approval_response(
approved=True,
)
run_id, tool_outputs, tool_approvals = chat_client._convert_required_action_to_tool_output([approval_response]) # type: ignore
run_id, tool_outputs, tool_approvals = chat_client._prepare_tool_outputs_for_azure_ai([approval_response]) # type: ignore
assert run_id == "run_123"
assert tool_outputs is None
@@ -1204,10 +1210,10 @@ async def test_azure_ai_chat_client_convert_required_action_approval_response(
assert tool_approvals[0].approve is True
async def test_azure_ai_chat_client_create_function_call_contents_approval_request(
async def test_azure_ai_chat_client_parse_function_calls_from_azure_ai_approval_request(
mock_agents_client: MagicMock,
) -> None:
"""Test _create_function_call_contents with approval action."""
"""Test _parse_function_calls_from_azure_ai with approval action."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
# Mock SubmitToolApprovalAction with RequiredMcpToolCall
@@ -1222,7 +1228,7 @@ async def test_azure_ai_chat_client_create_function_call_contents_approval_reque
mock_event_data = MagicMock(spec=ThreadRun)
mock_event_data.required_action = mock_approval_action
result = chat_client._create_function_call_contents(mock_event_data, "response_123") # type: ignore
result = chat_client._parse_function_calls_from_azure_ai(mock_event_data, "response_123") # type: ignore
assert len(result) == 1
assert isinstance(result[0], FunctionApprovalRequestContent)
@@ -1312,7 +1318,7 @@ async def test_azure_ai_chat_client_create_agent_stream_submit_tool_outputs(
with patch("azure.ai.agents.models.AsyncAgentEventHandler", return_value=mock_handler):
stream, final_thread_id = await chat_client._create_agent_stream( # type: ignore
thread_id="test-thread", agent_id="test-agent", run_options={}, required_action_results=[function_result]
agent_id="test-agent", run_options={"thread_id": "test-thread"}, required_action_results=[function_result]
)
# Should call submit_tool_outputs_stream since we have matching run ID
python/packages/azure-ai/tests/test_azure_ai_client.py
+28 -114
View File
@@ -249,10 +249,10 @@ async def test_azure_ai_client_get_agent_reference_missing_model(
await client._get_agent_reference_or_create({}, None) # type: ignore
async def test_azure_ai_client_prepare_input_with_system_messages(
async def test_azure_ai_client_prepare_messages_for_azure_ai_with_system_messages(
mock_project_client: MagicMock,
) -> None:
"""Test _prepare_input converts system/developer messages to instructions."""
"""Test _prepare_messages_for_azure_ai converts system/developer messages to instructions."""
client = create_test_azure_ai_client(mock_project_client)
messages = [
@@ -261,7 +261,7 @@ async def test_azure_ai_client_prepare_input_with_system_messages(
ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text="System response")]),
]
result_messages, instructions = client._prepare_input(messages) # type: ignore
result_messages, instructions = client._prepare_messages_for_azure_ai(messages) # type: ignore
assert len(result_messages) == 2
assert result_messages[0].role == Role.USER
@@ -269,10 +269,10 @@ async def test_azure_ai_client_prepare_input_with_system_messages(
assert instructions == "You are a helpful assistant."
async def test_azure_ai_client_prepare_input_no_system_messages(
async def test_azure_ai_client_prepare_messages_for_azure_ai_no_system_messages(
mock_project_client: MagicMock,
) -> None:
"""Test _prepare_input with no system/developer messages."""
"""Test _prepare_messages_for_azure_ai with no system/developer messages."""
client = create_test_azure_ai_client(mock_project_client)
messages = [
@@ -280,7 +280,7 @@ async def test_azure_ai_client_prepare_input_no_system_messages(
ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text="Hi there!")]),
]
result_messages, instructions = client._prepare_input(messages) # type: ignore
result_messages, instructions = client._prepare_messages_for_azure_ai(messages) # type: ignore
assert len(result_messages) == 2
assert instructions is None
@@ -294,14 +294,14 @@ async def test_azure_ai_client_prepare_options_basic(mock_project_client: MagicM
chat_options = ChatOptions()
with (
patch.object(client.__class__.__bases__[0], "prepare_options", return_value={"model": "test-model"}),
patch.object(client.__class__.__bases__[0], "_prepare_options", return_value={"model": "test-model"}),
patch.object(
client,
"_get_agent_reference_or_create",
return_value={"name": "test-agent", "version": "1.0", "type": "agent_reference"},
),
):
run_options = await client.prepare_options(messages, chat_options)
run_options = await client._prepare_options(messages, chat_options)
assert "extra_body" in run_options
assert run_options["extra_body"]["agent"]["name"] == "test-agent"
@@ -329,14 +329,14 @@ async def test_azure_ai_client_prepare_options_with_application_endpoint(
chat_options = ChatOptions()
with (
patch.object(client.__class__.__bases__[0], "prepare_options", return_value={"model": "test-model"}),
patch.object(client.__class__.__bases__[0], "_prepare_options", return_value={"model": "test-model"}),
patch.object(
client,
"_get_agent_reference_or_create",
return_value={"name": "test-agent", "version": "1", "type": "agent_reference"},
),
):
run_options = await client.prepare_options(messages, chat_options)
run_options = await client._prepare_options(messages, chat_options)
if expects_agent:
assert "extra_body" in run_options
@@ -369,14 +369,14 @@ async def test_azure_ai_client_prepare_options_with_application_project_client(
chat_options = ChatOptions()
with (
patch.object(client.__class__.__bases__[0], "prepare_options", return_value={"model": "test-model"}),
patch.object(client.__class__.__bases__[0], "_prepare_options", return_value={"model": "test-model"}),
patch.object(
client,
"_get_agent_reference_or_create",
return_value={"name": "test-agent", "version": "1", "type": "agent_reference"},
),
):
run_options = await client.prepare_options(messages, chat_options)
run_options = await client._prepare_options(messages, chat_options)
if expects_agent:
assert "extra_body" in run_options
@@ -386,13 +386,13 @@ async def test_azure_ai_client_prepare_options_with_application_project_client(
async def test_azure_ai_client_initialize_client(mock_project_client: MagicMock) -> None:
"""Test initialize_client method."""
"""Test _initialize_client method."""
client = create_test_azure_ai_client(mock_project_client)
mock_openai_client = MagicMock()
mock_project_client.get_openai_client = MagicMock(return_value=mock_openai_client)
await client.initialize_client()
await client._initialize_client()
assert client.client is mock_openai_client
mock_project_client.get_openai_client.assert_called_once()
@@ -727,7 +727,7 @@ async def test_azure_ai_client_prepare_options_excludes_response_format(
with (
patch.object(
client.__class__.__bases__[0],
"prepare_options",
"_prepare_options",
return_value={"model": "test-model", "response_format": ResponseFormatModel},
),
patch.object(
@@ -736,7 +736,7 @@ async def test_azure_ai_client_prepare_options_excludes_response_format(
return_value={"name": "test-agent", "version": "1.0", "type": "agent_reference"},
),
):
run_options = await client.prepare_options(messages, chat_options)
run_options = await client._prepare_options(messages, chat_options)
# response_format should be excluded from final run options
assert "response_format" not in run_options
@@ -745,94 +745,8 @@ async def test_azure_ai_client_prepare_options_excludes_response_format(
assert run_options["extra_body"]["agent"]["name"] == "test-agent"
async def test_azure_ai_client_prepare_options_with_resp_conversation_id(
mock_project_client: MagicMock,
) -> None:
"""Test prepare_options with conversation ID starting with 'resp_'."""
client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0")
messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
chat_options = ChatOptions(conversation_id="resp_12345")
with (
patch.object(
client.__class__.__bases__[0],
"prepare_options",
return_value={"model": "test-model", "previous_response_id": "old_value", "conversation": "old_conv"},
),
patch.object(
client,
"_get_agent_reference_or_create",
return_value={"name": "test-agent", "version": "1.0", "type": "agent_reference"},
),
):
run_options = await client.prepare_options(messages, chat_options)
# Should set previous_response_id and remove conversation property
assert run_options["previous_response_id"] == "resp_12345"
assert "conversation" not in run_options
async def test_azure_ai_client_prepare_options_with_conv_conversation_id(
mock_project_client: MagicMock,
) -> None:
"""Test prepare_options with conversation ID starting with 'conv_'."""
client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0")
messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
chat_options = ChatOptions(conversation_id="conv_67890")
with (
patch.object(
client.__class__.__bases__[0],
"prepare_options",
return_value={"model": "test-model", "previous_response_id": "old_value", "conversation": "old_conv"},
),
patch.object(
client,
"_get_agent_reference_or_create",
return_value={"name": "test-agent", "version": "1.0", "type": "agent_reference"},
),
):
run_options = await client.prepare_options(messages, chat_options)
# Should set conversation and remove previous_response_id property
assert run_options["conversation"] == "conv_67890"
assert "previous_response_id" not in run_options
async def test_azure_ai_client_prepare_options_with_client_conversation_id(
mock_project_client: MagicMock,
) -> None:
"""Test prepare_options using client's default conversation ID when chat options don't have one."""
client = create_test_azure_ai_client(
mock_project_client, agent_name="test-agent", agent_version="1.0", conversation_id="resp_client_default"
)
messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
chat_options = ChatOptions() # No conversation_id specified
with (
patch.object(
client.__class__.__bases__[0],
"prepare_options",
return_value={"model": "test-model", "previous_response_id": "old_value", "conversation": "old_conv"},
),
patch.object(
client,
"_get_agent_reference_or_create",
return_value={"name": "test-agent", "version": "1.0", "type": "agent_reference"},
),
):
run_options = await client.prepare_options(messages, chat_options)
# Should use client's default conversation_id and set previous_response_id
assert run_options["previous_response_id"] == "resp_client_default"
assert "conversation" not in run_options
def test_get_conversation_id_with_store_true_and_conversation_id() -> None:
"""Test get_conversation_id returns conversation ID when store is True and conversation exists."""
"""Test _get_conversation_id returns conversation ID when store is True and conversation exists."""
client = create_test_azure_ai_client(MagicMock())
# Mock OpenAI response with conversation
@@ -842,13 +756,13 @@ def test_get_conversation_id_with_store_true_and_conversation_id() -> None:
mock_conversation.id = "conv_67890"
mock_response.conversation = mock_conversation
result = client.get_conversation_id(mock_response, store=True)
result = client._get_conversation_id(mock_response, store=True)
assert result == "conv_67890"
def test_get_conversation_id_with_store_true_and_no_conversation() -> None:
"""Test get_conversation_id returns response ID when store is True and no conversation exists."""
"""Test _get_conversation_id returns response ID when store is True and no conversation exists."""
client = create_test_azure_ai_client(MagicMock())
# Mock OpenAI response without conversation
@@ -856,13 +770,13 @@ def test_get_conversation_id_with_store_true_and_no_conversation() -> None:
mock_response.id = "resp_12345"
mock_response.conversation = None
result = client.get_conversation_id(mock_response, store=True)
result = client._get_conversation_id(mock_response, store=True)
assert result == "resp_12345"
def test_get_conversation_id_with_store_true_and_empty_conversation_id() -> None:
"""Test get_conversation_id returns response ID when store is True and conversation ID is empty."""
"""Test _get_conversation_id returns response ID when store is True and conversation ID is empty."""
client = create_test_azure_ai_client(MagicMock())
# Mock OpenAI response with conversation but empty ID
@@ -872,13 +786,13 @@ def test_get_conversation_id_with_store_true_and_empty_conversation_id() -> None
mock_conversation.id = ""
mock_response.conversation = mock_conversation
result = client.get_conversation_id(mock_response, store=True)
result = client._get_conversation_id(mock_response, store=True)
assert result == "resp_12345"
def test_get_conversation_id_with_store_false() -> None:
"""Test get_conversation_id returns None when store is False."""
"""Test _get_conversation_id returns None when store is False."""
client = create_test_azure_ai_client(MagicMock())
# Mock OpenAI response with conversation
@@ -888,13 +802,13 @@ def test_get_conversation_id_with_store_false() -> None:
mock_conversation.id = "conv_67890"
mock_response.conversation = mock_conversation
result = client.get_conversation_id(mock_response, store=False)
result = client._get_conversation_id(mock_response, store=False)
assert result is None
def test_get_conversation_id_with_parsed_response_and_store_true() -> None:
"""Test get_conversation_id works with ParsedResponse when store is True."""
"""Test _get_conversation_id works with ParsedResponse when store is True."""
client = create_test_azure_ai_client(MagicMock())
# Mock ParsedResponse with conversation
@@ -904,13 +818,13 @@ def test_get_conversation_id_with_parsed_response_and_store_true() -> None:
mock_conversation.id = "conv_parsed_67890"
mock_response.conversation = mock_conversation
result = client.get_conversation_id(mock_response, store=True)
result = client._get_conversation_id(mock_response, store=True)
assert result == "conv_parsed_67890"
def test_get_conversation_id_with_parsed_response_no_conversation() -> None:
"""Test get_conversation_id returns response ID with ParsedResponse when no conversation exists."""
"""Test _get_conversation_id returns response ID with ParsedResponse when no conversation exists."""
client = create_test_azure_ai_client(MagicMock())
# Mock ParsedResponse without conversation
@@ -918,7 +832,7 @@ def test_get_conversation_id_with_parsed_response_no_conversation() -> None:
mock_response.id = "resp_parsed_12345"
mock_response.conversation = None
result = client.get_conversation_id(mock_response, store=True)
result = client._get_conversation_id(mock_response, store=True)
assert result == "resp_parsed_12345"