mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
test(harness): move live Azure chat tests to chat_live_azure suite
Per audit 5/8c: - test_azure_ai.py: 7 live tests (request_format, gpt5 reasoning, completion_azure, flexible api_base, model-router x3) moved with the AzureModelRouterStreamingCallback helper they alone used; the map_azure_model_group / image-url body / deepseek reasoning goldens stay - test_azure_openai.py: test_azure_safety_result, test_completion_azure_deployment_id, test_azure_openai_with_prompt_cache_key moved; header/url-builder/param-mapping goldens stay; TestAzureEmbedding stays (non-chat, Sameer) - test_azure_o_series.py: TestAzureOpenAIO3Mini moved; its keep-unit test_override_fake_stream stays behind as test_azure_o1_override_fake_stream (no network; router/model-info golden)
This commit is contained in:
mateo-berri
parent
2ab5a298b6
commit
b9a4d9bd41
6 changed files with 457 additions and 408 deletions
328
tests/harness_suites/chat_live_azure/test_azure_ai_chat_live.py
Normal file
328
tests/harness_suites/chat_live_azure/test_azure_ai_chat_live.py
Normal file
|
|
@ -0,0 +1,328 @@
|
|||
# What is this?
|
||||
## Unit tests for Azure AI integration
|
||||
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
import litellm.types
|
||||
import litellm.types.utils
|
||||
import json
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
# from base_rerank_unit_tests import BaseLLMRerankTest
|
||||
|
||||
load_dotenv()
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import completion
|
||||
|
||||
AZURE_AI_API_BASE = os.getenv("AZURE_AI_API_BASE")
|
||||
|
||||
class AzureModelRouterStreamingCallback(
|
||||
litellm.integrations.custom_logger.CustomLogger
|
||||
):
|
||||
"""
|
||||
Custom callback to capture streaming cost tracking for Azure Model Router.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.standard_logging_payload = None
|
||||
self.response_cost = None
|
||||
self.async_success_called = False
|
||||
self.complete_streaming_response = None
|
||||
super().__init__()
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
print(f"async_log_success_event called")
|
||||
self.async_success_called = True
|
||||
self.standard_logging_payload = kwargs.get("standard_logging_object")
|
||||
self.complete_streaming_response = kwargs.get("complete_streaming_response")
|
||||
|
||||
if self.standard_logging_payload:
|
||||
self.response_cost = self.standard_logging_payload.get("response_cost")
|
||||
print(
|
||||
f"standard_logging_payload model: {self.standard_logging_payload.get('model')}"
|
||||
)
|
||||
print(f"standard_logging_payload response_cost: {self.response_cost}")
|
||||
|
||||
if self.complete_streaming_response:
|
||||
print(
|
||||
f"complete_streaming_response model: {self.complete_streaming_response.model}"
|
||||
)
|
||||
print(
|
||||
f"complete_streaming_response usage: {self.complete_streaming_response.usage}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_ai_request_format():
|
||||
"""
|
||||
Test that Azure AI requests are formatted correctly with the proper endpoint and parameters
|
||||
for both synchronous and asynchronous calls
|
||||
"""
|
||||
|
||||
litellm._turn_on_debug()
|
||||
|
||||
# Set up the test parameters
|
||||
api_key = os.getenv("AZURE_AI_API_KEY")
|
||||
api_base = os.getenv("AZURE_AI_API_BASE")
|
||||
model = "azure_ai/gpt-4.1-mini"
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "Hello! How can I assist you today?"},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
|
||||
await litellm.acompletion(
|
||||
custom_llm_provider="azure_ai",
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", ["azure/gpt5_series/gpt-5-mini", "azure/gpt-5-mini"])
|
||||
async def test_azure_gpt5_reasoning(model):
|
||||
litellm._turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "What is the capital of France?"}],
|
||||
reasoning_effort="minimal",
|
||||
max_tokens=10,
|
||||
api_base=os.getenv("AZURE_AI_API_BASE"),
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
)
|
||||
print("response: ", response)
|
||||
assert response.choices[0].message.content is not None
|
||||
|
||||
|
||||
def test_completion_azure():
|
||||
try:
|
||||
from litellm import completion_cost
|
||||
|
||||
litellm.set_verbose = False
|
||||
## Test azure call
|
||||
response = completion(
|
||||
model="azure_ai/gpt-4.1-mini",
|
||||
api_base=os.getenv("AZURE_AI_API_BASE"),
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hello, how are you?",
|
||||
}
|
||||
],
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
)
|
||||
print(f"response: {response}")
|
||||
print(f"response hidden params: {response._hidden_params}")
|
||||
print(response)
|
||||
|
||||
cost = completion_cost(completion_response=response)
|
||||
assert cost > 0.0
|
||||
print("Cost for azure completion request", cost)
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[
|
||||
AZURE_AI_API_BASE,
|
||||
f"{AZURE_AI_API_BASE}/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2023-03-15-preview",
|
||||
],
|
||||
)
|
||||
def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base):
|
||||
try:
|
||||
litellm.set_verbose = True
|
||||
|
||||
response = completion(
|
||||
model="azure_ai/gpt-4.1-mini",
|
||||
api_base=api_base,
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
messages=[{"role": "user", "content": "What is the meaning of life?"}],
|
||||
)
|
||||
|
||||
print(response)
|
||||
except litellm.Timeout as e:
|
||||
pass
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_ai_model_router():
|
||||
"""
|
||||
Test Azure AI model router non-streaming response cost tracking.
|
||||
Verifies that the flat cost of $0.14 per M input tokens is applied.
|
||||
|
||||
Tests the pattern: azure_ai/model_router/<deployment-name>
|
||||
Where deployment-name is the Azure deployment (e.g., "azure-model-router").
|
||||
The model_router prefix is stripped before sending to Azure API.
|
||||
"""
|
||||
from litellm.llms.azure_ai.cost_calculator import (
|
||||
calculate_azure_model_router_flat_cost,
|
||||
)
|
||||
|
||||
litellm._turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model="azure_ai/model_router/azure-model-router",
|
||||
messages=[{"role": "user", "content": "hi who is this"}],
|
||||
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
|
||||
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
|
||||
)
|
||||
print("response: ", response)
|
||||
|
||||
# Check response cost
|
||||
tracked_cost = response._hidden_params["response_cost"]
|
||||
assert tracked_cost > 0
|
||||
print("Tracked cost: ", tracked_cost)
|
||||
|
||||
# Verify flat cost is included using the helper function
|
||||
usage = response.usage
|
||||
if usage and usage.prompt_tokens:
|
||||
expected_flat_cost = calculate_azure_model_router_flat_cost(
|
||||
model="model_router/azure-model-router", prompt_tokens=usage.prompt_tokens
|
||||
)
|
||||
print(f"Prompt tokens: {usage.prompt_tokens}")
|
||||
print(f"Expected flat cost: ${expected_flat_cost:.9f}")
|
||||
print(f"Total tracked cost: ${tracked_cost:.9f}")
|
||||
|
||||
# Total cost should be at least the flat cost
|
||||
assert (
|
||||
tracked_cost >= expected_flat_cost
|
||||
), f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}"
|
||||
|
||||
# Verify the flat cost is non-zero
|
||||
assert expected_flat_cost > 0, "Flat cost should be greater than 0"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_ai_model_router_streaming_model_in_chunk():
|
||||
"""
|
||||
Test that Azure AI model router streaming returns the actual model in each chunk.
|
||||
The response should contain the actual model used (e.g., gpt-4.1-nano) not the request model (azure-model-router).
|
||||
"""
|
||||
litellm._turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model="azure_ai/azure-model-router",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
|
||||
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Collect chunks and check model field
|
||||
chunks_with_model = []
|
||||
async for chunk in response:
|
||||
print(f"Chunk model: {chunk.model}")
|
||||
if chunk.model and chunk.model.strip():
|
||||
chunks_with_model.append(chunk.model)
|
||||
|
||||
print(f"All chunk models: {chunks_with_model}")
|
||||
|
||||
# At least some chunks should have a model
|
||||
assert len(chunks_with_model) > 0, "No chunks had a model field set"
|
||||
|
||||
# The model should NOT be azure-model-router (the request model)
|
||||
# It should be the actual model from the response (e.g., gpt-4.1-nano, gpt-5-nano, etc.)
|
||||
for model in chunks_with_model:
|
||||
assert (
|
||||
model != "azure-model-router"
|
||||
), f"Chunk model should be actual model, not request model. Got: {model}"
|
||||
# The actual model should be a real model name like gpt-4.1-nano, gpt-5-nano, etc.
|
||||
print(f"Verified chunk has actual model: {model}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_ai_model_router_streaming_cost_with_stream_options():
|
||||
"""
|
||||
Test Azure AI model router streaming cost tracking with stream_options include_usage=True.
|
||||
This tests the specific case where cost tracking fails with stream_options.
|
||||
"""
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
test_callback = AzureModelRouterStreamingCallback()
|
||||
litellm.callbacks = [test_callback]
|
||||
|
||||
try:
|
||||
litellm._turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model="azure_ai/azure-model-router",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
|
||||
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
|
||||
# Consume the stream and check chunks
|
||||
full_response = ""
|
||||
chunks_with_model = []
|
||||
async for chunk in response:
|
||||
print(
|
||||
f"Chunk: model={chunk.model}, choices={len(chunk.choices) if chunk.choices else 0}"
|
||||
)
|
||||
if chunk.model:
|
||||
chunks_with_model.append(chunk.model)
|
||||
if (
|
||||
chunk.choices
|
||||
and chunk.choices[0].delta
|
||||
and chunk.choices[0].delta.content
|
||||
):
|
||||
full_response += chunk.choices[0].delta.content
|
||||
|
||||
print(f"Full streamed response: {full_response}")
|
||||
print(f"Chunks with model: {chunks_with_model}")
|
||||
|
||||
# Give async logging time to complete
|
||||
import asyncio
|
||||
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Verify callback was called
|
||||
assert (
|
||||
test_callback.async_success_called is True
|
||||
), "async_log_success_event was not called"
|
||||
assert (
|
||||
test_callback.standard_logging_payload is not None
|
||||
), "standard_logging_payload is None"
|
||||
|
||||
# Check response cost
|
||||
print(f"Final response_cost: {test_callback.response_cost}")
|
||||
|
||||
# The first chunk may have the request model (azure-model-router) because it's created
|
||||
# before the API response is received. Subsequent chunks should have the actual model.
|
||||
# At least some chunks should have the actual model (not azure-model-router)
|
||||
actual_model_chunks = [
|
||||
m for m in chunks_with_model if m != "azure-model-router"
|
||||
]
|
||||
assert (
|
||||
len(actual_model_chunks) > 0
|
||||
), "No chunks had the actual model from the API response"
|
||||
print(f"Chunks with actual model: {actual_model_chunks}")
|
||||
|
||||
# Verify response cost is tracked - this is the main goal of this test
|
||||
assert (
|
||||
test_callback.response_cost is not None
|
||||
), "response_cost is None with stream_options"
|
||||
assert (
|
||||
test_callback.response_cost > 0
|
||||
), f"response_cost should be > 0, got {test_callback.response_cost}"
|
||||
print(
|
||||
f"Streaming cost tracking with stream_options passed. Cost: {test_callback.response_cost}"
|
||||
)
|
||||
|
||||
finally:
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
litellm.callbacks = []
|
||||
|
|
@ -0,0 +1,39 @@
|
|||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
|
||||
|
||||
import litellm
|
||||
from base_llm_unit_tests import BaseLLMChatTest
|
||||
|
||||
class TestAzureOpenAIO3Mini(BaseLLMChatTest):
|
||||
def get_base_completion_call_args(self):
|
||||
# Clear the LLM client cache to prevent test pollution from cached clients
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
return {
|
||||
"model": "azure/o3-mini",
|
||||
"api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
"api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
"api_version": "2024-12-01-preview",
|
||||
}
|
||||
|
||||
def get_client(self):
|
||||
from openai import AzureOpenAI
|
||||
|
||||
return AzureOpenAI(
|
||||
api_key="my-fake-o1-key",
|
||||
base_url="https://openai-prod-test.openai.azure.com",
|
||||
api_version="2024-02-15-preview",
|
||||
)
|
||||
|
||||
def test_basic_tool_calling(self):
|
||||
pass
|
||||
|
||||
def test_prompt_caching(self):
|
||||
"""Temporary override. o1 prompt caching is not working."""
|
||||
pass
|
||||
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
import sys
|
||||
import os
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
import pytest
|
||||
from litellm.llms.azure.common_utils import process_azure_headers
|
||||
from httpx import Headers
|
||||
from base_embedding_unit_tests import BaseLLMEmbeddingTest
|
||||
import litellm
|
||||
from litellm import completion
|
||||
|
||||
def test_azure_safety_result():
|
||||
"""Bubble up safety result from Azure OpenAI"""
|
||||
from litellm import completion
|
||||
|
||||
litellm._turn_on_debug()
|
||||
|
||||
response = completion(
|
||||
model="azure/gpt-4.1-mini",
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
api_base=os.getenv("AZURE_AI_API_BASE"),
|
||||
api_version="2024-12-01-preview",
|
||||
messages=[{"role": "user", "content": "Hello world"}],
|
||||
)
|
||||
print(f"response: {response}")
|
||||
assert response.choices[0].message.content is not None
|
||||
assert response.choices[0].provider_specific_fields is not None
|
||||
|
||||
|
||||
def test_completion_azure_deployment_id():
|
||||
"""
|
||||
Ensure deployment_id takes precedence over model.
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
response = completion(
|
||||
deployment_id="gpt-4.1-mini",
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hello, how are you?",
|
||||
}
|
||||
],
|
||||
)
|
||||
# Add any assertions here to check the response
|
||||
print(response)
|
||||
|
||||
|
||||
def test_azure_openai_with_prompt_cache_key():
|
||||
"""
|
||||
E2E test for Azure OpenAI with prompt cache key param on /chat/completions API.
|
||||
"""
|
||||
litellm._turn_on_debug()
|
||||
response = litellm.completion(
|
||||
model="azure/gpt-4.1-mini",
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
api_base=os.getenv("AZURE_AI_API_BASE"),
|
||||
api_version="2024-12-01-preview",
|
||||
messages=[{"role": "user", "content": "What is the weather in San Francisco?"}],
|
||||
prompt_cache_key="test_streaming_azure_openai",
|
||||
)
|
||||
|
|
@ -24,7 +24,6 @@ from unittest.mock import MagicMock, patch
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import completion
|
||||
|
||||
AZURE_AI_API_BASE = os.getenv("AZURE_AI_API_BASE")
|
||||
|
||||
|
|
@ -163,301 +162,3 @@ def test_azure_deepseek_reasoning_content():
|
|||
# }
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_ai_request_format():
|
||||
"""
|
||||
Test that Azure AI requests are formatted correctly with the proper endpoint and parameters
|
||||
for both synchronous and asynchronous calls
|
||||
"""
|
||||
|
||||
litellm._turn_on_debug()
|
||||
|
||||
# Set up the test parameters
|
||||
api_key = os.getenv("AZURE_AI_API_KEY")
|
||||
api_base = os.getenv("AZURE_AI_API_BASE")
|
||||
model = "azure_ai/gpt-4.1-mini"
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "Hello! How can I assist you today?"},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
|
||||
await litellm.acompletion(
|
||||
custom_llm_provider="azure_ai",
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", ["azure/gpt5_series/gpt-5-mini", "azure/gpt-5-mini"])
|
||||
async def test_azure_gpt5_reasoning(model):
|
||||
litellm._turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "What is the capital of France?"}],
|
||||
reasoning_effort="minimal",
|
||||
max_tokens=10,
|
||||
api_base=os.getenv("AZURE_AI_API_BASE"),
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
)
|
||||
print("response: ", response)
|
||||
assert response.choices[0].message.content is not None
|
||||
|
||||
|
||||
def test_completion_azure():
|
||||
try:
|
||||
from litellm import completion_cost
|
||||
|
||||
litellm.set_verbose = False
|
||||
## Test azure call
|
||||
response = completion(
|
||||
model="azure_ai/gpt-4.1-mini",
|
||||
api_base=os.getenv("AZURE_AI_API_BASE"),
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hello, how are you?",
|
||||
}
|
||||
],
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
)
|
||||
print(f"response: {response}")
|
||||
print(f"response hidden params: {response._hidden_params}")
|
||||
print(response)
|
||||
|
||||
cost = completion_cost(completion_response=response)
|
||||
assert cost > 0.0
|
||||
print("Cost for azure completion request", cost)
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[
|
||||
AZURE_AI_API_BASE,
|
||||
f"{AZURE_AI_API_BASE}/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2023-03-15-preview",
|
||||
],
|
||||
)
|
||||
def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base):
|
||||
try:
|
||||
litellm.set_verbose = True
|
||||
|
||||
response = completion(
|
||||
model="azure_ai/gpt-4.1-mini",
|
||||
api_base=api_base,
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
messages=[{"role": "user", "content": "What is the meaning of life?"}],
|
||||
)
|
||||
|
||||
print(response)
|
||||
except litellm.Timeout as e:
|
||||
pass
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_ai_model_router():
|
||||
"""
|
||||
Test Azure AI model router non-streaming response cost tracking.
|
||||
Verifies that the flat cost of $0.14 per M input tokens is applied.
|
||||
|
||||
Tests the pattern: azure_ai/model_router/<deployment-name>
|
||||
Where deployment-name is the Azure deployment (e.g., "azure-model-router").
|
||||
The model_router prefix is stripped before sending to Azure API.
|
||||
"""
|
||||
from litellm.llms.azure_ai.cost_calculator import (
|
||||
calculate_azure_model_router_flat_cost,
|
||||
)
|
||||
|
||||
litellm._turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model="azure_ai/model_router/azure-model-router",
|
||||
messages=[{"role": "user", "content": "hi who is this"}],
|
||||
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
|
||||
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
|
||||
)
|
||||
print("response: ", response)
|
||||
|
||||
# Check response cost
|
||||
tracked_cost = response._hidden_params["response_cost"]
|
||||
assert tracked_cost > 0
|
||||
print("Tracked cost: ", tracked_cost)
|
||||
|
||||
# Verify flat cost is included using the helper function
|
||||
usage = response.usage
|
||||
if usage and usage.prompt_tokens:
|
||||
expected_flat_cost = calculate_azure_model_router_flat_cost(
|
||||
model="model_router/azure-model-router", prompt_tokens=usage.prompt_tokens
|
||||
)
|
||||
print(f"Prompt tokens: {usage.prompt_tokens}")
|
||||
print(f"Expected flat cost: ${expected_flat_cost:.9f}")
|
||||
print(f"Total tracked cost: ${tracked_cost:.9f}")
|
||||
|
||||
# Total cost should be at least the flat cost
|
||||
assert (
|
||||
tracked_cost >= expected_flat_cost
|
||||
), f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}"
|
||||
|
||||
# Verify the flat cost is non-zero
|
||||
assert expected_flat_cost > 0, "Flat cost should be greater than 0"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_ai_model_router_streaming_model_in_chunk():
|
||||
"""
|
||||
Test that Azure AI model router streaming returns the actual model in each chunk.
|
||||
The response should contain the actual model used (e.g., gpt-4.1-nano) not the request model (azure-model-router).
|
||||
"""
|
||||
litellm._turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model="azure_ai/azure-model-router",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
|
||||
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Collect chunks and check model field
|
||||
chunks_with_model = []
|
||||
async for chunk in response:
|
||||
print(f"Chunk model: {chunk.model}")
|
||||
if chunk.model and chunk.model.strip():
|
||||
chunks_with_model.append(chunk.model)
|
||||
|
||||
print(f"All chunk models: {chunks_with_model}")
|
||||
|
||||
# At least some chunks should have a model
|
||||
assert len(chunks_with_model) > 0, "No chunks had a model field set"
|
||||
|
||||
# The model should NOT be azure-model-router (the request model)
|
||||
# It should be the actual model from the response (e.g., gpt-4.1-nano, gpt-5-nano, etc.)
|
||||
for model in chunks_with_model:
|
||||
assert (
|
||||
model != "azure-model-router"
|
||||
), f"Chunk model should be actual model, not request model. Got: {model}"
|
||||
# The actual model should be a real model name like gpt-4.1-nano, gpt-5-nano, etc.
|
||||
print(f"Verified chunk has actual model: {model}")
|
||||
|
||||
|
||||
class AzureModelRouterStreamingCallback(
|
||||
litellm.integrations.custom_logger.CustomLogger
|
||||
):
|
||||
"""
|
||||
Custom callback to capture streaming cost tracking for Azure Model Router.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.standard_logging_payload = None
|
||||
self.response_cost = None
|
||||
self.async_success_called = False
|
||||
self.complete_streaming_response = None
|
||||
super().__init__()
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
print(f"async_log_success_event called")
|
||||
self.async_success_called = True
|
||||
self.standard_logging_payload = kwargs.get("standard_logging_object")
|
||||
self.complete_streaming_response = kwargs.get("complete_streaming_response")
|
||||
|
||||
if self.standard_logging_payload:
|
||||
self.response_cost = self.standard_logging_payload.get("response_cost")
|
||||
print(
|
||||
f"standard_logging_payload model: {self.standard_logging_payload.get('model')}"
|
||||
)
|
||||
print(f"standard_logging_payload response_cost: {self.response_cost}")
|
||||
|
||||
if self.complete_streaming_response:
|
||||
print(
|
||||
f"complete_streaming_response model: {self.complete_streaming_response.model}"
|
||||
)
|
||||
print(
|
||||
f"complete_streaming_response usage: {self.complete_streaming_response.usage}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_ai_model_router_streaming_cost_with_stream_options():
|
||||
"""
|
||||
Test Azure AI model router streaming cost tracking with stream_options include_usage=True.
|
||||
This tests the specific case where cost tracking fails with stream_options.
|
||||
"""
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
test_callback = AzureModelRouterStreamingCallback()
|
||||
litellm.callbacks = [test_callback]
|
||||
|
||||
try:
|
||||
litellm._turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model="azure_ai/azure-model-router",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
|
||||
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
|
||||
# Consume the stream and check chunks
|
||||
full_response = ""
|
||||
chunks_with_model = []
|
||||
async for chunk in response:
|
||||
print(
|
||||
f"Chunk: model={chunk.model}, choices={len(chunk.choices) if chunk.choices else 0}"
|
||||
)
|
||||
if chunk.model:
|
||||
chunks_with_model.append(chunk.model)
|
||||
if (
|
||||
chunk.choices
|
||||
and chunk.choices[0].delta
|
||||
and chunk.choices[0].delta.content
|
||||
):
|
||||
full_response += chunk.choices[0].delta.content
|
||||
|
||||
print(f"Full streamed response: {full_response}")
|
||||
print(f"Chunks with model: {chunks_with_model}")
|
||||
|
||||
# Give async logging time to complete
|
||||
import asyncio
|
||||
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Verify callback was called
|
||||
assert (
|
||||
test_callback.async_success_called is True
|
||||
), "async_log_success_event was not called"
|
||||
assert (
|
||||
test_callback.standard_logging_payload is not None
|
||||
), "standard_logging_payload is None"
|
||||
|
||||
# Check response cost
|
||||
print(f"Final response_cost: {test_callback.response_cost}")
|
||||
|
||||
# The first chunk may have the request model (azure-model-router) because it's created
|
||||
# before the API response is received. Subsequent chunks should have the actual model.
|
||||
# At least some chunks should have the actual model (not azure-model-router)
|
||||
actual_model_chunks = [
|
||||
m for m in chunks_with_model if m != "azure-model-router"
|
||||
]
|
||||
assert (
|
||||
len(actual_model_chunks) > 0
|
||||
), "No chunks had the actual model from the API response"
|
||||
print(f"Chunks with actual model: {actual_model_chunks}")
|
||||
|
||||
# Verify response cost is tracked - this is the main goal of this test
|
||||
assert (
|
||||
test_callback.response_cost is not None
|
||||
), "response_cost is None with stream_options"
|
||||
assert (
|
||||
test_callback.response_cost > 0
|
||||
), f"response_cost should be > 0, got {test_callback.response_cost}"
|
||||
print(
|
||||
f"Streaming cost tracking with stream_options passed. Cost: {test_callback.response_cost}"
|
||||
)
|
||||
|
||||
finally:
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
litellm.callbacks = []
|
||||
|
|
|
|||
|
|
@ -5,67 +5,35 @@ sys.path.insert(
|
|||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
|
||||
|
||||
import litellm
|
||||
from base_llm_unit_tests import BaseLLMChatTest
|
||||
|
||||
|
||||
class TestAzureOpenAIO3Mini(BaseLLMChatTest):
|
||||
def get_base_completion_call_args(self):
|
||||
# Clear the LLM client cache to prevent test pollution from cached clients
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
return {
|
||||
"model": "azure/o3-mini",
|
||||
"api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
"api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
"api_version": "2024-12-01-preview",
|
||||
}
|
||||
def test_azure_o1_override_fake_stream():
|
||||
"""Test that native streaming is not supported for o1."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "azure/o1-preview",
|
||||
"litellm_params": {
|
||||
"model": "azure/o1-preview",
|
||||
"api_key": "my-fake-o1-key",
|
||||
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com",
|
||||
},
|
||||
"model_info": {
|
||||
"supports_native_streaming": True,
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
def get_client(self):
|
||||
from openai import AzureOpenAI
|
||||
|
||||
return AzureOpenAI(
|
||||
api_key="my-fake-o1-key",
|
||||
base_url="https://openai-prod-test.openai.azure.com",
|
||||
api_version="2024-02-15-preview",
|
||||
)
|
||||
|
||||
def test_basic_tool_calling(self):
|
||||
pass
|
||||
|
||||
def test_prompt_caching(self):
|
||||
"""Temporary override. o1 prompt caching is not working."""
|
||||
pass
|
||||
|
||||
def test_override_fake_stream(self):
|
||||
"""Test that native streaming is not supported for o1."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "azure/o1-preview",
|
||||
"litellm_params": {
|
||||
"model": "azure/o1-preview",
|
||||
"api_key": "my-fake-o1-key",
|
||||
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com",
|
||||
},
|
||||
"model_info": {
|
||||
"supports_native_streaming": True,
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
## check model info
|
||||
|
||||
model_info = litellm.get_model_info(
|
||||
model="azure/o1-preview", custom_llm_provider="azure"
|
||||
)
|
||||
assert model_info["supports_native_streaming"] is True
|
||||
|
||||
fake_stream = litellm.AzureOpenAIO1Config().should_fake_stream(
|
||||
model="azure/o1-preview", stream=True
|
||||
)
|
||||
assert fake_stream is False
|
||||
## check model info
|
||||
|
||||
model_info = litellm.get_model_info(
|
||||
model="azure/o1-preview", custom_llm_provider="azure"
|
||||
)
|
||||
assert model_info["supports_native_streaming"] is True
|
||||
|
||||
fake_stream = litellm.AzureOpenAIO1Config().should_fake_stream(
|
||||
model="azure/o1-preview", stream=True
|
||||
)
|
||||
assert fake_stream is False
|
||||
|
|
|
|||
|
|
@ -99,7 +99,6 @@ def test_process_azure_headers_with_dict_input():
|
|||
|
||||
from unittest.mock import MagicMock
|
||||
import litellm
|
||||
from litellm import completion
|
||||
import os
|
||||
|
||||
|
||||
|
|
@ -244,43 +243,6 @@ def test_map_openai_params():
|
|||
assert len(optional_params["tools"]) > 1
|
||||
|
||||
|
||||
def test_azure_safety_result():
|
||||
"""Bubble up safety result from Azure OpenAI"""
|
||||
from litellm import completion
|
||||
|
||||
litellm._turn_on_debug()
|
||||
|
||||
response = completion(
|
||||
model="azure/gpt-4.1-mini",
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
api_base=os.getenv("AZURE_AI_API_BASE"),
|
||||
api_version="2024-12-01-preview",
|
||||
messages=[{"role": "user", "content": "Hello world"}],
|
||||
)
|
||||
print(f"response: {response}")
|
||||
assert response.choices[0].message.content is not None
|
||||
assert response.choices[0].provider_specific_fields is not None
|
||||
|
||||
|
||||
def test_completion_azure_deployment_id():
|
||||
"""
|
||||
Ensure deployment_id takes precedence over model.
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
response = completion(
|
||||
deployment_id="gpt-4.1-mini",
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hello, how are you?",
|
||||
}
|
||||
],
|
||||
)
|
||||
# Add any assertions here to check the response
|
||||
print(response)
|
||||
|
||||
|
||||
def test_azure_with_content_safety_error():
|
||||
"""
|
||||
Verify user can access innererror from the Azure OpenAI exception
|
||||
|
|
@ -338,16 +300,3 @@ def test_azure_with_content_safety_error():
|
|||
)
|
||||
|
||||
|
||||
def test_azure_openai_with_prompt_cache_key():
|
||||
"""
|
||||
E2E test for Azure OpenAI with prompt cache key param on /chat/completions API.
|
||||
"""
|
||||
litellm._turn_on_debug()
|
||||
response = litellm.completion(
|
||||
model="azure/gpt-4.1-mini",
|
||||
api_key=os.getenv("AZURE_AI_API_KEY"),
|
||||
api_base=os.getenv("AZURE_AI_API_BASE"),
|
||||
api_version="2024-12-01-preview",
|
||||
messages=[{"role": "user", "content": "What is the weather in San Francisco?"}],
|
||||
prompt_cache_key="test_streaming_azure_openai",
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue