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

Python: [BREAKING] updated structure and samples (#875)

Browse Source 
* updated structure and samples

* updated names and removed cross tests

* updated projects etc

* updated tests

* updated test

* test fixes

* removed devui for now

* updated all-tests task

* removed old style configs

* remove coverage from tests

* updated to unit tests with all-tests

* updated foundry everywhere

* fix azure ai tests

* fix merge tests

* fix mypy
This commit is contained in:
Eduard van Valkenburg
2025-09-25 09:02:53 +02:00
committed by GitHub
Unverified
parent 366a7f7d47
commit 9355329dfd
169 changed files with 1159 additions and 1761 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/tests/azure/conftest.py
+62
View File
@@ -0,0 +1,62 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from pytest import fixture
from agent_framework import ChatMessage
# region: Connector Settings fixtures
@fixture
def exclude_list(request: Any) -> list[str]:
"""Fixture that returns a list of environment variables to exclude."""
return request.param if hasattr(request, "param") else []
@fixture
def override_env_param_dict(request: Any) -> dict[str, str]:
"""Fixture that returns a dict of environment variables to override."""
return request.param if hasattr(request, "param") else {}
# These two fixtures are used for multiple things, also non-connector tests
@fixture()
def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): # type: ignore
"""Fixture to set environment variables for AzureOpenAISettings."""
if exclude_list is None:
exclude_list = []
if override_env_param_dict is None:
override_env_param_dict = {}
env_vars = {
"AZURE_OPENAI_ENDPOINT": "https://test-endpoint.com",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "test_chat_deployment",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME": "test_chat_deployment",
"AZURE_OPENAI_TEXT_DEPLOYMENT_NAME": "test_text_deployment",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME": "test_embedding_deployment",
"AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME": "test_text_to_image_deployment",
"AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME": "test_audio_to_text_deployment",
"AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME": "test_text_to_audio_deployment",
"AZURE_OPENAI_REALTIME_DEPLOYMENT_NAME": "test_realtime_deployment",
"AZURE_OPENAI_API_KEY": "test_api_key",
"AZURE_OPENAI_API_VERSION": "2023-03-15-preview",
"AZURE_OPENAI_BASE_URL": "https://test_text_deployment.test-base-url.com",
"AZURE_OPENAI_TOKEN_ENDPOINT": "https://test-token-endpoint.com",
}
env_vars.update(override_env_param_dict) # type: ignore
for key, value in env_vars.items():
if key in exclude_list:
monkeypatch.delenv(key, raising=False) # type: ignore
continue
monkeypatch.setenv(key, value) # type: ignore
return env_vars
@fixture(scope="function")
def chat_history() -> list[ChatMessage]:
return []
python/packages/main/tests/azure/test_azure_assistants_client.py
+723
View File
@@ -0,0 +1,723 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Annotated
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from azure.identity import AzureCliCredential
from pydantic import Field
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
HostedCodeInterpreterTool,
TextContent,
)
from agent_framework.azure import AzureOpenAIAssistantsClient
from agent_framework.exceptions import ServiceInitializationError
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
)
def create_test_azure_assistants_client(
mock_async_azure_openai: MagicMock,
deployment_name: str | None = None,
assistant_id: str | None = None,
assistant_name: str | None = None,
thread_id: str | None = None,
should_delete_assistant: bool = False,
) -> AzureOpenAIAssistantsClient:
"""Helper function to create AzureOpenAIAssistantsClient instances for testing, bypassing Pydantic validation."""
return AzureOpenAIAssistantsClient.model_construct(
ai_model_id=deployment_name or "test_chat_deployment",
assistant_id=assistant_id,
assistant_name=assistant_name,
thread_id=thread_id,
api_key="test-api-key",
endpoint="https://test-endpoint.com",
client=mock_async_azure_openai,
_should_delete_assistant=should_delete_assistant,
)
@pytest.fixture
def mock_async_azure_openai() -> MagicMock:
"""Mock AsyncAzureOpenAI client."""
mock_client = MagicMock()
# Mock beta.assistants
mock_client.beta.assistants.create = AsyncMock(return_value=MagicMock(id="test-assistant-id"))
mock_client.beta.assistants.delete = AsyncMock()
# Mock beta.threads
mock_client.beta.threads.create = AsyncMock(return_value=MagicMock(id="test-thread-id"))
mock_client.beta.threads.delete = AsyncMock()
# Mock beta.threads.runs
mock_client.beta.threads.runs.create = AsyncMock(return_value=MagicMock(id="test-run-id"))
mock_client.beta.threads.runs.retrieve = AsyncMock()
mock_client.beta.threads.runs.submit_tool_outputs = AsyncMock()
# Mock beta.threads.messages
mock_client.beta.threads.messages.create = AsyncMock()
mock_client.beta.threads.messages.list = AsyncMock(return_value=MagicMock(data=[]))
return mock_client
def test_azure_assistants_client_init_with_client(mock_async_azure_openai: MagicMock) -> None:
"""Test AzureOpenAIAssistantsClient initialization with existing client."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai,
deployment_name="test_chat_deployment",
assistant_id="existing-assistant-id",
thread_id="test-thread-id",
)
assert chat_client.client is mock_async_azure_openai
assert chat_client.ai_model_id == "test_chat_deployment"
assert chat_client.assistant_id == "existing-assistant-id"
assert chat_client.thread_id == "test-thread-id"
assert not chat_client._should_delete_assistant # type: ignore
assert isinstance(chat_client, ChatClientProtocol)
def test_azure_assistants_client_init_auto_create_client(
azure_openai_unit_test_env: dict[str, str],
mock_async_azure_openai: MagicMock,
) -> None:
"""Test AzureOpenAIAssistantsClient initialization with auto-created client."""
chat_client = AzureOpenAIAssistantsClient.model_construct(
ai_model_id=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
assistant_id=None,
assistant_name="TestAssistant",
thread_id=None,
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
client=mock_async_azure_openai,
_should_delete_assistant=False,
)
assert chat_client.client is mock_async_azure_openai
assert chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert chat_client.assistant_id is None
assert chat_client.assistant_name == "TestAssistant"
assert not chat_client._should_delete_assistant # type: ignore
def test_azure_assistants_client_init_validation_fail() -> None:
"""Test AzureOpenAIAssistantsClient initialization with validation failure."""
with pytest.raises(ServiceInitializationError):
# Force failure by providing invalid deployment name type - this should cause validation to fail
AzureOpenAIAssistantsClient(deployment_name=123, api_key="valid-key") # type: ignore
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
def test_azure_assistants_client_init_missing_deployment_name(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test AzureOpenAIAssistantsClient initialization with missing deployment name."""
with pytest.raises(ServiceInitializationError):
AzureOpenAIAssistantsClient(
api_key=azure_openai_unit_test_env.get("AZURE_OPENAI_API_KEY", "test-key"), env_file_path="nonexistent.env"
)
def test_azure_assistants_client_init_with_default_headers(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test AzureOpenAIAssistantsClient initialization with default headers."""
default_headers = {"X-Unit-Test": "test-guid"}
chat_client = AzureOpenAIAssistantsClient(
deployment_name="test_chat_deployment",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
default_headers=default_headers,
)
assert chat_client.ai_model_id == "test_chat_deployment"
assert isinstance(chat_client, ChatClientProtocol)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
assert key in chat_client.client.default_headers
assert chat_client.client.default_headers[key] == value
async def test_azure_assistants_client_get_assistant_id_or_create_existing_assistant(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when assistant_id is already provided."""
chat_client = create_test_azure_assistants_client(mock_async_azure_openai, assistant_id="existing-assistant-id")
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "existing-assistant-id"
assert not chat_client._should_delete_assistant # type: ignore
mock_async_azure_openai.beta.assistants.create.assert_not_called()
async def test_azure_assistants_client_get_assistant_id_or_create_create_new(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when creating a new assistant."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai, deployment_name="test_chat_deployment", assistant_name="TestAssistant"
)
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "test-assistant-id"
assert chat_client._should_delete_assistant # type: ignore
mock_async_azure_openai.beta.assistants.create.assert_called_once()
async def test_azure_assistants_client_aclose_should_not_delete(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test close when assistant should not be deleted."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-keep", should_delete_assistant=False
)
await chat_client.close() # type: ignore
# Verify assistant deletion was not called
mock_async_azure_openai.beta.assistants.delete.assert_not_called()
assert not chat_client._should_delete_assistant # type: ignore
async def test_azure_assistants_client_aclose_should_delete(mock_async_azure_openai: MagicMock) -> None:
"""Test close method calls cleanup."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
await chat_client.close()
# Verify assistant deletion was called
mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
assert not chat_client._should_delete_assistant # type: ignore
async def test_azure_assistants_client_async_context_manager(mock_async_azure_openai: MagicMock) -> None:
"""Test async context manager functionality."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
# Test context manager
async with chat_client:
pass # Just test that we can enter and exit
# Verify cleanup was called on exit
mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
def test_azure_assistants_client_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test serialization of AzureOpenAIAssistantsClient."""
default_headers = {"X-Unit-Test": "test-guid"}
# Test basic initialization and to_dict
chat_client = AzureOpenAIAssistantsClient(
deployment_name="test_chat_deployment",
assistant_id="test-assistant-id",
assistant_name="TestAssistant",
thread_id="test-thread-id",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
default_headers=default_headers,
)
dumped_settings = chat_client.to_dict()
assert dumped_settings["ai_model_id"] == "test_chat_deployment"
assert dumped_settings["assistant_id"] == "test-assistant-id"
assert dumped_settings["assistant_name"] == "TestAssistant"
assert dumped_settings["thread_id"] == "test-thread-id"
assert dumped_settings["api_key"] == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
# Assert that the default header we added is present in the dumped_settings default headers
for key, value in default_headers.items():
assert key in dumped_settings["default_headers"]
assert dumped_settings["default_headers"][key] == value
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
assert "User-Agent" not in dumped_settings["default_headers"]
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
return f"The weather in {location} is sunny with a high of 25°C."
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_get_response() -> None:
"""Test Azure Assistants Client response."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="The weather in Seattle is currently sunny with a high of 25°C. "
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert any(word in response.text.lower() for word in ["sunny", "25", "weather", "seattle"])
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_get_response_tools() -> None:
"""Test Azure Assistants Client response with tools."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
# 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",
)
assert response is not None
assert isinstance(response, ChatResponse)
assert any(word in response.text.lower() for word in ["sunny", "25", "weather"])
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_streaming() -> None:
"""Test Azure Assistants Client streaming response."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="The weather in Seattle is currently sunny with a high of 25°C. "
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
# Test that the client can be used to get a response
response = azure_assistants_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 any(word in full_message.lower() for word in ["sunny", "25", "weather", "seattle"])
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_streaming_tools() -> None:
"""Test Azure Assistants Client streaming response with tools."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
# 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",
)
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 any(word in full_message.lower() for word in ["sunny", "25", "weather"])
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_with_existing_assistant() -> None:
"""Test Azure Assistants Client with existing assistant ID."""
# First create an assistant to use in the test
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as temp_client:
# Get the assistant ID by triggering assistant creation
messages = [ChatMessage(role="user", text="Hello")]
await temp_client.get_response(messages=messages)
assistant_id = temp_client.assistant_id
# Now test using the existing assistant
async with AzureOpenAIAssistantsClient(
assistant_id=assistant_id, credential=AzureCliCredential()
) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
assert azure_assistants_client.assistant_id == assistant_id
messages = [ChatMessage(role="user", text="What can you do?")]
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert len(response.text) > 0
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_basic_run():
"""Test ChatAgent basic run functionality with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
) as agent:
# Run a simple query
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
assert "Hello World" in response.text
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_basic_run_streaming():
"""Test ChatAgent basic streaming functionality with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
) as agent:
# Run streaming query
full_message: str = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert chunk is not None
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_message += chunk.text
# Validate streaming response
assert len(full_message) > 0
assert "streaming response test" in full_message.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_thread_persistence():
"""Test ChatAgent thread persistence across runs with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(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 message - establish context
first_response = await agent.run(
"Remember this number: 42. What number did I just tell you to remember?", thread=thread
)
assert isinstance(first_response, AgentRunResponse)
assert "42" in first_response.text
# Second message - test conversation memory
second_response = await agent.run(
"What number did I tell you to remember in my previous message?", thread=thread
)
assert isinstance(second_response, AgentRunResponse)
assert "42" in second_response.text
# Verify thread has been populated with conversation ID
assert thread.service_thread_id is not None
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_existing_thread_id():
"""Test ChatAgent with existing thread ID to continue conversations across agent instances."""
# First, create a conversation and capture the thread ID
existing_thread_id = None
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Start a conversation and get the thread ID
thread = agent.get_new_thread()
response1 = await agent.run("What's the weather in Paris?", thread=thread)
# Validate first response
assert isinstance(response1, AgentRunResponse)
assert response1.text is not None
assert any(word in response1.text.lower() for word in ["weather", "paris"])
# The thread ID is set after the first response
existing_thread_id = thread.service_thread_id
assert existing_thread_id is not None
# Now continue with the same thread ID in a new agent instance
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(thread_id=existing_thread_id, credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Create a thread with the existing ID
thread = AgentThread(service_thread_id=existing_thread_id)
# Ask about the previous conversation
response2 = await agent.run("What was the last city I asked about?", thread=thread)
# Validate that the agent remembers the previous conversation
assert isinstance(response2, AgentRunResponse)
assert response2.text is not None
# Should reference Paris from the previous conversation
assert "paris" in response2.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_code_interpreter():
"""Test ChatAgent with code interpreter through AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can write and execute Python code.",
tools=[HostedCodeInterpreterTool()],
) as agent:
# Request code execution
response = await agent.run("Write Python code to calculate the factorial of 5 and show the result.")
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.text is not None
# Factorial of 5 is 120
assert "120" in response.text or "factorial" in response.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Assistants Client."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(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"])
def test_azure_assistants_client_entra_id_authentication() -> None:
"""Test Entra ID authentication path with credential."""
mock_credential = MagicMock()
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key = None # No API key to trigger Entra ID path
mock_settings.token_endpoint = "https://login.microsoftonline.com/test"
mock_settings.get_azure_auth_token.return_value = "entra-token-12345"
mock_settings.api_version = "2024-05-01-preview"
mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
mock_settings.base_url = None
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
api_key="placeholder-key",
endpoint="https://test-endpoint.openai.azure.com",
credential=mock_credential,
token_endpoint="https://login.microsoftonline.com/test",
)
# Verify Entra ID token was requested
mock_settings.get_azure_auth_token.assert_called_once_with(mock_credential)
# Verify client was created with the token
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token"] == "entra-token-12345"
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_no_authentication_error() -> None:
"""Test authentication validation error when no auth provided."""
with patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class:
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key = None # No API key
mock_settings.token_endpoint = None # No token endpoint
mock_settings_class.return_value = mock_settings
# Test missing authentication raises error
with pytest.raises(ServiceInitializationError, match="API key, ad_token, or ad_token_provider is required"):
AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
# No authentication provided at all
)
def test_azure_assistants_client_ad_token_authentication() -> None:
"""Test ad_token authentication client parameter path."""
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key = None # No API key
mock_settings.api_version = "2024-05-01-preview"
mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
mock_settings.base_url = None
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
ad_token="test-ad-token-12345",
)
# ad_token path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token"] == "test-ad-token-12345"
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_ad_token_provider_authentication() -> None:
"""Test ad_token_provider authentication client parameter path."""
from openai.lib.azure import AsyncAzureADTokenProvider
mock_token_provider = MagicMock(spec=AsyncAzureADTokenProvider)
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key = None # No API key
mock_settings.api_version = "2024-05-01-preview"
mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
mock_settings.base_url = None
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
ad_token_provider=mock_token_provider,
)
# ad_token_provider path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token_provider"] is mock_token_provider
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_base_url_configuration() -> None:
"""Test base_url client parameter path."""
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key.get_secret_value.return_value = "test-api-key"
mock_settings.base_url = "https://custom-base-url.com"
mock_settings.endpoint = None # No endpoint, should use base_url
mock_settings.api_version = "2024-05-01-preview"
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment", api_key="test-api-key", base_url="https://custom-base-url.com"
)
# base_url path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["base_url"] == "https://custom-base-url.com"
assert "azure_endpoint" not in call_args
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_azure_endpoint_configuration() -> None:
"""Test azure_endpoint client parameter path."""
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key.get_secret_value.return_value = "test-api-key"
mock_settings.base_url = None # No base_url
mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
mock_settings.api_version = "2024-05-01-preview"
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
api_key="test-api-key",
endpoint="https://test-endpoint.openai.azure.com",
)
# azure_endpoint path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_endpoint"] == "https://test-endpoint.openai.azure.com"
assert "base_url" not in call_args
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
python/packages/main/tests/azure/test_azure_chat_client.py
+835
View File
@@ -0,0 +1,835 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import os
from unittest.mock import AsyncMock, MagicMock, patch
import openai
import pytest
from azure.identity import AzureCliCredential
from httpx import Request, Response
from openai import AsyncAzureOpenAI, AsyncStream
from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
from openai.types.chat import ChatCompletion, ChatCompletionChunk
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDelta
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
BaseChatClient,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
TextContent,
ai_function,
)
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
from agent_framework.openai import (
ContentFilterResultSeverity,
OpenAIContentFilterException,
)
# region Service Setup
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
)
def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
azure_chat_client = AzureOpenAIChatClient()
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, BaseChatClient)
def test_init_client(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization with client
client = MagicMock(spec=AsyncAzureOpenAI)
azure_chat_client = AzureOpenAIChatClient(async_client=client)
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
def test_init_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
# Custom header for testing
default_headers = {"X-Unit-Test": "test-guid"}
azure_chat_client = AzureOpenAIChatClient(
default_headers=default_headers,
)
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, BaseChatClient)
for key, value in default_headers.items():
assert key in azure_chat_client.client.default_headers
assert azure_chat_client.client.default_headers[key] == value
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
def test_init_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
azure_chat_client = AzureOpenAIChatClient()
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, BaseChatClient)
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
def test_init_with_empty_deployment_name(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ServiceInitializationError):
AzureOpenAIChatClient(
env_file_path="test.env",
)
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
def test_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ServiceInitializationError):
AzureOpenAIChatClient(
env_file_path="test.env",
)
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
def test_init_with_invalid_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ServiceInitializationError):
AzureOpenAIChatClient()
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Test": "test"}
settings = {
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
"default_headers": default_headers,
"env_file_path": "test.env",
}
azure_chat_client = AzureOpenAIChatClient.from_dict(settings)
dumped_settings = azure_chat_client.to_dict()
assert dumped_settings["ai_model_id"] == settings["deployment_name"]
assert str(settings["endpoint"]) in str(dumped_settings["base_url"])
assert str(settings["deployment_name"]) in str(dumped_settings["base_url"])
assert settings["api_key"] == dumped_settings["api_key"]
assert settings["api_version"] == dumped_settings["api_version"]
# Assert that the default header we added is present in the dumped_settings default headers
for key, value in default_headers.items():
assert key in dumped_settings["default_headers"]
assert dumped_settings["default_headers"][key] == value
# Assert that the 'User-agent' header is not present in the dumped_settings default headers
assert USER_AGENT_KEY not in dumped_settings["default_headers"]
# endregion
# region CMC
@pytest.fixture
def mock_chat_completion_response() -> ChatCompletion:
return ChatCompletion(
id="test_id",
choices=[
Choice(index=0, message=ChatCompletionMessage(content="test", role="assistant"), finish_reason="stop")
],
created=0,
model="test",
object="chat.completion",
)
@pytest.fixture
def mock_streaming_chat_completion_response() -> AsyncStream[ChatCompletionChunk]:
content = ChatCompletionChunk(
id="test_id",
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
created=0,
model="test",
object="chat.completion.chunk",
)
stream = MagicMock(spec=AsyncStream)
stream.__aiter__.return_value = [content]
return stream
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(text="hello world", role="user"))
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(
messages=chat_history,
)
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
stream=False,
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc_with_logit_bias(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
chat_history.append(ChatMessage(text=prompt, role="user"))
token_bias: dict[str | int, float] = {"1": -100}
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(messages=chat_history, logit_bias=token_bias)
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
stream=False,
logit_bias=token_bias,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc_with_stop(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
chat_history.append(ChatMessage(text=prompt, role="user"))
stop = ["!"]
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(messages=chat_history, stop=stop)
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
stream=False,
stop=stop,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_azure_on_your_data(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
Choice(
index=0,
message=ChatCompletionMessage(
content="test",
role="assistant",
context={ # type: ignore
"citations": [
{
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
}
],
"intent": "query used",
},
),
finish_reason="stop",
)
]
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
chat_history.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
{
"type": "AzureCognitiveSearch",
"parameters": {
"indexName": "test_index",
"endpoint": "https://test-endpoint-search.com",
"key": "test_key",
},
}
]
}
azure_chat_client = AzureOpenAIChatClient()
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
assert isinstance(content.messages[0].contents[0], TextContent)
assert len(content.messages[0].contents[0].annotations) == 1
assert content.messages[0].contents[0].annotations[0].title == "test title"
assert content.messages[0].contents[0].text == "test"
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
stream=False,
extra_body=expected_data_settings,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_azure_on_your_data_string(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
Choice(
index=0,
message=ChatCompletionMessage(
content="test",
role="assistant",
context=json.dumps({ # type: ignore
"citations": [
{
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
}
],
"intent": "query used",
}),
),
finish_reason="stop",
)
]
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
messages_in.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
{
"type": "AzureCognitiveSearch",
"parameters": {
"indexName": "test_index",
"endpoint": "https://test-endpoint-search.com",
"key": "test_key",
},
}
]
}
azure_chat_client = AzureOpenAIChatClient()
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
assert isinstance(content.messages[0].contents[0], TextContent)
assert len(content.messages[0].contents[0].annotations) == 1
assert content.messages[0].contents[0].annotations[0].title == "test title"
assert content.messages[0].contents[0].text == "test"
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
stream=False,
extra_body=expected_data_settings,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_azure_on_your_data_fail(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
Choice(
index=0,
message=ChatCompletionMessage(
content="test",
role="assistant",
context="not a dictionary", # type: ignore
),
finish_reason="stop",
)
]
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
messages_in.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
{
"type": "AzureCognitiveSearch",
"parameters": {
"indexName": "test_index",
"endpoint": "https://test-endpoint-search.com",
"key": "test_key",
},
}
]
}
azure_chat_client = AzureOpenAIChatClient()
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
assert isinstance(content.messages[0].contents[0], TextContent)
assert content.messages[0].contents[0].text == "test"
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
stream=False,
extra_body=expected_data_settings,
)
CONTENT_FILTERED_ERROR_MESSAGE = (
"The response was filtered due to the prompt triggering Azure OpenAI's content management policy. Please "
"modify your prompt and retry. To learn more about our content filtering policies please read our "
"documentation: https://go.microsoft.com/fwlink/?linkid=2198766"
)
CONTENT_FILTERED_ERROR_FULL_MESSAGE = (
"Error code: 400 - {'error': {'message': \"%s\", 'type': null, 'param': 'prompt', 'code': 'content_filter', "
"'status': 400, 'innererror': {'code': 'ResponsibleAIPolicyViolation', 'content_filter_result': {'hate': "
"{'filtered': True, 'severity': 'high'}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': "
"{'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}}}"
) % CONTENT_FILTERED_ERROR_MESSAGE
@patch.object(AsyncChatCompletions, "create")
async def test_content_filtering_raises_correct_exception(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
mock_create.side_effect = openai.BadRequestError(
CONTENT_FILTERED_ERROR_FULL_MESSAGE,
response=Response(400, request=Request("POST", test_endpoint)),
body={
"message": CONTENT_FILTERED_ERROR_MESSAGE,
"type": None,
"param": "prompt",
"code": "content_filter",
"status": 400,
"innererror": {
"code": "ResponsibleAIPolicyViolation",
"content_filter_result": {
"hate": {"filtered": True, "severity": "high"},
"self_harm": {"filtered": False, "severity": "safe"},
"sexual": {"filtered": False, "severity": "safe"},
"violence": {"filtered": False, "severity": "safe"},
},
},
},
)
azure_chat_client = AzureOpenAIChatClient()
with pytest.raises(OpenAIContentFilterException, match="service encountered a content error") as exc_info:
await azure_chat_client.get_response(
messages=chat_history,
)
content_filter_exc = exc_info.value
assert content_filter_exc.param == "prompt"
assert content_filter_exc.content_filter_result["hate"].filtered
assert content_filter_exc.content_filter_result["hate"].severity == ContentFilterResultSeverity.HIGH
@patch.object(AsyncChatCompletions, "create")
async def test_content_filtering_without_response_code_raises_with_default_code(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
mock_create.side_effect = openai.BadRequestError(
CONTENT_FILTERED_ERROR_FULL_MESSAGE,
response=Response(400, request=Request("POST", test_endpoint)),
body={
"message": CONTENT_FILTERED_ERROR_MESSAGE,
"type": None,
"param": "prompt",
"code": "content_filter",
"status": 400,
"innererror": {
"content_filter_result": {
"hate": {"filtered": True, "severity": "high"},
"self_harm": {"filtered": False, "severity": "safe"},
"sexual": {"filtered": False, "severity": "safe"},
"violence": {"filtered": False, "severity": "safe"},
},
},
},
)
azure_chat_client = AzureOpenAIChatClient()
with pytest.raises(OpenAIContentFilterException, match="service encountered a content error"):
await azure_chat_client.get_response(
messages=chat_history,
)
@patch.object(AsyncChatCompletions, "create")
async def test_bad_request_non_content_filter(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
mock_create.side_effect = openai.BadRequestError(
"The request was bad.", response=Response(400, request=Request("POST", test_endpoint)), body={}
)
azure_chat_client = AzureOpenAIChatClient()
with pytest.raises(ServiceResponseException, match="service failed to complete the prompt"):
await azure_chat_client.get_response(
messages=chat_history,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_get_streaming(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_streaming_chat_completion_response: AsyncStream[ChatCompletionChunk],
) -> None:
mock_create.return_value = mock_streaming_chat_completion_response
chat_history.append(ChatMessage(text="hello world", role="user"))
azure_chat_client = AzureOpenAIChatClient()
async for msg in azure_chat_client.get_streaming_response(
messages=chat_history,
):
assert msg is not None
assert msg.message_id is not None
assert msg.response_id is not None
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
stream=True,
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
# NOTE: The `stream_options={"include_usage": True}` is explicitly enforced in
# `OpenAIChatCompletionBase._inner_get_streaming_response`.
# To ensure consistency, we align the arguments here accordingly.
stream_options={"include_usage": True},
)
@ai_function
def get_story_text() -> str:
"""Returns a story about Emily and David."""
return (
"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."
)
@ai_function
def get_weather(location: str) -> str:
"""Get the current weather for a location."""
return f"The weather in {location} is sunny and 72°F."
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_response() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_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 azure_chat_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
# Check for any relevant keywords that indicate the AI understood the context
assert any(
word in response.text.lower() for word in ["scientists", "research", "antarctica", "glaciology", "climate"]
)
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_response_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_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 azure_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
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_streaming() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_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 = azure_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
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_streaming_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_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 = azure_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
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_basic_run():
"""Test Azure OpenAI chat client agent basic run functionality with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
) 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()
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_basic_run_streaming():
"""Test Azure OpenAI chat client agent basic streaming functionality with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(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()
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_thread_persistence():
"""Test Azure OpenAI chat client agent thread persistence across runs with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(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
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()
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_existing_thread():
"""Test Azure 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=AzureOpenAIChatClient(credential=AzureCliCredential()),
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=AzureOpenAIChatClient(credential=AzureCliCredential()),
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()
@skip_if_azure_integration_tests_disabled
async def test_azure_chat_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Chat Client."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(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"])
python/packages/main/tests/azure/test_azure_responses_client.py
+624
View File
@@ -0,0 +1,624 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Annotated
import pytest
from azure.identity import AzureCliCredential
from pydantic import BaseModel
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
HostedCodeInterpreterTool,
HostedFileSearchTool,
HostedMCPTool,
HostedVectorStoreContent,
TextContent,
ai_function,
)
from agent_framework.azure import AzureOpenAIResponsesClient
from agent_framework.exceptions import ServiceInitializationError
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
)
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
location: str
weather: str
@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 weather in {location} is sunny and 72°F."
async def create_vector_store(client: AzureOpenAIResponsesClient) -> tuple[str, HostedVectorStoreContent]:
"""Create a vector store with sample documents for testing."""
file = await client.client.files.create(
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="assistants"
)
vector_store = await client.client.vector_stores.create(
name="knowledge_base",
expires_after={"anchor": "last_active_at", "days": 1},
)
result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
if result.last_error is not None:
raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
return file.id, HostedVectorStoreContent(vector_store_id=vector_store.id)
async def delete_vector_store(client: AzureOpenAIResponsesClient, file_id: str, vector_store_id: str) -> None:
"""Delete the vector store after tests."""
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
await client.client.files.delete(file_id=file_id)
def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
azure_responses_client = AzureOpenAIResponsesClient()
assert azure_responses_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, ChatClientProtocol)
def test_init_validation_fail() -> None:
# Test successful initialization
with pytest.raises(ServiceInitializationError):
AzureOpenAIResponsesClient(api_key="34523", deployment_name={"test": "dict"}) # type: ignore
def test_init_ai_model_id_constructor(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
ai_model_id = "test_model_id"
azure_responses_client = AzureOpenAIResponsesClient(deployment_name=ai_model_id)
assert azure_responses_client.ai_model_id == ai_model_id
assert isinstance(azure_responses_client, ChatClientProtocol)
def test_init_with_default_header(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Unit-Test": "test-guid"}
# Test successful initialization
azure_responses_client = AzureOpenAIResponsesClient(
default_headers=default_headers,
)
assert azure_responses_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, ChatClientProtocol)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
assert key in azure_responses_client.client.default_headers
assert azure_responses_client.client.default_headers[key] == value
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]], indirect=True)
def test_init_with_empty_model_id(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ServiceInitializationError):
AzureOpenAIResponsesClient(
env_file_path="test.env",
)
def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Unit-Test": "test-guid"}
settings = {
"ai_model_id": azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"],
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
"default_headers": default_headers,
}
azure_responses_client = AzureOpenAIResponsesClient.from_dict(settings)
dumped_settings = azure_responses_client.to_dict()
assert dumped_settings["ai_model_id"] == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert dumped_settings["api_key"] == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
# Assert that the default header we added is present in the dumped_settings default headers
for key, value in default_headers.items():
assert key in dumped_settings["default_headers"]
assert dumped_settings["default_headers"][key] == value
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
assert "User-Agent" not in dumped_settings["default_headers"]
@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()
@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()
@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,
),
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()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_streaming_tools() -> None:
"""Test azure responses client streaming tools."""
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()
@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()
@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()
@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
@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,
)
# 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
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_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
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
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=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
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()
@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
@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"])
@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
@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
@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."""
mcp_tool = 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",
)
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=[mcp_tool],
) as agent:
response = await agent.run(
"How to create an Azure storage account using az cli?",
max_tokens=200,
)
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
@skip_if_azure_integration_tests_disabled
@pytest.mark.skip(reason="File search requires API key auth, subscription only allows token auth")
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 web 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
@skip_if_azure_integration_tests_disabled
@pytest.mark.skip(reason="File search requires API key auth, subscription only allows token auth")
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 web 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
python/packages/main/tests/azure/test_entra_id_authentication.py
+156
View File
@@ -0,0 +1,156 @@
# Copyright (c) Microsoft. All rights reserved.
from unittest.mock import AsyncMock, MagicMock
import pytest
from azure.core.exceptions import ClientAuthenticationError
from agent_framework.azure._entra_id_authentication import (
get_entra_auth_token,
get_entra_auth_token_async,
)
from agent_framework.exceptions import ServiceInvalidAuthError
@pytest.fixture
def mock_credential() -> MagicMock:
"""Mock synchronous TokenCredential."""
mock_cred = MagicMock()
# Create a mock token object with a .token attribute
mock_token = MagicMock()
mock_token.token = "test-access-token-12345"
mock_cred.get_token.return_value = mock_token
return mock_cred
@pytest.fixture
def mock_async_credential() -> MagicMock:
"""Mock asynchronous AsyncTokenCredential."""
mock_cred = MagicMock()
# Create a mock token object with a .token attribute
mock_token = MagicMock()
mock_token.token = "test-async-access-token-12345"
mock_cred.get_token = AsyncMock(return_value=mock_token)
return mock_cred
def test_get_entra_auth_token_success(mock_credential: MagicMock) -> None:
"""Test successful token retrieval with sync function."""
token_endpoint = "https://test-endpoint.com/.default"
result = get_entra_auth_token(mock_credential, token_endpoint)
# Assert - check the results
assert result == "test-access-token-12345"
mock_credential.get_token.assert_called_once_with(token_endpoint)
async def test_get_entra_auth_token_async_success(mock_async_credential: MagicMock) -> None:
"""Test successful token retrieval with async function."""
token_endpoint = "https://test-endpoint.com/.default"
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
# Assert - check the results
assert result == "test-async-access-token-12345"
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
def test_get_entra_auth_token_missing_endpoint(mock_credential: MagicMock) -> None:
"""Test that missing token endpoint raises ServiceInvalidAuthError."""
# Test with empty string
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
get_entra_auth_token(mock_credential, "")
# Test with None
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
get_entra_auth_token(mock_credential, None) # type: ignore
async def test_get_entra_auth_token_async_missing_endpoint(mock_async_credential: MagicMock) -> None:
"""Test that missing token endpoint raises ServiceInvalidAuthError in async function."""
# Test with empty string
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
await get_entra_auth_token_async(mock_async_credential, "")
# Test with None
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
await get_entra_auth_token_async(mock_async_credential, None) # type: ignore
def test_get_entra_auth_token_auth_failure(mock_credential: MagicMock) -> None:
"""Test that Azure authentication failure returns None."""
mock_credential.get_token.side_effect = ClientAuthenticationError("Auth failed")
token_endpoint = "https://test-endpoint.com/.default"
result = get_entra_auth_token(mock_credential, token_endpoint)
# Assert - should return None on auth failure
assert result is None
mock_credential.get_token.assert_called_once_with(token_endpoint)
async def test_get_entra_auth_token_async_auth_failure(mock_async_credential: MagicMock) -> None:
"""Test that Azure authentication failure returns None in async function."""
mock_async_credential.get_token.side_effect = ClientAuthenticationError("Auth failed")
token_endpoint = "https://test-endpoint.com/.default"
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
# Assert - should return None on auth failure
assert result is None
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
def test_get_entra_auth_token_none_token_response(mock_credential: MagicMock) -> None:
"""Test that None token response returns None."""
mock_credential.get_token.return_value = None
token_endpoint = "https://test-endpoint.com/.default"
result = get_entra_auth_token(mock_credential, token_endpoint)
# Assert
assert result is None
mock_credential.get_token.assert_called_once_with(token_endpoint)
async def test_get_entra_auth_token_async_none_token_response(mock_async_credential: MagicMock) -> None:
"""Test that None token response returns None in async function."""
mock_async_credential.get_token.return_value = None
token_endpoint = "https://test-endpoint.com/.default"
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
# Assert
assert result is None
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
def test_get_entra_auth_token_with_kwargs(mock_credential: MagicMock) -> None:
"""Test that kwargs are passed through to get_token."""
token_endpoint = "https://test-endpoint.com/.default"
extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"}
result = get_entra_auth_token(mock_credential, token_endpoint, **extra_kwargs)
# Assert
assert result == "test-access-token-12345"
mock_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)
async def test_get_entra_auth_token_async_with_kwargs(mock_async_credential: MagicMock) -> None:
"""Test that kwargs are passed through to async get_token."""
token_endpoint = "https://test-endpoint.com/.default"
extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"}
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint, **extra_kwargs)
# Assert
assert result == "test-async-access-token-12345"
mock_async_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)
python/packages/main/tests/openai/test_openai_responses_client.py
+1 -2
View File
@@ -1560,7 +1560,7 @@ async def test_openai_responses_client_agent_chat_options_agent_level() -> None:
@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."""
# Use the same MCP server as the Foundry example
mcp_tool = HostedMCPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
@@ -1573,7 +1573,6 @@ async def test_openai_responses_client_agent_hosted_mcp_tool() -> None:
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=[mcp_tool],
) as agent:
# Use the same query as the Foundry example
response = await agent.run(
"How to create an Azure storage account using az cli?",
max_tokens=200,