[BREAKING] Python: fix OpenAI Azure routing and provider samples (#4925)

* Python: fix OpenAI Azure routing and provider samples

Prefer OpenAI when OPENAI_API_KEY is present unless Azure is explicitly requested. Clarify constructor docs, keep deprecated Azure wrappers compatible with stricter settings validation, and refresh the provider samples and tests to use the current client patterns.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix bandit

* Python: align OpenAI embedding Azure routing

Extend the shared OpenAI-vs-Azure routing and credential behavior to the embedding client, add Azure embedding regression coverage, and refresh the embedding samples to use the generic client path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: fix embedding client pyright check

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: thin OpenAI embedding wrapper

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: document embedding overload routing

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: fix callable OpenAI key routing

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: fix Azure credential routing tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: address OpenAI review feedback

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: narrow Azure routing markers

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: refine OpenAI model fallback order

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: narrow Azure deployment docs

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: remove embedding routing wording

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: run embedding Azure integration tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* changed variable name

* Python: expand OpenAI package README

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* clarified readme

* Python: fix Azure OpenAI integration setup

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: correct Azure integration env mapping

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* updated code to fix int tests

* test updates

* test fix

* fix test setup

* updates to tests and setup

* remove openai assistants int tests

* improvements in int tests

* fix env var

* fix env vars

* fix azure responses test

* trigger actions

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-03-27 13:33:39 +00:00
committed by GitHub
co-authored by Copilot
parent 3611be82cf
commit cc0cfaaac8
103 changed files with 5451 additions and 4216 deletions
@@ -43,6 +43,8 @@ def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): #
"OPENAI_ORG_ID",
"OPENAI_MODEL",
"OPENAI_EMBEDDING_MODEL",
"OPENAI_CHAT_MODEL",
"OPENAI_RESPONSES_MODEL",
"OPENAI_TEXT_MODEL_ID",
"OPENAI_TEXT_TO_IMAGE_MODEL_ID",
"OPENAI_AUDIO_TO_TEXT_MODEL_ID",
@@ -53,6 +55,9 @@ def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): #
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_BASE_URL",
"AZURE_OPENAI_API_KEY",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME",
"AZURE_OPENAI_DEPLOYMENT_NAME",
"AZURE_OPENAI_API_VERSION",
],
@@ -97,6 +102,8 @@ def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dic
"OPENAI_ORG_ID",
"OPENAI_MODEL",
"OPENAI_EMBEDDING_MODEL",
"OPENAI_CHAT_MODEL",
"OPENAI_RESPONSES_MODEL",
"OPENAI_TEXT_MODEL_ID",
"OPENAI_TEXT_TO_IMAGE_MODEL_ID",
"OPENAI_AUDIO_TO_TEXT_MODEL_ID",
@@ -107,6 +114,9 @@ def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dic
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_BASE_URL",
"AZURE_OPENAI_API_KEY",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME",
"AZURE_OPENAI_DEPLOYMENT_NAME",
"AZURE_OPENAI_API_VERSION",
],
@@ -114,6 +124,9 @@ def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dic
env_vars = {
"AZURE_OPENAI_ENDPOINT": "https://test-endpoint.openai.azure.com",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "test_chat_deployment",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME": "test_responses_deployment",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME": "test_embedding_deployment",
"AZURE_OPENAI_DEPLOYMENT_NAME": "test_deployment",
"AZURE_OPENAI_API_KEY": "test_api_key",
"AZURE_OPENAI_API_VERSION": "2024-12-01-preview",
@@ -1,6 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Annotated, Any
from unittest.mock import AsyncMock, MagicMock
@@ -750,64 +749,3 @@ class TestToolMerging:
# endregion
# region Integration Tests
skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
reason="No real OPENAI_API_KEY provided; skipping integration tests.",
)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_openai_integration_tests_disabled
class TestOpenAIAssistantProviderIntegration:
"""Integration tests requiring real OpenAI API."""
async def test_create_and_run_agent(self) -> None:
"""End-to-end test of creating and running an agent."""
provider = OpenAIAssistantProvider()
agent = await provider.create_agent(
name="IntegrationTestAgent",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful assistant. Respond briefly.",
)
try:
result = await agent.run("Say 'hello' and nothing else.")
result_text = str(result)
assert "hello" in result_text.lower()
finally:
# Clean up the assistant
await provider._client.beta.assistants.delete(agent.id) # type: ignore[reportPrivateUsage, union-attr]
async def test_create_agent_with_function_tools_integration(self) -> None:
"""Integration test with function tools."""
provider = OpenAIAssistantProvider()
@tool(approval_mode="never_require")
def get_current_time() -> str:
"""Get the current time."""
from datetime import datetime
return datetime.now().strftime("%H:%M")
agent = await provider.create_agent(
name="TimeAgent",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful assistant.",
tools=[get_current_time],
)
try:
result = await agent.run("What time is it? Use the get_current_time function.")
result_text = str(result)
# The response should contain time information
assert ":" in result_text or "time" in result_text.lower()
finally:
await provider._client.beta.assistants.delete(agent.id) # type: ignore[reportPrivateUsage, union-attr]
# endregion
@@ -28,6 +28,7 @@ from agent_framework._sessions import (
from agent_framework.exceptions import (
ChatClientException,
ChatClientInvalidRequestException,
SettingNotFoundError,
)
from openai import BadRequestError
from openai.types.responses.response_reasoning_item import Summary
@@ -109,6 +110,14 @@ def test_init(openai_unit_test_env: dict[str, str]) -> None:
assert isinstance(openai_responses_client, SupportsChatGetResponse)
def test_init_prefers_openai_responses_model(monkeypatch, openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_RESPONSES_MODEL", "test_responses_model_id")
openai_responses_client = OpenAIChatClient()
assert openai_responses_client.model == "test_responses_model_id"
def test_init_validation_fail() -> None:
# Test successful initialization
with pytest.raises(ValueError):
@@ -143,7 +152,7 @@ def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None:
@pytest.mark.parametrize("exclude_list", [["OPENAI_MODEL"]], indirect=True)
def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ValueError):
with pytest.raises(SettingNotFoundError):
OpenAIChatClient()
@@ -151,7 +160,7 @@ def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
def test_init_with_empty_api_key(openai_unit_test_env: dict[str, str]) -> None:
model_id = "test_model_id"
with pytest.raises(ValueError):
with pytest.raises(SettingNotFoundError):
OpenAIChatClient(
model=model_id,
)
@@ -203,34 +212,56 @@ async def test_get_response_with_invalid_input() -> None:
async def test_get_response_with_all_parameters() -> None:
"""Test get_response with all possible parameters to cover parameter handling logic."""
"""Test request preparation with a comprehensive parameter set."""
client = OpenAIChatClient(model="test-model", api_key="test-key")
# Test with comprehensive parameter set - should fail due to invalid API key
with pytest.raises(ChatClientException):
await client.get_response(
messages=[Message(role="user", text="Test message")],
options={
"include": ["message.output_text.logprobs"],
"instructions": "You are a helpful assistant",
"max_tokens": 100,
"parallel_tool_calls": True,
"model": "gpt-4",
"previous_response_id": "prev-123",
"reasoning": {"chain_of_thought": "enabled"},
"service_tier": "auto",
"response_format": OutputStruct,
"seed": 42,
"store": True,
"temperature": 0.7,
"tool_choice": "auto",
"tools": [get_weather],
"top_p": 0.9,
"user": "test-user",
"truncation": "auto",
"timeout": 30.0,
"additional_properties": {"custom": "value"},
},
)
_, run_options, _ = await client._prepare_request(
messages=[Message(role="user", text="Test message")],
options={
"include": ["message.output_text.logprobs"],
"instructions": "You are a helpful assistant",
"max_tokens": 100,
"parallel_tool_calls": True,
"model": "gpt-4",
"previous_response_id": "prev-123",
"reasoning": {"chain_of_thought": "enabled"},
"service_tier": "auto",
"response_format": OutputStruct,
"seed": 42,
"store": True,
"temperature": 0.7,
"tool_choice": "auto",
"tools": [get_weather],
"top_p": 0.9,
"user": "test-user",
"truncation": "auto",
"timeout": 30.0,
"additional_properties": {"custom": "value"},
},
)
assert run_options["include"] == ["message.output_text.logprobs"]
assert run_options["max_output_tokens"] == 100
assert run_options["parallel_tool_calls"] is True
assert run_options["model"] == "gpt-4"
assert run_options["previous_response_id"] == "prev-123"
assert run_options["reasoning"] == {"chain_of_thought": "enabled"}
assert run_options["service_tier"] == "auto"
assert run_options["text_format"] is OutputStruct
assert run_options["store"] is True
assert run_options["temperature"] == 0.7
assert run_options["tool_choice"] == "auto"
assert run_options["top_p"] == 0.9
assert run_options["user"] == "test-user"
assert run_options["truncation"] == "auto"
assert run_options["timeout"] == 30.0
assert run_options["additional_properties"] == {"custom": "value"}
assert len(run_options["tools"]) == 1
assert run_options["tools"][0]["type"] == "function"
assert run_options["tools"][0]["name"] == "get_weather"
assert run_options["input"][0]["role"] == "system"
assert run_options["input"][0]["content"][0]["text"] == "You are a helpful assistant"
assert run_options["input"][1]["role"] == "user"
assert run_options["input"][1]["content"][0]["text"] == "Test message"
@pytest.mark.asyncio
@@ -248,12 +279,13 @@ async def test_web_search_tool_with_location() -> None:
}
)
# Should raise an authentication error due to invalid API key
with pytest.raises(ChatClientException):
await client.get_response(
messages=[Message(role="user", text="What's the weather?")],
options={"tools": [web_search_tool], "tool_choice": "auto"},
)
_, run_options, _ = await client._prepare_request(
messages=[Message(role="user", text="What's the weather?")],
options={"tools": [web_search_tool], "tool_choice": "auto"},
)
assert run_options["tools"] == [web_search_tool]
assert run_options["tool_choice"] == "auto"
async def test_code_interpreter_tool_variations() -> None:
@@ -263,20 +295,22 @@ async def test_code_interpreter_tool_variations() -> None:
# Test code interpreter using static method
code_tool = OpenAIChatClient.get_code_interpreter_tool()
with pytest.raises(ChatClientException):
await client.get_response(
messages=[Message("user", ["Run some code"])],
options={"tools": [code_tool]},
)
_, run_options, _ = await client._prepare_request(
messages=[Message("user", ["Run some code"])],
options={"tools": [code_tool]},
)
assert run_options["tools"] == [code_tool]
# Test code interpreter with files using static method
code_tool_with_files = OpenAIChatClient.get_code_interpreter_tool(file_ids=["file1", "file2"])
with pytest.raises(ChatClientException):
await client.get_response(
messages=[Message(role="user", text="Process these files")],
options={"tools": [code_tool_with_files]},
)
_, run_options, _ = await client._prepare_request(
messages=[Message(role="user", text="Process these files")],
options={"tools": [code_tool_with_files]},
)
assert run_options["tools"] == [code_tool_with_files]
async def test_content_filter_exception() -> None:
@@ -300,23 +334,23 @@ async def test_content_filter_exception() -> None:
@pytest.mark.asyncio
async def test_hosted_file_search_tool_validation() -> None:
"""Test get_response HostedFileSearchTool validation."""
"""Test HostedFileSearchTool validation and request preparation."""
client = OpenAIChatClient(model="test-model", api_key="test-key")
# Test file search tool with vector store IDs
file_search_tool = OpenAIChatClient.get_file_search_tool(vector_store_ids=["vs_123"])
# Test using file search tool - may raise various exceptions depending on API response
with pytest.raises((ValueError, ChatClientInvalidRequestException, ChatClientException)):
await client.get_response(
messages=[Message("user", ["Test"])],
options={"tools": [file_search_tool]},
)
_, run_options, _ = await client._prepare_request(
messages=[Message("user", ["Test"])],
options={"tools": [file_search_tool]},
)
assert run_options["tools"] == [file_search_tool]
async def test_chat_message_parsing_with_function_calls() -> None:
"""Test get_response message preparation with function call and result content types in conversation flow."""
"""Test message preparation with function call and function result content."""
client = OpenAIChatClient(model="test-model", api_key="test-key")
# Create messages with function call and result content
@@ -335,9 +369,27 @@ async def test_chat_message_parsing_with_function_calls() -> None:
Message(role="tool", contents=[function_result]),
]
# This should exercise the message parsing logic - will fail due to invalid API key
with pytest.raises(ChatClientException):
await client.get_response(messages=messages)
prepared_messages = client._prepare_messages_for_openai(messages)
assert prepared_messages == [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "Call a function"}],
},
{
"call_id": "test-call-id",
"id": "fc_test-fc-id",
"type": "function_call",
"name": "test_function",
"arguments": '{"param": "value"}',
},
{
"call_id": "test-call-id",
"type": "function_call_output",
"output": "Function executed successfully",
},
]
async def test_response_format_parse_path() -> None:
@@ -3043,8 +3095,6 @@ def test_with_callable_api_key() -> None:
"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"),
@@ -3057,7 +3107,6 @@ def test_with_callable_api_key() -> None:
# OpenAIChatOptions - 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
@@ -3113,70 +3162,56 @@ async def test_integration_options(
they don't cause failures. Options marked with needs_validation also
check that the feature actually works correctly.
"""
openai_responses_client = OpenAIChatClient()
client = OpenAIChatClient()
# Need at least 2 iterations for tool_choice tests: one to get function call, one to get final response
openai_responses_client.function_invocation_configuration["max_iterations"] = 2
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
# Prepare test message
# Prepare test message
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(Message(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [Message(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]
# Test streaming mode
response = await client.get_response(stream=True, messages=messages, options=options).get_final_response()
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"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
# 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"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(Message(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [Message(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_stream = openai_responses_client.get_response(
stream=True,
messages=messages,
options=options,
)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await openai_responses_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()
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.timeout(300)
@@ -3186,53 +3221,24 @@ async def test_integration_options(
async def test_integration_web_search() -> None:
client = OpenAIChatClient(model="gpt-5")
for streaming in [False, True]:
# Use static method for web search tool
web_search_tool = OpenAIChatClient.get_web_search_tool()
content = {
"messages": [
Message(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
"options": {
"tool_choice": "auto",
"tools": [web_search_tool],
},
}
if streaming:
response = await client.get_response(stream=True, **content).get_final_response()
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
web_search_tool_with_location = OpenAIChatClient.get_web_search_tool(
user_location={"country": "US", "city": "Seattle"},
)
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [web_search_tool_with_location],
},
}
if streaming:
response = await client.get_response(stream=True, **content).get_final_response()
else:
response = await client.get_response(**content)
assert response.text is not None
# Test that the client will use the web search tool with location
web_search_tool_with_location = OpenAIChatClient.get_web_search_tool(
user_location={"country": "US", "city": "Seattle"},
)
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [web_search_tool_with_location],
},
}
response = await client.get_response(stream=True, **content).get_final_response()
assert response.text is not None
@pytest.mark.skip(
@@ -3351,7 +3357,6 @@ async def test_integration_tool_rich_content_image() -> None:
assert "house" in response.text.lower(), f"Model did not describe the house image. Response: {response.text}"
@pytest.mark.timeout(300)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_openai_integration_tests_disabled
@@ -3363,14 +3368,11 @@ async def test_integration_agent_replays_local_tool_history_without_stale_fc_id(
async def search_hotels(city: Annotated[str, "The city to search for hotels in"]) -> str:
return f"The only hotel option in {city} is {hotel_code}."
client = OpenAIChatClient()
# override with model that does not do reasoning by default
client = OpenAIChatClient(model="gpt-5.4")
client.function_invocation_configuration["max_iterations"] = 2
agent = Agent(
client=client,
tools=[search_hotels],
default_options={"store": False},
)
agent = Agent(client=client, tools=[search_hotels], default_options={"store": False})
session = agent.create_session()
first_response = await agent.run(
@@ -4,12 +4,16 @@ from __future__ import annotations
import json
import os
from functools import wraps
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from agent_framework import Agent, AgentResponse, ChatResponse, Content, Message, SupportsChatGetResponse, tool
from azure.identity.aio import AzureCliCredential, get_bearer_token_provider
from agent_framework.exceptions import SettingNotFoundError
from azure.core.credentials_async import AsyncTokenCredential
from azure.identity.aio import AzureCliCredential
from openai import AsyncAzureOpenAI
from pydantic import BaseModel
from pytest import param
@@ -20,11 +24,40 @@ pytestmark = pytest.mark.azure
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.openai.azure.com")
or os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == "",
or (
os.getenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME", "") == ""
and os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == ""
),
reason="No real Azure OpenAI endpoint or responses deployment provided; skipping integration tests.",
)
def _with_azure_openai_debug() -> Any:
def decorator(func: Any) -> Any:
@wraps(func)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
try:
return await func(*args, **kwargs)
except Exception as exc:
model = os.getenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME") or os.getenv(
"AZURE_OPENAI_DEPLOYMENT_NAME", "<unset>"
)
api_version = os.getenv("AZURE_OPENAI_API_VERSION") or "preview"
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "<unset>")
debug_message = f"Azure OpenAI debug: endpoint={endpoint}, model={model}, api_version={api_version}"
if hasattr(exc, "add_note"):
exc.add_note(debug_message)
elif exc.args:
exc.args = (f"{exc.args[0]}\n{debug_message}", *exc.args[1:])
else:
exc.args = (debug_message,)
raise
return wrapper
return decorator
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
@@ -32,18 +65,6 @@ class OutputStruct(BaseModel):
weather: str | None = None
def _create_azure_openai_chat_client(
*,
api_key: Any = None,
) -> OpenAIChatClient:
return OpenAIChatClient(
model=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
api_key=api_key or os.environ["AZURE_OPENAI_API_KEY"],
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
)
async def create_vector_store(client: OpenAIChatClient) -> tuple[str, Content]:
"""Create a vector store with sample documents for testing."""
file = await client.client.files.create(
@@ -79,30 +100,117 @@ async def get_weather(location: str) -> str:
def test_init_with_azure_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_openai_chat_client()
client = OpenAIChatClient(credential=AzureCliCredential())
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(client, SupportsChatGetResponse)
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
assert client.api_version == azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"]
assert client.azure_endpoint.startswith(azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"])
def test_init_auto_detects_azure_env(azure_openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIChatClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
def test_init_uses_default_azure_api_version(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_openai_chat_client()
def test_openai_api_key_wins_over_azure_env(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatClient()
assert client.model == "gpt-5"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_api_version_alone_does_not_override_openai_api_key(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatClient(api_version="2024-10-21")
assert client.model == "gpt-5"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_explicit_credential_wins_over_openai_api_key(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatClient(credential=lambda: "token")
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
def test_init_falls_back_to_generic_azure_deployment_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME", raising=False)
client = OpenAIChatClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.api_version == "preview"
assert isinstance(client.client, AsyncAzureOpenAI)
def test_init_does_not_fall_back_to_openai_responses_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.setenv("OPENAI_RESPONSES_MODEL", "test_responses_model")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIChatClient()
def test_init_does_not_fall_back_to_openai_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("OPENAI_RESPONSES_MODEL", raising=False)
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIChatClient()
def test_init_with_credential_wraps_async_token_credential(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
class TestAsyncTokenCredential(AsyncTokenCredential):
async def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
credential = TestAsyncTokenCredential()
token_provider = MagicMock()
with patch("azure.identity.aio.get_bearer_token_provider", return_value=token_provider) as mock_provider:
client = OpenAIChatClient(credential=credential)
assert isinstance(client.client, AsyncAzureOpenAI)
mock_provider.assert_called_once_with(credential, "https://cognitiveservices.azure.com/.default")
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
def test_init_uses_default_azure_api_version(azure_openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIChatClient(credential=AzureCliCredential())
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert client.api_version is not None
def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
@@ -123,8 +231,6 @@ def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
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"),
@@ -136,7 +242,6 @@ def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_
param("tool_choice", "none", True, id="tool_choice_none"),
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"),
param("tools", [get_weather], True, id="tools_function"),
@@ -174,15 +279,14 @@ def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_
),
],
)
@_with_azure_openai_debug()
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
@@ -233,64 +337,34 @@ async def test_integration_options(
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_web_search() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
for streaming in [False, True]:
content = {
"messages": [
Message(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
"options": {
"tool_choice": "auto",
"tools": [OpenAIChatClient.get_web_search_tool()],
},
"stream": streaming,
}
if streaming:
response = await client.get_response(**content).get_final_response()
else:
response = await client.get_response(**content)
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [OpenAIChatClient.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})],
},
"stream": streaming,
}
if streaming:
response = await client.get_response(**content).get_final_response()
else:
response = await client.get_response(**content)
assert response.text is not None
response = await client.get_response(
messages=[
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
options={
"tools": [OpenAIChatClient.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})],
},
stream=True,
).get_final_response()
assert isinstance(response, ChatResponse)
assert response.text is not None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_file_search() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
file_id, vector_store = await create_vector_store(client)
try:
response = await client.get_response(
@@ -310,11 +384,10 @@ async def test_integration_client_file_search() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_file_search_streaming() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
file_id, vector_store = await create_vector_store(client)
try:
response_stream = client.get_response(
@@ -336,11 +409,10 @@ async def test_integration_client_file_search_streaming() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_hosted_mcp_tool() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
response = await client.get_response(
messages=[Message(role="user", text="How to create an Azure storage account using az cli?")],
options={
@@ -361,11 +433,10 @@ async def test_integration_client_agent_hosted_mcp_tool() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_hosted_code_interpreter_tool() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
response = await client.get_response(
messages=[Message(role="user", text="Calculate the sum of numbers from 1 to 10 using Python code.")],
@@ -381,14 +452,13 @@ async def test_integration_client_agent_hosted_code_interpreter_tool() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_existing_session() -> None:
async with AzureCliCredential() as credential:
preserved_session = None
async with Agent(
client=_create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
session = first_agent.create_session()
@@ -403,9 +473,7 @@ async def test_integration_client_agent_existing_session() -> None:
if preserved_session:
async with Agent(
client=_create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
second_response = await second_agent.run("What is my hobby?", session=preserved_session)
@@ -418,6 +486,7 @@ async def test_integration_client_agent_existing_session() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_tool_rich_content_image() -> None:
image_path = Path(__file__).parent.parent / "assets" / "sample_image.jpg"
image_bytes = image_path.read_bytes()
@@ -428,9 +497,7 @@ async def test_azure_openai_chat_client_tool_rich_content_image() -> None:
return Content.from_data(data=image_bytes, media_type="image/jpeg")
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
@@ -13,7 +13,7 @@ from agent_framework import (
SupportsChatGetResponse,
tool,
)
from agent_framework.exceptions import ChatClientException
from agent_framework.exceptions import ChatClientException, SettingNotFoundError
from openai import BadRequestError
from openai.types.chat.chat_completion import ChatCompletion, Choice
from openai.types.chat.chat_completion_message import ChatCompletionMessage
@@ -37,6 +37,14 @@ def test_init(openai_unit_test_env: dict[str, str]) -> None:
assert isinstance(open_ai_chat_completion, SupportsChatGetResponse)
def test_init_prefers_openai_chat_model(monkeypatch, openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_CHAT_MODEL", "test_chat_model_id")
open_ai_chat_completion = OpenAIChatCompletionClient()
assert open_ai_chat_completion.model == "test_chat_model_id"
def test_init_validation_fail() -> None:
# Test successful initialization
with pytest.raises(ValueError):
@@ -93,7 +101,7 @@ def test_init_base_url_from_settings_env() -> None:
@pytest.mark.parametrize("exclude_list", [["OPENAI_MODEL"]], indirect=True)
def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ValueError):
with pytest.raises(SettingNotFoundError):
OpenAIChatCompletionClient()
@@ -101,7 +109,7 @@ def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
def test_init_with_empty_api_key(openai_unit_test_env: dict[str, str]) -> None:
model_id = "test_model_id"
with pytest.raises(ValueError):
with pytest.raises(SettingNotFoundError):
OpenAIChatCompletionClient(
model=model_id,
)
@@ -1480,71 +1488,61 @@ async def test_integration_options(
# Need at least 2 iterations for tool_choice tests: one to get function call, one to get final response
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
# Prepare test message
# Prepare test message
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(Message(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [Message(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]
# Test streaming mode
response = await client.get_response(
messages=messages,
stream=True,
options=options,
).get_final_response()
assert response is not None
assert isinstance(response, ChatResponse)
assert response.messages is not None
if not option_name.startswith("tool_choice") and (
(isinstance(option_value, str) and option_value != "required")
or (isinstance(option_value, dict) and option_value.get("mode") != "required")
):
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"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
# 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"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(Message(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [Message(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_stream = client.get_response(
messages=messages,
stream=True,
options=options,
)
response = await response_stream.get_final_response()
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.messages is not None
if not option_name.startswith("tool_choice") and (
(isinstance(option_value, str) and option_value != "required")
or (isinstance(option_value, dict) and option_value.get("mode") != "required")
):
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()
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
@@ -3,7 +3,9 @@
from __future__ import annotations
import os
from collections.abc import Awaitable, Callable
from functools import wraps
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from agent_framework import (
@@ -16,7 +18,9 @@ from agent_framework import (
SupportsChatGetResponse,
tool,
)
from azure.identity.aio import AzureCliCredential, get_bearer_token_provider
from agent_framework.exceptions import SettingNotFoundError
from azure.core.credentials_async import AsyncTokenCredential
from azure.identity.aio import AzureCliCredential
from openai import AsyncAzureOpenAI
from agent_framework_openai import OpenAIChatCompletionClient
@@ -25,21 +29,37 @@ pytestmark = pytest.mark.azure
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.openai.azure.com")
or os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == "",
or (
os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", "") == "" and os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == ""
),
reason="No real Azure OpenAI endpoint or chat deployment provided; skipping integration tests.",
)
def _create_azure_chat_completion_client(
*,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
) -> OpenAIChatCompletionClient:
return OpenAIChatCompletionClient(
model=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
api_key=api_key or os.environ["AZURE_OPENAI_API_KEY"],
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
)
def _with_azure_openai_debug() -> Any:
def decorator(func: Any) -> Any:
@wraps(func)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
try:
return await func(*args, **kwargs)
except Exception as exc:
model = os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME") or os.getenv(
"AZURE_OPENAI_DEPLOYMENT_NAME", "<unset>"
)
api_version = os.getenv("AZURE_OPENAI_API_VERSION", "<unset>")
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "<unset>")
debug_message = f"Azure OpenAI debug: endpoint={endpoint}, model={model}, api_version={api_version}"
if hasattr(exc, "add_note"):
exc.add_note(debug_message)
elif exc.args:
exc.args = (f"{exc.args[0]}\n{debug_message}", *exc.args[1:])
else:
exc.args = (debug_message,)
raise
return wrapper
return decorator
@tool(approval_mode="never_require")
@@ -60,9 +80,9 @@ async def get_weather(location: str) -> str:
def test_init_with_azure_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_chat_completion_client()
client = OpenAIChatCompletionClient(azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"))
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(client, SupportsChatGetResponse)
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
@@ -73,18 +93,86 @@ def test_init_with_azure_endpoint(azure_openai_unit_test_env: dict[str, str]) ->
def test_init_auto_detects_azure_env(azure_openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIChatCompletionClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
def test_init_uses_default_azure_api_version(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_VERSION", "preview")
client = _create_azure_chat_completion_client()
def test_openai_api_key_wins_over_azure_env(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatCompletionClient()
assert client.model == "gpt-5"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_explicit_credential_wins_over_openai_api_key(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatCompletionClient(credential=lambda: "token")
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
def test_init_falls_back_to_generic_azure_deployment_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", raising=False)
client = OpenAIChatCompletionClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.api_version == "2024-10-21"
assert isinstance(client.client, AsyncAzureOpenAI)
def test_init_does_not_fall_back_to_openai_chat_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.setenv("OPENAI_CHAT_MODEL", "test_chat_model")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIChatCompletionClient()
def test_init_does_not_fall_back_to_openai_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("OPENAI_CHAT_MODEL", raising=False)
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIChatCompletionClient()
def test_init_with_credential_wraps_async_token_credential(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False)
class TestAsyncTokenCredential(AsyncTokenCredential):
async def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
credential = TestAsyncTokenCredential()
token_provider = MagicMock()
with patch("azure.identity.aio.get_bearer_token_provider", return_value=token_provider) as mock_provider:
client = OpenAIChatCompletionClient(credential=credential)
assert isinstance(client.client, AsyncAzureOpenAI)
mock_provider.assert_called_once_with(credential, "https://cognitiveservices.azure.com/.default")
def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
@@ -102,11 +190,10 @@ def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_response() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatCompletionClient(credential=credential)
assert isinstance(client, SupportsChatGetResponse)
messages = [
@@ -134,11 +221,10 @@ async def test_azure_openai_chat_completion_client_response() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_response_tools() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatCompletionClient(credential=credential)
response = await client.get_response(
messages=[Message(role="user", text="who are Emily and David?")],
@@ -153,11 +239,10 @@ async def test_azure_openai_chat_completion_client_response_tools() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_streaming() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatCompletionClient(credential=credential)
response = client.get_response(
messages=[
@@ -190,11 +275,10 @@ async def test_azure_openai_chat_completion_client_streaming() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_streaming_tools() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatCompletionClient(credential=credential)
response = client.get_response(
messages=[Message(role="user", text="who are Emily and David?")],
@@ -215,13 +299,12 @@ async def test_azure_openai_chat_completion_client_streaming_tools() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_agent_basic_run() -> None:
async with (
AzureCliCredential() as credential,
Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
) as agent,
):
response = await agent.run("Please respond with exactly: 'This is a response test.'")
@@ -234,20 +317,14 @@ async def test_azure_openai_chat_completion_client_agent_basic_run() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_agent_basic_run_streaming() -> None:
async with (
AzureCliCredential() as credential,
Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
) as agent,
Agent(client=OpenAIChatCompletionClient(credential=credential)) as agent,
):
full_text = ""
async for chunk in agent.run(
"Please respond with exactly: 'This is a streaming response test.'",
stream=True,
):
async for chunk in agent.run("Please respond with exactly: 'This is a streaming response test.'", stream=True):
assert isinstance(chunk, AgentResponseUpdate)
if chunk.text:
full_text += chunk.text
@@ -258,13 +335,12 @@ async def test_azure_openai_chat_completion_client_agent_basic_run_streaming() -
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_agent_session_persistence() -> None:
async with (
AzureCliCredential() as credential,
Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as agent,
):
@@ -281,14 +357,13 @@ async def test_azure_openai_chat_completion_client_agent_session_persistence() -
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_agent_existing_session() -> None:
async with AzureCliCredential() as credential:
preserved_session = None
async with Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
session = first_agent.create_session()
@@ -299,9 +374,7 @@ async def test_azure_openai_chat_completion_client_agent_existing_session() -> N
if preserved_session:
async with Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
second_response = await second_agent.run("What is my name?", session=preserved_session)
@@ -314,13 +387,12 @@ async def test_azure_openai_chat_completion_client_agent_existing_session() -> N
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_chat_completion_client_agent_level_tool_persistence() -> None:
async with (
AzureCliCredential() as credential,
Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather],
) as agent,
@@ -6,6 +6,7 @@ import os
from unittest.mock import AsyncMock, MagicMock
import pytest
from agent_framework.exceptions import SettingNotFoundError
from openai.types import CreateEmbeddingResponse
from openai.types import Embedding as OpenAIEmbedding
from openai.types.create_embedding_response import Usage
@@ -32,13 +33,6 @@ def _make_openai_response(
)
@pytest.fixture
def openai_unit_test_env(monkeypatch: pytest.MonkeyPatch) -> None:
"""Set up environment variables for OpenAI embedding client."""
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
# --- OpenAI unit tests ---
@@ -50,24 +44,39 @@ def test_openai_construction_with_explicit_params() -> None:
assert client.model == "text-embedding-3-small"
def test_openai_construction_from_env(openai_unit_test_env: None) -> None:
def test_openai_construction_from_env(openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIEmbeddingClient()
assert client.model == openai_unit_test_env["OPENAI_EMBEDDING_MODEL"]
def test_with_callable_api_key() -> None:
"""Test OpenAIEmbeddingClient initialization with callable API key."""
async def get_api_key() -> str:
return "test-api-key-123"
client = OpenAIEmbeddingClient(model="text-embedding-3-small", api_key=get_api_key)
assert client.model == "text-embedding-3-small"
assert client.client is not None
def test_openai_construction_missing_api_key_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
with pytest.raises(ValueError, match="API key is required"):
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
def test_openai_construction_without_openai_or_azure_config_raises_clear_error(
openai_unit_test_env: dict[str, str],
) -> None:
with pytest.raises(SettingNotFoundError):
OpenAIEmbeddingClient(model="text-embedding-3-small")
def test_openai_construction_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("OPENAI_EMBEDDING_MODEL", raising=False)
with pytest.raises(ValueError, match="embedding model is required"):
OpenAIEmbeddingClient(api_key="test-key")
@pytest.mark.parametrize("exclude_list", [["OPENAI_EMBEDDING_MODEL"]], indirect=True)
def test_openai_construction_falls_back_to_openai_model(openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIEmbeddingClient()
assert client.model == openai_unit_test_env["OPENAI_MODEL"]
async def test_openai_get_embeddings(openai_unit_test_env: None) -> None:
async def test_openai_get_embeddings(openai_unit_test_env: dict[str, str]) -> None:
mock_response = _make_openai_response(
embeddings=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]],
)
@@ -85,7 +94,7 @@ async def test_openai_get_embeddings(openai_unit_test_env: None) -> None:
assert result[0].dimensions == 3
async def test_openai_get_embeddings_usage(openai_unit_test_env: None) -> None:
async def test_openai_get_embeddings_usage(openai_unit_test_env: dict[str, str]) -> None:
mock_response = _make_openai_response(
embeddings=[[0.1]],
prompt_tokens=10,
@@ -103,7 +112,7 @@ async def test_openai_get_embeddings_usage(openai_unit_test_env: None) -> None:
assert result.usage["total_token_count"] == 10
async def test_openai_options_passthrough_dimensions(openai_unit_test_env: None) -> None:
async def test_openai_options_passthrough_dimensions(openai_unit_test_env: dict[str, str]) -> None:
mock_response = _make_openai_response(embeddings=[[0.1]])
client = OpenAIEmbeddingClient()
client.client = MagicMock()
@@ -118,7 +127,7 @@ async def test_openai_options_passthrough_dimensions(openai_unit_test_env: None)
assert result.options is options
async def test_openai_options_passthrough_encoding_format(openai_unit_test_env: None) -> None:
async def test_openai_options_passthrough_encoding_format(openai_unit_test_env: dict[str, str]) -> None:
mock_response = _make_openai_response(embeddings=[[0.1]])
client = OpenAIEmbeddingClient()
client.client = MagicMock()
@@ -132,7 +141,7 @@ async def test_openai_options_passthrough_encoding_format(openai_unit_test_env:
assert call_kwargs["encoding_format"] == "base64"
async def test_openai_base64_decoding(openai_unit_test_env: None) -> None:
async def test_openai_base64_decoding(openai_unit_test_env: dict[str, str]) -> None:
import base64
import struct
@@ -176,7 +185,7 @@ async def test_openai_error_when_no_model_id() -> None:
await client.get_embeddings(["test"])
async def test_openai_empty_values_returns_empty(openai_unit_test_env: None) -> None:
async def test_openai_empty_values_returns_empty(openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIEmbeddingClient()
client.client = MagicMock()
client.client.embeddings = MagicMock()
@@ -0,0 +1,250 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import os
from functools import wraps
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from agent_framework.exceptions import SettingNotFoundError
from azure.core.credentials_async import AsyncTokenCredential
from azure.identity.aio import AzureCliCredential
from openai import AsyncAzureOpenAI
from agent_framework_openai import OpenAIEmbeddingClient, OpenAIEmbeddingOptions
pytestmark = pytest.mark.azure
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.openai.azure.com")
or (
os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", "") == ""
and os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == ""
),
reason="No real Azure OpenAI endpoint or embedding deployment provided; skipping integration tests.",
)
def _with_azure_openai_debug() -> Any:
def decorator(func: Any) -> Any:
@wraps(func)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
try:
return await func(*args, **kwargs)
except Exception as exc:
model = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME") or os.getenv(
"AZURE_OPENAI_DEPLOYMENT_NAME", "<unset>"
)
api_version = os.getenv("AZURE_OPENAI_API_VERSION", "<unset>")
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "<unset>")
debug_message = f"Azure OpenAI debug: endpoint={endpoint}, model={model}, api_version={api_version}"
if hasattr(exc, "add_note"):
exc.add_note(debug_message)
elif exc.args:
exc.args = (f"{exc.args[0]}\n{debug_message}", *exc.args[1:])
else:
exc.args = (debug_message,)
raise
return wrapper
return decorator
def _get_azure_embedding_deployment_name() -> str:
return os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME") or os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"]
def _create_azure_embedding_client(
*,
api_key: str | None = None,
credential: AsyncTokenCredential | None = None,
) -> OpenAIEmbeddingClient:
resolved_api_key = (
api_key if api_key is not None else None if credential is not None else os.environ["AZURE_OPENAI_API_KEY"]
)
return OpenAIEmbeddingClient(
model=_get_azure_embedding_deployment_name(),
api_key=resolved_api_key,
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
credential=credential,
)
def test_init_with_azure_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_embedding_client()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
assert client.api_version == azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"]
def test_init_auto_detects_azure_embedding_env(azure_openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIEmbeddingClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
def test_init_falls_back_to_generic_azure_deployment_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", raising=False)
client = OpenAIEmbeddingClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
def test_init_does_not_fall_back_to_openai_embedding_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIEmbeddingClient()
def test_init_does_not_fall_back_to_openai_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("OPENAI_EMBEDDING_MODEL", raising=False)
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIEmbeddingClient()
def test_openai_api_key_wins_over_azure_env(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
client = OpenAIEmbeddingClient()
assert client.model == "text-embedding-3-small"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_api_version_alone_does_not_override_openai_api_key(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
client = OpenAIEmbeddingClient(api_version="2024-10-21")
assert client.model == "text-embedding-3-small"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_explicit_credential_wins_over_openai_api_key(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
client = OpenAIEmbeddingClient(credential=lambda: "token")
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
def test_init_with_credential_wraps_async_token_credential(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
class TestAsyncTokenCredential(AsyncTokenCredential):
async def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
credential = TestAsyncTokenCredential()
token_provider = MagicMock()
with patch("azure.identity.aio.get_bearer_token_provider", return_value=token_provider) as mock_provider:
client = OpenAIEmbeddingClient(credential=credential)
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
mock_provider.assert_called_once_with(credential, "https://cognitiveservices.azure.com/.default")
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
def test_init_uses_default_azure_api_version(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_embedding_client()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
assert client.api_version == "2024-10-21"
def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
monkeypatch.setenv("OPENAI_BASE_URL", "https://custom-openai-endpoint.com/v1")
client = OpenAIEmbeddingClient()
assert client.model == "text-embedding-3-small"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_get_embeddings() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_embedding_client(credential=credential)
result = await client.get_embeddings(["hello world"])
assert len(result) == 1
assert isinstance(result[0].vector, list)
assert len(result[0].vector) > 0
assert all(isinstance(v, float) for v in result[0].vector)
assert result[0].model is not None
assert result.usage is not None
assert result.usage["input_token_count"] > 0
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_get_embeddings_multiple() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_embedding_client(credential=credential)
result = await client.get_embeddings(["hello", "world", "test"])
assert len(result) == 3
dims = [len(embedding.vector) for embedding in result]
assert all(dimension == dims[0] for dimension in dims)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_get_embeddings_with_dimensions() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_embedding_client(credential=credential)
options: OpenAIEmbeddingOptions = {"dimensions": 256}
result = await client.get_embeddings(["hello world"], options=options)
assert len(result) == 1
assert len(result[0].vector) == 256
@@ -0,0 +1,54 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
from unittest.mock import MagicMock, patch
import pytest
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from agent_framework_openai._shared import AZURE_OPENAI_TOKEN_SCOPE, _resolve_azure_credential_to_token_provider
class _AsyncTokenCredentialStub(AsyncTokenCredential):
async def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
class _TokenCredentialStub(TokenCredential):
def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
def test_resolve_azure_async_credential_wraps_provider() -> None:
credential = _AsyncTokenCredentialStub()
token_provider = MagicMock()
with patch("azure.identity.aio.get_bearer_token_provider", return_value=token_provider) as mock_provider:
resolved = _resolve_azure_credential_to_token_provider(credential)
assert resolved is token_provider
mock_provider.assert_called_once_with(credential, AZURE_OPENAI_TOKEN_SCOPE)
def test_resolve_azure_sync_credential_wraps_provider() -> None:
credential = _TokenCredentialStub()
token_provider = MagicMock()
with patch("azure.identity.get_bearer_token_provider", return_value=token_provider) as mock_provider:
resolved = _resolve_azure_credential_to_token_provider(credential)
assert resolved is token_provider
mock_provider.assert_called_once_with(credential, AZURE_OPENAI_TOKEN_SCOPE)
def test_resolve_azure_callable_token_provider_passthrough() -> None:
token_provider = MagicMock()
assert _resolve_azure_credential_to_token_provider(token_provider) is token_provider
def test_resolve_azure_invalid_credential_raises() -> None:
with pytest.raises(ValueError, match="credential"):
_resolve_azure_credential_to_token_provider(object()) # type: ignore[arg-type]