mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(azure_ai): stamp the model router's selected model instead of matching on the model name
The model Azure Model Router served was recovered by checking whether the text "model_router" or "model-router" appeared in a model string. Spend logs applied that check to the litellm model path, where the route prefix guarantees a match, but the proxy applied it to the client's model group alias, which carries no prefix. A model group named anything else therefore lost the selected model in both the response and the spend row. AzureModelRouterConfig now stamps the served model onto _hidden_params, and the spend log payload and the proxy's response restamping read that stamp. The name heuristic survives as a fallback for callers with no response in hand, routed through get_azure_ai_route so it lives in one place.
This commit is contained in:
ryan-crabbe-berri
parent
1cef8823fa
commit
0ea6f5e159
7 changed files with 538 additions and 566 deletions
|
|
@ -5762,11 +5762,15 @@ def get_standard_logging_object_payload(
|
|||
response_model_name = final_response_obj.get("model")
|
||||
|
||||
# For Azure Model Router, preserve the actual model in the top-level standard
|
||||
# logging payload only when the user has opted in.
|
||||
# logging payload.
|
||||
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
|
||||
|
||||
requested_model: Final = kwargs.get("model")
|
||||
if (
|
||||
isinstance(requested_model, str)
|
||||
and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower())
|
||||
stamped_selected_model: Final = AzureFoundryModelInfo.get_model_router_selected_model(hidden_params)
|
||||
if stamped_selected_model is not None:
|
||||
model_name = stamped_selected_model
|
||||
elif (
|
||||
AzureFoundryModelInfo.is_model_router_call(model=requested_model, hidden_params=hidden_params)
|
||||
and isinstance(response_model_name, str)
|
||||
and response_model_name
|
||||
):
|
||||
|
|
|
|||
|
|
@ -65,15 +65,24 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
|
|||
|
||||
Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07)
|
||||
and returns it with the azure_ai/ prefix for proper display and cost tracking.
|
||||
|
||||
Also stamps that model onto ``_hidden_params`` so downstream consumers (spend logs,
|
||||
response restamping) can read it instead of guessing the route from the model string.
|
||||
"""
|
||||
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
|
||||
AzureFoundryModelInfo,
|
||||
)
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
get_hidden_params_dict,
|
||||
)
|
||||
|
||||
# Get base model for the parent call (strips routing prefixes for API compatibility)
|
||||
base_model: Final[str] = AzureFoundryModelInfo.get_base_model(model)
|
||||
|
||||
# Call parent transform_response first - this will extract the actual model
|
||||
# from the raw response (e.g., "gpt-5-nano-2025-08-07")
|
||||
model_response = super().transform_response(
|
||||
transformed_response: Final = super().transform_response(
|
||||
model=base_model,
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
|
|
@ -86,7 +95,15 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
|
|||
api_key=api_key,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
return model_response
|
||||
selected_model: Final = transformed_response.model
|
||||
if selected_model:
|
||||
# Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a
|
||||
# class-level dict, so an in-place write can bleed into unrelated responses.
|
||||
transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter
|
||||
**get_hidden_params_dict(transformed_response),
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model,
|
||||
}
|
||||
return transformed_response
|
||||
|
||||
def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> dict | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final, Literal
|
||||
|
||||
import litellm
|
||||
|
|
@ -5,6 +6,8 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
|
||||
|
||||
|
||||
class AzureFoundryModelInfo(BaseLLMModelInfo):
|
||||
"""Model info for Azure AI / Azure Foundry models."""
|
||||
|
|
@ -37,6 +40,39 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
|
|||
return "model_router"
|
||||
return "default"
|
||||
|
||||
@staticmethod
|
||||
def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None:
|
||||
"""The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``.
|
||||
|
||||
Reading this beats re-deriving the route from a model string: the stamp is set on the
|
||||
code path that was actually taken, so it holds no matter what the caller named the model.
|
||||
"""
|
||||
if not hidden_params:
|
||||
return None
|
||||
selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY)
|
||||
if isinstance(selected, str) and selected:
|
||||
return selected
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def is_model_router_call(
|
||||
model: str | None = None,
|
||||
hidden_params: Mapping[str, object] | None = None,
|
||||
) -> bool:
|
||||
"""Whether a request went down the Azure Model Router route.
|
||||
|
||||
Prefers the response stamp, then the deployment's litellm model path, and only then the
|
||||
caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router
|
||||
name heuristic lives in exactly one place.
|
||||
"""
|
||||
if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None:
|
||||
return True
|
||||
deployment_model: Final = (hidden_params or {}).get("litellm_model_name") or (hidden_params or {}).get("model")
|
||||
return any(
|
||||
isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router"
|
||||
for candidate in (deployment_model, model)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_api_base(api_base: str | None = None) -> str | None:
|
||||
return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE")
|
||||
|
|
|
|||
|
|
@ -1136,24 +1136,25 @@ async def open_sse_before_first_byte(
|
|||
)
|
||||
|
||||
|
||||
def _is_azure_model_router_request(model: str) -> bool:
|
||||
def _is_azure_model_router_request(model: str, hidden_params: Mapping[str, object] | None = None) -> bool:
|
||||
"""
|
||||
Check if the requested model is an Azure Model Router.
|
||||
Check if a request went down the Azure Model Router route.
|
||||
|
||||
Azure Model Router models follow the pattern:
|
||||
- azure_ai/model_router/<deployment-name>
|
||||
- azure_ai/model-router
|
||||
- model_router/<deployment-name>
|
||||
- model-router
|
||||
``model`` here is what the *client* sent, a model group alias with no ``model_router/``
|
||||
prefix, so matching on it alone only works when the operator happened to put "model-router"
|
||||
in the alias. Where the response is in hand its stamp answers this outright, so callers
|
||||
should pass ``hidden_params``.
|
||||
|
||||
Args:
|
||||
model: The requested model name
|
||||
hidden_params: ``_hidden_params`` from the response, when the caller has it
|
||||
|
||||
Returns:
|
||||
bool: True if this is an Azure Model Router request
|
||||
"""
|
||||
model_lower: Final = model.lower()
|
||||
return "model-router" in model_lower or "model_router" in model_lower
|
||||
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
|
||||
|
||||
return AzureFoundryModelInfo.is_model_router_call(model=model, hidden_params=hidden_params)
|
||||
|
||||
|
||||
def _override_openai_response_model(
|
||||
|
|
@ -1221,7 +1222,7 @@ def _override_openai_response_model(
|
|||
return
|
||||
|
||||
# Check if this is an Azure Model Router request - if so, preserve the actual model used
|
||||
if _is_azure_model_router_request(requested_model):
|
||||
if _is_azure_model_router_request(requested_model, hidden_params):
|
||||
verbose_proxy_logger.debug(
|
||||
"%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.",
|
||||
log_context,
|
||||
|
|
|
|||
|
|
@ -6,9 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
|
||||
|
||||
import time
|
||||
|
||||
|
|
@ -277,9 +275,7 @@ def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata():
|
|||
|
||||
assert cost is not None, "Cost should not be None"
|
||||
expected_cost = (10 * custom_input_cost) + (5 * custom_output_cost)
|
||||
assert cost == pytest.approx(
|
||||
expected_cost
|
||||
), f"Expected {expected_cost}, got {cost}"
|
||||
assert cost == pytest.approx(expected_cost), f"Expected {expected_cost}, got {cost}"
|
||||
finally:
|
||||
litellm.model_cost.pop(custom_model_id, None)
|
||||
|
||||
|
|
@ -876,13 +872,8 @@ async def test_datadog_logger_not_shadowed_by_llm_obs(monkeypatch):
|
|||
|
||||
# Regression check: we expect a distinct DataDogLogger, not the LLM Obs logger
|
||||
assert type(datadog_logger) is DataDogLogger
|
||||
assert any(
|
||||
isinstance(cb, DataDogLLMObsLogger)
|
||||
for cb in logging_module._in_memory_loggers
|
||||
)
|
||||
assert any(
|
||||
type(cb) is DataDogLogger for cb in logging_module._in_memory_loggers
|
||||
)
|
||||
assert any(isinstance(cb, DataDogLLMObsLogger) for cb in logging_module._in_memory_loggers)
|
||||
assert any(type(cb) is DataDogLogger for cb in logging_module._in_memory_loggers)
|
||||
finally:
|
||||
logging_module._in_memory_loggers.clear()
|
||||
|
||||
|
|
@ -893,9 +884,7 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch):
|
|||
|
||||
# Required env vars for Logfire integration
|
||||
monkeypatch.setenv("LOGFIRE_TOKEN", "test-token")
|
||||
monkeypatch.setenv(
|
||||
"LOGFIRE_BASE_URL", "https://logfire-api-custom.pydantic.dev"
|
||||
) # no trailing slash on purpose
|
||||
monkeypatch.setenv("LOGFIRE_BASE_URL", "https://logfire-api-custom.pydantic.dev") # no trailing slash on purpose
|
||||
|
||||
# Import after env vars are set (important if module-level caching exists)
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry # logger class
|
||||
|
|
@ -914,9 +903,7 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch):
|
|||
|
||||
# Sanity: we got the right logger type and it is cached
|
||||
assert type(logger) is OpenTelemetry
|
||||
assert any(
|
||||
type(cb) is OpenTelemetry for cb in logging_module._in_memory_loggers
|
||||
)
|
||||
assert any(type(cb) is OpenTelemetry for cb in logging_module._in_memory_loggers)
|
||||
|
||||
# Core regression check: base URL env var should influence the exporter endpoint.
|
||||
#
|
||||
|
|
@ -927,9 +914,7 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch):
|
|||
or getattr(logger, "config", None)
|
||||
or getattr(logger, "_otel_config", None)
|
||||
)
|
||||
assert (
|
||||
cfg is not None
|
||||
), "Expected OpenTelemetry logger to keep an otel config on the instance"
|
||||
assert cfg is not None, "Expected OpenTelemetry logger to keep an otel config on the instance"
|
||||
|
||||
endpoint = getattr(cfg, "endpoint", None) or getattr(cfg, "otlp_endpoint", None)
|
||||
assert endpoint is not None, "Expected otel config to expose the OTLP endpoint"
|
||||
|
|
@ -1087,9 +1072,7 @@ async def test_logging_non_streaming_request():
|
|||
|
||||
# Use the filtered call for assertions
|
||||
call_args = calls_with_expected_input[0]
|
||||
standard_logging_object = call_args.kwargs["kwargs"][
|
||||
"standard_logging_object"
|
||||
]
|
||||
standard_logging_object = call_args.kwargs["kwargs"]["standard_logging_object"]
|
||||
assert standard_logging_object["stream"] is not True
|
||||
finally:
|
||||
# Restore original callbacks to ensure test isolation
|
||||
|
|
@ -1107,18 +1090,14 @@ async def test_logging_non_streaming_request():
|
|||
"agenerate_content_stream",
|
||||
],
|
||||
)
|
||||
def test_success_handler_skips_sync_callbacks_for_async_requests(
|
||||
logging_obj, async_flag
|
||||
):
|
||||
def test_success_handler_skips_sync_callbacks_for_async_requests(logging_obj, async_flag):
|
||||
"""Ensure sync success callbacks are skipped when async call type flags are set."""
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
class DummyLogger(CustomLogger):
|
||||
pass
|
||||
|
||||
logging_obj.stream = (
|
||||
False # simulate non-streaming request where sync callbacks would normally run
|
||||
)
|
||||
logging_obj.stream = False # simulate non-streaming request where sync callbacks would normally run
|
||||
logging_obj.model_call_details["litellm_params"] = {async_flag: True}
|
||||
logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"]
|
||||
|
||||
|
|
@ -1194,21 +1173,11 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call
|
|||
def test_is_sync_litellm_request():
|
||||
assert LitellmLogging._is_sync_litellm_request({}) is True
|
||||
assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False
|
||||
assert (
|
||||
LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True})
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False
|
||||
)
|
||||
assert LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) is False
|
||||
assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False
|
||||
assert LitellmLogging._is_sync_litellm_request({"agenerate_content": True}) is False
|
||||
assert (
|
||||
LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True})
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
LitellmLogging._is_sync_litellm_request({"aanthropic_messages": False}) is True
|
||||
)
|
||||
assert LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) is False
|
||||
assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": False}) is True
|
||||
|
||||
|
||||
def test_get_litellm_params_propagates_allm_passthrough_route():
|
||||
|
|
@ -1255,9 +1224,7 @@ async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream
|
|||
logging_obj.model_call_details["litellm_params"] = {"acompletion": True}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
mock_callback, "async_log_success_event", new_callable=AsyncMock
|
||||
) as mock_async_log,
|
||||
patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log,
|
||||
patch.object(mock_callback, "log_success_event") as mock_sync_log,
|
||||
patch.object(
|
||||
logging_obj,
|
||||
|
|
@ -1318,9 +1285,7 @@ async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_fin
|
|||
|
||||
with (
|
||||
patch.object(mock_callback, "log_success_event") as mock_sync_log,
|
||||
patch.object(
|
||||
mock_callback, "async_log_success_event", new_callable=AsyncMock
|
||||
) as mock_async_log,
|
||||
patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log,
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_success_handler_helper_fn",
|
||||
|
|
@ -1362,20 +1327,14 @@ async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks(
|
|||
logging_obj.model_call_details["litellm_params"] = {}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
logging_obj, "async_success_handler", new_callable=AsyncMock
|
||||
) as mock_async,
|
||||
patch.object(
|
||||
logging_obj, "success_handler", new_callable=MagicMock
|
||||
) as mock_sync,
|
||||
patch.object(logging_obj, "async_success_handler", new_callable=AsyncMock) as mock_async,
|
||||
patch.object(logging_obj, "success_handler", new_callable=MagicMock) as mock_sync,
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_should_run_sync_callbacks_for_async_calls",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.executor.submit"
|
||||
) as mock_submit,
|
||||
patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit,
|
||||
):
|
||||
await logging_obj.dispatch_success_handlers(
|
||||
result=result,
|
||||
|
|
@ -1409,9 +1368,7 @@ async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through
|
|||
|
||||
try:
|
||||
with (
|
||||
patch.object(
|
||||
mock_callback, "async_log_success_event", new_callable=AsyncMock
|
||||
) as mock_async_log,
|
||||
patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log,
|
||||
patch.object(mock_callback, "log_success_event") as mock_sync_log,
|
||||
):
|
||||
await logging_obj.dispatch_success_handlers(result={"id": "pt-1"})
|
||||
|
|
@ -1438,20 +1395,14 @@ async def test_dispatch_failure_handlers_prefer_async_does_not_submit_sync_handl
|
|||
logging_obj.model_call_details["litellm_params"] = {}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
logging_obj, "async_failure_handler", new_callable=AsyncMock
|
||||
) as mock_async,
|
||||
patch.object(
|
||||
logging_obj, "failure_handler", new_callable=MagicMock
|
||||
) as mock_sync,
|
||||
patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock) as mock_async,
|
||||
patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync,
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_should_run_sync_failure_callbacks_for_async_calls",
|
||||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.executor.submit"
|
||||
) as mock_submit,
|
||||
patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit,
|
||||
):
|
||||
await logging_obj.dispatch_failure_handlers(
|
||||
exception,
|
||||
|
|
@ -1534,12 +1485,8 @@ async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_c
|
|||
patch.object(litellm, "success_callback", []),
|
||||
patch.object(litellm, "failure_callback", [_sync_failure_callback]),
|
||||
patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock),
|
||||
patch.object(
|
||||
logging_obj, "failure_handler", new_callable=MagicMock
|
||||
) as mock_sync,
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.executor.submit"
|
||||
) as mock_submit,
|
||||
patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync,
|
||||
patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit,
|
||||
):
|
||||
await logging_obj.dispatch_failure_handlers(
|
||||
exception,
|
||||
|
|
@ -1566,15 +1513,9 @@ async def test_dispatch_failure_handlers_sync_sdk_shortcut_runs_sync_handler_inl
|
|||
logging_obj.model_call_details["litellm_params"] = {}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
logging_obj, "async_failure_handler", new_callable=AsyncMock
|
||||
) as mock_async,
|
||||
patch.object(
|
||||
logging_obj, "failure_handler", new_callable=MagicMock
|
||||
) as mock_sync,
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.executor.submit"
|
||||
) as mock_submit,
|
||||
patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock) as mock_async,
|
||||
patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync,
|
||||
patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit,
|
||||
):
|
||||
await logging_obj.dispatch_failure_handlers(
|
||||
exception,
|
||||
|
|
@ -1621,14 +1562,10 @@ def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj)
|
|||
event_hook=GuardrailEventHooks.logging_only,
|
||||
)
|
||||
guardrail.should_run_guardrail = MagicMock(return_value=False)
|
||||
guardrail.logging_hook = MagicMock(
|
||||
return_value=(logging_obj.model_call_details, model_response)
|
||||
)
|
||||
guardrail.logging_hook = MagicMock(return_value=(logging_obj.model_call_details, model_response))
|
||||
|
||||
dummy_logger = DummyLogger()
|
||||
dummy_logger.logging_hook = MagicMock(
|
||||
return_value=(logging_obj.model_call_details, model_response)
|
||||
)
|
||||
dummy_logger.logging_hook = MagicMock(return_value=(logging_obj.model_call_details, model_response))
|
||||
|
||||
with patch.object(
|
||||
logging_obj,
|
||||
|
|
@ -1762,11 +1699,7 @@ def test_get_request_tags_from_metadata_and_litellm_metadata():
|
|||
|
||||
# Test case 2: Tags in litellm_metadata only
|
||||
tags = StandardLoggingPayloadSetup._get_request_tags(
|
||||
litellm_params={
|
||||
"litellm_metadata": {
|
||||
"tags": ["litellm-metadata-tag-1", "litellm-metadata-tag-2"]
|
||||
}
|
||||
},
|
||||
litellm_params={"litellm_metadata": {"tags": ["litellm-metadata-tag-1", "litellm-metadata-tag-2"]}},
|
||||
proxy_server_request={},
|
||||
)
|
||||
assert "litellm-metadata-tag-1" in tags
|
||||
|
|
@ -1871,15 +1804,9 @@ def test_get_request_tags_does_not_mutate_original_tags():
|
|||
user_agent_count_2 = len([t for t in tags2 if t.startswith("User-Agent:")])
|
||||
user_agent_count_3 = len([t for t in tags3 if t.startswith("User-Agent:")])
|
||||
|
||||
assert (
|
||||
user_agent_count_1 == 2
|
||||
), f"Expected 2 User-Agent tags, got {user_agent_count_1}"
|
||||
assert (
|
||||
user_agent_count_2 == 2
|
||||
), f"Expected 2 User-Agent tags, got {user_agent_count_2}"
|
||||
assert (
|
||||
user_agent_count_3 == 2
|
||||
), f"Expected 2 User-Agent tags, got {user_agent_count_3}"
|
||||
assert user_agent_count_1 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_1}"
|
||||
assert user_agent_count_2 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_2}"
|
||||
assert user_agent_count_3 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_3}"
|
||||
|
||||
# Verify all returned lists are independent (different objects)
|
||||
assert tags1 is not tags2
|
||||
|
|
@ -1912,9 +1839,7 @@ def test_get_extra_header_tags():
|
|||
|
||||
# Test case 3: Extra headers configured but request has no headers dict
|
||||
litellm.extra_spend_tag_headers = ["x-custom", "x-tenant"]
|
||||
result = StandardLoggingPayloadSetup._get_extra_header_tags(
|
||||
proxy_server_request={"headers": "not-a-dict"}
|
||||
)
|
||||
result = StandardLoggingPayloadSetup._get_extra_header_tags(proxy_server_request={"headers": "not-a-dict"})
|
||||
assert result is None
|
||||
|
||||
# Test case 4: Extra headers configured but none match request headers
|
||||
|
|
@ -2215,9 +2140,7 @@ def test_get_masked_values():
|
|||
"presidio_anonymizer_api_base": None,
|
||||
"vertex_credentials": "{sensitive_api_key}",
|
||||
}
|
||||
masked_values = _get_masked_values(
|
||||
sensitive_object, unmasked_length=4, number_of_asterisks=4
|
||||
)
|
||||
masked_values = _get_masked_values(sensitive_object, unmasked_length=4, number_of_asterisks=4)
|
||||
assert masked_values["presidio_anonymizer_api_base"] is None
|
||||
assert masked_values["vertex_credentials"] == "{s****y}"
|
||||
|
||||
|
|
@ -2242,9 +2165,7 @@ async def test_e2e_generate_cold_storage_object_key_successful():
|
|||
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key,
|
||||
):
|
||||
# Mock the S3 object key generation to return a predictable result
|
||||
mock_get_s3_key.return_value = (
|
||||
"2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
)
|
||||
mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
|
||||
# Call the function
|
||||
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
|
||||
|
|
@ -2285,16 +2206,12 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
|
|||
|
||||
with (
|
||||
patch("litellm.cold_storage_custom_logger", "s3_v2"),
|
||||
patch(
|
||||
"litellm.logging_callback_manager.get_active_custom_logger_for_callback_name"
|
||||
) as mock_get_logger,
|
||||
patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger,
|
||||
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key,
|
||||
):
|
||||
# Setup mocks
|
||||
mock_get_logger.return_value = mock_custom_logger
|
||||
mock_get_s3_key.return_value = (
|
||||
"storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
)
|
||||
mock_get_s3_key.return_value = "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
|
||||
# Call the function
|
||||
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
|
||||
|
|
@ -2313,9 +2230,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
|
|||
)
|
||||
|
||||
# Verify the result
|
||||
assert (
|
||||
result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
)
|
||||
assert result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -2338,16 +2253,12 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
|
|||
|
||||
with (
|
||||
patch("litellm.cold_storage_custom_logger", "s3_v2"),
|
||||
patch(
|
||||
"litellm.logging_callback_manager.get_active_custom_logger_for_callback_name"
|
||||
) as mock_get_logger,
|
||||
patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger,
|
||||
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key,
|
||||
):
|
||||
# Setup mocks
|
||||
mock_get_logger.return_value = mock_custom_logger
|
||||
mock_get_s3_key.return_value = (
|
||||
"2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
)
|
||||
mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
|
||||
# Call the function
|
||||
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
|
||||
|
|
@ -2463,9 +2374,7 @@ def test_get_usage_as_dict():
|
|||
assert result == {"prompt_tokens": 20, "completion_tokens": 30}
|
||||
|
||||
# Test case 5: response_obj with no usage key returns empty
|
||||
result = StandardLoggingPayloadSetup.get_usage_as_dict(
|
||||
response_obj={"id": "resp-1", "choices": []}
|
||||
)
|
||||
result = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj={"id": "resp-1", "choices": []})
|
||||
assert result == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
|
||||
|
||||
|
||||
|
|
@ -2478,26 +2387,20 @@ def test_append_system_prompt_messages():
|
|||
# Test case 1: system in kwargs with existing messages
|
||||
kwargs = {"system": "You are a helpful assistant"}
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(
|
||||
kwargs=kwargs, messages=messages
|
||||
)
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages)
|
||||
assert len(result) == 2
|
||||
assert result[0] == {"role": "system", "content": "You are a helpful assistant"}
|
||||
assert result[1] == {"role": "user", "content": "Hello"}
|
||||
|
||||
# Test case 2: system in kwargs with None messages
|
||||
kwargs = {"system": "You are a helpful assistant"}
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(
|
||||
kwargs=kwargs, messages=None
|
||||
)
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=None)
|
||||
assert len(result) == 1
|
||||
assert result[0] == {"role": "system", "content": "You are a helpful assistant"}
|
||||
|
||||
# Test case 3: system in kwargs with empty messages list
|
||||
kwargs = {"system": "You are a helpful assistant"}
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(
|
||||
kwargs=kwargs, messages=[]
|
||||
)
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=[])
|
||||
assert len(result) == 1
|
||||
assert result[0] == {"role": "system", "content": "You are a helpful assistant"}
|
||||
|
||||
|
|
@ -2507,24 +2410,18 @@ def test_append_system_prompt_messages():
|
|||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(
|
||||
kwargs=kwargs, messages=messages
|
||||
)
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages)
|
||||
assert len(result) == 2
|
||||
assert result[0] == {"role": "system", "content": "You are a helpful assistant"}
|
||||
|
||||
# Test case 5: no system in kwargs returns messages unchanged
|
||||
kwargs = {}
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(
|
||||
kwargs=kwargs, messages=messages
|
||||
)
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages)
|
||||
assert result == messages
|
||||
|
||||
# Test case 6: None kwargs returns messages unchanged
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(
|
||||
kwargs=None, messages=messages
|
||||
)
|
||||
result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=None, messages=messages)
|
||||
assert result == messages
|
||||
|
||||
|
||||
|
|
@ -2585,12 +2482,11 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu
|
|||
|
||||
# Verify that standard_logging_object was set
|
||||
assert "standard_logging_object" in logging_obj.model_call_details, (
|
||||
"standard_logging_object should be set for pass-through endpoints "
|
||||
"even when complete_streaming_response is None"
|
||||
"standard_logging_object should be set for pass-through endpoints even when complete_streaming_response is None"
|
||||
)
|
||||
assert logging_obj.model_call_details["standard_logging_object"] is not None, (
|
||||
"standard_logging_object should not be None for pass-through endpoints"
|
||||
)
|
||||
assert (
|
||||
logging_obj.model_call_details["standard_logging_object"] is not None
|
||||
), "standard_logging_object should not be None for pass-through endpoints"
|
||||
|
||||
# Verify that async_complete_streaming_response was set to prevent re-processing
|
||||
# This is consistent with the existing code pattern for regular streaming
|
||||
|
|
@ -2598,15 +2494,13 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu
|
|||
"async_complete_streaming_response should be set to prevent re-processing, "
|
||||
"consistent with the existing code pattern"
|
||||
)
|
||||
assert (
|
||||
logging_obj.model_call_details["async_complete_streaming_response"] is result
|
||||
), "async_complete_streaming_response should be set to the result"
|
||||
assert logging_obj.model_call_details["async_complete_streaming_response"] is result, (
|
||||
"async_complete_streaming_response should be set to the result"
|
||||
)
|
||||
|
||||
# Verify that response_cost is set to None (cost calculation not possible for pass-through)
|
||||
# This is consistent with the error handling in the non-pass-through code path
|
||||
assert (
|
||||
"response_cost" in logging_obj.model_call_details
|
||||
), "response_cost should be set for pass-through endpoints"
|
||||
assert "response_cost" in logging_obj.model_call_details, "response_cost should be set for pass-through endpoints"
|
||||
assert logging_obj.model_call_details["response_cost"] is None, (
|
||||
"response_cost should be None for pass-through endpoints since "
|
||||
"StandardPassThroughResponseObject doesn't have standard usage info"
|
||||
|
|
@ -2665,14 +2559,10 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp
|
|||
# Verify first call set the values
|
||||
assert "standard_logging_object" in logging_obj.model_call_details
|
||||
assert "async_complete_streaming_response" in logging_obj.model_call_details
|
||||
first_standard_logging_object = logging_obj.model_call_details[
|
||||
"standard_logging_object"
|
||||
]
|
||||
first_standard_logging_object = logging_obj.model_call_details["standard_logging_object"]
|
||||
|
||||
# Second call - should return early due to async_complete_streaming_response guard
|
||||
with patch.object(
|
||||
logging_obj, "get_combined_callback_list", return_value=[]
|
||||
) as mock_callbacks:
|
||||
with patch.object(logging_obj, "get_combined_callback_list", return_value=[]) as mock_callbacks:
|
||||
await logging_obj.async_success_handler(
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
|
|
@ -2683,10 +2573,9 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp
|
|||
mock_callbacks.assert_not_called()
|
||||
|
||||
# Verify standard_logging_object wasn't modified by second call
|
||||
assert (
|
||||
logging_obj.model_call_details["standard_logging_object"]
|
||||
is first_standard_logging_object
|
||||
), "standard_logging_object should not be modified on re-processing"
|
||||
assert logging_obj.model_call_details["standard_logging_object"] is first_standard_logging_object, (
|
||||
"standard_logging_object should not be modified on re-processing"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -2725,9 +2614,7 @@ async def test_async_success_handler_sets_standard_logging_object_for_streaming_
|
|||
}
|
||||
|
||||
# Create a pass-through response object (simulating unparseable streaming response)
|
||||
result = StandardPassThroughResponseObject(
|
||||
response='data: {"chunk": 1}\ndata: {"chunk": 2}\ndata: [DONE]'
|
||||
)
|
||||
result = StandardPassThroughResponseObject(response='data: {"chunk": 1}\ndata: {"chunk": 2}\ndata: [DONE]')
|
||||
|
||||
start_time = datetime.now()
|
||||
end_time = datetime.now()
|
||||
|
|
@ -2747,9 +2634,9 @@ async def test_async_success_handler_sets_standard_logging_object_for_streaming_
|
|||
"standard_logging_object should be set for streaming pass-through endpoints "
|
||||
"even when the response cannot be parsed into a ModelResponse"
|
||||
)
|
||||
assert (
|
||||
logging_obj.model_call_details["standard_logging_object"] is not None
|
||||
), "standard_logging_object should not be None for streaming pass-through endpoints"
|
||||
assert logging_obj.model_call_details["standard_logging_object"] is not None, (
|
||||
"standard_logging_object should not be None for streaming pass-through endpoints"
|
||||
)
|
||||
|
||||
|
||||
def test_get_error_information_error_code_priority():
|
||||
|
|
@ -2791,30 +2678,22 @@ def test_get_error_information_error_code_priority():
|
|||
self.message = message
|
||||
super().__init__(message)
|
||||
|
||||
both_exception = BothAttributesException(
|
||||
code="400", status_code=500, message="Bad Request"
|
||||
)
|
||||
both_exception = BothAttributesException(code="400", status_code=500, message="Bad Request")
|
||||
result = StandardLoggingPayloadSetup.get_error_information(both_exception)
|
||||
assert result["error_code"] == "400" # Should prefer 'code' over 'status_code'
|
||||
|
||||
# Test case 4: Exception with 'code' as empty string - should fall back to 'status_code'
|
||||
empty_code_exception = BothAttributesException(
|
||||
code="", status_code=404, message="Not Found"
|
||||
)
|
||||
empty_code_exception = BothAttributesException(code="", status_code=404, message="Not Found")
|
||||
result = StandardLoggingPayloadSetup.get_error_information(empty_code_exception)
|
||||
assert result["error_code"] == "404" # Should fall back to status_code
|
||||
|
||||
# Test case 5: Exception with 'code' as "None" string - should fall back to 'status_code'
|
||||
none_string_exception = BothAttributesException(
|
||||
code="None", status_code=503, message="Service Unavailable"
|
||||
)
|
||||
none_string_exception = BothAttributesException(code="None", status_code=503, message="Service Unavailable")
|
||||
result = StandardLoggingPayloadSetup.get_error_information(none_string_exception)
|
||||
assert result["error_code"] == "503" # Should fall back to status_code
|
||||
|
||||
# Test case 6: Exception with 'code' as None - should fall back to 'status_code'
|
||||
none_code_exception = BothAttributesException(
|
||||
code=None, status_code=401, message="Unauthorized"
|
||||
)
|
||||
none_code_exception = BothAttributesException(code=None, status_code=401, message="Unauthorized")
|
||||
result = StandardLoggingPayloadSetup.get_error_information(none_code_exception)
|
||||
assert result["error_code"] == "401" # Should fall back to status_code
|
||||
|
||||
|
|
@ -2863,9 +2742,7 @@ def test_get_error_information_prefers_message_attribute_over_str():
|
|||
)
|
||||
|
||||
result = StandardLoggingPayloadSetup.get_error_information(exc)
|
||||
assert (
|
||||
result["error_message"] == msg
|
||||
), f"expected message from .message attribute, got {result['error_message']!r}"
|
||||
assert result["error_message"] == msg, f"expected message from .message attribute, got {result['error_message']!r}"
|
||||
assert result["error_code"] == "401"
|
||||
assert result["error_class"] == "ProxyExceptionLike"
|
||||
|
||||
|
|
@ -2940,8 +2817,7 @@ def test_get_error_information_preserves_explicit_empty_message():
|
|||
exc = ProxyExceptionLike(message="", code=500)
|
||||
result = StandardLoggingPayloadSetup.get_error_information(exc)
|
||||
assert result["error_message"] == "", (
|
||||
"explicit empty .message must survive verbatim; got "
|
||||
f"{result['error_message']!r}"
|
||||
f"explicit empty .message must survive verbatim; got {result['error_message']!r}"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -3204,9 +3080,7 @@ def test_process_hidden_params_recalculates_cost_after_failure_handler_zero():
|
|||
choices=[{"message": {"role": "assistant", "content": "ok"}}],
|
||||
usage=Usage(prompt_tokens=9698, completion_tokens=30, total_tokens=9728),
|
||||
)
|
||||
logging_obj._process_hidden_params_and_response_cost(
|
||||
result, datetime.now(), datetime.now()
|
||||
)
|
||||
logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now())
|
||||
|
||||
cost = logging_obj.model_call_details.get("response_cost")
|
||||
assert cost is not None and cost > 0
|
||||
|
|
@ -3230,9 +3104,7 @@ def test_process_hidden_params_preserves_zero_cost_in_hidden_params():
|
|||
litellm_call_id="test-hidden-zero-cost",
|
||||
function_id="test-hidden-zero-cost",
|
||||
)
|
||||
logging_obj.model_call_details["litellm_params"] = {
|
||||
"model": "gemini-2.5-flash-lite"
|
||||
}
|
||||
logging_obj.model_call_details["litellm_params"] = {"model": "gemini-2.5-flash-lite"}
|
||||
logging_obj.optional_params = {}
|
||||
|
||||
result = ModelResponse(
|
||||
|
|
@ -3242,9 +3114,7 @@ def test_process_hidden_params_preserves_zero_cost_in_hidden_params():
|
|||
)
|
||||
result._hidden_params = {"response_cost": 0.0}
|
||||
|
||||
logging_obj._process_hidden_params_and_response_cost(
|
||||
result, datetime.now(), datetime.now()
|
||||
)
|
||||
logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now())
|
||||
|
||||
assert logging_obj.model_call_details.get("response_cost") == 0.0
|
||||
slo = logging_obj.model_call_details.get("standard_logging_object") or {}
|
||||
|
|
@ -3293,9 +3163,7 @@ def test_process_hidden_params_uses_hidden_params_cost_after_failure_handler_zer
|
|||
)
|
||||
result._hidden_params = {"response_cost": passthrough_cost}
|
||||
|
||||
logging_obj._process_hidden_params_and_response_cost(
|
||||
result, datetime.now(), datetime.now()
|
||||
)
|
||||
logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now())
|
||||
|
||||
assert logging_obj.model_call_details.get("response_cost") == passthrough_cost
|
||||
slo = logging_obj.model_call_details.get("standard_logging_object") or {}
|
||||
|
|
@ -3352,9 +3220,7 @@ def test_function_setup_litellm_metadata_populates_metadata():
|
|||
assert litellm_metadata.get("user_api_key_hash") == test_api_key_hash
|
||||
|
||||
# metadata should be a COPY, not an alias — mutating one must not affect the other
|
||||
assert (
|
||||
metadata is not litellm_metadata
|
||||
), "litellm_params['metadata'] should be a copy, not the same object"
|
||||
assert metadata is not litellm_metadata, "litellm_params['metadata'] should be a copy, not the same object"
|
||||
|
||||
|
||||
def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup():
|
||||
|
|
@ -3399,9 +3265,9 @@ def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup():
|
|||
litellm_params = logging_obj.model_call_details.get("litellm_params", {})
|
||||
litellm_metadata = litellm_params.get("litellm_metadata")
|
||||
assert litellm_metadata is not None
|
||||
assert litellm_metadata.get("standard_logging_guardrail_information") == [
|
||||
guardrail_entry
|
||||
], "guardrail writes after function_setup must be visible to the logging object"
|
||||
assert litellm_metadata.get("standard_logging_guardrail_information") == [guardrail_entry], (
|
||||
"guardrail writes after function_setup must be visible to the logging object"
|
||||
)
|
||||
assert litellm_metadata.get("applied_guardrails") == ["pam-ethical-request"]
|
||||
|
||||
merged = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params)
|
||||
|
|
@ -3570,9 +3436,7 @@ def test_failure_handler_skips_sync_callbacks_for_pass_through_requests(logging_
|
|||
|
||||
|
||||
@pytest.mark.parametrize("call_type", ["completion", "acompletion"])
|
||||
def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests(
|
||||
logging_obj, call_type
|
||||
):
|
||||
def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests(logging_obj, call_type):
|
||||
"""Ensure sync failure callbacks still fire for normal (non-pass-through) requests."""
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
|
@ -3733,9 +3597,7 @@ def test_standard_logging_hidden_params_backfills_response_cost_without_mutating
|
|||
)
|
||||
response._hidden_params = {"response_cost": None, "model_id": "mid-test"}
|
||||
|
||||
payload = logging_obj._build_standard_logging_payload(
|
||||
response, datetime.now(), datetime.now()
|
||||
)
|
||||
payload = logging_obj._build_standard_logging_payload(response, datetime.now(), datetime.now())
|
||||
|
||||
assert payload is not None
|
||||
assert payload["hidden_params"]["response_cost"] == 0.002
|
||||
|
|
@ -3789,10 +3651,7 @@ def test_merge_hidden_params_from_response_into_metadata_no_op_when_empty():
|
|||
_hidden_params = {}
|
||||
|
||||
logging_obj._merge_hidden_params_from_response_into_metadata(_NoHp())
|
||||
assert (
|
||||
"hidden_params"
|
||||
not in logging_obj.model_call_details["litellm_params"]["metadata"]
|
||||
)
|
||||
assert "hidden_params" not in logging_obj.model_call_details["litellm_params"]["metadata"]
|
||||
|
||||
|
||||
# ── StandardLoggingPayloadSetup.get_additional_headers ───────────────────────
|
||||
|
|
@ -3870,6 +3729,82 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
|
|||
assert payload["litellm_call_id"] == call_id
|
||||
|
||||
|
||||
# ── Azure Model Router selected-model attribution ────────────────────────────
|
||||
|
||||
|
||||
def _model_router_response(selected_model: str, stamp: bool):
|
||||
"""A ModelResponse as AzureModelRouterConfig hands it back, with or without the stamp."""
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
response = ModelResponse(model=selected_model)
|
||||
response._hidden_params = {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {}
|
||||
return response
|
||||
|
||||
|
||||
def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj):
|
||||
"""
|
||||
The selected model must win off the stamp, not off "model-router" appearing in the
|
||||
requested model. An operator whose model group is named anything else was invisible
|
||||
to the name check, so their logs and spend rows named the router instead.
|
||||
"""
|
||||
import datetime
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
get_standard_logging_object_payload,
|
||||
)
|
||||
|
||||
now = datetime.datetime.now()
|
||||
payload = get_standard_logging_object_payload(
|
||||
kwargs={
|
||||
"model": "azure_ai/smart-pick",
|
||||
"custom_llm_provider": "azure_ai",
|
||||
"messages": [],
|
||||
"litellm_params": {"metadata": {}},
|
||||
},
|
||||
init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=True),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
logging_obj=logging_obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["model"] == "azure_ai/grok-4-1-fast-reasoning"
|
||||
|
||||
|
||||
def test_standard_logging_payload_keeps_requested_model_without_router_stamp(logging_obj):
|
||||
"""
|
||||
Control for the test above: an ordinary azure_ai deployment is unaffected, so the stamp
|
||||
is what redirects attribution rather than the response model winning unconditionally.
|
||||
"""
|
||||
import datetime
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
get_standard_logging_object_payload,
|
||||
)
|
||||
|
||||
now = datetime.datetime.now()
|
||||
payload = get_standard_logging_object_payload(
|
||||
kwargs={
|
||||
"model": "azure_ai/smart-pick",
|
||||
"custom_llm_provider": "azure_ai",
|
||||
"messages": [],
|
||||
"litellm_params": {"metadata": {}},
|
||||
},
|
||||
init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=False),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
logging_obj=logging_obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["model"] == "azure_ai/smart-pick"
|
||||
|
||||
|
||||
def _make_dict_logging_obj():
|
||||
"""Build a Logging instance configured for a non-streaming dict result."""
|
||||
obj = LitellmLogging(
|
||||
|
|
@ -3905,9 +3840,7 @@ def test_success_handler_computes_cost_for_dict_response():
|
|||
"_build_standard_logging_payload",
|
||||
return_value={"response_cost": expected_cost},
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
|
||||
),
|
||||
patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_is_recognized_call_type_for_logging",
|
||||
|
|
@ -3944,9 +3877,7 @@ def test_success_handler_preserves_precomputed_cost_for_dict_response():
|
|||
"_build_standard_logging_payload",
|
||||
return_value={"response_cost": precomputed_cost},
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
|
||||
),
|
||||
patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_is_recognized_call_type_for_logging",
|
||||
|
|
@ -3985,9 +3916,7 @@ def test_success_handler_unified_helper_runs_for_typed_results():
|
|||
"_build_standard_logging_payload",
|
||||
return_value={"response_cost": expected_cost},
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
|
||||
),
|
||||
patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_is_recognized_call_type_for_logging",
|
||||
|
|
@ -4042,9 +3971,7 @@ class TestFirstApiCallStartTimeSetOnce:
|
|||
assert first == obj.model_call_details["api_call_start_time"]
|
||||
# Set on the logging object only — user metadata untouched.
|
||||
assert user_meta == {}
|
||||
assert (
|
||||
"first_api_call_start_time" not in obj.model_call_details["litellm_params"]
|
||||
)
|
||||
assert "first_api_call_start_time" not in obj.model_call_details["litellm_params"]
|
||||
|
||||
time.sleep(0.002) # ensure a distinct retry timestamp
|
||||
obj.pre_call(input="hi", api_key="sk-test")
|
||||
|
|
@ -4061,18 +3988,16 @@ def test_get_error_information_for_logging_payload_ignores_spoofed_disconnect_wi
|
|||
baseline = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=ValueError("provider failure"),
|
||||
)
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={
|
||||
"error_information": {
|
||||
"error_code": "499",
|
||||
"error_message": "Client disconnected the request",
|
||||
"error_class": "ClientDisconnected",
|
||||
}
|
||||
},
|
||||
original_exception=ValueError("provider failure"),
|
||||
error_str="provider failure",
|
||||
)
|
||||
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={
|
||||
"error_information": {
|
||||
"error_code": "499",
|
||||
"error_message": "Client disconnected the request",
|
||||
"error_class": "ClientDisconnected",
|
||||
}
|
||||
},
|
||||
original_exception=ValueError("provider failure"),
|
||||
error_str="provider failure",
|
||||
)
|
||||
assert error_information == baseline
|
||||
assert error_str == "provider failure"
|
||||
|
|
@ -4086,22 +4011,18 @@ def test_get_error_information_for_logging_payload_client_disconnect():
|
|||
"error_message": "Client disconnected the request",
|
||||
"error_class": "ClientDisconnected",
|
||||
}
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={"client_disconnected": True, "error_information": custom_error},
|
||||
original_exception=None,
|
||||
error_str=None,
|
||||
)
|
||||
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={"client_disconnected": True, "error_information": custom_error},
|
||||
original_exception=None,
|
||||
error_str=None,
|
||||
)
|
||||
assert error_information == custom_error
|
||||
assert error_str == "Client disconnected the request"
|
||||
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={"client_disconnected": True},
|
||||
original_exception=None,
|
||||
error_str="existing error",
|
||||
)
|
||||
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={"client_disconnected": True},
|
||||
original_exception=None,
|
||||
error_str="existing error",
|
||||
)
|
||||
assert error_information["error_code"] == "499"
|
||||
assert error_str == "existing error"
|
||||
|
|
@ -4109,12 +4030,10 @@ def test_get_error_information_for_logging_payload_client_disconnect():
|
|||
baseline = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=None,
|
||||
)
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={},
|
||||
original_exception=None,
|
||||
error_str=None,
|
||||
)
|
||||
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={},
|
||||
original_exception=None,
|
||||
error_str=None,
|
||||
)
|
||||
assert error_information == baseline
|
||||
assert error_str is None
|
||||
|
|
@ -4149,9 +4068,7 @@ def test_get_error_information_prefers_message_attribute_over_empty_str():
|
|||
def __str__(self):
|
||||
return ""
|
||||
|
||||
info = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=_SilentExc()
|
||||
)
|
||||
info = StandardLoggingPayloadSetup.get_error_information(original_exception=_SilentExc())
|
||||
assert info["error_message"] == "real failure detail"
|
||||
assert info["error_code"] == "401"
|
||||
|
||||
|
|
@ -4182,9 +4099,7 @@ def _responses_api_response_with_text(text="hello world"):
|
|||
type="message",
|
||||
role="assistant",
|
||||
status="completed",
|
||||
content=[
|
||||
ResponseOutputText(annotations=[], text=text, type="output_text")
|
||||
],
|
||||
content=[ResponseOutputText(annotations=[], text=text, type="output_text")],
|
||||
)
|
||||
],
|
||||
usage=ResponseAPIUsage(input_tokens=11, output_tokens=7, total_tokens=18),
|
||||
|
|
@ -4199,9 +4114,7 @@ def _responses_api_response_with_text(text="hello world"):
|
|||
("ResponseFailedEvent", "response.failed"),
|
||||
],
|
||||
)
|
||||
def test_handle_anthropic_messages_response_logging_translates_terminal_responses_api_event(
|
||||
event_cls, event_type
|
||||
):
|
||||
def test_handle_anthropic_messages_response_logging_translates_terminal_responses_api_event(event_cls, event_type):
|
||||
"""Regression for #28595 / #28943. When anthropic_messages routes to the OpenAI
|
||||
Responses backend and stream=True, success_handler receives a terminal Responses
|
||||
API event. The handler must translate it to a ModelResponse whose choices carry
|
||||
|
|
@ -4240,10 +4153,7 @@ def test_handle_anthropic_messages_response_logging_passes_model_response_throug
|
|||
"""Anthropic-native path already yields a ModelResponse; it must be returned unchanged."""
|
||||
logging_obj = _anthropic_messages_logging_obj()
|
||||
model_response = ModelResponse()
|
||||
assert (
|
||||
logging_obj._handle_anthropic_messages_response_logging(result=model_response)
|
||||
is model_response
|
||||
)
|
||||
assert logging_obj._handle_anthropic_messages_response_logging(result=model_response) is model_response
|
||||
|
||||
|
||||
def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_responses_payload():
|
||||
|
|
@ -4539,9 +4449,7 @@ def test_non_image_response_has_no_output_image_count(logging_obj):
|
|||
|
||||
def test_zero_token_video_usage_preserves_duration_seconds(logging_obj):
|
||||
"""Video usage bills by duration; the payload must keep duration_seconds even with zero tokens."""
|
||||
payload = _build_payload_for_media_response(
|
||||
logging_obj, {"id": "video-1", "usage": {"duration_seconds": 4.0}}
|
||||
)
|
||||
payload = _build_payload_for_media_response(logging_obj, {"id": "video-1", "usage": {"duration_seconds": 4.0}})
|
||||
|
||||
assert payload is not None
|
||||
assert payload["metadata"]["usage_object"]["duration_seconds"] == 4.0
|
||||
|
|
|
|||
|
|
@ -5,9 +5,7 @@ from unittest.mock import MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path
|
||||
from litellm.llms.azure_ai.azure_model_router.transformation import (
|
||||
AzureModelRouterConfig,
|
||||
)
|
||||
|
|
@ -120,9 +118,7 @@ def test_azure_ai_grok_stop_parameter_handling():
|
|||
# Test supported parameters for Grok models
|
||||
for model in ("grok-4-fast", "grok-4.3"):
|
||||
grok_params = config.get_supported_openai_params(model)
|
||||
assert (
|
||||
"stop" not in grok_params
|
||||
), "Grok models should not support stop parameter"
|
||||
assert "stop" not in grok_params, "Grok models should not support stop parameter"
|
||||
|
||||
# Test supported parameters for non-Grok models
|
||||
gpt_params = config.get_supported_openai_params("gpt-4")
|
||||
|
|
@ -201,11 +197,84 @@ def test_azure_model_router_response_shows_actual_model():
|
|||
|
||||
# Verify that the response contains the actual model used, not the router model
|
||||
assert result.model == "azure_ai/gpt-5-nano-2025-08-07", (
|
||||
f"Expected model to be 'azure_ai/gpt-5-nano-2025-08-07' (actual model used), "
|
||||
f"but got '{result.model}'"
|
||||
f"Expected model to be 'azure_ai/gpt-5-nano-2025-08-07' (actual model used), but got '{result.model}'"
|
||||
)
|
||||
|
||||
|
||||
def test_azure_model_router_stamps_selected_model_on_hidden_params():
|
||||
"""
|
||||
The selected model must be stamped on _hidden_params, not left for downstream code to
|
||||
re-derive by looking for "model-router" in the model string. Deployments whose alias
|
||||
does not contain that text are invisible to the string check.
|
||||
"""
|
||||
from httpx import Response
|
||||
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
|
||||
AzureFoundryModelInfo,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
raw_response_json = {
|
||||
"id": "chatcmpl-test456",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "grok-4-1-fast-reasoning",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "pong"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
|
||||
mock_response = MagicMock(spec=Response)
|
||||
mock_response.json.return_value = raw_response_json
|
||||
mock_response.text = json.dumps(raw_response_json)
|
||||
mock_response.headers = {}
|
||||
|
||||
logging_obj = MagicMock(spec=LiteLLMLoggingObj)
|
||||
logging_obj.post_call = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
result = AzureModelRouterConfig().transform_response(
|
||||
model="smart-pick",
|
||||
raw_response=mock_response,
|
||||
model_response=ModelResponse(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={},
|
||||
messages=[{"role": "user", "content": "Reply with just pong"}],
|
||||
optional_params={},
|
||||
litellm_params={"model": "azure_ai/model_router/smart-pick"},
|
||||
encoding=None,
|
||||
api_key="test-key",
|
||||
json_mode=False,
|
||||
)
|
||||
|
||||
assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == result.model
|
||||
assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == "azure_ai/grok-4-1-fast-reasoning"
|
||||
assert AzureFoundryModelInfo.get_model_router_selected_model(result._hidden_params) == (
|
||||
"azure_ai/grok-4-1-fast-reasoning"
|
||||
)
|
||||
assert AzureFoundryModelInfo.is_model_router_call(model="smart-pick", hidden_params=result._hidden_params) is True
|
||||
|
||||
|
||||
def test_azure_model_router_stamp_does_not_leak_across_responses():
|
||||
"""
|
||||
ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written
|
||||
as a fresh dict. Mutating in place would bleed the selected model into unrelated responses.
|
||||
"""
|
||||
from litellm.llms.azure_ai.common_utils import AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
untouched = ModelResponse()
|
||||
|
||||
assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {})
|
||||
|
||||
|
||||
def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name():
|
||||
"""
|
||||
Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name.
|
||||
|
|
@ -226,14 +295,10 @@ def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name():
|
|||
mock_response.text = error_text
|
||||
mock_response.json.return_value = json.loads(error_text)
|
||||
mock_response.status_code = 400
|
||||
e = httpx.HTTPStatusError(
|
||||
message="400", request=MagicMock(), response=mock_response
|
||||
)
|
||||
e = httpx.HTTPStatusError(message="400", request=MagicMock(), response=mock_response)
|
||||
|
||||
assert config._error_has_tool_level_extra_fields(error_text) is True
|
||||
assert (
|
||||
config.should_retry_llm_api_inside_llm_translation_on_http_error(e, {}) is True
|
||||
)
|
||||
assert config.should_retry_llm_api_inside_llm_translation_on_http_error(e, {}) is True
|
||||
|
||||
request_data = {
|
||||
"model": "FW-Kimi-K2.6",
|
||||
|
|
@ -354,9 +419,7 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages():
|
|||
{
|
||||
"role": "assistant",
|
||||
"content": "I can help.",
|
||||
"thinking_blocks": [
|
||||
{"type": "thinking", "thinking": "Reading the file.", "signature": "sig"}
|
||||
],
|
||||
"thinking_blocks": [{"type": "thinking", "thinking": "Reading the file.", "signature": "sig"}],
|
||||
"provider_specific_fields": {"thought_signature": "sig-top"},
|
||||
"tool_calls": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -126,16 +126,12 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
assert json.loads(result.body) == guardrailed_body
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch):
|
||||
"""The guardrail JSON path must forward upstream response headers (e.g.
|
||||
x-amzn-requestid) alongside the x-litellm-* headers, matching the
|
||||
non-guardrail passthrough path, while dropping length headers that no
|
||||
longer match the rewritten body."""
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(
|
||||
data={"custom_llm_provider": "bedrock"}
|
||||
)
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
|
|
@ -175,14 +171,10 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
assert result.headers["content-length"] == str(len(result.body))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch):
|
||||
"""The guardrail event-stream branch must also forward upstream response
|
||||
headers alongside the x-litellm-* headers."""
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(
|
||||
data={"custom_llm_provider": "bedrock"}
|
||||
)
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
|
|
@ -224,15 +216,11 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
assert result.headers["x-litellm-call-id"] == "test-call-id"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(self, monkeypatch):
|
||||
"""Guardrailed non-streaming passthrough responses must include headers
|
||||
injected by post_call_response_headers_hook, matching the headers a
|
||||
non-guardrailed passthrough response would carry."""
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(
|
||||
data={"custom_llm_provider": "bedrock"}
|
||||
)
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
|
|
@ -251,9 +239,7 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
return kwargs["response"]
|
||||
|
||||
proxy_logging_obj.post_call_success_hook = fake_post_call_success_hook
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value={"x-litellm-custom": "from-hook"}
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-litellm-custom": "from-hook"})
|
||||
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=upstream,
|
||||
|
|
@ -2221,6 +2207,52 @@ class TestOverrideOpenAIResponseModel:
|
|||
assert response_obj.model == actual_model_used
|
||||
assert response_obj.model != requested_model
|
||||
|
||||
def test_override_model_preserves_model_router_model_for_alias_without_router_in_name(self):
|
||||
"""
|
||||
The client sends a model group alias, which carries no model_router/ prefix, so the
|
||||
name check alone only fires when the operator happened to put "model-router" in the
|
||||
alias. With the stamp on the response the actual model survives whatever it is named.
|
||||
"""
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
|
||||
)
|
||||
|
||||
requested_model = "smart-pick"
|
||||
actual_model_used = "azure_ai/grok-4-1-fast-reasoning"
|
||||
|
||||
response_obj = MagicMock()
|
||||
response_obj.model = actual_model_used
|
||||
response_obj._hidden_params = {
|
||||
"additional_headers": {},
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: actual_model_used,
|
||||
}
|
||||
|
||||
_override_openai_response_model(
|
||||
response_obj=response_obj,
|
||||
requested_model=requested_model,
|
||||
log_context="test_context",
|
||||
)
|
||||
assert response_obj.model == actual_model_used
|
||||
|
||||
def test_override_model_still_restamps_non_router_alias_without_stamp(self):
|
||||
"""
|
||||
Control for the test above: absent the stamp, an ordinary deployment keeps being
|
||||
restamped to the requested model, so the stamp is doing the work rather than the
|
||||
preserve branch having gone unconditional.
|
||||
"""
|
||||
requested_model = "smart-pick"
|
||||
|
||||
response_obj = MagicMock()
|
||||
response_obj.model = "azure_ai/grok-4-1-fast-reasoning"
|
||||
response_obj._hidden_params = {"additional_headers": {}}
|
||||
|
||||
_override_openai_response_model(
|
||||
response_obj=response_obj,
|
||||
requested_model=requested_model,
|
||||
log_context="test_context",
|
||||
)
|
||||
assert response_obj.model == requested_model
|
||||
|
||||
def test_override_model_uses_winning_model_for_fastest_response(self):
|
||||
"""
|
||||
Test that when fastest_response batch completion is used with a
|
||||
|
|
@ -2793,9 +2825,7 @@ class TestStreamCloseOnDisconnect:
|
|||
finally:
|
||||
closed.set()
|
||||
|
||||
response = _UpstreamClosingStreamingResponse(
|
||||
body(), media_type="text/event-stream"
|
||||
)
|
||||
response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream")
|
||||
|
||||
async def receive():
|
||||
await asyncio.Event().wait()
|
||||
|
|
@ -2826,9 +2856,7 @@ class TestStreamCloseOnDisconnect:
|
|||
finally:
|
||||
closed.set()
|
||||
|
||||
response = _UpstreamClosingStreamingResponse(
|
||||
body(), media_type="text/event-stream"
|
||||
)
|
||||
response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream")
|
||||
|
||||
async def receive():
|
||||
await disconnected.wait()
|
||||
|
|
@ -2899,9 +2927,7 @@ class TestStreamCloseOnDisconnect:
|
|||
finally:
|
||||
inner_closed.set()
|
||||
|
||||
response = await create_response(
|
||||
generator=wrapped(), media_type="text/event-stream", headers={}
|
||||
)
|
||||
response = await create_response(generator=wrapped(), media_type="text/event-stream", headers={})
|
||||
|
||||
async def receive():
|
||||
await asyncio.Event().wait()
|
||||
|
|
@ -3097,9 +3123,7 @@ class TestStreamCloseOnDisconnect:
|
|||
|
||||
with pytest.raises(_ClientDisconnectedBeforeFirstChunk):
|
||||
await asyncio.wait_for(
|
||||
_buffer_first_chunk_honoring_disconnect(
|
||||
AcloseRaises(), request=self._request_that_disconnects()
|
||||
),
|
||||
_buffer_first_chunk_honoring_disconnect(AcloseRaises(), request=self._request_that_disconnects()),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
|
@ -3115,9 +3139,7 @@ class TestStreamCloseOnDisconnect:
|
|||
|
||||
with pytest.raises(_ClientDisconnectedBeforeFirstChunk):
|
||||
await asyncio.wait_for(
|
||||
_buffer_first_chunk_honoring_disconnect(
|
||||
blocking_gen(), request=self._request_that_disconnects()
|
||||
),
|
||||
_buffer_first_chunk_honoring_disconnect(blocking_gen(), request=self._request_that_disconnects()),
|
||||
timeout=5,
|
||||
)
|
||||
assert closed.is_set()
|
||||
|
|
@ -3133,9 +3155,7 @@ class TestHandleLLMApiExceptionRetryAfter:
|
|||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test")
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value=callback_headers or {}
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {})
|
||||
|
||||
try:
|
||||
await processor._handle_llm_api_exception(
|
||||
|
|
@ -3187,9 +3207,7 @@ class TestHandleLLMApiExceptionRetryAfter:
|
|||
enable_pre_call_checks=False,
|
||||
cooldown_list=[],
|
||||
)
|
||||
proxy_exc = await self._invoke(
|
||||
exc, callback_headers={"retry-after": "", "x-custom": "1"}
|
||||
)
|
||||
proxy_exc = await self._invoke(exc, callback_headers={"retry-after": "", "x-custom": "1"})
|
||||
assert proxy_exc.headers["retry-after"] == "43"
|
||||
assert proxy_exc.headers["x-custom"] == "1"
|
||||
|
||||
|
|
@ -3385,9 +3403,7 @@ class TestDisconnectGatherCleanup:
|
|||
return Request(scope={"type": "http", "headers": []}, receive=receive)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_raises_499_on_client_disconnect(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_base_process_llm_request_raises_499_on_client_disconnect(self, monkeypatch):
|
||||
"""With cancel_on_disconnect enabled, base_process_llm_request returns 499."""
|
||||
import asyncio
|
||||
|
||||
|
|
@ -3416,9 +3432,7 @@ class TestDisconnectGatherCleanup:
|
|||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await processing_obj.base_process_llm_request(
|
||||
|
|
@ -3436,9 +3450,7 @@ class TestDisconnectGatherCleanup:
|
|||
assert "disconnected" in exc_info.value.detail.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(self, monkeypatch):
|
||||
import asyncio
|
||||
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
|
|
@ -3463,9 +3475,7 @@ class TestDisconnectGatherCleanup:
|
|||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
monkeypatch.setattr(
|
||||
cpr,
|
||||
"route_request",
|
||||
|
|
@ -3526,9 +3536,7 @@ class TestDisconnectGatherCleanup:
|
|||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await processing_obj.base_process_llm_request(
|
||||
|
|
@ -3579,9 +3587,7 @@ class TestDisconnectGatherCleanup:
|
|||
assert task.done()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_preserves_llm_error_after_gather(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_base_process_llm_request_preserves_llm_error_after_gather(self, monkeypatch):
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
|
|
@ -3610,9 +3616,7 @@ class TestDisconnectGatherCleanup:
|
|||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.is_disconnected = AsyncMock(return_value=False)
|
||||
|
|
@ -3649,19 +3653,13 @@ class TestStreamingClientDisconnectLogging:
|
|||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
|
||||
|
||||
assert recorded is True
|
||||
assert request_data["metadata"]["client_disconnected"] is True
|
||||
assert request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
assert (
|
||||
request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
)
|
||||
assert (
|
||||
mock_logging_obj.model_call_details["litellm_params"]["metadata"][
|
||||
"error_information"
|
||||
]["error_code"]
|
||||
mock_logging_obj.model_call_details["litellm_params"]["metadata"]["error_information"]["error_code"]
|
||||
== "499"
|
||||
)
|
||||
|
||||
|
|
@ -3675,9 +3673,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
mock_request.is_disconnected = AsyncMock(return_value=False)
|
||||
request_data = {"metadata": {}}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
|
||||
|
||||
assert recorded is False
|
||||
assert "client_disconnected" not in request_data["metadata"]
|
||||
|
|
@ -3702,22 +3698,12 @@ class TestStreamingClientDisconnectLogging:
|
|||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
|
||||
|
||||
assert recorded is True
|
||||
assert request_data["metadata"]["client_disconnected"] is True
|
||||
assert (
|
||||
mock_logging_obj.model_call_details["litellm_params"]["metadata"][
|
||||
"client_disconnected"
|
||||
]
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
mock_logging_obj.model_call_details["metadata"]["client_disconnected"]
|
||||
is True
|
||||
)
|
||||
assert mock_logging_obj.model_call_details["litellm_params"]["metadata"]["client_disconnected"] is True
|
||||
assert mock_logging_obj.model_call_details["metadata"]["client_disconnected"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_streaming_client_disconnect_handles_none_request_data_metadata(self):
|
||||
|
|
@ -3733,15 +3719,11 @@ class TestStreamingClientDisconnectLogging:
|
|||
"litellm_params": {"metadata": None},
|
||||
}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
|
||||
|
||||
assert recorded is True
|
||||
assert request_data["metadata"]["client_disconnected"] is True
|
||||
assert (
|
||||
request_data["litellm_params"]["metadata"]["client_disconnected"] is True
|
||||
)
|
||||
assert request_data["litellm_params"]["metadata"]["client_disconnected"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_client_disconnect_metadata_none_returns_early(self):
|
||||
|
|
@ -3752,9 +3734,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
_apply_client_disconnect_metadata(None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(self, monkeypatch):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
|
@ -3786,9 +3766,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
assert request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(self, monkeypatch):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
|
@ -3818,9 +3796,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
assert "client_disconnected" not in request_data["metadata"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_streaming_data_generator_records_499_on_early_aclose(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_async_streaming_data_generator_records_499_on_early_aclose(self, monkeypatch):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
|
@ -3835,9 +3811,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
yield {"choices": [{"delta": {"content": " there"}}]}
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.async_post_call_streaming_iterator_hook = (
|
||||
mock_streaming_iterator
|
||||
)
|
||||
mock_proxy_logging.async_post_call_streaming_iterator_hook = mock_streaming_iterator
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
|
|
@ -3848,9 +3822,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
"model": "gemini-2.0-flash",
|
||||
"metadata": {},
|
||||
"litellm_params": {"metadata": {}},
|
||||
"litellm_logging_obj": MagicMock(
|
||||
model_call_details={"metadata": {}, "litellm_params": {}}
|
||||
),
|
||||
"litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}),
|
||||
}
|
||||
|
||||
gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
|
|
@ -3869,6 +3841,8 @@ class TestStreamingClientDisconnectLogging:
|
|||
assert request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
|
||||
class TestCancelOnDisconnect:
|
||||
"""
|
||||
Coverage for the opt-in `general_settings.cancel_on_disconnect` flag:
|
||||
|
|
@ -3895,23 +3869,17 @@ class TestCancelOnDisconnect:
|
|||
llm_call = asyncio.get_running_loop().create_future()
|
||||
disconnect_event = asyncio.Event()
|
||||
|
||||
await _cancel_llm_call_on_client_disconnect(
|
||||
request, llm_call, disconnect_event
|
||||
)
|
||||
await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
|
||||
|
||||
assert llm_call.cancelled()
|
||||
assert disconnect_event.is_set()
|
||||
|
||||
async def test_monitor_is_noop_while_client_stays_connected(self):
|
||||
request = self._request(
|
||||
[{"type": "http.request", "body": b"", "more_body": False}]
|
||||
)
|
||||
request = self._request([{"type": "http.request", "body": b"", "more_body": False}])
|
||||
llm_call = asyncio.get_running_loop().create_future()
|
||||
disconnect_event = asyncio.Event()
|
||||
|
||||
monitor = asyncio.create_task(
|
||||
_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
|
||||
)
|
||||
monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert not monitor.done()
|
||||
|
|
@ -3930,9 +3898,7 @@ class TestCancelOnDisconnect:
|
|||
llm_call = asyncio.get_running_loop().create_future()
|
||||
disconnect_event = asyncio.Event()
|
||||
|
||||
await _cancel_llm_call_on_client_disconnect(
|
||||
request, llm_call, disconnect_event
|
||||
)
|
||||
await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
|
||||
|
||||
assert not llm_call.cancelled()
|
||||
assert not disconnect_event.is_set()
|
||||
|
|
@ -3947,9 +3913,7 @@ class TestCancelOnDisconnect:
|
|||
with pytest.raises(asyncio.CancelledError):
|
||||
await _await_llm_call_cancelling_on_disconnect(request, llm_call)
|
||||
|
||||
async def _drive_base_process_llm_request(
|
||||
self, monkeypatch, general_settings: dict, llm_call, request: Request
|
||||
):
|
||||
async def _drive_base_process_llm_request(self, monkeypatch, general_settings: dict, llm_call, request: Request):
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
|
@ -3958,9 +3922,7 @@ class TestCancelOnDisconnect:
|
|||
logging_obj._on_deferred_stream_complete = None
|
||||
logging_obj.cost_breakdown = None
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(
|
||||
data={"model": "fake-model", "litellm_logging_obj": logging_obj}
|
||||
)
|
||||
processor = ProxyBaseLLMRequestProcessing(data={"model": "fake-model", "litellm_logging_obj": logging_obj})
|
||||
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
|
||||
|
|
@ -3968,9 +3930,7 @@ class TestCancelOnDisconnect:
|
|||
proxy_logging_obj.post_call_success_hook = AsyncMock(
|
||||
side_effect=lambda data, user_api_key_dict, response: response
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def fake_route_request(**kwargs):
|
||||
return llm_call()
|
||||
|
|
@ -4049,9 +4009,7 @@ class TestCancelOnDisconnect:
|
|||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await processor._handle_llm_api_exception(
|
||||
e=HTTPException(
|
||||
status_code=499, detail="Client disconnected the request"
|
||||
),
|
||||
e=HTTPException(status_code=499, detail="Client disconnected the request"),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
|
@ -4117,7 +4075,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
|
|||
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
|
||||
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", capture_hook)
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=httpx_response,
|
||||
|
|
@ -4167,7 +4127,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
|
|||
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
|
||||
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", non_dict_hook)
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=httpx_response,
|
||||
|
|
@ -4205,7 +4167,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
|
|||
hook_spy = AsyncMock()
|
||||
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy)
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=httpx_response,
|
||||
|
|
@ -4246,7 +4210,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
|
|||
hook_spy = AsyncMock()
|
||||
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy)
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=httpx_response,
|
||||
|
|
@ -4358,7 +4324,9 @@ class TestEventStreamAllmPassthroughRoute:
|
|||
"content-length": "99",
|
||||
}
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=mock_response,
|
||||
|
|
@ -4389,9 +4357,7 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
de-anonymized.
|
||||
"""
|
||||
|
||||
def _build_processing_obj(
|
||||
self, custom_llm_provider: str, endpoint: str = ""
|
||||
) -> ProxyBaseLLMRequestProcessing:
|
||||
def _build_processing_obj(self, custom_llm_provider: str, endpoint: str = "") -> ProxyBaseLLMRequestProcessing:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "call-123"
|
||||
logging_obj.cost_breakdown = None
|
||||
|
|
@ -4442,14 +4408,17 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
processing_obj = self._build_processing_obj("anthropic")
|
||||
chunks = [b"chunk-1", b"chunk-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
),
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
|
|
@ -4458,27 +4427,27 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
assert streamed == chunks
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_converse_stream_is_buffered_through_handler(
|
||||
self, monkeypatch
|
||||
):
|
||||
processing_obj = self._build_processing_obj(
|
||||
"bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream"
|
||||
)
|
||||
async def test_bedrock_converse_stream_is_buffered_through_handler(self, monkeypatch):
|
||||
processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream")
|
||||
chunks = [b"raw-1", b"raw-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
), patch(
|
||||
"litellm.llms.bedrock.passthrough.guardrail_translation.handler."
|
||||
"BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
|
||||
new=AsyncMock(return_value=b"modified-body"),
|
||||
) as mock_handler:
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"litellm.llms.bedrock.passthrough.guardrail_translation.handler."
|
||||
"BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
|
||||
new=AsyncMock(return_value=b"modified-body"),
|
||||
) as mock_handler,
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
assert isinstance(result, Response)
|
||||
|
|
@ -4494,19 +4463,23 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
)
|
||||
chunks = [b"raw-1", b"raw-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
), patch(
|
||||
"litellm.llms.bedrock.passthrough.guardrail_translation.handler."
|
||||
"BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
|
||||
new=AsyncMock(return_value=b"modified-body"),
|
||||
) as mock_handler:
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"litellm.llms.bedrock.passthrough.guardrail_translation.handler."
|
||||
"BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
|
||||
new=AsyncMock(return_value=b"modified-body"),
|
||||
) as mock_handler,
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
assert isinstance(result, StreamingResponse)
|
||||
|
|
@ -4528,14 +4501,17 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
)
|
||||
chunks = [b"raw-1", b"raw-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=False,
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=False,
|
||||
),
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
|
|
@ -4554,14 +4530,17 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
processing_obj = self._build_processing_obj("anthropic")
|
||||
chunks = [b"chunk-1", b"chunk-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=False,
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=False,
|
||||
),
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
|
|
@ -4902,7 +4881,6 @@ class TestResponseCostHeaderForTypedDictResponses:
|
|||
|
||||
|
||||
class TestPreCallWithFallbacksOnLocalRateLimit:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_triggered_on_local_rate_limit(self):
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
|
|
@ -5054,9 +5032,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}]
|
||||
|
||||
user_api_key_dict = MagicMock()
|
||||
user_api_key_dict.router_settings = {
|
||||
"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]
|
||||
}
|
||||
user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]}
|
||||
|
||||
with patch.object(
|
||||
processor,
|
||||
|
|
@ -5087,9 +5063,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(
|
||||
data={"model": "gpt-4", "disable_fallbacks": True}
|
||||
)
|
||||
processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True})
|
||||
|
||||
async def mock_pre_call_logic(**kwargs):
|
||||
raise ProxyRateLimitError(
|
||||
|
|
@ -5215,9 +5189,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
|
||||
# Real per-key per-model TPM limiter + a key carrying the customer's
|
||||
# `model_tpm_limit` metadata (only the primary is capped).
|
||||
limiter = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache())
|
||||
)
|
||||
limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="sk-lit3890",
|
||||
metadata={"model_tpm_limit": {primary_model: 100}},
|
||||
|
|
@ -5225,10 +5197,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
|
||||
# Pre-seed the primary's per-model token counter at the cap so the very
|
||||
# next request trips it. The counter key uses the *hashed* api_key.
|
||||
counter_key = (
|
||||
f"{user_api_key_dict.api_key}::{primary_model}"
|
||||
f"::{precise_minute}::request_count"
|
||||
)
|
||||
counter_key = f"{user_api_key_dict.api_key}::{primary_model}::{precise_minute}::request_count"
|
||||
await limiter.internal_usage_cache.async_set_cache(
|
||||
key=counter_key,
|
||||
value={"current_requests": 0, "current_tpm": 100, "current_rpm": 0},
|
||||
|
|
@ -5259,9 +5228,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
mock_router = MagicMock()
|
||||
mock_router.fallbacks = [{primary_model: [fallback_model]}]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock
|
||||
):
|
||||
with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock):
|
||||
with patch.object(
|
||||
processor,
|
||||
"common_processing_pre_call_logic",
|
||||
|
|
@ -5291,9 +5258,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
|
||||
# Sanity-check the premise: the limiter genuinely raises a
|
||||
# ProxyRateLimitError for the capped primary under the frozen clock.
|
||||
with patch(
|
||||
"litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock
|
||||
):
|
||||
with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock):
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await limiter.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -5654,16 +5619,12 @@ class TestStreamingClientDisconnectBilling:
|
|||
prompt_tokens=1000,
|
||||
completion_tokens=10,
|
||||
total_tokens=1010,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=500
|
||||
),
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
event = await self._bill_and_collect_success_event(
|
||||
append_openai_style_cached_usage_chunk
|
||||
)
|
||||
event = await self._bill_and_collect_success_event(append_openai_style_cached_usage_chunk)
|
||||
|
||||
usage = event["response_obj"].usage
|
||||
assert getattr(usage, "cache_read_input_tokens", None) == 500
|
||||
|
|
@ -6433,9 +6394,7 @@ class TestInjectCostIntoUsageDict:
|
|||
logging_obj.model_call_details["custom_llm_provider"] = "anthropic"
|
||||
assert logging_obj.cost_breakdown is None
|
||||
|
||||
model_response = ModelResponse(
|
||||
usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)
|
||||
)
|
||||
model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224))
|
||||
cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj)
|
||||
|
||||
assert cost is not None and cost > 0
|
||||
|
|
@ -6464,9 +6423,7 @@ class TestInjectCostIntoUsageDict:
|
|||
)
|
||||
existing = logging_obj.cost_breakdown
|
||||
|
||||
model_response = ModelResponse(
|
||||
usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)
|
||||
)
|
||||
model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224))
|
||||
ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj)
|
||||
|
||||
assert logging_obj.cost_breakdown is existing
|
||||
|
|
@ -6761,9 +6718,7 @@ def test_ttft_keepalive_interval_only_arms_for_a_streaming_request(request_data,
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream_requested, expect_ping", [(True, True), (False, False)])
|
||||
async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(
|
||||
stream_requested, expect_ping
|
||||
):
|
||||
async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(stream_requested, expect_ping):
|
||||
"""The wiring, not the helper: every route funnels through this method, and the
|
||||
whole time-to-first-token is spent inside the call it wraps."""
|
||||
|
||||
|
|
@ -6909,9 +6864,7 @@ async def test_a_late_failure_is_reported_to_the_failure_hook():
|
|||
async def record(exc):
|
||||
audited.append(exc)
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=record
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=record)
|
||||
collected = await _drain(response)
|
||||
|
||||
assert [type(exc).__name__ for exc in audited] == ["HTTPException"]
|
||||
|
|
@ -6928,9 +6881,7 @@ async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame():
|
|||
async def broken_hook(exc):
|
||||
raise RuntimeError("the audit backend is down")
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook)
|
||||
collected = await _drain(response)
|
||||
|
||||
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
|
||||
|
|
@ -6984,9 +6935,7 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke
|
|||
[(0, False), (None, True)],
|
||||
ids=["operator-hard-disabled-this-deployment", "deployment-says-nothing"],
|
||||
)
|
||||
async def test_base_process_llm_request_honours_a_deployment_hard_disable(
|
||||
deployment_keepalive, expect_ping
|
||||
):
|
||||
async def test_base_process_llm_request_honours_a_deployment_hard_disable(deployment_keepalive, expect_ping):
|
||||
"""`keepalive_seconds: 0` is documented as a disable a request cannot lift. The
|
||||
funnel has to hand its router to the gate for that to hold before the upstream
|
||||
has answered, since no deployment has served the request yet."""
|
||||
|
|
@ -7032,9 +6981,7 @@ async def test_a_hook_returning_a_replacement_decides_what_the_client_sees():
|
|||
async def sanitize(exc):
|
||||
return HTTPException(status_code=502, detail="upstream unavailable")
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize)
|
||||
collected = await _drain(response)
|
||||
|
||||
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
|
||||
|
|
@ -7073,9 +7020,7 @@ async def test_a_hook_that_returns_nothing_leaves_the_real_error_intact():
|
|||
async def audit_only(exc):
|
||||
return None
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only)
|
||||
collected = await _drain(response)
|
||||
|
||||
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
|
||||
|
|
@ -7092,9 +7037,7 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug():
|
|||
async def broken_hook(exc):
|
||||
raise RuntimeError("the audit backend is down")
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook)
|
||||
collected = await _drain(response)
|
||||
|
||||
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue