Python: [BREAKING]: Introducing Options as TypedDict and Generic (#3140)

* WIP typeddict for options

* updated all clients and ChatAgents

* updated everything

* added ADR

* fix mypy

* proper typevar imports

* fixed import

* fixed other imports

* slight update in the sample

* updated from feedback

* fixes

* fixed missing covariants and test fixes

* fixed typing

* updated anthropic thinking config

* ruff fixes

* fixed int tests

* fix tests and mypy

* updated integration tests

* updated docstring and test fix

* improved options handling in obser

* mypy fix

* updated a host of integration tests

* fix tests

* bedrock fix
This commit is contained in:
Eduard van Valkenburg
2026-01-13 17:41:05 +01:00
committed by GitHub
Unverified
parent 5faa2851bb
commit 3e97425245
111 changed files with 6141 additions and 4715 deletions
@@ -299,8 +299,7 @@ async def test_azure_assistants_client_get_response_tools() -> None:
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
options={"tools": [get_weather], "tool_choice": "auto"},
)
assert response is not None
@@ -352,8 +351,7 @@ async def test_azure_assistants_client_streaming_tools() -> None:
# Test that the client can be used to get a response
response = azure_assistants_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
options={"tools": [get_weather], "tool_choice": "auto"},
)
full_message: str = ""
async for chunk in response:
@@ -212,7 +212,7 @@ async def test_cmc_with_logit_bias(
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(messages=chat_history, logit_bias=token_bias)
await azure_chat_client.get_response(messages=chat_history, options={"logit_bias": token_bias})
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
@@ -237,7 +237,7 @@ async def test_cmc_with_stop(
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(messages=chat_history, stop=stop)
await azure_chat_client.get_response(messages=chat_history, options={"stop": stop})
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
@@ -300,7 +300,7 @@ async def test_azure_on_your_data(
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
options={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
@@ -370,7 +370,7 @@ async def test_azure_on_your_data_string(
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
options={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
@@ -429,7 +429,7 @@ async def test_azure_on_your_data_fail(
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
options={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
@@ -652,8 +652,7 @@ async def test_azure_openai_chat_client_response_tools() -> None:
# Test that the client can be used to get a response
response = await azure_chat_client.get_response(
messages=messages,
tools=[get_story_text],
tool_choice="auto",
options={"tools": [get_story_text], "tool_choice": "auto"},
)
assert response is not None
@@ -709,8 +708,7 @@ async def test_azure_openai_chat_client_streaming_tools() -> None:
# Test that the client can be used to get a response
response = azure_chat_client.get_streaming_response(
messages=messages,
tools=[get_story_text],
tool_choice="auto",
options={"tools": [get_story_text], "tool_choice": "auto"},
)
full_message: str = ""
async for chunk in response:
@@ -1,26 +1,25 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import os
from typing import Annotated
from typing import Annotated, Any
import pytest
from azure.identity import AzureCliCredential
from pydantic import BaseModel
from pytest import param
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
HostedCodeInterpreterTool,
HostedFileSearchTool,
HostedMCPTool,
HostedVectorStoreContent,
TextContent,
HostedWebSearchTool,
ai_function,
)
from agent_framework.azure import AzureOpenAIResponsesClient
@@ -74,7 +73,7 @@ async def delete_vector_store(client: AzureOpenAIResponsesClient, file_id: str,
def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
azure_responses_client = AzureOpenAIResponsesClient()
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert azure_responses_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, ChatClientProtocol)
@@ -141,283 +140,286 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
assert "User-Agent" not in dumped_settings["default_headers"]
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_response() -> None:
"""Test azure responses client responses."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await azure_responses_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert "scientists" in response.text
messages.clear()
messages.append(ChatMessage(role="user", text="The weather in New York is sunny"))
messages.append(ChatMessage(role="user", text="What is the weather in New York?"))
# Test that the client can be used to get a structured response
structured_response = await azure_responses_client.get_response( # type: ignore[reportAssignmentType]
messages=messages,
response_format=OutputStruct,
)
assert structured_response is not None
assert isinstance(structured_response, ChatResponse)
assert isinstance(structured_response.value, OutputStruct)
assert structured_response.value.location == "New York"
assert "sunny" in structured_response.value.weather.lower()
# region Integration Tests
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_response_tools() -> None:
"""Test azure responses client tools."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What is the weather in New York?"))
# Test that the client can be used to get a response
response = await azure_responses_client.get_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
)
assert response is not None
assert isinstance(response, ChatResponse)
assert "sunny" in response.text
messages.clear()
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
# Test that the client can be used to get a response
structured_response: ChatResponse = await azure_responses_client.get_response( # type: ignore[reportAssignmentType]
messages=messages,
tools=[get_weather],
tool_choice="auto",
response_format=OutputStruct,
)
assert structured_response is not None
assert isinstance(structured_response, ChatResponse)
assert isinstance(structured_response.value, OutputStruct)
assert "Seattle" in structured_response.value.location
assert "sunny" in structured_response.value.weather.lower()
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_streaming() -> None:
"""Test Azure azure responses client streaming responses."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = azure_responses_client.get_streaming_response(messages=messages)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "scientists" in full_message
messages.clear()
messages.append(ChatMessage(role="user", text="The weather in Seattle is sunny"))
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
structured_response = await ChatResponse.from_chat_response_generator(
azure_responses_client.get_streaming_response(
messages=messages,
response_format=OutputStruct,
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
# Simple ChatOptions - just verify they don't fail
param("temperature", 0.7, False, id="temperature"),
param("top_p", 0.9, False, id="top_p"),
param("max_tokens", 500, False, id="max_tokens"),
param("seed", 123, False, id="seed"),
param("user", "test-user-id", False, id="user"),
param("metadata", {"test_key": "test_value"}, False, id="metadata"),
param("frequency_penalty", 0.5, False, id="frequency_penalty"),
param("presence_penalty", 0.3, False, id="presence_penalty"),
param("stop", ["END"], False, id="stop"),
param("allow_multiple_tool_calls", True, False, id="allow_multiple_tool_calls"),
param("tool_choice", "none", True, id="tool_choice_none"),
# OpenAIResponsesOptions - just verify they don't fail
param("safety_identifier", "user-hash-abc123", False, id="safety_identifier"),
param("truncation", "auto", False, id="truncation"),
param("top_logprobs", 5, False, id="top_logprobs"),
param("prompt_cache_key", "test-cache-key", False, id="prompt_cache_key"),
param("max_tool_calls", 3, False, id="max_tool_calls"),
# Complex options requiring output validation
param("tools", [get_weather], True, id="tools_function"),
param("tool_choice", "auto", True, id="tool_choice_auto"),
param(
"tool_choice",
{"mode": "required", "required_function_name": "get_weather"},
True,
id="tool_choice_required",
),
output_format_type=OutputStruct,
)
assert structured_response is not None
assert isinstance(structured_response, ChatResponse)
assert isinstance(structured_response.value, OutputStruct)
assert "Seattle" in structured_response.value.location
assert "sunny" in structured_response.value.weather.lower()
param("response_format", OutputStruct, True, id="response_format_pydantic"),
param(
"response_format",
{
"type": "json_schema",
"json_schema": {
"name": "WeatherDigest",
"strict": True,
"schema": {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
"temperature_c": {"type": "number"},
"advisory": {"type": "string"},
},
"required": ["location", "conditions", "temperature_c", "advisory"],
"additionalProperties": False,
},
},
},
True,
id="response_format_runtime_json_schema",
),
],
)
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
"""Parametrized test covering all ChatOptions and OpenAIResponsesOptions.
Tests both streaming and non-streaming modes for each option to ensure
they don't cause failures. Options marked with needs_validation also
check that the feature actually works correctly.
"""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
# to ensure toolmode required does not endlessly loop
client.function_invocation_configuration.max_iterations = 1
for streaming in [False, True]:
# Prepare test message
if option_name == "tools" or option_name == "tool_choice":
# Use weather-related prompt for tool tests
messages = [ChatMessage(role="user", text="What is the weather in Seattle?")]
elif option_name == "response_format":
# Use prompt that works well with structured output
messages = [ChatMessage(role="user", text="The weather in Seattle is sunny")]
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
# Add tools if testing tool_choice to avoid errors
if option_name == "tool_choice":
options["tools"] = [get_weather]
if streaming:
# Test streaming mode
response_gen = client.get_streaming_response(
messages=messages,
options=options,
)
output_format = option_value if option_name == "response_format" else None
response = await ChatResponse.from_chat_response_generator(response_gen, output_format_type=output_format)
else:
# Test non-streaming mode
response = await client.get_response(
messages=messages,
options=options,
)
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation:
if option_name == "tools" or option_name == "tool_choice":
# Should have called the weather function
text = response.text.lower()
assert "sunny" in text or "seattle" in text, f"Tool not invoked for {option_name}"
elif option_name == "response_format":
if option_value == OutputStruct:
# Should have structured output
assert response.value is not None, "No structured output"
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
# Runtime JSON schema
assert response.value is None, "No structured output, can't parse any json."
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
assert "seattle" in response_value["location"].lower()
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_streaming_tools() -> None:
"""Test azure responses client streaming tools."""
async def test_integration_web_search() -> None:
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
for streaming in [False, True]:
content = {
"messages": "Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
"options": {
"tool_choice": "auto",
"tools": [HostedWebSearchTool()],
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
assert response is not None
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
# Test that the client will use the web search tool with location
additional_properties = {
"user_location": {
"country": "US",
"city": "Seattle",
}
}
content = {
"messages": "What is the current weather? Do not ask for my current location.",
"options": {
"tool_choice": "auto",
"tools": [HostedWebSearchTool(additional_properties=additional_properties)],
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
assert response.text is not None
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_integration_client_file_search() -> None:
"""Test Azure responses client with file search tool."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = [ChatMessage(role="user", text="What is the weather in Seattle?")]
# Test that the client can be used to get a response
response = azure_responses_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "sunny" in full_message
messages.clear()
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
structured_response = azure_responses_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
response_format=OutputStruct,
)
full_message = ""
async for chunk in structured_response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
output = OutputStruct.model_validate_json(full_message)
assert "Seattle" in output.location
assert "sunny" in output.weather.lower()
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_basic_run():
"""Test Azure Responses Client agent basic run functionality with AzureOpenAIResponsesClient."""
agent = AzureOpenAIResponsesClient(credential=AzureCliCredential()).create_agent(
instructions="You are a helpful assistant.",
)
# Test basic run
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
assert "hello world" in response.text.lower()
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_basic_run_streaming():
"""Test Azure Responses Client agent basic streaming functionality with AzureOpenAIResponsesClient."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
) as agent:
# Test streaming run
full_text = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_text += chunk.text
assert len(full_text) > 0
assert "streaming response test" in full_text.lower()
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_thread_persistence():
"""Test Azure Responses Client agent thread persistence across runs with AzureOpenAIResponsesClient."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# First interaction
first_response = await agent.run("My favorite programming language is Python. Remember this.", thread=thread)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Second interaction - test memory
second_response = await agent.run("What is my favorite programming language?", thread=thread)
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_thread_storage_with_store_true():
"""Test Azure Responses Client agent with store=True to verify service_thread_id is returned."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
) as agent:
# Create a new thread
thread = AgentThread()
# Initially, service_thread_id should be None
assert thread.service_thread_id is None
# Run with store=True to store messages on Azure/OpenAI side
response = await agent.run(
"Hello! Please remember that my name is Alex.",
thread=thread,
store=True,
file_id, vector_store = await create_vector_store(azure_responses_client)
try:
# Test that the client will use the file search tool
response = await azure_responses_client.get_response(
messages=[
ChatMessage(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
options={"tools": [HostedFileSearchTool(inputs=vector_store)], "tool_choice": "auto"},
)
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# After store=True, service_thread_id should be populated
assert thread.service_thread_id is not None
assert isinstance(thread.service_thread_id, str)
assert len(thread.service_thread_id) > 0
assert "sunny" in response.text.lower()
assert "75" in response.text
finally:
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_existing_thread():
async def test_integration_client_file_search_streaming() -> None:
"""Test Azure responses client with file search tool and streaming."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
file_id, vector_store = await create_vector_store(azure_responses_client)
# Test that the client will use the file search tool
try:
response = azure_responses_client.get_streaming_response(
messages=[
ChatMessage(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
options={"tools": [HostedFileSearchTool(inputs=vector_store)], "tool_choice": "auto"},
)
assert response is not None
full_response = await ChatResponse.from_chat_response_generator(response)
assert "sunny" in full_response.text.lower()
assert "75" in full_response.text
finally:
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_hosted_mcp_tool() -> None:
"""Integration test for HostedMCPTool with Azure Response Agent using Microsoft Learn MCP."""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
response = await client.get_response(
"How to create an Azure storage account using az cli?",
options={
# this needs to be high enough to handle the full MCP tool response.
"max_tokens": 5000,
"tools": HostedMCPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
description="A Microsoft Learn MCP server for documentation questions",
approval_mode="never_require",
),
},
)
assert isinstance(response, ChatResponse)
assert response.text
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_hosted_code_interpreter_tool():
"""Test Azure Responses Client agent with HostedCodeInterpreterTool through AzureOpenAIResponsesClient."""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
response = await client.get_response(
"Calculate the sum of numbers from 1 to 10 using Python code.",
options={
"tools": [HostedCodeInterpreterTool()],
},
)
# Should contain calculation result (sum of 1-10 = 55) or code execution content
contains_relevant_content = any(
term in response.text.lower() for term in ["55", "sum", "code", "python", "calculate", "10"]
)
assert contains_relevant_content or len(response.text.strip()) > 10
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_existing_thread():
"""Test Azure Responses Client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
@@ -428,7 +430,7 @@ async def test_azure_responses_client_agent_existing_thread():
) as first_agent:
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My hobby is photography. Remember this.", thread=thread)
first_response = await first_agent.run("My hobby is photography. Remember this.", thread=thread, store=True)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
@@ -448,189 +450,3 @@ async def test_azure_responses_client_agent_existing_thread():
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
assert "photography" in second_response.text.lower()
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_hosted_code_interpreter_tool():
"""Test Azure Responses Client agent with HostedCodeInterpreterTool through AzureOpenAIResponsesClient."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can execute Python code.",
tools=[HostedCodeInterpreterTool()],
) as agent:
# Test code interpreter functionality
response = await agent.run("Calculate the sum of numbers from 1 to 10 using Python code.")
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# Should contain calculation result (sum of 1-10 = 55) or code execution content
contains_relevant_content = any(
term in response.text.lower() for term in ["55", "sum", "code", "python", "calculate", "10"]
)
assert contains_relevant_content or len(response.text.strip()) > 10
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Responses Client."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
# First run - agent-level tool should be available
first_response = await agent.run("What's the weather like in Chicago?")
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Should use the agent-level weather tool
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - agent-level tool should still be available (persistence test)
second_response = await agent.run("What's the weather in Miami?")
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
# Should use the agent-level weather tool again
assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_chat_options_run_level() -> None:
"""Integration test for comprehensive ChatOptions parameter coverage with Azure Response Agent."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
) as agent:
response = await agent.run(
"Provide a brief, helpful response.",
max_tokens=100,
temperature=0.7,
top_p=0.9,
seed=123,
user="comprehensive-test-user",
tools=[get_weather],
tool_choice="auto",
)
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_chat_options_agent_level() -> None:
"""Integration test for comprehensive ChatOptions parameter coverage with Azure Response Agent."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
max_tokens=100,
temperature=0.7,
top_p=0.9,
seed=123,
user="comprehensive-test-user",
tools=[get_weather],
tool_choice="auto",
) as agent:
response = await agent.run(
"Provide a brief, helpful response.",
)
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_hosted_mcp_tool() -> None:
"""Integration test for HostedMCPTool with Azure Response Agent using Microsoft Learn MCP."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=HostedMCPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
description="A Microsoft Learn MCP server for documentation questions",
approval_mode="never_require",
),
) as agent:
response = await agent.run(
"How to create an Azure storage account using az cli?",
# this needs to be high enough to handle the full MCP tool response.
max_tokens=5000,
)
assert isinstance(response, AgentRunResponse)
assert response.text
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_file_search() -> None:
"""Test Azure responses client with file search tool."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
file_id, vector_store = await create_vector_store(azure_responses_client)
# Test that the client will use the file search tool
response = await azure_responses_client.get_response(
messages=[
ChatMessage(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
tools=[HostedFileSearchTool(inputs=vector_store)],
tool_choice="auto",
)
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
assert "sunny" in response.text.lower()
assert "75" in response.text
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_file_search_streaming() -> None:
"""Test Azure responses client with file search tool and streaming."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
file_id, vector_store = await create_vector_store(azure_responses_client)
# Test that the client will use the file search tool
response = azure_responses_client.get_streaming_response(
messages=[
ChatMessage(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
tools=[HostedFileSearchTool(inputs=vector_store)],
tool_choice="auto",
)
assert response is not None
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
assert "sunny" in full_message.lower()
assert "75" in full_message
+10 -10
View File
@@ -4,7 +4,7 @@ import asyncio
import logging
import sys
from collections.abc import AsyncIterable, MutableSequence
from typing import Any
from typing import Any, Generic
from unittest.mock import patch
from uuid import uuid4
@@ -18,7 +18,6 @@ from agent_framework import (
AgentThread,
BaseChatClient,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Role,
@@ -28,6 +27,7 @@ from agent_framework import (
use_chat_middleware,
use_function_invocation,
)
from agent_framework._clients import TOptions_co
if sys.version_info >= (3, 12):
from typing import override # type: ignore
@@ -113,7 +113,7 @@ class MockChatClient:
@use_chat_middleware
class MockBaseChatClient(BaseChatClient):
class MockBaseChatClient(BaseChatClient[TOptions_co], Generic[TOptions_co]):
"""Mock implementation of the BaseChatClient."""
def __init__(self, **kwargs: Any):
@@ -127,27 +127,27 @@ class MockBaseChatClient(BaseChatClient):
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
options: dict[str, Any],
**kwargs: Any,
) -> ChatResponse:
"""Send a chat request to the AI service.
Args:
messages: The chat messages to send.
chat_options: The options for the request.
options: The options dict for the request.
kwargs: Any additional keyword arguments.
Returns:
The chat response contents representing the response(s).
"""
logger.debug(f"Running base chat client inner, with: {messages=}, {chat_options=}, {kwargs=}")
logger.debug(f"Running base chat client inner, with: {messages=}, {options=}, {kwargs=}")
self.call_count += 1
if not self.run_responses:
return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[-1].text}"))
response = self.run_responses.pop(0)
if chat_options.tool_choice == "none":
if options.get("tool_choice") == "none":
return ChatResponse(
messages=ChatMessage(
role="assistant",
@@ -163,14 +163,14 @@ class MockBaseChatClient(BaseChatClient):
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
options: dict[str, Any],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
logger.debug(f"Running base chat client inner stream, with: {messages=}, {chat_options=}, {kwargs=}")
logger.debug(f"Running base chat client inner stream, with: {messages=}, {options=}, {kwargs=}")
if not self.streaming_responses:
yield ChatResponseUpdate(text=f"update - {messages[0].text}", role="assistant")
return
if chat_options.tool_choice == "none":
if options.get("tool_choice") == "none":
yield ChatResponseUpdate(text="I broke out of the function invocation loop...", role="assistant")
return
response = self.streaming_responses.pop(0)
+39 -28
View File
@@ -118,8 +118,8 @@ async def test_prepare_thread_does_not_mutate_agent_chat_options(chat_client: Ch
tool = HostedCodeInterpreterTool()
agent = ChatAgent(chat_client=chat_client, tools=[tool])
assert agent.chat_options.tools is not None
base_tools = agent.chat_options.tools
assert agent.default_options.get("tools") is not None
base_tools = agent.default_options["tools"]
thread = agent.get_new_thread()
_, prepared_chat_options, _ = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
@@ -127,11 +127,11 @@ async def test_prepare_thread_does_not_mutate_agent_chat_options(chat_client: Ch
input_messages=[ChatMessage(role=Role.USER, text="Test")],
)
assert prepared_chat_options.tools is not None
assert base_tools is not prepared_chat_options.tools
assert prepared_chat_options.get("tools") is not None
assert base_tools is not prepared_chat_options["tools"]
prepared_chat_options.tools.append(HostedCodeInterpreterTool()) # type: ignore[arg-type]
assert len(agent.chat_options.tools) == 1
prepared_chat_options["tools"].append(HostedCodeInterpreterTool()) # type: ignore[arg-type]
assert len(agent.default_options["tools"]) == 1
async def test_chat_client_agent_update_thread_id(chat_client_base: ChatClientProtocol) -> None:
@@ -597,61 +597,68 @@ async def test_chat_agent_tool_choice_run_level_overrides_agent_level(
chat_client_base: Any, ai_function_tool: Any
) -> None:
"""Verify that tool_choice passed to run() overrides agent-level tool_choice."""
from agent_framework import ChatOptions, ToolMode
captured_options: list[ChatOptions] = []
captured_options: list[dict[str, Any]] = []
# Store the original inner method
original_inner = chat_client_base._inner_get_response
async def capturing_inner(
*, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
*, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> ChatResponse:
captured_options.append(chat_options)
return await original_inner(messages=messages, chat_options=chat_options, **kwargs)
captured_options.append(options)
return await original_inner(messages=messages, options=options, **kwargs)
chat_client_base._inner_get_response = capturing_inner
# Create agent with agent-level tool_choice="auto" and a tool (tools required for tool_choice to be meaningful)
agent = ChatAgent(chat_client=chat_client_base, tool_choice="auto", tools=[ai_function_tool])
agent = ChatAgent(
chat_client=chat_client_base,
tools=[ai_function_tool],
options={"tool_choice": "auto"},
)
# Run with run-level tool_choice="required"
await agent.run("Hello", tool_choice="required")
await agent.run("Hello", options={"tool_choice": "required"})
# Verify the client received tool_choice="required", not "auto"
assert len(captured_options) >= 1
assert captured_options[0].tool_choice == "required"
assert captured_options[0].tool_choice == ToolMode.REQUIRED_ANY
assert captured_options[0]["tool_choice"] == "required"
async def test_chat_agent_tool_choice_agent_level_used_when_run_level_not_specified(
chat_client_base: Any, ai_function_tool: Any
) -> None:
"""Verify that agent-level tool_choice is used when run() doesn't specify one."""
from agent_framework import ChatOptions, ToolMode
from agent_framework import ChatOptions
captured_options: list[ChatOptions] = []
original_inner = chat_client_base._inner_get_response
async def capturing_inner(
*, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
*, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> ChatResponse:
captured_options.append(chat_options)
return await original_inner(messages=messages, chat_options=chat_options, **kwargs)
captured_options.append(options)
return await original_inner(messages=messages, options=options, **kwargs)
chat_client_base._inner_get_response = capturing_inner
# Create agent with agent-level tool_choice="required" and a tool
agent = ChatAgent(chat_client=chat_client_base, tool_choice="required", tools=[ai_function_tool])
agent = ChatAgent(
chat_client=chat_client_base,
tools=[ai_function_tool],
default_options={"tool_choice": "required"},
)
# Run without specifying tool_choice
await agent.run("Hello")
# Verify the client received tool_choice="required" from agent-level
assert len(captured_options) >= 1
assert captured_options[0].tool_choice == "required"
assert captured_options[0].tool_choice == ToolMode.REQUIRED_ANY
assert captured_options[0]["tool_choice"] == "required"
# older code compared to ToolMode constants; ensure value is 'required'
assert captured_options[0]["tool_choice"] == "required"
async def test_chat_agent_tool_choice_none_at_run_preserves_agent_level(
@@ -665,19 +672,23 @@ async def test_chat_agent_tool_choice_none_at_run_preserves_agent_level(
original_inner = chat_client_base._inner_get_response
async def capturing_inner(
*, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
*, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> ChatResponse:
captured_options.append(chat_options)
return await original_inner(messages=messages, chat_options=chat_options, **kwargs)
captured_options.append(options)
return await original_inner(messages=messages, options=options, **kwargs)
chat_client_base._inner_get_response = capturing_inner
# Create agent with agent-level tool_choice="auto" and a tool
agent = ChatAgent(chat_client=chat_client_base, tool_choice="auto", tools=[ai_function_tool])
agent = ChatAgent(
chat_client=chat_client_base,
tools=[ai_function_tool],
default_options={"tool_choice": "auto"},
)
# Run with explicitly passing None (same as not specifying)
await agent.run("Hello", tool_choice=None)
await agent.run("Hello", options={"tool_choice": None})
# Verify the client received tool_choice="auto" from agent-level
assert len(captured_options) >= 1
assert captured_options[0].tool_choice == "auto"
assert captured_options[0]["tool_choice"] == "auto"
@@ -0,0 +1,433 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import os
from typing import Annotated
import pytest
from pydantic import BaseModel
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatAgent,
HostedCodeInterpreterTool,
HostedImageGenerationTool,
HostedMCPTool,
MCPStreamableHTTPTool,
ai_function,
)
from agent_framework.openai import OpenAIResponsesClient
skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
or os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
reason="No real OPENAI_API_KEY provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
)
@ai_function
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
# Implementation of the tool to get weather
return f"The current weather in {location} is sunny."
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_basic_run_streaming():
"""Test OpenAI Responses Client agent basic streaming functionality with OpenAIResponsesClient."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
) as agent:
# Test streaming run
full_text = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_text += chunk.text
assert len(full_text) > 0
assert "streaming response test" in full_text.lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_thread_persistence():
"""Test OpenAI Responses Client agent thread persistence across runs with OpenAIResponsesClient."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# First interaction
first_response = await agent.run("My favorite programming language is Python. Remember this.", thread=thread)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Second interaction - test memory
second_response = await agent.run("What is my favorite programming language?", thread=thread)
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_thread_storage_with_store_true():
"""Test OpenAI Responses Client agent with store=True to verify service_thread_id is returned."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant.",
) as agent:
# Create a new thread
thread = AgentThread()
# Initially, service_thread_id should be None
assert thread.service_thread_id is None
# Run with store=True to store messages on OpenAI side
response = await agent.run(
"Hello! Please remember that my name is Alex.",
thread=thread,
options={"store": True},
)
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# After store=True, service_thread_id should be populated
assert thread.service_thread_id is not None
assert isinstance(thread.service_thread_id, str)
assert len(thread.service_thread_id) > 0
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_existing_thread():
"""Test OpenAI Responses Client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My hobby is photography. Remember this.", thread=thread)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Preserve the thread for reuse
preserved_thread = thread
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved thread
second_response = await second_agent.run("What is my hobby?", thread=preserved_thread)
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
assert "photography" in second_response.text.lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_hosted_code_interpreter_tool():
"""Test OpenAI Responses Client agent with HostedCodeInterpreterTool through OpenAIResponsesClient."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can execute Python code.",
tools=[HostedCodeInterpreterTool()],
) as agent:
# Test code interpreter functionality
response = await agent.run("Calculate the sum of numbers from 1 to 10 using Python code.")
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# Should contain calculation result (sum of 1-10 = 55) or code execution content
contains_relevant_content = any(
term in response.text.lower() for term in ["55", "sum", "code", "python", "calculate", "10"]
)
assert contains_relevant_content or len(response.text.strip()) > 10
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_image_generation_tool():
"""Test OpenAI Responses Client agent with raw image_generation tool through OpenAIResponsesClient."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can generate images.",
tools=HostedImageGenerationTool(options={"image_size": "1024x1024", "media_type": "png"}),
) as agent:
# Test image generation functionality
response = await agent.run("Generate an image of a cute red panda sitting on a tree branch in a forest.")
assert isinstance(response, AgentRunResponse)
assert response.messages
# Verify we got image content - look for ImageGenerationToolResultContent
image_content_found = False
for message in response.messages:
for content in message.contents:
if content.type == "image_generation_tool_result" and content.outputs:
image_content_found = True
break
if image_content_found:
break
# The test passes if we got image content
assert image_content_found, "Expected to find image content in response"
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with OpenAI Responses Client."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
# First run - agent-level tool should be available
first_response = await agent.run("What's the weather like in Chicago?")
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Should use the agent-level weather tool
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - agent-level tool should still be available (persistence test)
second_response = await agent.run("What's the weather in Miami?")
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
# Should use the agent-level weather tool again
assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_run_level_tool_isolation():
"""Test that run-level tools are isolated to specific runs and don't persist with OpenAI Responses Client."""
# Counter to track how many times the weather tool is called
call_count = 0
@ai_function
async def get_weather_with_counter(
location: Annotated[str, "The location as a city name"],
) -> str:
"""Get the current weather in a given location."""
nonlocal call_count
call_count += 1
return f"The weather in {location} is sunny and 72°F."
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant.",
) as agent:
# First run - use run-level tool
first_response = await agent.run(
"What's the weather like in Chicago?",
tools=[get_weather_with_counter], # Run-level tool
)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Should use the run-level weather tool (call count should be 1)
assert call_count == 1
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - run-level tool should NOT persist (key isolation test)
second_response = await agent.run("What's the weather like in Miami?")
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
# Should NOT use the weather tool since it was only run-level in previous call
# Call count should still be 1 (no additional calls)
assert call_count == 1
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_chat_options_agent_level() -> None:
"""Integration test for comprehensive ChatOptions parameter coverage with OpenAI Response Agent."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant.",
tools=[get_weather],
default_options={
"max_tokens": 100,
"temperature": 0.7,
"top_p": 0.9,
"seed": 123,
"user": "comprehensive-test-user",
"tool_choice": "auto",
},
) as agent:
response = await agent.run(
"Provide a brief, helpful response.",
)
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_hosted_mcp_tool() -> None:
"""Integration test for HostedMCPTool with OpenAI Response Agent using Microsoft Learn MCP."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=HostedMCPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
description="A Microsoft Learn MCP server for documentation questions",
approval_mode="never_require",
),
) as agent:
response = await agent.run(
"How to create an Azure storage account using az cli?",
# this needs to be high enough to handle the full MCP tool response.
options={"max_tokens": 5000},
)
assert isinstance(response, AgentRunResponse)
assert response.text
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_local_mcp_tool() -> None:
"""Integration test for MCPStreamableHTTPTool with OpenAI Response Agent using Microsoft Learn MCP."""
mcp_tool = MCPStreamableHTTPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
)
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=[mcp_tool],
) as agent:
response = await agent.run(
"How to create an Azure storage account using az cli?",
options={"max_tokens": 200},
)
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
class ReleaseBrief(BaseModel):
"""Structured output model for release brief testing."""
title: str
summary: str
highlights: list[str]
model_config = {"extra": "forbid"}
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_with_response_format_pydantic() -> None:
"""Integration test for response_format with Pydantic model using OpenAI Responses Client."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that returns structured JSON responses.",
) as agent:
response = await agent.run(
"Summarize the following release notes into a ReleaseBrief:\n\n"
"Version 2.0 Release Notes:\n"
"- Added new streaming API for real-time responses\n"
"- Improved error handling with detailed messages\n"
"- Performance boost of 50% in batch processing\n"
"- Fixed memory leak in connection pooling",
options={
"response_format": ReleaseBrief,
},
)
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.value is not None
assert isinstance(response.value, ReleaseBrief)
# Validate structured output fields
brief = response.value
assert len(brief.title) > 0
assert len(brief.summary) > 0
assert len(brief.highlights) > 0
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_with_runtime_json_schema() -> None:
"""Integration test for response_format with runtime JSON schema using OpenAI Responses Client."""
runtime_schema = {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
"temperature_c": {"type": "number"},
"advisory": {"type": "string"},
},
"required": ["location", "conditions", "temperature_c", "advisory"],
"additionalProperties": False,
}
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="Return only JSON that matches the provided schema. Do not add commentary.",
) as agent:
response = await agent.run(
"Give a brief weather digest for Seattle.",
options={
"response_format": {
"type": "json_schema",
"json_schema": {
"name": runtime_schema["title"],
"strict": True,
"schema": runtime_schema,
},
},
},
)
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.text is not None
# Parse JSON and validate structure
parsed = json.loads(response.text)
assert "location" in parsed
assert "conditions" in parsed
assert "temperature_c" in parsed
assert "advisory" in parsed
@@ -7,7 +7,6 @@ from agent_framework import (
BaseChatClient,
ChatClientProtocol,
ChatMessage,
ChatOptions,
Role,
)
@@ -50,12 +49,22 @@ async def test_chat_client_instructions_handling(chat_client_base: ChatClientPro
chat_client_base,
"_inner_get_response",
) as mock_inner_get_response:
await chat_client_base.get_response("hello", chat_options=ChatOptions(instructions=instructions))
await chat_client_base.get_response("hello", options={"instructions": instructions})
mock_inner_get_response.assert_called_once()
_, kwargs = mock_inner_get_response.call_args
messages = kwargs.get("messages", [])
assert len(messages) == 2
assert messages[0].role == Role.SYSTEM
assert messages[0].text == instructions
assert messages[1].role == Role.USER
assert messages[1].text == "hello"
assert len(messages) == 1
assert messages[0].role == Role.USER
assert messages[0].text == "hello"
from agent_framework._types import prepend_instructions_to_messages
appended_messages = prepend_instructions_to_messages(
[ChatMessage(role=Role.USER, text="hello")],
instructions,
)
assert len(appended_messages) == 2
assert appended_messages[0].role == Role.SYSTEM
assert appended_messages[0].text == "You are a helpful assistant."
assert appended_messages[1].role == Role.USER
assert appended_messages[1].text == "hello"
@@ -1,6 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Awaitable, Callable
from typing import Any
import pytest
@@ -8,7 +9,6 @@ from agent_framework import (
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
FunctionApprovalRequestContent,
@@ -39,7 +39,7 @@ async def test_base_client_with_function_calling(chat_client_base: ChatClientPro
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]})
assert exec_counter == 1
assert len(response.messages) == 3
assert response.messages[0].role == Role.ASSISTANT
@@ -79,7 +79,7 @@ async def test_base_client_with_function_calling_resets(chat_client_base: ChatCl
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]})
assert exec_counter == 2
assert len(response.messages) == 5
assert response.messages[0].role == Role.ASSISTANT
@@ -121,7 +121,9 @@ async def test_base_client_with_streaming_function_calling(chat_client_base: Cha
],
]
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[ai_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [ai_func]}
):
updates.append(update)
assert len(updates) == 4 # two updates with the function call, the function result and the final text
assert updates[0].contents[0].call_id == "1"
@@ -371,18 +373,18 @@ async def test_function_invocation_scenarios(
]
# Execute the test
chat_options = ChatOptions(tool_choice="auto", tools=tools)
options: dict[str, Any] = {"tool_choice": "auto", "tools": tools}
if thread_type == "service":
# For service threads, we need to pass conversation_id via ChatOptions
chat_options.store = True
chat_options.conversation_id = conversation_id
# For service threads, we need to pass conversation_id via options
options["store"] = True
options["conversation_id"] = conversation_id
if not streaming:
response = await chat_client_base.get_response("hello", chat_options=chat_options)
response = await chat_client_base.get_response("hello", options=options)
messages = response.messages
else:
updates = []
async for update in chat_client_base.get_streaming_response("hello", chat_options=chat_options):
async for update in chat_client_base.get_streaming_response("hello", options=options):
updates.append(update)
messages = updates
@@ -492,7 +494,9 @@ async def test_rejected_approval(chat_client_base: ChatClientProtocol):
]
# Get the response with approval requests
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[func_approved, func_rejected])
response = await chat_client_base.get_response(
"hello", options={"tool_choice": "auto", "tools": [func_approved, func_rejected]}
)
# Approval requests are now added to the assistant message, not a separate message
assert len(response.messages) == 1
# Assistant message should have: 2 FunctionCallContent + 2 FunctionApprovalRequestContent
@@ -519,7 +523,9 @@ async def test_rejected_approval(chat_client_base: ChatClientProtocol):
all_messages = response.messages + [ChatMessage(role="user", contents=[approved_response, rejected_response])]
# Call get_response which will process the approvals
await chat_client_base.get_response(all_messages, tool_choice="auto", tools=[func_approved, func_rejected])
await chat_client_base.get_response(
all_messages, options={"tool_choice": "auto", "tools": [func_approved, func_rejected]}
)
# Verify the approval/rejection was processed correctly
# Find the results in the input messages (modified in-place)
@@ -574,7 +580,9 @@ async def test_approval_requests_in_assistant_message(chat_client_base: ChatClie
),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[func_with_approval])
response = await chat_client_base.get_response(
"hello", options={"tool_choice": "auto", "tools": [func_with_approval]}
)
# Should have one assistant message containing both the call and approval request
assert len(response.messages) == 1
@@ -610,7 +618,9 @@ async def test_persisted_approval_messages_replay_correctly(chat_client_base: Ch
]
# Get approval request
response1 = await chat_client_base.get_response("hello", tool_choice="auto", tools=[func_with_approval])
response1 = await chat_client_base.get_response(
"hello", options={"tool_choice": "auto", "tools": [func_with_approval]}
)
# Store messages (like a thread would)
persisted_messages = [
@@ -628,7 +638,9 @@ async def test_persisted_approval_messages_replay_correctly(chat_client_base: Ch
persisted_messages.append(ChatMessage(role="user", contents=[approval_response]))
# Continue with all persisted messages
response2 = await chat_client_base.get_response(persisted_messages, tool_choice="auto", tools=[func_with_approval])
response2 = await chat_client_base.get_response(
persisted_messages, options={"tool_choice": "auto", "tools": [func_with_approval]}
)
# Should execute successfully
assert response2 is not None
@@ -656,7 +668,9 @@ async def test_no_duplicate_function_calls_after_approval_processing(chat_client
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response1 = await chat_client_base.get_response("hello", tool_choice="auto", tools=[func_with_approval])
response1 = await chat_client_base.get_response(
"hello", options={"tool_choice": "auto", "tools": [func_with_approval]}
)
approval_req = [c for c in response1.messages[0].contents if isinstance(c, FunctionApprovalRequestContent)][0]
approval_response = FunctionApprovalResponseContent(
@@ -666,7 +680,7 @@ async def test_no_duplicate_function_calls_after_approval_processing(chat_client
)
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
await chat_client_base.get_response(all_messages, tool_choice="auto", tools=[func_with_approval])
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [func_with_approval]})
# Count function calls with the same call_id
function_call_count = sum(
@@ -699,7 +713,9 @@ async def test_rejection_result_uses_function_call_id(chat_client_base: ChatClie
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response1 = await chat_client_base.get_response("hello", tool_choice="auto", tools=[func_with_approval])
response1 = await chat_client_base.get_response(
"hello", options={"tool_choice": "auto", "tools": [func_with_approval]}
)
approval_req = [c for c in response1.messages[0].contents if isinstance(c, FunctionApprovalRequestContent)][0]
rejection_response = FunctionApprovalResponseContent(
@@ -709,7 +725,7 @@ async def test_rejection_result_uses_function_call_id(chat_client_base: ChatClie
)
all_messages = response1.messages + [ChatMessage(role="user", contents=[rejection_response])]
await chat_client_base.get_response(all_messages, tool_choice="auto", tools=[func_with_approval])
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [func_with_approval]})
# Find the rejection result
rejection_result = next(
@@ -753,7 +769,7 @@ async def test_max_iterations_limit(chat_client_base: ChatClientProtocol):
# Set max_iterations to 1 in additional_properties
chat_client_base.function_invocation_configuration.max_iterations = 1
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]})
# With max_iterations=1, we should:
# 1. Execute first function call (exec_counter=1)
@@ -780,7 +796,7 @@ async def test_function_invocation_config_enabled_false(chat_client_base: ChatCl
# Disable function invocation
chat_client_base.function_invocation_configuration.enabled = False
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]})
# Function should not be executed - when enabled=False, the loop doesn't run
assert exec_counter == 0
@@ -827,7 +843,7 @@ async def test_function_invocation_config_max_consecutive_errors(chat_client_bas
# Set max_consecutive_errors to 2
chat_client_base.function_invocation_configuration.max_consecutive_errors_per_request = 2
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[error_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]})
# Should stop after 2 consecutive errors and force a non-tool response
error_results = [
@@ -870,7 +886,7 @@ async def test_function_invocation_config_terminate_on_unknown_calls_false(chat_
# Set terminate_on_unknown_calls to False (default)
chat_client_base.function_invocation_configuration.terminate_on_unknown_calls = False
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[known_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [known_func]})
# Should have a result message indicating the tool wasn't found
assert len(response.messages) == 3
@@ -904,7 +920,7 @@ async def test_function_invocation_config_terminate_on_unknown_calls_true(chat_c
# Should raise an exception when encountering an unknown function
with pytest.raises(KeyError, match='Error: Requested function "unknown_function" not found'):
await chat_client_base.get_response("hello", tool_choice="auto", tools=[known_func])
await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [known_func]})
assert exec_counter == 0
@@ -940,7 +956,7 @@ async def test_function_invocation_config_additional_tools(chat_client_base: Cha
chat_client_base.function_invocation_configuration.additional_tools = [hidden_func]
# Only pass visible_func in the tools parameter
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[visible_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [visible_func]})
# Additional tools are treated as declaration_only, so not executed
# The function call should be in the messages but not executed
@@ -976,7 +992,7 @@ async def test_function_invocation_config_include_detailed_errors_false(chat_cli
# Set include_detailed_errors to False (default)
chat_client_base.function_invocation_configuration.include_detailed_errors = False
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[error_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]})
# Should have a generic error message
error_result = next(
@@ -1008,7 +1024,7 @@ async def test_function_invocation_config_include_detailed_errors_true(chat_clie
# Set include_detailed_errors to True
chat_client_base.function_invocation_configuration.include_detailed_errors = True
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[error_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]})
# Should have detailed error message
error_result = next(
@@ -1076,7 +1092,7 @@ async def test_argument_validation_error_with_detailed_errors(chat_client_base:
# Set include_detailed_errors to True
chat_client_base.function_invocation_configuration.include_detailed_errors = True
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[typed_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [typed_func]})
# Should have detailed validation error
error_result = next(
@@ -1108,7 +1124,7 @@ async def test_argument_validation_error_without_detailed_errors(chat_client_bas
# Set include_detailed_errors to False (default)
chat_client_base.function_invocation_configuration.include_detailed_errors = False
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[typed_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [typed_func]})
# Should have generic validation error
error_result = next(
@@ -1175,7 +1191,7 @@ async def test_unapproved_tool_execution_raises_exception(chat_client_base: Chat
]
# Get approval request
response1 = await chat_client_base.get_response("hello", tool_choice="auto", tools=[test_func])
response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [test_func]})
approval_req = [c for c in response1.messages[0].contents if isinstance(c, FunctionApprovalRequestContent)][0]
@@ -1190,7 +1206,7 @@ async def test_unapproved_tool_execution_raises_exception(chat_client_base: Chat
all_messages = response1.messages + [ChatMessage(role="user", contents=[rejection_response])]
# This should handle the rejection gracefully (not raise ToolException to user)
await chat_client_base.get_response(all_messages, tool_choice="auto", tools=[test_func])
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [test_func]})
# Should have a rejection result
rejection_result = next(
@@ -1235,7 +1251,7 @@ async def test_approved_function_call_with_error_without_detailed_errors(chat_cl
chat_client_base.function_invocation_configuration.include_detailed_errors = False
# Get approval request
response1 = await chat_client_base.get_response("hello", tool_choice="auto", tools=[error_func])
response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]})
approval_req = [c for c in response1.messages[0].contents if isinstance(c, FunctionApprovalRequestContent)][0]
@@ -1249,7 +1265,7 @@ async def test_approved_function_call_with_error_without_detailed_errors(chat_cl
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
# Execute the approved function (which will error)
await chat_client_base.get_response(all_messages, tool_choice="auto", tools=[error_func])
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [error_func]})
# Should have executed the function
assert exec_counter == 1
@@ -1299,7 +1315,7 @@ async def test_approved_function_call_with_error_with_detailed_errors(chat_clien
chat_client_base.function_invocation_configuration.include_detailed_errors = True
# Get approval request
response1 = await chat_client_base.get_response("hello", tool_choice="auto", tools=[error_func])
response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]})
approval_req = [c for c in response1.messages[0].contents if isinstance(c, FunctionApprovalRequestContent)][0]
@@ -1313,7 +1329,7 @@ async def test_approved_function_call_with_error_with_detailed_errors(chat_clien
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
# Execute the approved function (which will error)
await chat_client_base.get_response(all_messages, tool_choice="auto", tools=[error_func])
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [error_func]})
# Should have executed the function
assert exec_counter == 1
@@ -1361,7 +1377,7 @@ async def test_approved_function_call_with_validation_error(chat_client_base: Ch
chat_client_base.function_invocation_configuration.include_detailed_errors = True
# Get approval request
response1 = await chat_client_base.get_response("hello", tool_choice="auto", tools=[typed_func])
response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [typed_func]})
approval_req = [c for c in response1.messages[0].contents if isinstance(c, FunctionApprovalRequestContent)][0]
@@ -1375,7 +1391,7 @@ async def test_approved_function_call_with_validation_error(chat_client_base: Ch
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
# Execute the approved function (which will fail validation)
await chat_client_base.get_response(all_messages, tool_choice="auto", tools=[typed_func])
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [typed_func]})
# Should NOT have executed the function (validation failed before execution)
assert exec_counter == 0
@@ -1418,7 +1434,7 @@ async def test_approved_function_call_successful_execution(chat_client_base: Cha
]
# Get approval request
response1 = await chat_client_base.get_response("hello", tool_choice="auto", tools=[success_func])
response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [success_func]})
approval_req = [c for c in response1.messages[0].contents if isinstance(c, FunctionApprovalRequestContent)][0]
@@ -1432,7 +1448,7 @@ async def test_approved_function_call_successful_execution(chat_client_base: Cha
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
# Execute the approved function
await chat_client_base.get_response(all_messages, tool_choice="auto", tools=[success_func])
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [success_func]})
# Should have executed successfully
assert exec_counter == 1
@@ -1476,7 +1492,9 @@ async def test_declaration_only_tool(chat_client_base: ChatClientProtocol):
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[declaration_func])
response = await chat_client_base.get_response(
"hello", options={"tool_choice": "auto", "tools": [declaration_func]}
)
# Should have the function call in messages but not a result
function_calls = [
@@ -1530,7 +1548,7 @@ async def test_multiple_function_calls_parallel_execution(chat_client_base: Chat
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[func1, func2])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [func1, func2]})
# Both functions should have been executed
assert "func1_start" in exec_order
@@ -1566,7 +1584,7 @@ async def test_callable_function_converted_to_ai_function(chat_client_base: Chat
]
# Pass plain function (will be auto-converted)
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[plain_function])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [plain_function]})
# Function should be executed
assert exec_counter == 1
@@ -1598,7 +1616,7 @@ async def test_conversation_id_handling(chat_client_base: ChatClientProtocol):
),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[test_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [test_func]})
# Should have executed the function
results = [
@@ -1625,7 +1643,7 @@ async def test_function_result_appended_to_existing_assistant_message(chat_clien
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[test_func])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [test_func]})
# Should have messages with both function call and function result
assert len(response.messages) >= 2
@@ -1667,7 +1685,7 @@ async def test_error_recovery_resets_counter(chat_client_base: ChatClientProtoco
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[sometimes_fails])
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [sometimes_fails]})
# Should have both an error and a success
error_results = [
@@ -1714,7 +1732,7 @@ async def test_streaming_approval_request_generated(chat_client_base: ChatClient
# Get the streaming response with approval request
updates = []
async for update in chat_client_base.get_streaming_response(
"hello", tool_choice="auto", tools=[func_with_approval]
"hello", options={"tool_choice": "auto", "tools": [func_with_approval]}
):
updates.append(update)
@@ -1770,7 +1788,9 @@ async def test_streaming_max_iterations_limit(chat_client_base: ChatClientProtoc
chat_client_base.function_invocation_configuration.max_iterations = 1
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[ai_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [ai_func]}
):
updates.append(update)
# With max_iterations=1, we should only execute first function
@@ -1798,7 +1818,9 @@ async def test_streaming_function_invocation_config_enabled_false(chat_client_ba
chat_client_base.function_invocation_configuration.enabled = False
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[ai_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [ai_func]}
):
updates.append(update)
# Function should not be executed - when enabled=False, the loop doesn't run
@@ -1841,7 +1863,9 @@ async def test_streaming_function_invocation_config_max_consecutive_errors(chat_
chat_client_base.function_invocation_configuration.max_consecutive_errors_per_request = 2
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[error_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [error_func]}
):
updates.append(update)
# Should stop after 2 consecutive errors
@@ -1887,7 +1911,9 @@ async def test_streaming_function_invocation_config_terminate_on_unknown_calls_f
chat_client_base.function_invocation_configuration.terminate_on_unknown_calls = False
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[known_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [known_func]}
):
updates.append(update)
# Should have a result message indicating the tool wasn't found
@@ -1926,7 +1952,9 @@ async def test_streaming_function_invocation_config_terminate_on_unknown_calls_t
# Should raise an exception when encountering an unknown function
with pytest.raises(KeyError, match='Error: Requested function "unknown_function" not found'):
async for _ in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[known_func]):
async for _ in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [known_func]}
):
pass
assert exec_counter == 0
@@ -1953,7 +1981,9 @@ async def test_streaming_function_invocation_config_include_detailed_errors_true
chat_client_base.function_invocation_configuration.include_detailed_errors = True
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[error_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [error_func]}
):
updates.append(update)
# Should have detailed error message
@@ -1989,7 +2019,9 @@ async def test_streaming_function_invocation_config_include_detailed_errors_fals
chat_client_base.function_invocation_configuration.include_detailed_errors = False
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[error_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [error_func]}
):
updates.append(update)
# Should have a generic error message
@@ -2023,7 +2055,9 @@ async def test_streaming_argument_validation_error_with_detailed_errors(chat_cli
chat_client_base.function_invocation_configuration.include_detailed_errors = True
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[typed_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [typed_func]}
):
updates.append(update)
# Should have detailed validation error
@@ -2057,7 +2091,9 @@ async def test_streaming_argument_validation_error_without_detailed_errors(chat_
chat_client_base.function_invocation_configuration.include_detailed_errors = False
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[typed_func]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [typed_func]}
):
updates.append(update)
# Should have generic validation error
@@ -2105,7 +2141,9 @@ async def test_streaming_multiple_function_calls_parallel_execution(chat_client_
]
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[func1, func2]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [func1, func2]}
):
updates.append(update)
# Both functions should have been executed
@@ -2144,7 +2182,7 @@ async def test_streaming_approval_requests_in_assistant_message(chat_client_base
updates = []
async for update in chat_client_base.get_streaming_response(
"hello", tool_choice="auto", tools=[func_with_approval]
"hello", options={"tool_choice": "auto", "tools": [func_with_approval]}
):
updates.append(update)
@@ -2189,7 +2227,9 @@ async def test_streaming_error_recovery_resets_counter(chat_client_base: ChatCli
]
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[sometimes_fails]):
async for update in chat_client_base.get_streaming_response(
"hello", options={"tool_choice": "auto", "tools": [sometimes_fails]}
):
updates.append(update)
# Should have both an error and a success
@@ -2246,8 +2286,7 @@ async def test_terminate_loop_single_function_call(chat_client_base: ChatClientP
response = await chat_client_base.get_response(
"hello",
tool_choice="auto",
tools=[ai_func],
options={"tool_choice": "auto", "tools": [ai_func]},
middleware=[TerminateLoopMiddleware()],
)
@@ -2314,8 +2353,7 @@ async def test_terminate_loop_multiple_function_calls_one_terminates(chat_client
response = await chat_client_base.get_response(
"hello",
tool_choice="auto",
tools=[normal_func, terminating_func],
options={"tool_choice": "auto", "tools": [normal_func, terminating_func]},
middleware=[SelectiveTerminateMiddleware()],
)
@@ -2366,8 +2404,7 @@ async def test_terminate_loop_streaming_single_function_call(chat_client_base: C
updates = []
async for update in chat_client_base.get_streaming_response(
"hello",
tool_choice="auto",
tools=[ai_func],
options={"tool_choice": "auto", "tools": [ai_func]},
middleware=[TerminateLoopMiddleware()],
):
updates.append(update)
@@ -0,0 +1,218 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for kwargs propagation from get_response() to @ai_function tools."""
from typing import Any
from agent_framework import (
ChatMessage,
ChatResponse,
ChatResponseUpdate,
FunctionCallContent,
TextContent,
ai_function,
)
from agent_framework._tools import _handle_function_calls_response, _handle_function_calls_streaming_response
class TestKwargsPropagationToAIFunction:
"""Test cases for kwargs flowing from get_response() to @ai_function tools."""
async def test_kwargs_propagate_to_ai_function_with_kwargs(self) -> None:
"""Test that kwargs passed to get_response() are available in @ai_function **kwargs."""
captured_kwargs: dict[str, Any] = {}
@ai_function
def capture_kwargs_tool(x: int, **kwargs: Any) -> str:
"""A tool that captures kwargs for testing."""
captured_kwargs.update(kwargs)
return f"result: x={x}"
# Create a mock client
mock_client = type("MockClient", (), {})()
call_count = [0]
async def mock_get_response(self, messages, **kwargs):
call_count[0] += 1
if call_count[0] == 1:
# First call: return a function call
return ChatResponse(
messages=[
ChatMessage(
role="assistant",
contents=[
FunctionCallContent(call_id="call_1", name="capture_kwargs_tool", arguments='{"x": 42}')
],
)
]
)
# Second call: return final response
return ChatResponse(messages=[ChatMessage(role="assistant", text="Done!")])
# Wrap the function with function invocation decorator
wrapped = _handle_function_calls_response(mock_get_response)
# Call with custom kwargs that should propagate to the tool
# Note: tools are passed in options dict, custom kwargs are passed separately
result = await wrapped(
mock_client,
messages=[],
options={"tools": [capture_kwargs_tool]},
user_id="user-123",
session_token="secret-token",
custom_data={"key": "value"},
)
# Verify the tool was called and received the kwargs
assert "user_id" in captured_kwargs, f"Expected 'user_id' in captured kwargs: {captured_kwargs}"
assert captured_kwargs["user_id"] == "user-123"
assert "session_token" in captured_kwargs
assert captured_kwargs["session_token"] == "secret-token"
assert "custom_data" in captured_kwargs
assert captured_kwargs["custom_data"] == {"key": "value"}
# Verify result
assert result.messages[-1].text == "Done!"
async def test_kwargs_not_forwarded_to_ai_function_without_kwargs(self) -> None:
"""Test that kwargs are NOT forwarded to @ai_function that doesn't accept **kwargs."""
@ai_function
def simple_tool(x: int) -> str:
"""A simple tool without **kwargs."""
# This should not receive any extra kwargs
return f"result: x={x}"
mock_client = type("MockClient", (), {})()
call_count = [0]
async def mock_get_response(self, messages, **kwargs):
call_count[0] += 1
if call_count[0] == 1:
return ChatResponse(
messages=[
ChatMessage(
role="assistant",
contents=[FunctionCallContent(call_id="call_1", name="simple_tool", arguments='{"x": 99}')],
)
]
)
return ChatResponse(messages=[ChatMessage(role="assistant", text="Completed!")])
wrapped = _handle_function_calls_response(mock_get_response)
# Call with kwargs - the tool should work but not receive them
result = await wrapped(
mock_client,
messages=[],
options={"tools": [simple_tool]},
user_id="user-123", # This kwarg should be ignored by the tool
)
# Verify the tool was called successfully (no error from extra kwargs)
assert result.messages[-1].text == "Completed!"
async def test_kwargs_isolated_between_function_calls(self) -> None:
"""Test that kwargs don't leak between different function call invocations."""
invocation_kwargs: list[dict[str, Any]] = []
@ai_function
def tracking_tool(name: str, **kwargs: Any) -> str:
"""A tool that tracks kwargs from each invocation."""
invocation_kwargs.append(dict(kwargs))
return f"called with {name}"
mock_client = type("MockClient", (), {})()
call_count = [0]
async def mock_get_response(self, messages, **kwargs):
call_count[0] += 1
if call_count[0] == 1:
# Two function calls in one response
return ChatResponse(
messages=[
ChatMessage(
role="assistant",
contents=[
FunctionCallContent(
call_id="call_1", name="tracking_tool", arguments='{"name": "first"}'
),
FunctionCallContent(
call_id="call_2", name="tracking_tool", arguments='{"name": "second"}'
),
],
)
]
)
return ChatResponse(messages=[ChatMessage(role="assistant", text="All done!")])
wrapped = _handle_function_calls_response(mock_get_response)
# Call with kwargs
result = await wrapped(
mock_client,
messages=[],
options={"tools": [tracking_tool]},
request_id="req-001",
trace_context={"trace_id": "abc"},
)
# Both invocations should have received the same kwargs
assert len(invocation_kwargs) == 2
for kwargs in invocation_kwargs:
assert kwargs.get("request_id") == "req-001"
assert kwargs.get("trace_context") == {"trace_id": "abc"}
assert result.messages[-1].text == "All done!"
async def test_streaming_response_kwargs_propagation(self) -> None:
"""Test that kwargs propagate to @ai_function in streaming mode."""
captured_kwargs: dict[str, Any] = {}
@ai_function
def streaming_capture_tool(value: str, **kwargs: Any) -> str:
"""A tool that captures kwargs during streaming."""
captured_kwargs.update(kwargs)
return f"processed: {value}"
mock_client = type("MockClient", (), {})()
call_count = [0]
async def mock_get_streaming_response(self, messages, **kwargs):
call_count[0] += 1
if call_count[0] == 1:
# First call: return function call update
yield ChatResponseUpdate(
role="assistant",
contents=[
FunctionCallContent(
call_id="stream_call_1",
name="streaming_capture_tool",
arguments='{"value": "streaming-test"}',
)
],
is_finished=True,
)
else:
# Second call: return final response
yield ChatResponseUpdate(text=TextContent(text="Stream complete!"), role="assistant", is_finished=True)
wrapped = _handle_function_calls_streaming_response(mock_get_streaming_response)
# Collect streaming updates
updates: list[ChatResponseUpdate] = []
async for update in wrapped(
mock_client,
messages=[],
options={"tools": [streaming_capture_tool]},
streaming_session="session-xyz",
correlation_id="corr-123",
):
updates.append(update)
# Verify kwargs were captured by the tool
assert "streaming_session" in captured_kwargs, f"Expected 'streaming_session' in {captured_kwargs}"
assert captured_kwargs["streaming_session"] == "session-xyz"
assert captured_kwargs["correlation_id"] == "corr-123"
@@ -29,7 +29,6 @@ from agent_framework._middleware import (
FunctionMiddlewarePipeline,
)
from agent_framework._tools import AIFunction
from agent_framework._types import ChatOptions
class TestAgentRunContext:
@@ -100,12 +99,12 @@ class TestChatContext:
def test_init_with_defaults(self, mock_chat_client: Any) -> None:
"""Test ChatContext initialization with default values."""
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
assert context.chat_client is mock_chat_client
assert context.messages == messages
assert context.chat_options is chat_options
assert context.options is chat_options
assert context.is_streaming is False
assert context.metadata == {}
assert context.result is None
@@ -114,13 +113,13 @@ class TestChatContext:
def test_init_with_custom_values(self, mock_chat_client: Any) -> None:
"""Test ChatContext initialization with custom values."""
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions(temperature=0.5)
chat_options: dict[str, Any] = {"temperature": 0.5}
metadata = {"key": "value"}
context = ChatContext(
chat_client=mock_chat_client,
messages=messages,
chat_options=chat_options,
options=chat_options,
is_streaming=True,
metadata=metadata,
terminate=True,
@@ -128,7 +127,7 @@ class TestChatContext:
assert context.chat_client is mock_chat_client
assert context.messages == messages
assert context.chat_options is chat_options
assert context.options is chat_options
assert context.is_streaming is True
assert context.metadata == metadata
assert context.terminate is True
@@ -562,8 +561,8 @@ class TestChatMiddlewarePipeline:
"""Test pipeline execution with no middleware."""
pipeline = ChatMiddlewarePipeline()
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
expected_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
@@ -589,8 +588,8 @@ class TestChatMiddlewarePipeline:
middleware = OrderTrackingChatMiddleware("test")
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
expected_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
@@ -606,8 +605,8 @@ class TestChatMiddlewarePipeline:
"""Test pipeline streaming execution with no middleware."""
pipeline = ChatMiddlewarePipeline()
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
async def final_handler(ctx: ChatContext) -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="chunk1")])
@@ -637,10 +636,8 @@ class TestChatMiddlewarePipeline:
middleware = StreamOrderTrackingChatMiddleware("test")
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(
chat_client=mock_chat_client, messages=messages, chat_options=chat_options, is_streaming=True
)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
async def final_handler(ctx: ChatContext) -> AsyncIterable[ChatResponseUpdate]:
execution_order.append("handler_start")
@@ -662,8 +659,8 @@ class TestChatMiddlewarePipeline:
middleware = self.PreNextTerminateChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
execution_order: list[str] = []
async def final_handler(ctx: ChatContext) -> ChatResponse:
@@ -682,8 +679,8 @@ class TestChatMiddlewarePipeline:
middleware = self.PostNextTerminateChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
execution_order: list[str] = []
async def final_handler(ctx: ChatContext) -> ChatResponse:
@@ -702,10 +699,8 @@ class TestChatMiddlewarePipeline:
middleware = self.PreNextTerminateChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(
chat_client=mock_chat_client, messages=messages, chat_options=chat_options, is_streaming=True
)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
execution_order: list[str] = []
async def final_handler(ctx: ChatContext) -> AsyncIterable[ChatResponseUpdate]:
@@ -729,10 +724,8 @@ class TestChatMiddlewarePipeline:
middleware = self.PostNextTerminateChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(
chat_client=mock_chat_client, messages=messages, chat_options=chat_options, is_streaming=True
)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
execution_order: list[str] = []
async def final_handler(ctx: ChatContext) -> AsyncIterable[ChatResponseUpdate]:
@@ -962,8 +955,8 @@ class TestMixedMiddleware:
pipeline = ChatMiddlewarePipeline([ClassChatMiddleware(), function_chat_middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
async def final_handler(ctx: ChatContext) -> ChatResponse:
execution_order.append("handler")
@@ -1093,8 +1086,8 @@ class TestMultipleMiddlewareOrdering:
middleware = [FirstChatMiddleware(), SecondChatMiddleware(), ThirdChatMiddleware()]
pipeline = ChatMiddlewarePipeline(middleware) # type: ignore
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
async def final_handler(ctx: ChatContext) -> ChatResponse:
execution_order.append("handler")
@@ -1203,7 +1196,7 @@ class TestContextContentValidation:
# Verify context has all expected attributes
assert hasattr(context, "chat_client")
assert hasattr(context, "messages")
assert hasattr(context, "chat_options")
assert hasattr(context, "options")
assert hasattr(context, "is_streaming")
assert hasattr(context, "metadata")
assert hasattr(context, "result")
@@ -1216,8 +1209,8 @@ class TestContextContentValidation:
assert context.messages[0].text == "test"
assert context.is_streaming is False
assert isinstance(context.metadata, dict)
assert isinstance(context.chat_options, ChatOptions)
assert context.chat_options.temperature == 0.5
assert isinstance(context.options, dict)
assert context.options.get("temperature") == 0.5
# Add custom metadata
context.metadata["validated"] = True
@@ -1227,8 +1220,8 @@ class TestContextContentValidation:
middleware = ChatContextValidationMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions(temperature=0.5)
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {"temperature": 0.5}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
async def final_handler(ctx: ChatContext) -> ChatResponse:
# Verify metadata was set by middleware
@@ -1331,10 +1324,10 @@ class TestStreamingScenarios:
middleware = ChatStreamingFlagMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
chat_options: dict[str, Any] = {}
# Test non-streaming
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
async def final_handler(ctx: ChatContext) -> ChatResponse:
streaming_flags.append(ctx.is_streaming)
@@ -1344,7 +1337,7 @@ class TestStreamingScenarios:
# Test streaming
context_stream = ChatContext(
chat_client=mock_chat_client, messages=messages, chat_options=chat_options, is_streaming=True
chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True
)
async def final_stream_handler(ctx: ChatContext) -> AsyncIterable[ChatResponseUpdate]:
@@ -1373,10 +1366,8 @@ class TestStreamingScenarios:
middleware = ChatStreamProcessingMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(
chat_client=mock_chat_client, messages=messages, chat_options=chat_options, is_streaming=True
)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
async def final_stream_handler(ctx: ChatContext) -> AsyncIterable[ChatResponseUpdate]:
chunks_processed.append("stream_start")
@@ -1590,8 +1581,8 @@ class TestMiddlewareExecutionControl:
middleware = NoNextChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
handler_called = False
@@ -1618,10 +1609,8 @@ class TestMiddlewareExecutionControl:
middleware = NoNextStreamingChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(
chat_client=mock_chat_client, messages=messages, chat_options=chat_options, is_streaming=True
)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
handler_called = False
@@ -1656,8 +1645,8 @@ class TestMiddlewareExecutionControl:
pipeline = ChatMiddlewarePipeline([FirstChatMiddleware(), SecondChatMiddleware()])
messages = [ChatMessage(role=Role.USER, text="test")]
chat_options = ChatOptions()
context = ChatContext(chat_client=mock_chat_client, messages=messages, chat_options=chat_options)
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
handler_called = False
@@ -734,7 +734,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
async def test_function_middleware_can_access_and_override_custom_kwargs(
self, chat_client: "MockChatClient"
) -> None:
"""Test that function middleware can access and override custom parameters like temperature."""
"""Test that function middleware can access and override custom parameters."""
captured_kwargs: dict[str, Any] = {}
modified_kwargs: dict[str, Any] = {}
middleware_called = False
@@ -747,38 +747,20 @@ class TestChatAgentFunctionMiddlewareWithTools:
middleware_called = True
# Capture the original kwargs
captured_kwargs["has_chat_options"] = "chat_options" in context.kwargs
captured_kwargs["has_custom_param"] = "custom_param" in context.kwargs
captured_kwargs["custom_param"] = context.kwargs.get("custom_param")
# Capture original chat_options values if present
if "chat_options" in context.kwargs:
chat_options = context.kwargs["chat_options"]
captured_kwargs["original_temperature"] = getattr(chat_options, "temperature", None)
captured_kwargs["original_max_tokens"] = getattr(chat_options, "max_tokens", None)
# Modify some kwargs
context.kwargs["temperature"] = 0.9
context.kwargs["max_tokens"] = 500
context.kwargs["new_param"] = "added_by_middleware"
# Also modify chat_options if present
if "chat_options" in context.kwargs:
context.kwargs["chat_options"].temperature = 0.9
context.kwargs["chat_options"].max_tokens = 500
# Store modified kwargs for verification
modified_kwargs["temperature"] = context.kwargs.get("temperature")
modified_kwargs["max_tokens"] = context.kwargs.get("max_tokens")
modified_kwargs["new_param"] = context.kwargs.get("new_param")
modified_kwargs["custom_param"] = context.kwargs.get("custom_param")
# Capture modified chat_options values if present
if "chat_options" in context.kwargs:
chat_options = context.kwargs["chat_options"]
modified_kwargs["chat_options_temperature"] = getattr(chat_options, "temperature", None)
modified_kwargs["chat_options_max_tokens"] = getattr(chat_options, "max_tokens", None)
await next(context)
chat_client.responses = [
@@ -800,9 +782,9 @@ class TestChatAgentFunctionMiddlewareWithTools:
# Create ChatAgent with function middleware
agent = ChatAgent(chat_client=chat_client, middleware=[kwargs_middleware], tools=[sample_tool_function])
# Execute the agent with custom parameters
# Execute the agent with custom parameters passed as kwargs
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages, temperature=0.7, max_tokens=100, custom_param="test_value")
response = await agent.run(messages, custom_param="test_value")
# Verify response
assert response is not None
@@ -812,19 +794,14 @@ class TestChatAgentFunctionMiddlewareWithTools:
assert middleware_called, "Function middleware was not called"
# Verify middleware captured the original kwargs
assert captured_kwargs["has_chat_options"] is True
assert captured_kwargs["has_custom_param"] is True
assert captured_kwargs["custom_param"] == "test_value"
assert captured_kwargs["original_temperature"] == 0.7
assert captured_kwargs["original_max_tokens"] == 100
# Verify middleware could modify the kwargs
assert modified_kwargs["temperature"] == 0.9
assert modified_kwargs["max_tokens"] == 500
assert modified_kwargs["new_param"] == "added_by_middleware"
assert modified_kwargs["custom_param"] == "test_value"
assert modified_kwargs["chat_options_temperature"] == 0.9
assert modified_kwargs["chat_options_max_tokens"] == 500
class TestMiddlewareDynamicRebuild:
@@ -366,7 +366,7 @@ class TestChatMiddleware:
# Execute the chat client directly with tools - this should trigger function invocation and middleware
messages = [ChatMessage(role=Role.USER, text="What's the weather in San Francisco?")]
response = await chat_client.get_response(messages, tools=[sample_tool])
response = await chat_client.get_response(messages, options={"tools": [sample_tool]})
# Verify response
assert response is not None
@@ -423,7 +423,7 @@ class TestChatMiddleware:
# Execute the chat client directly with run-level middleware and tools
messages = [ChatMessage(role=Role.USER, text="What's the weather in New York?")]
response = await chat_client.get_response(
messages, tools=[sample_tool], middleware=[run_level_function_middleware]
messages, options={"tools": [sample_tool]}, middleware=[run_level_function_middleware]
)
# Verify response
@@ -17,7 +17,6 @@ from agent_framework import (
AgentThread,
BaseChatClient,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Role,
@@ -215,7 +214,7 @@ def mock_chat_client():
return "https://test.example.com"
async def _inner_get_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
):
return ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
@@ -224,7 +223,7 @@ def mock_chat_client():
)
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
):
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
@@ -405,7 +404,7 @@ def mock_chat_agent():
self.id = "test_agent_id"
self.name = "test_agent"
self.description = "Test agent description"
self.chat_options = ChatOptions(model_id="TestModel")
self.default_options: dict[str, Any] = {"model_id": "TestModel"}
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentRunResponse(
+13 -11
View File
@@ -429,7 +429,7 @@ async def test_ai_function_invoke_ignores_additional_kwargs() -> None:
result = await simple_tool.invoke(
arguments=args,
api_token="secret-token",
chat_options={"model_id": "dummy"},
options={"model_id": "dummy"},
)
assert result == "HELLO WORLD"
@@ -1035,7 +1035,7 @@ async def test_non_streaming_single_function_no_approval():
wrapped = _handle_function_calls_response(mock_get_response)
# Execute
result = await wrapped(mock_client, messages=[], tools=[no_approval_tool])
result = await wrapped(mock_client, messages=[], options={"tools": [no_approval_tool]})
# Verify: should have 3 messages: function call, function result, final answer
assert len(result.messages) == 3
@@ -1075,7 +1075,7 @@ async def test_non_streaming_single_function_requires_approval():
wrapped = _handle_function_calls_response(mock_get_response)
# Execute
result = await wrapped(mock_client, messages=[], tools=[requires_approval_tool])
result = await wrapped(mock_client, messages=[], options={"tools": [requires_approval_tool]})
# Verify: should return 1 message with function call and approval request
from agent_framework import FunctionApprovalRequestContent
@@ -1121,7 +1121,7 @@ async def test_non_streaming_two_functions_both_no_approval():
wrapped = _handle_function_calls_response(mock_get_response)
# Execute
result = await wrapped(mock_client, messages=[], tools=[no_approval_tool])
result = await wrapped(mock_client, messages=[], options={"tools": [no_approval_tool]})
# Verify: should have function calls, results, and final answer
from agent_framework import FunctionResultContent
@@ -1167,7 +1167,7 @@ async def test_non_streaming_two_functions_both_require_approval():
wrapped = _handle_function_calls_response(mock_get_response)
# Execute
result = await wrapped(mock_client, messages=[], tools=[requires_approval_tool])
result = await wrapped(mock_client, messages=[], options={"tools": [requires_approval_tool]})
# Verify: should return 1 message with function calls and approval requests
from agent_framework import FunctionApprovalRequestContent
@@ -1213,7 +1213,7 @@ async def test_non_streaming_two_functions_mixed_approval():
wrapped = _handle_function_calls_response(mock_get_response)
# Execute
result = await wrapped(mock_client, messages=[], tools=[no_approval_tool, requires_approval_tool])
result = await wrapped(mock_client, messages=[], options={"tools": [no_approval_tool, requires_approval_tool]})
# Verify: should return approval requests for both (when one needs approval, all are sent for approval)
from agent_framework import FunctionApprovalRequestContent
@@ -1253,7 +1253,7 @@ async def test_streaming_single_function_no_approval():
# Execute and collect updates
updates = []
async for update in wrapped(mock_client, messages=[], tools=[no_approval_tool]):
async for update in wrapped(mock_client, messages=[], options={"tools": [no_approval_tool]}):
updates.append(update)
# Verify: should have function call update, tool result update (injected), and final update
@@ -1298,7 +1298,7 @@ async def test_streaming_single_function_requires_approval():
# Execute and collect updates
updates = []
async for update in wrapped(mock_client, messages=[], tools=[requires_approval_tool]):
async for update in wrapped(mock_client, messages=[], options={"tools": [requires_approval_tool]}):
updates.append(update)
# Verify: should yield function call and then approval request
@@ -1343,7 +1343,7 @@ async def test_streaming_two_functions_both_no_approval():
# Execute and collect updates
updates = []
async for update in wrapped(mock_client, messages=[], tools=[no_approval_tool]):
async for update in wrapped(mock_client, messages=[], options={"tools": [no_approval_tool]}):
updates.append(update)
# Verify: should have both function calls, one tool result update with both results, and final message
@@ -1392,7 +1392,7 @@ async def test_streaming_two_functions_both_require_approval():
# Execute and collect updates
updates = []
async for update in wrapped(mock_client, messages=[], tools=[requires_approval_tool]):
async for update in wrapped(mock_client, messages=[], options={"tools": [requires_approval_tool]}):
updates.append(update)
# Verify: should yield both function calls and then approval requests
@@ -1439,7 +1439,9 @@ async def test_streaming_two_functions_mixed_approval():
# Execute and collect updates
updates = []
async for update in wrapped(mock_client, messages=[], tools=[no_approval_tool, requires_approval_tool]):
async for update in wrapped(
mock_client, messages=[], options={"tools": [no_approval_tool, requires_approval_tool]}
):
updates.append(update)
# Verify: should yield both function calls and then approval requests (when one needs approval, all wait)
+96 -111
View File
@@ -43,6 +43,7 @@ from agent_framework import (
UsageContent,
UsageDetails,
ai_function,
merge_chat_options,
prepare_function_call_results,
)
from agent_framework.exceptions import AdditionItemMismatch, ContentError
@@ -866,117 +867,149 @@ async def test_chat_response_from_async_generator_output_format_in_method():
def test_chat_tool_mode():
"""Test the ToolMode class to ensure it initializes correctly."""
# Create instances of ToolMode
auto_mode = ToolMode.AUTO
required_any = ToolMode.REQUIRED_ANY
required_mode = ToolMode.REQUIRED("example_function")
none_mode = ToolMode.NONE
auto_mode: ToolMode = {"mode": "auto"}
required_any: ToolMode = {"mode": "required"}
required_mode: ToolMode = {"mode": "required", "required_function_name": "example_function"}
none_mode: ToolMode = {"mode": "none"}
# Check the type and content
assert auto_mode.mode == "auto"
assert auto_mode.required_function_name is None
assert required_any.mode == "required"
assert required_any.required_function_name is None
assert required_mode.mode == "required"
assert required_mode.required_function_name == "example_function"
assert none_mode.mode == "none"
assert none_mode.required_function_name is None
assert auto_mode["mode"] == "auto"
assert "required_function_name" not in auto_mode
assert required_any["mode"] == "required"
assert "required_function_name" not in required_any
assert required_mode["mode"] == "required"
assert required_mode["required_function_name"] == "example_function"
assert none_mode["mode"] == "none"
assert "required_function_name" not in none_mode
# Ensure the instances are of type ToolMode
assert isinstance(auto_mode, ToolMode)
assert isinstance(required_any, ToolMode)
assert isinstance(required_mode, ToolMode)
assert isinstance(none_mode, ToolMode)
assert ToolMode.REQUIRED("example_function") == ToolMode.REQUIRED("example_function")
# serializer returns just the mode
assert ToolMode.REQUIRED_ANY.serialize_model() == "required"
# equality of dicts
assert {"mode": "required", "required_function_name": "example_function"} == {
"mode": "required",
"required_function_name": "example_function",
}
def test_chat_tool_mode_from_dict():
"""Test creating ToolMode from a dictionary."""
mode_dict = {"mode": "required", "required_function_name": "example_function"}
mode = ToolMode(**mode_dict)
mode: ToolMode = {"mode": "required", "required_function_name": "example_function"}
# Check the type and content
assert mode.mode == "required"
assert mode.required_function_name == "example_function"
# Ensure the instance is of type ToolMode
assert isinstance(mode, ToolMode)
assert mode["mode"] == "required"
assert mode["required_function_name"] == "example_function"
# region ChatOptions
def test_chat_options_init() -> None:
options = ChatOptions()
assert options.model_id is None
"""Test that ChatOptions can be created as a TypedDict."""
options: ChatOptions = {}
assert options.get("model_id") is None
# With values
options_with_model: ChatOptions = {"model_id": "gpt-4o", "temperature": 0.7}
assert options_with_model.get("model_id") == "gpt-4o"
assert options_with_model.get("temperature") == 0.7
def test_chat_options_tool_choice_validation_errors():
with raises((ContentError, TypeError)):
ChatOptions(tool_choice="invalid-choice")
def test_chat_options_tool_choice_validation():
"""Test validate_tool_mode utility function."""
from agent_framework._types import validate_tool_mode
# Valid string values
assert validate_tool_mode("auto") == {"mode": "auto"}
assert validate_tool_mode("required") == {"mode": "required"}
assert validate_tool_mode("none") == {"mode": "none"}
# Valid ToolMode dict values
assert validate_tool_mode({"mode": "auto"}) == {"mode": "auto"}
assert validate_tool_mode({"mode": "required"}) == {"mode": "required"}
assert validate_tool_mode({"mode": "required", "required_function_name": "example_function"}) == {
"mode": "required",
"required_function_name": "example_function",
}
assert validate_tool_mode({"mode": "none"}) == {"mode": "none"}
# None should return mode==none
assert validate_tool_mode(None) == {"mode": "none"}
with raises(ContentError):
validate_tool_mode("invalid_mode")
with raises(ContentError):
validate_tool_mode({"mode": "invalid_mode"})
with raises(ContentError):
validate_tool_mode({"mode": "auto", "required_function_name": "should_not_be_here"})
def test_chat_options_and(ai_function_tool, ai_tool) -> None:
options1 = ChatOptions(model_id="gpt-4o", tools=[ai_function_tool], logit_bias={"x": 1}, metadata={"a": "b"})
options2 = ChatOptions(model_id="gpt-4.1", tools=[ai_tool], additional_properties={"p": 1})
def test_chat_options_merge(ai_function_tool, ai_tool) -> None:
"""Test merge_chat_options utility function."""
from agent_framework import merge_chat_options
options1: ChatOptions = {
"model_id": "gpt-4o",
"tools": [ai_function_tool],
"logit_bias": {"x": 1},
"metadata": {"a": "b"},
}
options2: ChatOptions = {"model_id": "gpt-4.1", "tools": [ai_tool]}
assert options1 != options2
options3 = options1 & options2
assert options3.model_id == "gpt-4.1"
assert options3.tools == [ai_function_tool, ai_tool]
assert options3.logit_bias == {"x": 1}
assert options3.metadata == {"a": "b"}
assert options3.additional_properties.get("p") == 1
# Merge options - override takes precedence for non-collection fields
options3 = merge_chat_options(options1, options2)
assert options3.get("model_id") == "gpt-4.1"
assert options3.get("tools") == [ai_function_tool, ai_tool] # tools are combined
assert options3.get("logit_bias") == {"x": 1} # base value preserved
assert options3.get("metadata") == {"a": "b"} # base value preserved
def test_chat_options_and_tool_choice_override() -> None:
"""Test that tool_choice from other takes precedence in ChatOptions merge."""
# Agent-level defaults to "auto"
agent_options = ChatOptions(model_id="gpt-4o", tool_choice="auto")
agent_options: ChatOptions = {"model_id": "gpt-4o", "tool_choice": "auto"}
# Run-level specifies "required"
run_options = ChatOptions(tool_choice="required")
run_options: ChatOptions = {"tool_choice": "required"}
merged = agent_options & run_options
merged = merge_chat_options(agent_options, run_options)
# Run-level should override agent-level
assert merged.tool_choice == "required"
assert merged.model_id == "gpt-4o" # Other fields preserved
assert merged.get("tool_choice") == "required"
assert merged.get("model_id") == "gpt-4o" # Other fields preserved
def test_chat_options_and_tool_choice_none_in_other_uses_self() -> None:
"""Test that when other.tool_choice is None, self.tool_choice is used."""
agent_options = ChatOptions(tool_choice="auto")
run_options = ChatOptions(model_id="gpt-4.1") # tool_choice is None
agent_options: ChatOptions = {"tool_choice": "auto"}
run_options: ChatOptions = {"model_id": "gpt-4.1"} # tool_choice is None
merged = agent_options & run_options
merged = merge_chat_options(agent_options, run_options)
# Should keep agent-level tool_choice since run-level is None
assert merged.tool_choice == "auto"
assert merged.model_id == "gpt-4.1"
assert merged.get("tool_choice") == "auto"
assert merged.get("model_id") == "gpt-4.1"
def test_chat_options_and_tool_choice_with_tool_mode() -> None:
"""Test ChatOptions merge with ToolMode objects."""
agent_options = ChatOptions(tool_choice=ToolMode.AUTO)
run_options = ChatOptions(tool_choice=ToolMode.REQUIRED_ANY)
agent_options: ChatOptions = {"tool_choice": "auto"}
run_options: ChatOptions = {"tool_choice": "required"}
merged = agent_options & run_options
merged = merge_chat_options(agent_options, run_options)
assert merged.tool_choice == ToolMode.REQUIRED_ANY
assert merged.tool_choice == "required" # ToolMode equality with string
assert merged.get("tool_choice") == "required"
assert merged.get("tool_choice") == "required"
def test_chat_options_and_tool_choice_required_specific_function() -> None:
"""Test ChatOptions merge with required specific function."""
agent_options = ChatOptions(tool_choice="auto")
run_options = ChatOptions(tool_choice=ToolMode.REQUIRED(function_name="get_weather"))
agent_options: ChatOptions = {"tool_choice": "auto"}
run_options: ChatOptions = {"tool_choice": {"mode": "required", "required_function_name": "get_weather"}}
merged = agent_options & run_options
merged = merge_chat_options(agent_options, run_options)
assert merged.tool_choice == "required"
assert merged.tool_choice.required_function_name == "get_weather"
tool_choice = merged.get("tool_choice")
assert tool_choice == {"mode": "required", "required_function_name": "get_weather"}
assert tool_choice["required_function_name"] == "get_weather"
# region Agent Response Fixtures
@@ -1249,7 +1282,7 @@ def test_function_call_content_parse_numeric_or_list():
def test_chat_tool_mode_eq_with_string():
assert ToolMode.AUTO == "auto"
assert {"mode": "auto"} == {"mode": "auto"}
# region AgentRunResponse
@@ -1437,30 +1470,6 @@ def test_chat_message_from_dict_with_mixed_content():
assert len(message_dict["contents"]) == 3
def test_chat_options_edge_cases():
"""Test ChatOptions with edge cases for better coverage."""
# Test with tools conversion
def sample_tool():
return "test"
options = ChatOptions(tools=[sample_tool], tool_choice="auto")
assert options.tool_choice == ToolMode.AUTO
# Test to_dict with ToolMode
options_dict = options.to_dict()
assert "tool_choice" in options_dict
# Test from_dict with tool_choice dict
data_with_dict_tool_choice = {
"model_id": "gpt-4",
"tool_choice": {"mode": "required", "required_function_name": "test_func"},
}
options_from_dict = ChatOptions.from_dict(data_with_dict_tool_choice)
assert options_from_dict.tool_choice.mode == "required"
assert options_from_dict.tool_choice.required_function_name == "test_func"
def test_text_content_add_type_error():
"""Test TextContent __add__ raises TypeError for incompatible types."""
t1 = TextContent("Hello")
@@ -1501,30 +1510,6 @@ def test_comprehensive_serialization_methods():
assert result_content.result == "success"
def test_chat_options_tool_choice_variations():
"""Test ChatOptions from_dict and to_dict with various tool_choice values."""
# Test with string tool_choice
data = {"model_id": "gpt-4", "tool_choice": "auto", "temperature": 0.7}
options = ChatOptions.from_dict(data)
assert options.tool_choice == ToolMode.AUTO
# Test with dict tool_choice
data_dict = {
"model_id": "gpt-4",
"tool_choice": {"mode": "required", "required_function_name": "test_func"},
"temperature": 0.7,
}
options_dict = ChatOptions.from_dict(data_dict)
assert options_dict.tool_choice.mode == "required"
assert options_dict.tool_choice.required_function_name == "test_func"
# Test to_dict with ToolMode
options_dict_serialized = options_dict.to_dict()
assert "tool_choice" in options_dict_serialized
assert isinstance(options_dict_serialized["tool_choice"], dict)
def test_chat_message_complex_content_serialization():
"""Test ChatMessage serialization with various content types."""
@@ -17,7 +17,6 @@ from agent_framework import (
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
FunctionCallContent,
@@ -27,7 +26,6 @@ from agent_framework import (
HostedVectorStoreContent,
Role,
TextContent,
ToolMode,
UriContent,
UsageContent,
ai_function,
@@ -43,6 +41,8 @@ skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
else "Integration tests are disabled.",
)
INTEGRATION_TEST_MODEL = "gpt-4.1-nano"
def create_test_openai_assistants_client(
mock_async_openai: MagicMock,
@@ -117,7 +117,7 @@ def mock_async_openai() -> MagicMock:
return mock_client
def test_openai_assistants_client_init_with_client(mock_async_openai: MagicMock) -> None:
def test_init_with_client(mock_async_openai: MagicMock) -> None:
"""Test OpenAIAssistantsClient initialization with existing client."""
chat_client = create_test_openai_assistants_client(
mock_async_openai, model_id="gpt-4", assistant_id="existing-assistant-id", thread_id="test-thread-id"
@@ -131,7 +131,7 @@ def test_openai_assistants_client_init_with_client(mock_async_openai: MagicMock)
assert isinstance(chat_client, ChatClientProtocol)
def test_openai_assistants_client_init_auto_create_client(
def test_init_auto_create_client(
openai_unit_test_env: dict[str, str],
mock_async_openai: MagicMock,
) -> None:
@@ -151,7 +151,7 @@ def test_openai_assistants_client_init_auto_create_client(
assert not chat_client._should_delete_assistant # type: ignore
def test_openai_assistants_client_init_validation_fail() -> None:
def test_init_validation_fail() -> None:
"""Test OpenAIAssistantsClient initialization with validation failure."""
with pytest.raises(ServiceInitializationError):
# Force failure by providing invalid model ID type - this should cause validation to fail
@@ -159,7 +159,7 @@ def test_openai_assistants_client_init_validation_fail() -> None:
@pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True)
def test_openai_assistants_client_init_missing_model_id(openai_unit_test_env: dict[str, str]) -> None:
def test_init_missing_model_id(openai_unit_test_env: dict[str, str]) -> None:
"""Test OpenAIAssistantsClient initialization with missing model ID."""
with pytest.raises(ServiceInitializationError):
OpenAIAssistantsClient(
@@ -168,13 +168,13 @@ def test_openai_assistants_client_init_missing_model_id(openai_unit_test_env: di
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
def test_openai_assistants_client_init_missing_api_key(openai_unit_test_env: dict[str, str]) -> None:
def test_init_missing_api_key(openai_unit_test_env: dict[str, str]) -> None:
"""Test OpenAIAssistantsClient initialization with missing API key."""
with pytest.raises(ServiceInitializationError):
OpenAIAssistantsClient(model_id="gpt-4", env_file_path="nonexistent.env")
def test_openai_assistants_client_init_with_default_headers(openai_unit_test_env: dict[str, str]) -> None:
def test_init_with_default_headers(openai_unit_test_env: dict[str, str]) -> None:
"""Test OpenAIAssistantsClient initialization with default headers."""
default_headers = {"X-Unit-Test": "test-guid"}
@@ -193,7 +193,7 @@ def test_openai_assistants_client_init_with_default_headers(openai_unit_test_env
assert chat_client.client.default_headers[key] == value
async def test_openai_assistants_client_get_assistant_id_or_create_existing_assistant(
async def test_get_assistant_id_or_create_existing_assistant(
mock_async_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when assistant_id is already provided."""
@@ -206,7 +206,7 @@ async def test_openai_assistants_client_get_assistant_id_or_create_existing_assi
mock_async_openai.beta.assistants.create.assert_not_called()
async def test_openai_assistants_client_get_assistant_id_or_create_create_new(
async def test_get_assistant_id_or_create_create_new(
mock_async_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when creating a new assistant."""
@@ -221,7 +221,7 @@ async def test_openai_assistants_client_get_assistant_id_or_create_create_new(
mock_async_openai.beta.assistants.create.assert_called_once()
async def test_openai_assistants_client_aclose_should_not_delete(
async def test_aclose_should_not_delete(
mock_async_openai: MagicMock,
) -> None:
"""Test close when assistant should not be deleted."""
@@ -236,7 +236,7 @@ async def test_openai_assistants_client_aclose_should_not_delete(
assert not chat_client._should_delete_assistant # type: ignore
async def test_openai_assistants_client_aclose_should_delete(mock_async_openai: MagicMock) -> None:
async def test_aclose_should_delete(mock_async_openai: MagicMock) -> None:
"""Test close method calls cleanup."""
chat_client = create_test_openai_assistants_client(
mock_async_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
@@ -249,7 +249,7 @@ async def test_openai_assistants_client_aclose_should_delete(mock_async_openai:
assert not chat_client._should_delete_assistant # type: ignore
async def test_openai_assistants_client_async_context_manager(mock_async_openai: MagicMock) -> None:
async def test_async_context_manager(mock_async_openai: MagicMock) -> None:
"""Test async context manager functionality."""
chat_client = create_test_openai_assistants_client(
mock_async_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
@@ -263,7 +263,7 @@ async def test_openai_assistants_client_async_context_manager(mock_async_openai:
mock_async_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
def test_openai_assistants_client_serialize(openai_unit_test_env: dict[str, str]) -> None:
def test_serialize(openai_unit_test_env: dict[str, str]) -> None:
"""Test serialization of OpenAIAssistantsClient."""
default_headers = {"X-Unit-Test": "test-guid"}
@@ -294,7 +294,7 @@ def test_openai_assistants_client_serialize(openai_unit_test_env: dict[str, str]
assert "User-Agent" not in dumped_settings["default_headers"]
async def test_openai_assistants_client_get_active_thread_run_none_thread_id(mock_async_openai: MagicMock) -> None:
async def test_get_active_thread_run_none_thread_id(mock_async_openai: MagicMock) -> None:
"""Test _get_active_thread_run with None thread_id returns None."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -305,7 +305,7 @@ async def test_openai_assistants_client_get_active_thread_run_none_thread_id(moc
mock_async_openai.beta.threads.runs.list.assert_not_called()
async def test_openai_assistants_client_get_active_thread_run_with_active_run(mock_async_openai: MagicMock) -> None:
async def test_get_active_thread_run_with_active_run(mock_async_openai: MagicMock) -> None:
"""Test _get_active_thread_run finds an active run."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -326,7 +326,7 @@ async def test_openai_assistants_client_get_active_thread_run_with_active_run(mo
mock_async_openai.beta.threads.runs.list.assert_called_once_with(thread_id="thread-123", limit=1, order="desc")
async def test_openai_assistants_client_prepare_thread_create_new(mock_async_openai: MagicMock) -> None:
async def test_prepare_thread_create_new(mock_async_openai: MagicMock) -> None:
"""Test _prepare_thread creates new thread when thread_id is None."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -353,7 +353,7 @@ async def test_openai_assistants_client_prepare_thread_create_new(mock_async_ope
)
async def test_openai_assistants_client_prepare_thread_cancel_existing_run(mock_async_openai: MagicMock) -> None:
async def test_prepare_thread_cancel_existing_run(mock_async_openai: MagicMock) -> None:
"""Test _prepare_thread cancels existing run when provided."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -369,7 +369,7 @@ async def test_openai_assistants_client_prepare_thread_cancel_existing_run(mock_
mock_async_openai.beta.threads.runs.cancel.assert_called_once_with(run_id="run-456", thread_id="thread-123")
async def test_openai_assistants_client_prepare_thread_existing_no_run(mock_async_openai: MagicMock) -> None:
async def test_prepare_thread_existing_no_run(mock_async_openai: MagicMock) -> None:
"""Test _prepare_thread with existing thread_id but no active run."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -382,7 +382,7 @@ async def test_openai_assistants_client_prepare_thread_existing_no_run(mock_asyn
mock_async_openai.beta.threads.runs.cancel.assert_not_called()
async def test_openai_assistants_client_process_stream_events_thread_run_created(mock_async_openai: MagicMock) -> None:
async def test_process_stream_events_thread_run_created(mock_async_openai: MagicMock) -> None:
"""Test _process_stream_events with thread.run.created event."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -415,7 +415,7 @@ async def test_openai_assistants_client_process_stream_events_thread_run_created
assert update.raw_representation == mock_response.data
async def test_openai_assistants_client_process_stream_events_message_delta_text(mock_async_openai: MagicMock) -> None:
async def test_process_stream_events_message_delta_text(mock_async_openai: MagicMock) -> None:
"""Test _process_stream_events with thread.message.delta event containing text."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -459,7 +459,7 @@ async def test_openai_assistants_client_process_stream_events_message_delta_text
assert update.raw_representation == mock_message_delta
async def test_openai_assistants_client_process_stream_events_requires_action(mock_async_openai: MagicMock) -> None:
async def test_process_stream_events_requires_action(mock_async_openai: MagicMock) -> None:
"""Test _process_stream_events with thread.run.requires_action event."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -502,7 +502,7 @@ async def test_openai_assistants_client_process_stream_events_requires_action(mo
chat_client._parse_function_calls_from_assistants.assert_called_once_with(mock_run, None) # type: ignore
async def test_openai_assistants_client_process_stream_events_run_step_created(mock_async_openai: MagicMock) -> None:
async def test_process_stream_events_run_step_created(mock_async_openai: MagicMock) -> None:
"""Test _process_stream_events with thread.run.step.created event."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -534,7 +534,7 @@ async def test_openai_assistants_client_process_stream_events_run_step_created(m
assert len(updates) == 0
async def test_openai_assistants_client_process_stream_events_run_completed_with_usage(
async def test_process_stream_events_run_completed_with_usage(
mock_async_openai: MagicMock,
) -> None:
"""Test _process_stream_events with thread.run.completed event containing usage."""
@@ -585,7 +585,7 @@ async def test_openai_assistants_client_process_stream_events_run_completed_with
assert update.raw_representation == mock_run
def test_openai_assistants_client_parse_function_calls_from_assistants_basic(mock_async_openai: MagicMock) -> None:
def test_parse_function_calls_from_assistants_basic(mock_async_openai: MagicMock) -> None:
"""Test _parse_function_calls_from_assistants with a simple function call."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -614,22 +614,22 @@ def test_openai_assistants_client_parse_function_calls_from_assistants_basic(moc
assert contents[0].arguments == {"location": "Seattle"}
def test_openai_assistants_client_prepare_options_basic(mock_async_openai: MagicMock) -> None:
def test_prepare_options_basic(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with basic chat options."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
# Create basic chat options
chat_options = ChatOptions(
max_tokens=100,
model_id="gpt-4",
temperature=0.7,
top_p=0.9,
)
# Create basic chat options as a dict
options = {
"max_tokens": 100,
"model_id": "gpt-4",
"temperature": 0.7,
"top_p": 0.9,
}
messages = [ChatMessage(role=Role.USER, text="Hello")]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, chat_options) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
# Check basic options were set
assert run_options["max_completion_tokens"] == 100
@@ -639,7 +639,7 @@ def test_openai_assistants_client_prepare_options_basic(mock_async_openai: Magic
assert tool_results is None
def test_openai_assistants_client_prepare_options_with_ai_function_tool(mock_async_openai: MagicMock) -> None:
def test_prepare_options_with_ai_function_tool(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with AIFunction tool."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -650,15 +650,15 @@ def test_openai_assistants_client_prepare_options_with_ai_function_tool(mock_asy
"""A test function."""
return f"Result for {query}"
chat_options = ChatOptions(
tools=[test_function],
tool_choice="auto",
)
options = {
"tools": [test_function],
"tool_choice": "auto",
}
messages = [ChatMessage(role=Role.USER, text="Hello")]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, chat_options) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
# Check tools were set correctly
assert "tools" in run_options
@@ -668,22 +668,22 @@ def test_openai_assistants_client_prepare_options_with_ai_function_tool(mock_asy
assert run_options["tool_choice"] == "auto"
def test_openai_assistants_client_prepare_options_with_code_interpreter(mock_async_openai: MagicMock) -> None:
def test_prepare_options_with_code_interpreter(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with HostedCodeInterpreterTool."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
# Create a real HostedCodeInterpreterTool
code_tool = HostedCodeInterpreterTool()
chat_options = ChatOptions(
tools=[code_tool],
tool_choice="auto",
)
options = {
"tools": [code_tool],
"tool_choice": "auto",
}
messages = [ChatMessage(role=Role.USER, text="Calculate something")]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, chat_options) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
# Check code interpreter tool was set correctly
assert "tools" in run_options
@@ -692,39 +692,39 @@ def test_openai_assistants_client_prepare_options_with_code_interpreter(mock_asy
assert run_options["tool_choice"] == "auto"
def test_openai_assistants_client_prepare_options_tool_choice_none(mock_async_openai: MagicMock) -> None:
def test_prepare_options_tool_choice_none(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with tool_choice set to 'none'."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
chat_options = ChatOptions(
tool_choice="none",
)
options = {
"tool_choice": "none",
}
messages = [ChatMessage(role=Role.USER, text="Hello")]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, chat_options) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
# Should set tool_choice to none and not include tools
assert run_options["tool_choice"] == "none"
assert "tools" not in run_options
def test_openai_assistants_client_prepare_options_required_function(mock_async_openai: MagicMock) -> None:
def test_prepare_options_required_function(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with required function tool choice."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
# Create a required function tool choice
tool_choice = ToolMode(mode="required", required_function_name="specific_function")
# Create a required function tool choice as dict
tool_choice = {"mode": "required", "required_function_name": "specific_function"}
chat_options = ChatOptions(
tool_choice=tool_choice,
)
options = {
"tool_choice": tool_choice,
}
messages = [ChatMessage(role=Role.USER, text="Hello")]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, chat_options) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
# Check required function tool choice was set correctly
expected_tool_choice = {
@@ -734,7 +734,7 @@ def test_openai_assistants_client_prepare_options_required_function(mock_async_o
assert run_options["tool_choice"] == expected_tool_choice
def test_openai_assistants_client_prepare_options_with_file_search_tool(mock_async_openai: MagicMock) -> None:
def test_prepare_options_with_file_search_tool(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with HostedFileSearchTool."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -742,15 +742,15 @@ def test_openai_assistants_client_prepare_options_with_file_search_tool(mock_asy
# Create a HostedFileSearchTool with max_results
file_search_tool = HostedFileSearchTool(max_results=10)
chat_options = ChatOptions(
tools=[file_search_tool],
tool_choice="auto",
)
options = {
"tools": [file_search_tool],
"tool_choice": "auto",
}
messages = [ChatMessage(role=Role.USER, text="Search for information")]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, chat_options) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
# Check file search tool was set correctly
assert "tools" in run_options
@@ -760,22 +760,22 @@ def test_openai_assistants_client_prepare_options_with_file_search_tool(mock_asy
assert run_options["tool_choice"] == "auto"
def test_openai_assistants_client_prepare_options_with_mapping_tool(mock_async_openai: MagicMock) -> None:
def test_prepare_options_with_mapping_tool(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with MutableMapping tool."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
# Create a tool as a MutableMapping (dict)
mapping_tool = {"type": "custom_tool", "parameters": {"setting": "value"}}
chat_options = ChatOptions(
tools=[mapping_tool], # type: ignore
tool_choice="auto",
)
options = {
"tools": [mapping_tool], # type: ignore
"tool_choice": "auto",
}
messages = [ChatMessage(role=Role.USER, text="Use custom tool")]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, chat_options) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
# Check mapping tool was set correctly
assert "tools" in run_options
@@ -784,7 +784,7 @@ def test_openai_assistants_client_prepare_options_with_mapping_tool(mock_async_o
assert run_options["tool_choice"] == "auto"
def test_openai_assistants_client_prepare_options_with_system_message(mock_async_openai: MagicMock) -> None:
def test_prepare_options_with_system_message(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with system message converted to instructions."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -794,7 +794,7 @@ def test_openai_assistants_client_prepare_options_with_system_message(mock_async
]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, None) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, {}) # type: ignore
# Check that additional_messages only contains the user message
# System message should be converted to instructions (though this is handled internally)
@@ -803,7 +803,7 @@ def test_openai_assistants_client_prepare_options_with_system_message(mock_async
assert run_options["additional_messages"][0]["role"] == "user"
def test_openai_assistants_client_prepare_options_with_image_content(mock_async_openai: MagicMock) -> None:
def test_prepare_options_with_image_content(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with image content."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -813,7 +813,7 @@ def test_openai_assistants_client_prepare_options_with_image_content(mock_async_
messages = [ChatMessage(role=Role.USER, contents=[image_content])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, None) # type: ignore
run_options, tool_results = chat_client._prepare_options(messages, {}) # type: ignore
# Check that image content was processed
assert "additional_messages" in run_options
@@ -825,7 +825,7 @@ def test_openai_assistants_client_prepare_options_with_image_content(mock_async_
assert message["content"][0]["image_url"]["url"] == "https://example.com/image.jpg"
def test_openai_assistants_client_prepare_tool_outputs_for_assistants_empty(mock_async_openai: MagicMock) -> None:
def test_prepare_tool_outputs_for_assistants_empty(mock_async_openai: MagicMock) -> None:
"""Test _prepare_tool_outputs_for_assistants with empty list."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -835,7 +835,7 @@ def test_openai_assistants_client_prepare_tool_outputs_for_assistants_empty(mock
assert tool_outputs is None
def test_openai_assistants_client_prepare_tool_outputs_for_assistants_valid(mock_async_openai: MagicMock) -> None:
def test_prepare_tool_outputs_for_assistants_valid(mock_async_openai: MagicMock) -> None:
"""Test _prepare_tool_outputs_for_assistants with valid function results."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
@@ -851,7 +851,7 @@ def test_openai_assistants_client_prepare_tool_outputs_for_assistants_valid(mock
assert tool_outputs[0].get("output") == "Function executed successfully"
def test_openai_assistants_client_prepare_tool_outputs_for_assistants_mismatched_run_ids(
def test_prepare_tool_outputs_for_assistants_mismatched_run_ids(
mock_async_openai: MagicMock,
) -> None:
"""Test _prepare_tool_outputs_for_assistants with mismatched run IDs."""
@@ -872,7 +872,7 @@ def test_openai_assistants_client_prepare_tool_outputs_for_assistants_mismatched
assert tool_outputs[0].get("tool_call_id") == "call-456"
def test_openai_assistants_client_update_agent_name_and_description(mock_async_openai: MagicMock) -> None:
def test_update_agent_name_and_description(mock_async_openai: MagicMock) -> None:
"""Test _update_agent_name_and_description method updates assistant_name when not already set."""
# Test updating agent name when assistant_name is None
chat_client = create_test_openai_assistants_client(mock_async_openai, assistant_name=None)
@@ -883,7 +883,7 @@ def test_openai_assistants_client_update_agent_name_and_description(mock_async_o
assert chat_client.assistant_name == "New Assistant Name"
def test_openai_assistants_client_update_agent_name_and_description_existing(mock_async_openai: MagicMock) -> None:
def test_update_agent_name_and_description_existing(mock_async_openai: MagicMock) -> None:
"""Test _update_agent_name_and_description method doesn't override existing assistant_name."""
# Test that existing assistant_name is not overridden
chat_client = create_test_openai_assistants_client(mock_async_openai, assistant_name="Existing Assistant")
@@ -895,7 +895,7 @@ def test_openai_assistants_client_update_agent_name_and_description_existing(moc
assert chat_client.assistant_name == "Existing Assistant"
def test_openai_assistants_client_update_agent_name_and_description_none(mock_async_openai: MagicMock) -> None:
def test_update_agent_name_and_description_none(mock_async_openai: MagicMock) -> None:
"""Test _update_agent_name_and_description method with None agent_name parameter."""
# Test that None agent_name doesn't change anything
chat_client = create_test_openai_assistants_client(mock_async_openai, assistant_name=None)
@@ -916,9 +916,9 @@ def get_weather(
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_client_get_response() -> None:
async def test_get_response() -> None:
"""Test OpenAI Assistants Client response."""
async with OpenAIAssistantsClient() as openai_assistants_client:
async with OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL) as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
@@ -941,9 +941,9 @@ async def test_openai_assistants_client_get_response() -> None:
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_client_get_response_tools() -> None:
async def test_get_response_tools() -> None:
"""Test OpenAI Assistants Client response with tools."""
async with OpenAIAssistantsClient() as openai_assistants_client:
async with OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL) as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
@@ -952,8 +952,7 @@ async def test_openai_assistants_client_get_response_tools() -> None:
# Test that the client can be used to get a response
response = await openai_assistants_client.get_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
options={"tools": [get_weather], "tool_choice": "auto"},
)
assert response is not None
@@ -963,9 +962,9 @@ async def test_openai_assistants_client_get_response_tools() -> None:
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_client_streaming() -> None:
async def test_streaming() -> None:
"""Test OpenAI Assistants Client streaming response."""
async with OpenAIAssistantsClient() as openai_assistants_client:
async with OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL) as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
@@ -994,9 +993,9 @@ async def test_openai_assistants_client_streaming() -> None:
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_client_streaming_tools() -> None:
async def test_streaming_tools() -> None:
"""Test OpenAI Assistants Client streaming response with tools."""
async with OpenAIAssistantsClient() as openai_assistants_client:
async with OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL) as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
@@ -1005,8 +1004,10 @@ async def test_openai_assistants_client_streaming_tools() -> None:
# Test that the client can be used to get a response
response = openai_assistants_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
options={
"tools": [get_weather],
"tool_choice": "auto",
},
)
full_message: str = ""
async for chunk in response:
@@ -1021,10 +1022,10 @@ async def test_openai_assistants_client_streaming_tools() -> None:
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_client_with_existing_assistant() -> None:
async def test_with_existing_assistant() -> None:
"""Test OpenAI Assistants Client with existing assistant ID."""
# First create an assistant to use in the test
async with OpenAIAssistantsClient() as temp_client:
async with OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL) as temp_client:
# Get the assistant ID by triggering assistant creation
messages = [ChatMessage(role="user", text="Hello")]
await temp_client.get_response(messages=messages)
@@ -1032,7 +1033,7 @@ async def test_openai_assistants_client_with_existing_assistant() -> None:
# Now test using the existing assistant
async with OpenAIAssistantsClient(
model_id="gpt-4o-mini", assistant_id=assistant_id
model_id=INTEGRATION_TEST_MODEL, assistant_id=assistant_id
) as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClientProtocol)
assert openai_assistants_client.assistant_id == assistant_id
@@ -1050,9 +1051,9 @@ async def test_openai_assistants_client_with_existing_assistant() -> None:
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
@pytest.mark.skip(reason="OpenAI file search functionality is currently broken - tracked in GitHub issue")
async def test_openai_assistants_client_file_search() -> None:
async def test_file_search() -> None:
"""Test OpenAI Assistants Client response."""
async with OpenAIAssistantsClient() as openai_assistants_client:
async with OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL) as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
@@ -1061,8 +1062,10 @@ async def test_openai_assistants_client_file_search() -> None:
file_id, vector_store = await create_vector_store(openai_assistants_client)
response = await openai_assistants_client.get_response(
messages=messages,
tools=[HostedFileSearchTool()],
tool_resources={"file_search": {"vector_store_ids": [vector_store.vector_store_id]}},
options={
"tools": [HostedFileSearchTool()],
"tool_resources": {"file_search": {"vector_store_ids": [vector_store.vector_store_id]}},
},
)
await delete_vector_store(openai_assistants_client, file_id, vector_store.vector_store_id)
@@ -1074,9 +1077,9 @@ async def test_openai_assistants_client_file_search() -> None:
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
@pytest.mark.skip(reason="OpenAI file search functionality is currently broken - tracked in GitHub issue")
async def test_openai_assistants_client_file_search_streaming() -> None:
async def test_file_search_streaming() -> None:
"""Test OpenAI Assistants Client response."""
async with OpenAIAssistantsClient() as openai_assistants_client:
async with OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL) as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
@@ -1085,8 +1088,10 @@ async def test_openai_assistants_client_file_search_streaming() -> None:
file_id, vector_store = await create_vector_store(openai_assistants_client)
response = openai_assistants_client.get_streaming_response(
messages=messages,
tools=[HostedFileSearchTool()],
tool_resources={"file_search": {"vector_store_ids": [vector_store.vector_store_id]}},
options={
"tools": [HostedFileSearchTool()],
"tool_resources": {"file_search": {"vector_store_ids": [vector_store.vector_store_id]}},
},
)
assert response is not None
@@ -1107,7 +1112,7 @@ async def test_openai_assistants_client_file_search_streaming() -> None:
async def test_openai_assistants_agent_basic_run():
"""Test ChatAgent basic run functionality with OpenAIAssistantsClient."""
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
chat_client=OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL),
) as agent:
# Run a simple query
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
@@ -1124,7 +1129,7 @@ async def test_openai_assistants_agent_basic_run():
async def test_openai_assistants_agent_basic_run_streaming():
"""Test ChatAgent basic streaming functionality with OpenAIAssistantsClient."""
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
chat_client=OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL),
) as agent:
# Run streaming query
full_message: str = ""
@@ -1144,7 +1149,7 @@ async def test_openai_assistants_agent_basic_run_streaming():
async def test_openai_assistants_agent_thread_persistence():
"""Test ChatAgent thread persistence across runs with OpenAIAssistantsClient."""
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
chat_client=OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
@@ -1176,7 +1181,7 @@ async def test_openai_assistants_agent_existing_thread_id():
existing_thread_id = None
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
chat_client=OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
@@ -1219,7 +1224,7 @@ async def test_openai_assistants_agent_code_interpreter():
"""Test ChatAgent with code interpreter through OpenAIAssistantsClient."""
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
chat_client=OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL),
instructions="You are a helpful assistant that can write and execute Python code.",
tools=[HostedCodeInterpreterTool()],
) as agent:
@@ -1235,11 +1240,11 @@ async def test_openai_assistants_agent_code_interpreter():
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_client_agent_level_tool_persistence():
async def test_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with OpenAI Assistants Client."""
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
chat_client=OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
@@ -1261,7 +1266,7 @@ async def test_openai_assistants_client_agent_level_tool_persistence():
# Callable API Key Tests
def test_openai_assistants_client_with_callable_api_key() -> None:
def test_with_callable_api_key() -> None:
"""Test OpenAIAssistantsClient initialization with callable API key."""
async def get_api_key() -> str:
@@ -1,25 +1,22 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import os
from typing import Annotated
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from openai import BadRequestError
from pydantic import BaseModel
from pytest import param
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
DataContent,
FunctionResultContent,
HostedWebSearchTool,
TextContent,
ToolProtocol,
ai_function,
prepare_function_call_results,
@@ -170,7 +167,7 @@ async def test_content_filter_exception_handling(openai_unit_test_env: dict[str,
patch.object(client.client.chat.completions, "create", side_effect=mock_error),
pytest.raises(OpenAIContentFilterException),
):
await client._inner_get_response(messages=messages, chat_options=ChatOptions()) # type: ignore
await client._inner_get_response(messages=messages, options={}) # type: ignore
def test_unsupported_tool_handling(openai_unit_test_env: dict[str, str]) -> None:
@@ -183,12 +180,12 @@ def test_unsupported_tool_handling(openai_unit_test_env: dict[str, str]) -> None
# This should ignore the unsupported ToolProtocol and return empty list
result = client._prepare_tools_for_openai([unsupported_tool]) # type: ignore
assert result == []
assert result == {}
# Also test with a non-ToolProtocol that should be converted to dict
dict_tool = {"type": "function", "name": "test"}
result = client._prepare_tools_for_openai([dict_tool]) # type: ignore
assert result == [dict_tool]
assert result["tools"] == [dict_tool]
@ai_function
@@ -208,407 +205,6 @@ def get_weather(location: str) -> str:
return f"The weather in {location} is sunny and 72°F."
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_completion_response() -> None:
"""Test OpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await openai_chat_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert "scientists" in response.text
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_completion_response_params() -> None:
"""Test OpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await openai_chat_client.get_response(
messages=messages, chat_options=ChatOptions(max_tokens=150, temperature=0.7, top_p=0.9)
)
assert response is not None
assert isinstance(response, ChatResponse)
assert "scientists" in response.text
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_completion_response_tools() -> None:
"""Test OpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await openai_chat_client.get_response(
messages=messages,
tools=[get_story_text],
tool_choice="auto",
)
assert response is not None
assert isinstance(response, ChatResponse)
assert "scientists" in response.text
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_streaming() -> None:
"""Test Azure OpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = openai_chat_client.get_streaming_response(messages=messages)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
assert chunk.message_id is not None
assert chunk.response_id is not None
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "scientists" in full_message
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_streaming_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = openai_chat_client.get_streaming_response(
messages=messages,
tools=[get_story_text],
tool_choice="auto",
)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "scientists" in full_message
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_web_search() -> None:
# Currently only a select few models support web search tool calls
openai_chat_client = OpenAIChatClient(model_id="gpt-4o-search-preview")
assert isinstance(openai_chat_client, ChatClientProtocol)
# Test that the client will use the web search tool
response = await openai_chat_client.get_response(
messages=[
ChatMessage(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
tools=[HostedWebSearchTool()],
tool_choice="auto",
)
assert response is not None
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
# Test that the client will use the web search tool with location
additional_properties = {
"user_location": {
"country": "US",
"city": "Seattle",
}
}
response = await openai_chat_client.get_response(
messages=[ChatMessage(role="user", text="What is the current weather? Do not ask for my current location.")],
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
tool_choice="auto",
)
assert response.text is not None
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_web_search_streaming() -> None:
openai_chat_client = OpenAIChatClient(model_id="gpt-4o-search-preview")
assert isinstance(openai_chat_client, ChatClientProtocol)
# Test that the client will use the web search tool
response = openai_chat_client.get_streaming_response(
messages=[
ChatMessage(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
tools=[HostedWebSearchTool()],
tool_choice="auto",
)
assert response is not None
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "Rumi" in full_message
assert "Mira" in full_message
assert "Zoey" in full_message
# Test that the client will use the web search tool with location
additional_properties = {
"user_location": {
"country": "US",
"city": "Seattle",
}
}
response = openai_chat_client.get_streaming_response(
messages=[ChatMessage(role="user", text="What is the current weather? Do not ask for my current location.")],
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
tool_choice="auto",
)
assert response is not None
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert full_message is not None
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_agent_basic_run():
"""Test OpenAI chat client agent basic run functionality with OpenAIChatClient."""
async with ChatAgent(
chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"),
) as agent:
# Test basic run
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
assert "hello world" in response.text.lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_agent_basic_run_streaming():
"""Test OpenAI chat client agent basic streaming functionality with OpenAIChatClient."""
async with ChatAgent(
chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"),
) as agent:
# Test streaming run
full_text = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_text += chunk.text
assert len(full_text) > 0
assert "streaming response test" in full_text.lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_agent_thread_persistence():
"""Test OpenAI chat client agent thread persistence across runs with OpenAIChatClient."""
async with ChatAgent(
chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# First interaction
response1 = await agent.run("My name is Alice. Remember this.", thread=thread)
assert isinstance(response1, AgentRunResponse)
assert response1.text is not None
# Second interaction - test memory
response2 = await agent.run("What is my name?", thread=thread)
assert isinstance(response2, AgentRunResponse)
assert response2.text is not None
assert "alice" in response2.text.lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_agent_existing_thread():
"""Test OpenAI chat client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
async with ChatAgent(
chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My name is Alice. Remember this.", thread=thread)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Preserve the thread for reuse
preserved_thread = thread
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
async with ChatAgent(
chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved thread
second_response = await second_agent.run("What is my name?", thread=preserved_thread)
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
assert "alice" in second_response.text.lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with OpenAI Chat Client."""
async with ChatAgent(
chat_client=OpenAIChatClient(model_id="gpt-4.1"),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
# First run - agent-level tool should be available
first_response = await agent.run("What's the weather like in Chicago?")
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Should use the agent-level weather tool
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - agent-level tool should still be available (persistence test)
second_response = await agent.run("What's the weather in Miami?")
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
# Should use the agent-level weather tool again
assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_run_level_tool_isolation():
"""Test that run-level tools are isolated to specific runs and don't persist with OpenAI Chat Client."""
# Counter to track how many times the weather tool is called
call_count = 0
@ai_function
async def get_weather_with_counter(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
nonlocal call_count
call_count += 1
return f"The weather in {location} is sunny and 72°F."
async with ChatAgent(
chat_client=OpenAIChatClient(model_id="gpt-4.1"),
instructions="You are a helpful assistant.",
) as agent:
# First run - use run-level tool
first_response = await agent.run(
"What's the weather like in Chicago?",
tools=[get_weather_with_counter], # Run-level tool
)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Should use the run-level weather tool (call count should be 1)
assert call_count == 1
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - run-level tool should NOT persist (key isolation test)
second_response = await agent.run("What's the weather like in Miami?")
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
# Should NOT use the weather tool since it was only run-level in previous call
# Call count should still be 1 (no additional calls)
assert call_count == 1
async def test_exception_message_includes_original_error_details() -> None:
"""Test that exception messages include original error details in the new format."""
client = OpenAIChatClient(model_id="test-model", api_key="test-key")
@@ -627,7 +223,7 @@ async def test_exception_message_includes_original_error_details() -> None:
patch.object(client.client.chat.completions, "create", side_effect=mock_error),
pytest.raises(ServiceResponseException) as exc_info,
):
await client._inner_get_response(messages=messages, chat_options=ChatOptions()) # type: ignore
await client._inner_get_response(messages=messages, options={}) # type: ignore
exception_message = str(exc_info.value)
assert "service failed to complete the prompt:" in exception_message
@@ -667,7 +263,7 @@ def test_chat_response_content_order_text_before_tool_calls(openai_unit_test_env
)
client = OpenAIChatClient()
response = client._parse_response_from_openai(mock_response, ChatOptions())
response = client._parse_response_from_openai(mock_response, {})
# Verify we have both text and tool call content
assert len(response.messages) == 1
@@ -894,3 +490,191 @@ def test_prepare_content_for_openai_document_file_mapping(openai_unit_test_env:
assert result["type"] == "file"
assert "filename" not in result["file"] # None filename should be omitted
# region Integration Tests
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
location: str
weather: str | None = None
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
# Simple ChatOptions - just verify they don't fail
param("temperature", 0.7, False, id="temperature"),
param("top_p", 0.9, False, id="top_p"),
param("max_tokens", 500, False, id="max_tokens"),
param("seed", 123, False, id="seed"),
param("user", "test-user-id", False, id="user"),
param("frequency_penalty", 0.5, False, id="frequency_penalty"),
param("presence_penalty", 0.3, False, id="presence_penalty"),
param("stop", ["END"], False, id="stop"),
param("allow_multiple_tool_calls", True, False, id="allow_multiple_tool_calls"),
# OpenAIChatOptions - just verify they don't fail
param("logit_bias", {"50256": -1}, False, id="logit_bias"),
param("prediction", {"type": "content", "content": "hello world"}, False, id="prediction"),
# Complex options requiring output validation
param("tools", [get_weather], True, id="tools_function"),
param("tool_choice", "auto", True, id="tool_choice_auto"),
param("tool_choice", "none", True, id="tool_choice_none"),
param("tool_choice", "required", True, id="tool_choice_required_any"),
param(
"tool_choice",
{"mode": "required", "required_function_name": "get_weather"},
True,
id="tool_choice_required",
),
param("response_format", OutputStruct, True, id="response_format_pydantic"),
param(
"response_format",
{
"type": "json_schema",
"json_schema": {
"name": "WeatherDigest",
"strict": True,
"schema": {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
"temperature_c": {"type": "number"},
"advisory": {"type": "string"},
},
"required": ["location", "conditions", "temperature_c", "advisory"],
"additionalProperties": False,
},
},
},
True,
id="response_format_runtime_json_schema",
),
],
)
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
"""Parametrized test covering all ChatOptions and OpenAIChatOptions.
Tests both streaming and non-streaming modes for each option to ensure
they don't cause failures. Options marked with needs_validation also
check that the feature actually works correctly.
"""
client = OpenAIChatClient()
# to ensure toolmode required does not endlessly loop
client.function_invocation_configuration.max_iterations = 1
for streaming in [False, True]:
# Prepare test message
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [ChatMessage(role="user", text="What is the weather in Seattle?")]
elif option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [ChatMessage(role="user", text="The weather in Seattle is sunny")]
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
# Add tools if testing tool_choice to avoid errors
if option_name.startswith("tool_choice"):
options["tools"] = [get_weather]
if streaming:
# Test streaming mode
response_gen = client.get_streaming_response(
messages=messages,
options=options,
)
output_format = option_value if option_name.startswith("response_format") else None
response = await ChatResponse.from_chat_response_generator(response_gen, output_format_type=output_format)
else:
# Test non-streaming mode
response = await client.get_response(
messages=messages,
options=options,
)
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation:
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
# Should have called the weather function
text = response.text.lower()
assert "sunny" in text or "seattle" in text, f"Tool not invoked for {option_name}"
elif option_name.startswith("response_format"):
if option_value == OutputStruct:
# Should have structured output
assert response.value is not None, "No structured output"
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
# Runtime JSON schema
assert response.value is None, "No structured output, can't parse any json."
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
assert "seattle" in response_value["location"].lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_integration_web_search() -> None:
client = OpenAIChatClient(model_id="gpt-4o-search-preview")
for streaming in [False, True]:
content = {
"messages": "Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
"options": {
"tool_choice": "auto",
"tools": [HostedWebSearchTool()],
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
assert response is not None
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
# Test that the client will use the web search tool with location
additional_properties = {
"user_location": {
"country": "US",
"city": "Seattle",
}
}
content = {
"messages": "What is the current weather? Do not ask for my current location.",
"options": {
"tool_choice": "auto",
"tools": [HostedWebSearchTool(additional_properties=additional_properties)],
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
assert response.text is not None
@@ -115,7 +115,6 @@ async def test_cmc_no_fcc_in_response(
openai_chat_completion = OpenAIChatClient()
await openai_chat_completion.get_response(
messages=chat_history,
arguments={},
)
mock_create.assert_awaited_once_with(
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
@@ -199,7 +198,7 @@ async def test_cmc_additional_properties(
chat_history.append(ChatMessage(role="user", text="hello world"))
openai_chat_completion = OpenAIChatClient()
await openai_chat_completion.get_response(messages=chat_history, additional_properties={"reasoning_effort": "low"})
await openai_chat_completion.get_response(messages=chat_history, options={"reasoning_effort": "low"})
mock_create.assert_awaited_once_with(
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
stream=False,
@@ -382,8 +381,6 @@ def test_chat_response_created_at_uses_utc(openai_unit_test_env: dict[str, str])
This is a regression test for the issue where created_at was using local time
but labeling it as UTC (with 'Z' suffix).
"""
from agent_framework import ChatOptions
# Use a specific Unix timestamp: 1733011890 = 2024-12-01T00:31:30Z (UTC)
# This ensures we test that the timestamp is actually converted to UTC
utc_timestamp = 1733011890
@@ -399,7 +396,7 @@ def test_chat_response_created_at_uses_utc(openai_unit_test_env: dict[str, str])
)
client = OpenAIChatClient()
response = client._parse_response_from_openai(mock_response, ChatOptions())
response = client._parse_response_from_openai(mock_response, {})
# Verify that created_at is correctly formatted as UTC
assert response.created_at is not None
File diff suppressed because it is too large Load Diff
@@ -785,13 +785,13 @@ class TestAgentManagerConfiguration:
chat_client = MagicMock()
manager_agent = ChatAgent(chat_client=chat_client, name="Coordinator")
assert manager_agent.chat_options.response_format is None
assert manager_agent.default_options.get("response_format") is None
worker = StubAgent("worker", "response")
builder = GroupChatBuilder().set_manager(manager_agent).participants([worker])
assert manager_agent.chat_options.response_format is ManagerSelectionResponse
assert manager_agent.default_options.get("response_format") is ManagerSelectionResponse
assert builder._manager_participant is manager_agent # type: ignore[attr-defined]
async def test_set_manager_accepts_agent_manager(self) -> None:
@@ -820,13 +820,15 @@ class TestAgentManagerConfiguration:
value: str
chat_client = MagicMock()
manager_agent = ChatAgent(chat_client=chat_client, name="Coordinator", response_format=CustomResponse)
manager_agent = ChatAgent(
chat_client=chat_client, name="Coordinator", default_options={"response_format": CustomResponse}
)
worker = StubAgent("worker", "response")
with pytest.raises(ValueError, match="response_format must be ManagerSelectionResponse"):
GroupChatBuilder().set_manager(manager_agent).participants([worker])
assert manager_agent.chat_options.response_format is CustomResponse
assert manager_agent.default_options.get("response_format") is CustomResponse
class TestFactoryFunctions:
@@ -504,8 +504,8 @@ async def test_clone_chat_agent_preserves_mcp_tools() -> None:
assert hasattr(cloned_agent, "_local_mcp_tools")
assert len(cloned_agent._local_mcp_tools) == 1 # type: ignore[reportPrivateUsage]
assert cloned_agent._local_mcp_tools[0] == mock_mcp_tool # type: ignore[reportPrivateUsage]
assert cloned_agent.chat_options.tools is not None
assert len(cloned_agent.chat_options.tools) == 1
assert cloned_agent.default_options.get("tools") is not None
assert len(cloned_agent.default_options.get("tools")) == 1
async def test_return_to_previous_routing():
@@ -658,15 +658,14 @@ async def test_tool_choice_preserved_from_agent_config():
"""Verify that agent-level tool_choice configuration is preserved and not overridden."""
from unittest.mock import AsyncMock
from agent_framework import ChatResponse, ToolMode
from agent_framework import ChatResponse
# Create a mock chat client that records the tool_choice used
recorded_tool_choices: list[Any] = []
async def mock_get_response(messages: Any, **kwargs: Any) -> ChatResponse:
chat_options = kwargs.get("chat_options")
if chat_options:
recorded_tool_choices.append(chat_options.tool_choice)
async def mock_get_response(messages: Any, options: dict[str, Any] | None = None, **kwargs: Any) -> ChatResponse:
if options:
recorded_tool_choices.append(options.get("tool_choice"))
return ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Response")],
response_id="test_response",
@@ -675,11 +674,11 @@ async def test_tool_choice_preserved_from_agent_config():
mock_client = MagicMock()
mock_client.get_response = AsyncMock(side_effect=mock_get_response)
# Create agent with specific tool_choice configuration
# Create agent with specific tool_choice configuration via default_options
agent = ChatAgent(
chat_client=mock_client,
name="test_agent",
tool_choice=ToolMode(mode="required"), # type: ignore[arg-type]
default_options={"tool_choice": {"mode": "required"}},
)
# Run the agent
@@ -689,7 +688,7 @@ async def test_tool_choice_preserved_from_agent_config():
assert len(recorded_tool_choices) > 0, "No tool_choice recorded"
last_tool_choice = recorded_tool_choices[-1]
assert last_tool_choice is not None, "tool_choice should not be None"
assert str(last_tool_choice) == "required", f"Expected 'required', got {last_tool_choice}"
assert last_tool_choice == {"mode": "required"}, f"Expected 'required', got {last_tool_choice}"
async def test_handoff_builder_with_request_info():