scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166
litellm/tests/unit/llms/azure/test_azure.py
tessellark c9a5777739
fix(azure): share one Entra token request per worker and audit #47348 (#48769) 
* test(azure): audit #47348

Add proxy integration tests for #47348's cached Entra token providers and
off-loop token refresh, and fix the token herd the audit found: once the
fetch moved to a worker thread, concurrent cold requests each fetched their
own token. Every cached bearer token provider now serializes its calls so
the first fetch fills the cache and the rest reuse it

* fix(azure): build each cached Entra token provider once under concurrent first calls

lru_cache does not hold back concurrent misses, so worker threads that miss together each built their own credential and token provider, and each fetched its own token. A lock around each provider cache makes the first caller build it and the rest reuse it

* fix(azure): share Entra token fetches only while no valid token is held

Concurrent callers now share one result (or one error) from the cold
fetch, while callers inside the refresh window keep the still-valid
token. The oidc/azure provider build and the client-assertion exchange
are single-flighted per audience, and the shared helpers no longer call
subscripted generic aliases, which crashed on Python 3.10.

* fix(azure): refresh Entra tokens once while other requests keep the valid token

MSAL treats a token with less than five minutes left as expired, so every
request in that window fetched a new one. Now one request refreshes while
the others keep using the still-valid token.

Requests that waited on a fetch that failed retry once through a fresh
fetch instead of all failing with the first error.

Credentials that only implement get_token keep working, and OIDC token
exchanges are cached per scope.

* fix(azure): bound Entra token waits during an outage

Requests waiting on a failed token fetch now retry only when the fetch
before it succeeded, so during an Entra outage later bursts fail after one
fetch instead of two.

Requests during a refresh only get the held token while it has more than
60 seconds left, and a failed refresh falls back to that token instead of
failing its request.

* fix(azure): share one retry per failed Entra token fetch and raise on a failed refresh

Waiters of a failed fetch now retry the flight they joined, so a fast-failing Entra makes one retry instead of one per waiter. A failed refresh raises again instead of returning the held token, which Postgres token auth read as a fresh one. The DefaultAzureCredential fallback reuses one provider per scope.

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-10-28 12:24:20 -07:00

793 lines
27 KiB
Python
Raw Permalink Blame History

"""Tests for litellm/llms/azure/azure.py AzureChatCompletion handler behaviour."""
import asyncio
import json
import os
import time
from collections.abc import Mapping
from typing import Final
from unittest.mock import ANY, MagicMock, patch
import httpx
import pytest
import respx
from httpx import Client, Headers
from openai import AsyncAzureOpenAI, AzureOpenAI
import litellm
from litellm import completion
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.azure.azure import AzureChatCompletion
from litellm.llms.azure.common_utils import process_azure_headers
class _FakeRawResponse:
headers: Final = {"x-ms-is-spilled-over": "true"}
def parse(self):
return iter(())
class _FakeRawCompletions:
def create(self, **kwargs):
return _FakeRawResponse()
def test_sync_streaming_stamps_response_headers_on_the_logging_obj() -> None:
"""Sync streaming must mirror async_streaming and record the provider response
headers on model_call_details, or downstream consumers (spillover-aware cost
calculation) cannot see them."""
client = AzureOpenAI(api_key="fake", api_version="2024-02-01", azure_endpoint="https://fake.openai.azure.com")
client.chat.completions.with_raw_response = _FakeRawCompletions()
logging_obj = LiteLLMLoggingObj(
model="azure/gpt-4o-spill-test",
messages=[{"role": "user", "content": "Hi"}],
stream=True,
call_type="completion",
start_time=time.time(),
litellm_call_id="spill-sync-1",
function_id="f",
)
AzureChatCompletion().streaming(
logging_obj=logging_obj,
api_base="https://fake.openai.azure.com",
api_key="fake",
api_version="2024-02-01",
dynamic_params=False,
data={"messages": [{"role": "user", "content": "Hi"}], "stream": True},
model="gpt-4o-spill-test",
timeout=30.0,
max_retries=0,
client=client,
)
assert logging_obj.model_call_details["response_headers"] == {"x-ms-is-spilled-over": "true"}
class _CancelledRawCompletions:
async def create(self, **kwargs):
raise asyncio.CancelledError()
@pytest.mark.asyncio
async def test_acompletion_propagates_cancelled_error() -> None:
client = AsyncAzureOpenAI(
api_key="fake-key",
api_version="2024-02-01",
azure_endpoint="https://fake-resource.openai.azure.com",
)
client.chat.completions.with_raw_response = _CancelledRawCompletions()
with pytest.raises(asyncio.CancelledError):
await litellm.acompletion(
model="azure/fake-deployment",
messages=[{"role": "user", "content": "hi"}],
client=client,
)
def test_process_azure_headers_empty():
result = process_azure_headers({})
assert result == {}, "Expected empty dictionary for no input"
def test_process_azure_headers_with_all_headers():
input_headers = Headers(
{
"x-ratelimit-limit-requests": "100",
"x-ratelimit-remaining-requests": "90",
"x-ratelimit-limit-tokens": "10000",
"x-ratelimit-remaining-tokens": "9000",
"other-header": "value",
}
)
expected_output = {
"x-ratelimit-limit-requests": "100",
"x-ratelimit-remaining-requests": "90",
"x-ratelimit-limit-tokens": "10000",
"x-ratelimit-remaining-tokens": "9000",
"llm_provider-x-ratelimit-limit-requests": "100",
"llm_provider-x-ratelimit-remaining-requests": "90",
"llm_provider-x-ratelimit-limit-tokens": "10000",
"llm_provider-x-ratelimit-remaining-tokens": "9000",
"llm_provider-other-header": "value",
}
result = process_azure_headers(input_headers)
assert result == expected_output, "Unexpected output for all Azure headers"
def test_process_azure_headers_with_partial_headers():
input_headers = Headers(
{
"x-ratelimit-limit-requests": "100",
"x-ratelimit-remaining-tokens": "9000",
"other-header": "value",
}
)
expected_output = {
"x-ratelimit-limit-requests": "100",
"x-ratelimit-remaining-tokens": "9000",
"llm_provider-x-ratelimit-limit-requests": "100",
"llm_provider-x-ratelimit-remaining-tokens": "9000",
"llm_provider-other-header": "value",
}
result = process_azure_headers(input_headers)
assert result == expected_output, "Unexpected output for partial Azure headers"
def test_process_azure_headers_with_no_matching_headers():
input_headers = Headers({"unrelated-header-1": "value1", "unrelated-header-2": "value2"})
expected_output = {
"llm_provider-unrelated-header-1": "value1",
"llm_provider-unrelated-header-2": "value2",
}
result = process_azure_headers(input_headers)
assert result == expected_output, "Unexpected output for non-matching headers"
def test_process_azure_headers_with_dict_input():
input_headers = {
"x-ratelimit-limit-requests": "100",
"x-ratelimit-remaining-requests": "90",
"other-header": "value",
}
expected_output = {
"x-ratelimit-limit-requests": "100",
"x-ratelimit-remaining-requests": "90",
"llm_provider-x-ratelimit-limit-requests": "100",
"llm_provider-x-ratelimit-remaining-requests": "90",
"llm_provider-other-header": "value",
}
result = process_azure_headers(input_headers)
assert result == expected_output, "Unexpected output for dict input"
@pytest.mark.parametrize(
"input, call_type",
[
({"messages": [{"role": "user", "content": "Hello world"}]}, "completion"),
({"input": "Hello world"}, "embedding"),
({"prompt": "Hello world"}, "image_generation"),
],
)
@pytest.mark.parametrize(
"header_value",
[
"headers",
"extra_headers",
],
)
def test_azure_extra_headers(input, call_type, header_value):
from litellm import embedding, image_generation
# Clear the LLM clients cache to ensure the new http_client is used
litellm.in_memory_llm_clients_cache.flush_cache()
http_client = Client()
messages = [{"role": "user", "content": "Hello world"}]
with patch.object(http_client, "send", new=MagicMock()) as mock_client:
litellm.client_session = http_client
try:
if call_type == "completion":
func = completion
elif call_type == "embedding":
func = embedding
elif call_type == "image_generation":
func = image_generation
data = {
"model": "azure/gpt-4.1-mini",
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com",
"api_version": "2023-07-01-preview",
"api_key": "my-azure-api-key",
header_value: {
"Authorization": "my-bad-key",
"Ocp-Apim-Subscription-Key": "hello-world-testing",
},
**input,
}
response = func(**data)
print(response)
except Exception as e:
print(e)
mock_client.assert_called()
print(f"mock_client.call_args: {mock_client.call_args}")
request = mock_client.call_args[0][0]
print(request.method) # This will print 'POST'
print(request.url) # This will print the full URL
print(request.headers) # This will print the full URL
auth_header = request.headers.get("Authorization")
apim_key = request.headers.get("Ocp-Apim-Subscription-Key")
print(auth_header)
assert auth_header == "my-bad-key"
assert apim_key == "hello-world-testing"
@pytest.mark.parametrize(
"api_base, model, expected_endpoint",
[
(
"https://fake-azure-endpoint.invalid",
"dall-e-3-test",
"https://fake-azure-endpoint.invalid/openai/deployments/dall-e-3-test/images/generations?api-version=2023-12-01-preview",
),
(
"https://fake-azure-endpoint.invalid/openai/deployments/my-custom-deployment",
"dall-e-3",
"https://fake-azure-endpoint.invalid/openai/deployments/my-custom-deployment/images/generations?api-version=2023-12-01-preview",
),
],
)
def test_process_azure_endpoint_url(api_base, model, expected_endpoint):
from litellm.llms.azure.azure import AzureChatCompletion
azure_chat_completion = AzureChatCompletion()
input_args = {
"azure_client_params": {
"api_version": "2023-12-01-preview",
"azure_endpoint": api_base,
"azure_deployment": model,
"max_retries": 2,
"timeout": 600,
"api_key": "sk-test-mock-key-505",
},
"model": model,
}
result = azure_chat_completion.create_azure_base_url(**input_args)
assert result == expected_endpoint, "Unexpected endpoint"
@patch("azure.identity.UsernamePasswordCredential")
@patch("azure.identity.get_bearer_token_provider")
def test_get_azure_ad_token_from_username_password(mock_get_bearer_token_provider, mock_credential):
from litellm.llms.azure.common_utils import (
get_azure_ad_token_from_username_password,
)
client_id = "test-client-id"
username = "test-username"
password = "test-password"
mock_token_provider = lambda: "mock-token"
mock_get_bearer_token_provider.return_value = mock_token_provider
result = get_azure_ad_token_from_username_password(
client_id=client_id, azure_username=username, azure_password=password
)
mock_credential.assert_called_once_with(client_id=client_id, username=username, password=password)
mock_get_bearer_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default")
assert result() == "mock-token"
def test_azure_openai_gpt_4o_naming(monkeypatch):
from pydantic import BaseModel, Field
monkeypatch.setenv("AZURE_API_VERSION", "2024-10-21")
client = AzureOpenAI(
api_key="test-api-key",
base_url="https://fake-azure-endpoint.invalid",
api_version="2023-12-01-preview",
)
class ResponseFormat(BaseModel):
number: str = Field(description="total number of days in a week")
days: list[str] = Field(description="name of days in a week")
with patch.object(client, "post") as mock_post:
try:
completion(
model="azure/gpt4o",
messages=[{"role": "user", "content": "Hello world"}],
response_format=ResponseFormat,
client=client,
)
except Exception as e:
print(e)
mock_post.assert_called_once()
print(mock_post.call_args.kwargs["body"])
assert "tool_calls" not in mock_post.call_args.kwargs["body"]
@pytest.mark.parametrize(
"api_version",
[
"2024-10-21",
# "2024-02-15-preview",
],
)
def test_azure_gpt_4o_with_tool_call_and_response_format(api_version):
from litellm import completion
from typing import Optional
from pydantic import BaseModel
import litellm
client = AzureOpenAI(
api_key="fake-key",
base_url="https://fake-azure.openai.azure.com",
api_version=api_version,
)
class InvestigationOutput(BaseModel):
alert_explanation: Optional[str] = None
investigation: Optional[str] = None
conclusions_and_possible_root_causes: Optional[str] = None
next_steps: Optional[str] = None
related_logs: Optional[str] = None
app_or_infra: Optional[str] = None
external_links: Optional[str] = None
tools = [
{
"type": "function",
"function": {
"name": "get_current_time",
"description": "Returns the current date and time",
"strict": True,
"parameters": {
"properties": {
"timezone": {
"type": "string",
"description": "The timezone to get the current time for (e.g., 'UTC', 'America/New_York')",
}
},
"required": ["timezone"],
"type": "object",
"additionalProperties": False,
},
},
}
]
with patch.object(client, "post") as mock_post:
mock_post.return_value.headers = {}
mock_post.return_value.parse.return_value = litellm.ModelResponse(
choices=[{"message": {"role": "assistant", "content": InvestigationOutput().model_dump_json()}}]
)
response = litellm.completion(
model="azure/gpt-4.1-mini",
messages=[
{
"role": "system",
"content": "You are a tool-calling AI assist provided with common devops and IT tools that you can use to troubleshoot problems or answer questions.\nWhenever possible you MUST first use tools to investigate then answer the question.",
},
{
"role": "user",
"content": "What is the current date and time in NYC?",
},
],
drop_params=True,
temperature=0.00000001,
tools=tools,
tool_choice="auto",
response_format=InvestigationOutput, # commenting this line will cause the output to be correct
api_version=api_version,
client=client,
)
mock_post.assert_called_once()
if api_version == "2024-10-21":
assert "response_format" in mock_post.call_args.kwargs["body"]
else:
assert "response_format" not in mock_post.call_args.kwargs["body"]
assert response.choices[0].message.content == InvestigationOutput().model_dump_json()
def test_map_openai_params():
"""
Ensure response_format does not override tools
"""
from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig
azure_openai_config = AzureOpenAIConfig()
tools = [
{
"type": "function",
"function": {
"name": "get_current_time",
"description": "Returns the current date and time",
"strict": True,
"parameters": {
"properties": {
"timezone": {
"type": "string",
"description": "The timezone to get the current time for (e.g., 'UTC', 'America/New_York')",
}
},
"required": ["timezone"],
"type": "object",
"additionalProperties": False,
},
},
}
]
received_args = {
"non_default_params": {
"temperature": 1e-08,
"response_format": {
"type": "json_schema",
"json_schema": {
"schema": {
"properties": {
"alert_explanation": {
"anyOf": [{"type": "string"}, {"type": "null"}],
"title": "Alert Explanation",
},
"investigation": {
"anyOf": [{"type": "string"}, {"type": "null"}],
"title": "Investigation",
},
"conclusions_and_possible_root_causes": {
"anyOf": [{"type": "string"}, {"type": "null"}],
"title": "Conclusions And Possible Root Causes",
},
"next_steps": {
"anyOf": [{"type": "string"}, {"type": "null"}],
"title": "Next Steps",
},
"related_logs": {
"anyOf": [{"type": "string"}, {"type": "null"}],
"title": "Related Logs",
},
"app_or_infra": {
"anyOf": [{"type": "string"}, {"type": "null"}],
"title": "App Or Infra",
},
"external_links": {
"anyOf": [{"type": "string"}, {"type": "null"}],
"title": "External Links",
},
},
"title": "InvestigationOutput",
"type": "object",
"additionalProperties": False,
"required": [
"alert_explanation",
"investigation",
"conclusions_and_possible_root_causes",
"next_steps",
"related_logs",
"app_or_infra",
"external_links",
],
},
"name": "InvestigationOutput",
"strict": True,
},
},
"tools": tools,
"tool_choice": "auto",
},
"optional_params": {},
"model": "gpt-4o",
"drop_params": True,
"api_version": "2024-02-15-preview",
}
optional_params = azure_openai_config.map_openai_params(**received_args)
assert "tools" in optional_params
assert len(optional_params["tools"]) > 1
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.parametrize("max_retries", [0, 4])
@pytest.mark.parametrize("stream", [True, False])
@patch(
"litellm.main.azure_chat_completions.make_sync_azure_openai_chat_completion_request"
)
def test_azure_max_retries_0(
mock_make_sync_azure_openai_chat_completion_request, max_retries, stream
):
import litellm
from litellm import completion
# Clear the LLM clients cache to ensure max_retries is set correctly
litellm.in_memory_llm_clients_cache.flush_cache()
try:
completion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "Hello world"}],
max_retries=max_retries,
stream=stream,
)
except Exception as e:
print(e)
mock_make_sync_azure_openai_chat_completion_request.assert_called_once()
assert (
mock_make_sync_azure_openai_chat_completion_request.call_args.kwargs[
"azure_client"
].max_retries
== max_retries
)
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.parametrize("max_retries", [0, 4])
@pytest.mark.parametrize("stream", [True, False])
@patch("litellm.main.azure_chat_completions.make_azure_openai_chat_completion_request")
@pytest.mark.asyncio
async def test_async_azure_max_retries_0(
make_azure_openai_chat_completion_request, max_retries, stream
):
import litellm
from litellm import acompletion
# Clear the LLM clients cache to ensure max_retries is set correctly
litellm.in_memory_llm_clients_cache.flush_cache()
try:
await acompletion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "Hello world"}],
max_retries=max_retries,
stream=stream,
)
except Exception as e:
print(e)
make_azure_openai_chat_completion_request.assert_called_once()
assert (
make_azure_openai_chat_completion_request.call_args.kwargs[
"azure_client"
].max_retries
== max_retries
)
def test_azure_openai_responses_bridge():
from litellm import completion
import litellm
litellm.turn_on_debug()
with patch.object(litellm, "responses") as mock_responses:
try:
response = completion(
model="azure/responses/test-azure-computer-use-preview",
messages=[{"role": "user", "content": "Hello world"}],
api_base=os.getenv("AZURE_COMPUTER_USE_API_BASE"),
api_version="2025-04-01-preview",
api_key=os.getenv("AZURE_COMPUTER_USE_API_KEY"),
)
except Exception as e:
print(e)
mock_responses.assert_called_once()
assert (
mock_responses.call_args.kwargs["model"]
== "azure/test-azure-computer-use-preview"
)
assert mock_responses.call_args.kwargs["custom_llm_provider"] == "azure"
def test_azure_with_content_safety_error():
"""
Verify user can access innererror from the Azure OpenAI exception
"""
from litellm import completion
from litellm.exceptions import ContentPolicyViolationError
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
from unittest.mock import MagicMock
mock_exception = Exception(
"The response was filtered due to the prompt triggering Azure OpenAI's content management policy"
)
mock_exception.body = {
"innererror": {
"code": "ResponsibleAIPolicyViolation",
"content_filter_result": {
"hate": {"filtered": False, "severity": "safe"},
"jailbreak": {"filtered": False, "detected": False},
"self_harm": {"filtered": False, "severity": "safe"},
"sexual": {"filtered": False, "severity": "safe"},
"violence": {"filtered": True, "severity": "high"},
},
}
}
mock_response = MagicMock()
mock_response.status_code = 400
mock_exception.response = mock_response
with pytest.raises(ContentPolicyViolationError) as exc_info:
exception_type(
model="azure/gpt-4o-new-test",
original_exception=mock_exception,
custom_llm_provider="azure",
)
e = exc_info.value
print("got exception=", e)
assert e.provider_specific_fields is not None
print("got provider_specific_fields=", e.provider_specific_fields)
assert e.provider_specific_fields.get("innererror") is not None
assert (
e.provider_specific_fields["innererror"]["code"]
== "ResponsibleAIPolicyViolation"
)
assert (
e.provider_specific_fields["innererror"]["content_filter_result"]["violence"][
"filtered"
]
is True
)
assert (
e.provider_specific_fields["innererror"]["content_filter_result"]["violence"][
"severity"
]
== "high"
)
def _azure_sse_body(chunks: tuple[Mapping[str, object], ...]) -> str:
return "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
@respx.mock
def test_azure_chat_stream_preserves_text_and_special_characters():
response_chunks: Final = (
{
"id": "chatcmpl-azure",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [{"index": 0, "delta": {"content": "<xml>café ☃</xml>"}, "finish_reason": None}],
},
{
"id": "chatcmpl-azure",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
},
)
route: Final = respx.post(
url__regex=r"https://azure-stream\.openai\.azure\.com/.*/chat/completions.*"
).mock(
return_value=httpx.Response(
200,
text=_azure_sse_body(response_chunks),
headers={"content-type": "text/event-stream"},
)
)
stream: Final = litellm.completion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "Respond with the <xml> tag only"}],
api_base="https://azure-stream.openai.azure.com",
api_version="2024-02-15-preview",
api_key="test-azure-key",
max_retries=0,
stream=True,
)
chunks: Final = list(stream)
assert route.call_count == 1
assert [chunk._hidden_params["custom_llm_provider"] for chunk in chunks] == ["azure"] * len(chunks)
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "<xml>café ☃</xml>"
assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"]
@respx.mock
def test_azure_chat_stream_reassembles_tool_call_arguments():
response_chunks: Final = (
{
"id": "chatcmpl-azure-tools",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [
{
"index": 0,
"delta": {
"tool_calls": [
{
"index": 0,
"id": "call-weather",
"type": "function",
"function": {"name": "get_current_weather", "arguments": '{"city":'},
}
]
},
"finish_reason": None,
}
],
},
{
"id": "chatcmpl-azure-tools",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [
{
"index": 0,
"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '"Boston"}'}}]},
"finish_reason": "tool_calls",
}
],
},
)
route: Final = respx.post(
url__regex=r"https://azure-tool-stream\.openai\.azure\.com/.*/chat/completions.*"
).mock(
return_value=httpx.Response(
200,
text=_azure_sse_body(response_chunks),
headers={"content-type": "text/event-stream"},
)
)
messages: Final = [{"role": "user", "content": "What is the weather in Boston?"}]
stream: Final = litellm.completion(
model="azure/gpt-4.1-mini",
messages=messages,
tools=[
{
"type": "function",
"function": {
"name": "get_current_weather",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"],
},
},
}
],
tool_choice="auto",
api_base="https://azure-tool-stream.openai.azure.com",
api_version="2024-02-15-preview",
api_key="test-azure-key",
max_retries=0,
stream=True,
)
chunks: Final = list(stream)
response: Final = litellm.stream_chunk_builder(chunks, messages=messages)
tool_call: Final = response.choices[0].message.tool_calls[0]
assert route.call_count == 1
assert tool_call.function.name == "get_current_weather"
assert json.loads(tool_call.function.arguments) == {"city": "Boston"}
assert response.choices[0].finish_reason == "tool_calls"
Reference in a new issue View git blame Copy permalink