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/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"