From 512e730f2fbbb4a832a42ca65b8b29c41c77508d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 21:30:48 -0700 Subject: [PATCH 01/24] fix(azure_ai): add passthrough config so router-model relays reach the deployment's own endpoint Every /azure_ai// relay failed with HTTP 500 because azure_ai had no passthrough config. The new AzureAIPassthroughConfig strips the router-model prefix from the relayed path, forwards to the deployment's api_base with its own credential (api-key on Foundry and Azure OpenAI hosts, Bearer elsewhere, Entra as the fallback), and delegates chat/completions cost logging to the Azure passthrough config. The router's provider inference now receives the deployment's api_base so an OpenAI-family model on a Foundry resource stays azure_ai instead of flipping to azure through the AZURE_AI_API_BASE env var. --- basedpyright-code-budget.json | 6 +- litellm/llms/azure_ai/chat/transformation.py | 12 +- litellm/llms/azure_ai/common_utils.py | 7 + .../azure_ai/passthrough/transformation.py | 106 +++++++++++ .../base_llm/passthrough/transformation.py | 3 +- litellm/router.py | 2 + litellm/utils.py | 6 + ruff-strict-budget.json | 2 +- ...est_azure_ai_passthrough_transformation.py | 175 ++++++++++++++++++ .../passthrough/test_passthrough_main.py | 115 ++++++++++++ type-discipline-budget.json | 2 +- 11 files changed, 423 insertions(+), 13 deletions(-) create mode 100644 litellm/llms/azure_ai/passthrough/transformation.py create mode 100644 tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 669107bb5b1..f31d620557d 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15284 + "limit": 15283 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,10 +105,10 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38309 + "limit": 38308 }, "reportUnknownParameterType": { - "limit": 19621 + "limit": 19620 }, "reportUnknownVariableType": { "limit": 29844 diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index f2d405e9a17..4c0a2e1f546 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -2,7 +2,6 @@ import copy import enum import re from typing import TYPE_CHECKING, Final, cast -from urllib.parse import urlparse import httpx from httpx import Response @@ -15,7 +14,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base +from litellm.llms.azure_ai.common_utils import ( + api_key_header_for_base, + is_foundry_model_inference_base, +) from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig @@ -99,11 +101,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Returns True if the request should use `api-key` header for authentication. """ - parsed_url: Final = urlparse(api_base) - host: Final = parsed_url.hostname - if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): - return True - return False + return api_key_header_for_base(api_base) == "api-key" def get_complete_url( self, diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index aa34bab5b2e..0665b3f64c5 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -19,6 +19,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool: return "/openai/deployments" not in parsed.path +def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader: + host: Final = urlparse(api_base).hostname if api_base else None + if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): + return "api-key" + return "Authorization" + + def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: """ Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py new file mode 100644 index 00000000000..afe0b814524 --- /dev/null +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + api_key_header_for_base, + get_azure_ai_auth_headers, +) +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.types.llms.openai import AllMessageValues + +if TYPE_CHECKING: + from httpx import URL, Response + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CostResponseTypes + + +def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: + path: Final = endpoint.lstrip("/") + for model_name in model_names: + if not model_name: + continue + if path == model_name: + return "" + if path.startswith(f"{model_name}/"): + return path[len(model_name) + 1 :] + return path + + +class PassthroughMetadata(BaseModel): + model_config = ConfigDict(extra="ignore") + + model_group: str = "" + + +def model_group_from(litellm_params: Mapping[str, object]) -> str: + try: + return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group + except ValidationError: + return "" + + +class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): + def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: + return request_data.get("stream") is True + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + endpoint: str, + request_query_params: Mapping[str, object] | None, + litellm_params: Mapping[str, object], + ) -> tuple[URL, str]: + base_target_url: Final = self.get_api_base(api_base) + if base_target_url is None: + raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") + + native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) + return ( + self.format_url(native_endpoint, base_target_url, request_query_params), + base_target_url, + ) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: base class contract returns dict for httpx + auth_headers: Final = get_azure_ai_auth_headers( + api_key=api_key, + litellm_params=litellm_params, + api_key_header=api_key_header_for_base(api_base), + ) + return {**headers, **auth_headers} # mutable-ok: base class contract returns dict for httpx + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: Mapping[str, object], + logging_obj: Logging, + endpoint: str, + ) -> CostResponseTypes | None: + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig + + return AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict + model=model, + custom_llm_provider=custom_llm_provider, + httpx_response=httpx_response, + request_data=dict(request_data), # mutable-ok: AzurePassthroughConfig wants a dict + logging_obj=logging_obj, + endpoint=endpoint, + ) diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index adbf2e126fb..34533e4c603 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,4 +1,5 @@ from abc import abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Union from ..base_utils import BaseLLMModelInfo @@ -23,7 +24,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, endpoint: str, base_target_url: str, - request_query_params: dict | None, + request_query_params: Mapping[str, object] | None, ) -> "URL": """ Helper function to add query params to the url diff --git a/litellm/router.py b/litellm/router.py index 6c7611c6236..b84f790287f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5051,6 +5051,7 @@ class Router: _, inferred_custom_llm_provider, _, _ = get_llm_provider( model=data["model"], custom_llm_provider=custom_llm_provider, + api_base=data.get("api_base"), ) custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider except Exception: @@ -5570,6 +5571,7 @@ class Router: _, inferred_custom_llm_provider, _, _ = get_llm_provider( model=data["model"], custom_llm_provider=custom_llm_provider, + api_base=data.get("api_base"), ) custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider except Exception: diff --git a/litellm/utils.py b/litellm/utils.py index 9d20d32d147..0f055308e3d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8831,6 +8831,12 @@ class ProviderConfigManager: ) return AzurePassthroughConfig() + elif LlmProviders.AZURE_AI == provider: + from litellm.llms.azure_ai.passthrough.transformation import ( + AzureAIPassthroughConfig, + ) + + return AzureAIPassthroughConfig() elif LlmProviders.GIGACHAT == provider: from litellm.llms.gigachat.passthrough.transformation import ( GigaChatPassthroughConfig, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 70408ea022b..a011230c2d3 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 119 + "limit": 118 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py new file mode 100644 index 00000000000..e1abbe1a7e1 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -0,0 +1,175 @@ +import json +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig +from litellm.types.utils import LlmProviders, ModelResponse +from litellm.utils import ProviderConfigManager + +FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" + + +@pytest.fixture(autouse=True) +def clear_azure_ai_env(monkeypatch): + for env_var in ("AZURE_AI_API_BASE", "AZURE_AI_API_KEY", "AZURE_AD_TOKEN", "AZURE_API_KEY"): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setattr(litellm, "api_key", None) + + +def test_provider_config_manager_resolves_azure_ai_passthrough_config(): + config = ProviderConfigManager.get_provider_passthrough_config(model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI) + + assert isinstance(config, AzureAIPassthroughConfig) + + +def test_router_model_prefix_is_stripped_and_native_path_kept_verbatim(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key=None, + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert base == FOUNDRY_BASE + + +def test_model_group_prefix_is_stripped_when_router_metadata_names_it(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key=None, + model="Cohere-parse-v5", + endpoint="/parse-alias/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={"litellm_metadata": {"model_group": "parse-alias"}}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + + +def test_model_inside_the_path_stays_and_query_params_are_forwarded(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/", + api_key=None, + model="gpt-5.4-mini", + endpoint="openai/deployments/gpt-5.4-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/openai/deployments/gpt-5.4-mini/chat/completions?api-version=2024-10-21" + + +def test_missing_api_base_raises_instead_of_building_a_relative_url(): + with pytest.raises(ValueError, match="AZURE_AI_API_BASE"): + AzureAIPassthroughConfig().get_complete_url( + api_base=None, + api_key=None, + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + +def _auth_headers(api_key: str | None, api_base: str, litellm_params: dict | None = None) -> dict: + return AzureAIPassthroughConfig().validate_environment( + headers={"content-type": "application/json"}, + model="Cohere-parse-v5", + messages=[], + optional_params={}, + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + ) + + +def test_foundry_host_gets_the_api_key_header(): + headers = _auth_headers(api_key="deployment-key", api_base=FOUNDRY_BASE) + + assert headers == {"content-type": "application/json", "api-key": "deployment-key"} + + +def test_serverless_host_gets_a_bearer_token(): + headers = _auth_headers(api_key="deployment-key", api_base="https://cohere-parse.eastus.models.ai.azure.com") + + assert headers["Authorization"] == "Bearer deployment-key" + assert "api-key" not in headers + + +def test_entra_token_is_used_when_the_deployment_has_no_api_key(): + headers = _auth_headers(api_key=None, api_base=FOUNDRY_BASE, litellm_params={"azure_ad_token": "entra-token"}) + + assert headers["Authorization"] == "Bearer entra-token" + + +def test_no_credentials_at_all_raises(): + with pytest.raises(ValueError, match="Missing Azure AI credentials"): + _auth_headers(api_key=None, api_base=FOUNDRY_BASE) + + +@pytest.mark.parametrize( + "request_data, expected", + [({"stream": True}, True), ({"stream": False}, False), ({}, False)], +) +def test_is_streaming_request_reads_the_stream_flag(request_data, expected): + assert AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) is expected + + +def _chat_completion_response() -> httpx.Response: + body = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, + } + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", f"{FOUNDRY_BASE}/models/chat/completions"), + ) + + +def test_chat_completions_relay_yields_a_model_response_for_cost_tracking(): + result = AzureAIPassthroughConfig().logging_non_streaming_response( + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + httpx_response=_chat_completion_response(), + request_data={"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "hi"}]}, + logging_obj=MagicMock(), + endpoint="models/chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "hi" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 8 + + +def test_non_chat_relay_yields_no_cost_response(): + parse_response = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id":"parse-1","pages":[]}', + request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), + ) + + result = AzureAIPassthroughConfig().logging_non_streaming_response( + model="Cohere-parse-v5", + custom_llm_provider="azure_ai", + httpx_response=parse_response, + request_data={"model": "Cohere-parse-v5"}, + logging_obj=MagicMock(), + endpoint="providers/cohere/v2/parse", + ) + + assert result is None diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 1950c37a12e..546cff18b5d 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -873,3 +873,118 @@ def test_llm_passthrough_route_propagates_allm_passthrough_route_to_logging_obj( assert captured_litellm_params.get("allm_passthrough_route") is True assert LitellmLogging._is_sync_litellm_request(captured_litellm_params) is False + + +FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" + + +def _foundry_parse_response() -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id":"parse-1","pages":[]}', + request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), + ) + + +def test_azure_ai_relay_reaches_the_deployment_with_its_own_credential(): + """ + Regression for LIT-7022: azure_ai had no passthrough config, so every + /azure_ai// relay raised "Provider azure_ai not found" + before a request was built. + """ + client = HTTPHandler() + + with patch.object(client.client, "send", return_value=_foundry_parse_response()) as mock_send: + response = llm_passthrough_route( + model="azure_ai/Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + method="POST", + custom_llm_provider="azure_ai", + api_base=FOUNDRY_BASE, + api_key="deployment-key", + json={"model": "Cohere-parse-v5", "document": {"type": "image_url", "image_url": "https://x/y.png"}}, + client=client, + litellm_logging_obj=MagicMock(), + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "Cohere-parse-v5" + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_router_relays_azure_ai_model_through_the_deployment_api_base(): + router = litellm.Router( + model_list=[ + { + "model_name": "foundry-parse", + "litellm_params": { + "model": "azure_ai/Cohere-parse-v5", + "api_base": FOUNDRY_BASE, + "api_key": "deployment-key", + }, + } + ] + ) + async_client = AsyncHTTPHandler() + + with patch.object(async_client.client, "send", AsyncMock(return_value=_foundry_parse_response())) as mock_send: + response = await router.allm_passthrough_route( + model="foundry-parse", + method="POST", + endpoint="foundry-parse/providers/cohere/v2/parse", + json={"model": "foundry-parse", "document": {"type": "image_url", "image_url": "https://x/y.png"}}, + client=async_client, + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "Cohere-parse-v5" + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_router_relays_an_openai_model_on_a_foundry_base_as_azure_ai(monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") + router = litellm.Router( + model_list=[ + { + "model_name": "foundry-gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": FOUNDRY_BASE, + "api_key": "deployment-key", + }, + } + ] + ) + async_client = AsyncHTTPHandler() + upstream = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=( + b'{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-5.4-mini",' + b'"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"hi"}}],' + b'"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' + ), + request=httpx.Request("POST", f"{FOUNDRY_BASE}/models/chat/completions"), + ) + + with patch.object(async_client.client, "send", AsyncMock(return_value=upstream)) as mock_send: + await router.allm_passthrough_route( + model="foundry-gpt", + method="POST", + endpoint="foundry-gpt/models/chat/completions", + request_query_params={"api-version": "2024-05-01-preview"}, + json={"model": "foundry-gpt", "messages": [{"role": "user", "content": "hi"}]}, + client=async_client, + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "gpt-5.4-mini" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3d01c08e8eb..7d12721175f 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22326 + "limit": 22325 }, "LIT002": { "limit": 26748 From a276690ce200bebe3547c110ac9a830b32052173 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:32:26 -0700 Subject: [PATCH 02/24] fix(router): keep a model's own provider prefix for generic SDK calls Generic passthrough calls inferred the provider from the bare model name, so an azure_ai/gpt-* deployment on an Azure OpenAI host flipped to azure and get_llm_provider re-prefixed the deployment name into azure_ai/gpt-5.4-mini, a 404 DeploymentNotFound. provider_for_generic_call takes the declared custom_llm_provider first, then the model's own prefix, and only infers for unprefixed models --- litellm/router.py | 24 ++--------- litellm/router_utils/common_utils.py | 27 ++++++++++++ .../test_router_utils_common_utils.py | 18 ++++++++ tests/test_litellm/test_router.py | 43 +++++++++++++++++++ 4 files changed, 91 insertions(+), 21 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b84f790287f..9b7a7ee7ce8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -135,6 +135,7 @@ from litellm.router_utils.common_utils import ( _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, + provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, warn_on_provider_credential_mismatch, @@ -5045,17 +5046,7 @@ class Router: kwargs=kwargs, model=model, model_name=model_name ) - # Get custom_llm_provider from deployment params - try: - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - api_base=data.get("api_base"), - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - except Exception: - custom_llm_provider = None + custom_llm_provider: Final = provider_for_generic_call(data) response_kwargs: Final = { **data, @@ -5566,16 +5557,7 @@ class Router: # Perform pre-call checks for routing strategy self.routing_strategy_pre_call_checks(deployment=deployment) - try: - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - api_base=data.get("api_base"), - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - except Exception: - custom_llm_provider = None + custom_llm_provider: Final = provider_for_generic_call(data) response: Final = original_function( **{ diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 280a7defcf8..ec3ffa1247e 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject +import litellm from litellm._logging import verbose_logger, verbose_router_logger from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError @@ -244,6 +245,32 @@ PROVIDER_SCOPED_CREDENTIAL_PARAMS: Final[Mapping[str, frozenset[str]]] = Mapping ) +def provider_for_generic_call(litellm_params: Mapping[str, object]) -> str | None: + """ + The provider the router hands a deployment's generic SDK call, or None when it cannot be resolved. + + A model that carries its own provider prefix keeps that prefix even where get_llm_provider + would resolve it to a sibling provider (azure_ai/ on an Azure OpenAI host + resolves to azure): the SDK call still receives the prefixed model, and an explicit provider + that contradicts the prefix makes get_llm_provider re-prefix it into a deployment name that + does not exist upstream. + """ + declared: Final = litellm_params.get("custom_llm_provider") + if isinstance(declared, str) and declared: + return declared + model: Final = litellm_params.get("model") + if not isinstance(model, str) or not model: + return None + prefix: Final = model.split("/", 1)[0] + if "/" in model and prefix in litellm.provider_list: + return prefix + try: + _, inferred, _, _ = get_llm_provider(model=model) + except BadRequestError: + return None + return inferred + + def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mapping[str, object]) -> str | None: """ Warn when a deployment carries one provider's credentials but resolves to another. diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index 30f658d7ea2..ac18b4889dd 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -12,6 +12,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, PROVIDER_SCOPED_CREDENTIAL_PARAMS, @@ -756,3 +757,20 @@ class TestWarnOnProviderCredentialMismatch: ) is None ) + + +@pytest.mark.parametrize( + ("litellm_params", "expected"), + [ + ({"model": "azure_ai/gpt-5.4-mini", "custom_llm_provider": "azure"}, "azure"), + ({"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.openai.azure.com"}, "azure_ai"), + ({"model": "cohere/command-r"}, "cohere"), + ({"model": "gpt-5.4-mini"}, "openai"), + ({"model": "no-provider-knows-this-model"}, None), + ({"api_base": "https://my-resource.openai.azure.com"}, None), + ], + ids=["declared_wins", "prefix_beats_host_flip", "prefix_beats_cohere_chat_flip", "unprefixed_inferred", "unknown", "no_model"], +) +def test_provider_for_generic_call(litellm_params, expected, monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") + assert provider_for_generic_call(litellm_params) == expected diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 31eb46f1458..b220b23c338 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12938,3 +12938,46 @@ async def test_router_retry_policy_controls_upstream_attempt_count( await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) assert upstream.call_count == expected_upstream_calls + + +@pytest.mark.asyncio +async def test_generic_call_keeps_the_deployment_name_of_an_azure_ai_model_on_an_azure_openai_host(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + { + "model_name": "aoai-gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.openai.azure.com", + "api_key": "deployment-key", + }, + } + ] + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upstream = respx_mock.post(host="my-resource.openai.azure.com", path__regex=r"^/openai/.*responses$").mock( + return_value=httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + }, + ) + ) + await router.aresponses(model="aoai-gpt", input="hi") + + assert json.loads(upstream.calls.last.request.content)["model"] == "gpt-5.4-mini" From cd25eb918995ef115aa67e8fd39337e4786bdcd2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:32:27 -0700 Subject: [PATCH 03/24] fix(azure_ai): cost streaming relays and return upstream errors from router relays Streaming chat relays on Azure and azure_ai deployments rebuild the response from the SSE chunks through the OpenAI passthrough assembler, so the spend log carries usage. The router relays keep the JSON body when the Content-Type carries a charset, return the upstream status and body instead of a 500 when the deployment rejects the call, and fall back to the caller's api-version when the deployment sets none. Lint budgets ratcheted to the measured totals --- basedpyright-code-budget.json | 8 +- .../llms/azure/passthrough/transformation.py | 24 ++++- .../azure_ai/passthrough/transformation.py | 18 ++++ .../proxy/common_utils/http_parsing_utils.py | 4 +- .../llm_passthrough_endpoints.py | 53 ++++++----- .../openai_passthrough_logging_handler.py | 3 +- ruff-strict-budget.json | 4 +- .../test_azure_passthrough_transformation.py | 69 +++++++++++++++ ...est_azure_ai_passthrough_transformation.py | 21 +++++ .../test_llm_pass_through_endpoints.py | 87 +++++++++++++++++++ type-discipline-budget.json | 4 +- 11 files changed, 260 insertions(+), 35 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index b446c79f49b..ef1bd675c5e 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15279 + "limit": 15278 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38281 + "limit": 38276 }, "reportUnknownParameterType": { - "limit": 19582 + "limit": 19581 }, "reportUnknownVariableType": { - "limit": 29829 + "limit": 29825 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 898852e645f..7654958fc0b 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from typing import TYPE_CHECKING, Final, Optional import httpx @@ -43,7 +44,7 @@ class AzurePassthroughConfig(BasePassthroughConfig): api_base=base_target_url, litellm_params=litellm_params, route=endpoint, - default_api_version=litellm_params.get("api_version"), + default_api_version=request_query_params.get("api-version") if request_query_params else None, ) return ( httpx.URL(complete_url), @@ -116,3 +117,24 @@ class AzurePassthroughConfig(BasePassthroughConfig): ) return litellm_model_response + + def handle_logging_collected_chunks( + self, + all_chunks: Sequence[str], + litellm_logging_obj: Logging, + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, + ) + + if "chat/completions" not in endpoint: + return None + + return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + ) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index afe0b814524..5948c865bad 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -104,3 +104,21 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): logging_obj=logging_obj, endpoint=endpoint, ) + + def handle_logging_collected_chunks( + self, + all_chunks: Sequence[str], + litellm_logging_obj: Logging, + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> CostResponseTypes | None: + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig + + return AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + custom_llm_provider=custom_llm_provider, + endpoint=endpoint, + ) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 54a0f18fd63..635cbf3ba9b 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -39,7 +39,7 @@ def _is_form_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) in _FORM_CONTENT_TYPES -def _is_json_content_type(content_type: str) -> bool: +def is_json_content_type(content_type: str) -> bool: """True iff the body should be parsed as JSON.""" return _normalize_media_type(content_type) == "application/json" @@ -406,7 +406,7 @@ async def get_request_body(request: Request) -> dict[str, Any]: """ if request.method == "POST": content_type: Final = request.headers.get("content-type", "") - if _is_json_content_type(content_type): + if is_json_content_type(content_type): return await _read_request_body(request) elif _is_form_content_type(content_type): return await get_form_data(request) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6b1d6405a6a..cf7b7fdf07a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_set_request_parsed_body, get_form_data, get_request_body, + is_json_content_type, ) from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, @@ -408,7 +409,7 @@ async def vllm_proxy_route( content=None, data=None, files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), params=None, headers=None, cookies=None, @@ -1492,6 +1493,14 @@ async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> Async await upstream.aclose() +async def _relay_upstream_response(upstream: httpx.Response) -> Response: + return Response( + content=await upstream.aread(), + status_code=upstream.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), + ) + + async def _relay_azure_router_model( llm_router: litellm.Router, model: str, @@ -1501,30 +1510,28 @@ async def _relay_azure_router_model( is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, ) -> Response: - result: Final = await llm_router.allm_passthrough_route( - model=model, - method=request.method, - endpoint=endpoint, - request_query_params=request.query_params, - request_headers=_safe_get_request_headers(request), - stream=is_streaming_request, - content=None, - data=None, - files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), - params=None, - headers=None, - cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), - ) + try: + result: Final = await llm_router.allm_passthrough_route( + model=model, + method=request.method, + endpoint=endpoint, + request_query_params=request.query_params, + request_headers=_safe_get_request_headers(request), + stream=is_streaming_request, + content=None, + data=None, + files=None, + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), + params=None, + headers=None, + cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + ) + except httpx.HTTPStatusError as upstream_error: + return await _relay_upstream_response(upstream_error.response) if not is_streaming_request: - upstream: Final = cast(httpx.Response, result) - return Response( - content=await upstream.aread(), - status_code=upstream.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), - ) + return await _relay_upstream_response(cast(httpx.Response, result)) if inspect.isasyncgen(result): sse_headers: Final = {"content-type": "text/event-stream"} diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 5f6489a69ca..a39ea342f0a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -4,6 +4,7 @@ OpenAI Passthrough Logging Handler Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. """ +from collections.abc import Sequence from datetime import datetime from typing import Final from urllib.parse import urlparse @@ -512,7 +513,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): def _build_complete_streaming_response( self, - all_chunks: list[str], + all_chunks: Sequence[str], litellm_logging_obj: LiteLLMLoggingObj, model: str, ) -> ModelResponse | TextCompletionResponse | None: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 694ce6e86e0..3354610bfcf 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2916 + "limit": 2914 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 116 + "limit": 115 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 29b74c2ee4a..68fbc99b5d5 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -92,3 +92,72 @@ def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_retur ) assert result is None + + +def _sse_line(payload: dict) -> str: + return "data: " + json.dumps(payload) + + +def _azure_chat_completion_chunks() -> list[str]: + head = {"id": "chatcmpl-abc123", "object": "chat.completion.chunk", "created": 1700000000, "model": "gpt-4.1-mini"} + return [ + _sse_line({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}]}), + _sse_line({**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]}), + _sse_line({**head, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), + _sse_line({**head, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}}), + "data: [DONE]", + ] + + +def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response(): + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_chat_completion_chunks(), + litellm_logging_obj=MagicMock(), + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens == 10 + assert response.usage.completion_tokens == 8 + + +def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_chat_completion_chunks(), + litellm_logging_obj=MagicMock(), + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/embeddings", + ) + + assert response is None + + +def _complete_url(request_query_params: dict, litellm_params: dict) -> httpx.URL: + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params=request_query_params, + litellm_params=litellm_params, + ) + return url + + +def test_azure_passthrough_url_falls_back_to_the_callers_api_version(): + url = _complete_url(request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={}) + + assert url.path == "/openai/deployments/gpt-4.1-mini/chat/completions" + assert url.params["api-version"] == "2025-04-01-preview" + + +def test_azure_passthrough_url_prefers_the_deployments_api_version(): + url = _complete_url( + request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={"api_version": "2024-10-21"} + ) + + assert url.params["api-version"] == "2024-10-21" diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index e1abbe1a7e1..67851d2d58e 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -173,3 +173,24 @@ def test_non_chat_relay_yields_no_cost_response(): ) assert result is None + + +def test_streaming_chat_completion_chunks_are_costed_like_azure(): + head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} + chunks = [ + "data: " + json.dumps({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]}), + "data: " + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), + "data: [DONE]", + ] + + response = AzureAIPassthroughConfig().handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=MagicMock(), + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + endpoint="chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "hi" + assert response.usage.total_tokens == 4 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index acb45038df0..a2043808b6b 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -5206,3 +5206,90 @@ class TestAzureRouterModelStreamingKeepalive: assert result.headers["x-upstream"] == "kept" assert chunks == [b"data: hello\n\n"] + + +class TestRouterModelRelayUpstreamContract: + def _request(self, content_type: str) -> MagicMock: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": content_type} + request.query_params = {} + return request + + def _install_router(self, monkeypatch, router, body: dict) -> None: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + + def _recording_router(self, captured: list[dict]): + class RecordingRouter: + async def allm_passthrough_route(self, **kwargs): + captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + return RecordingRouter() + + @pytest.mark.asyncio + async def test_azure_relay_keeps_the_json_body_when_the_content_type_carries_a_charset(self, monkeypatch): + body = {"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}]} + captured: list[dict] = [] + self._install_router(monkeypatch, self._recording_router(captured), body) + + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json; charset=utf-8"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert captured[0]["json"] == body + + @pytest.mark.asyncio + async def test_vllm_relay_keeps_the_json_body_when_the_content_type_carries_a_charset(self, monkeypatch): + body = {"model": "router-model", "messages": [{"role": "user", "content": "hi"}]} + captured: list[dict] = [] + self._install_router(monkeypatch, self._recording_router(captured), body) + + await vllm_proxy_route( + endpoint="/chat/completions", + request=self._request("application/json; charset=utf-8"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert captured[0]["json"] == body + + @pytest.mark.asyncio + async def test_azure_relay_returns_the_upstream_status_and_body_when_the_deployment_rejects_the_call( + self, monkeypatch + ): + upstream_body = {"error": {"code": "DeploymentNotFound", "message": "The API deployment does not exist."}} + + class RejectingRouter: + async def allm_passthrough_route(self, **kwargs): + upstream_request = httpx.Request( + "POST", "https://my-azure.openai.azure.com/openai/deployments/gpt-5/chat/completions" + ) + upstream = httpx.Response( + 404, json=upstream_body, headers={"x-ms-request-id": "req-1"}, request=upstream_request + ) + raise httpx.HTTPStatusError("404", request=upstream_request, response=upstream) + + self._install_router(monkeypatch, RejectingRouter(), {"model": "gpt-5", "stream": False}) + + result = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert result.status_code == 404 + assert json.loads(result.body) == upstream_body + assert result.headers["x-ms-request-id"] == "req-1" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index bfa055b43e5..30bd580c4e2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22180 + "limit": 22178 }, "LIT002": { "limit": 26745 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16464 + "limit": 16458 }, "LIT011": { "limit": 5506 From 9ea9aa2e7b1ef97f80e2141f0bc299498e351143 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:00:23 -0700 Subject: [PATCH 04/24] fix(azure_ai): count prompt tokens for streaming relays that carry no usage chunk --- basedpyright-code-budget.json | 8 ++--- .../streaming_chunk_builder_utils.py | 4 +-- .../llms/azure/passthrough/transformation.py | 20 ++++++++++- litellm/main.py | 4 +-- .../openai_passthrough_logging_handler.py | 7 ++-- ruff-strict-budget.json | 4 +-- .../test_azure_passthrough_transformation.py | 33 +++++++++++++++++-- type-discipline-budget.json | 4 +-- 8 files changed, 67 insertions(+), 17 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ef1bd675c5e..4dc7a0d8671 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15278 + "limit": 15277 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38276 + "limit": 38271 }, "reportUnknownParameterType": { - "limit": 19581 + "limit": 19580 }, "reportUnknownVariableType": { - "limit": 29825 + "limit": 29821 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0e01577b20e..1beb06f52e7 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -209,7 +209,7 @@ def apply_grounding_request_counts( class ChunkProcessor: - def __init__(self, chunks: list, messages: list | None = None): + def __init__(self, chunks: list, messages: Sequence | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages self.first_chunk = chunks[0] @@ -992,7 +992,7 @@ class ChunkProcessor: chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, - messages: list | None = None, + messages: Sequence | None = None, reasoning_tokens: int | None = None, ) -> Usage: """ diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 7654958fc0b..0b8b216faae 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,8 +1,9 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional import httpx from httpx import Response +from pydantic import BaseModel, ValidationError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM @@ -17,6 +18,22 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes +class RelayedChatRequest(BaseModel): + messages: Sequence[Mapping[str, object]] | None = None + + +class RelayedCallDetails(BaseModel): + request_data: RelayedChatRequest | None = None + + +def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None: + try: + details: Final = RelayedCallDetails.model_validate(litellm_logging_obj.model_call_details) + except ValidationError: + return None + return details.request_data.messages if details.request_data else None + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in request_data @@ -137,4 +154,5 @@ class AzurePassthroughConfig(BasePassthroughConfig): all_chunks=all_chunks, litellm_logging_obj=litellm_logging_obj, model=model, + messages=_relayed_messages(litellm_logging_obj), ) diff --git a/litellm/main.py b/litellm/main.py index 2929790f2bd..a7fa5560d30 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8553,7 +8553,7 @@ def config_completion(**kwargs): ) -def stream_chunk_builder_text_completion(chunks: list, messages: list | None = None) -> TextCompletionResponse: +def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse: id: Final = chunks[0]["id"] object: Final = chunks[0]["object"] created: Final = chunks[0]["created"] @@ -8670,7 +8670,7 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o def stream_chunk_builder( chunks: list, - messages: list | None = None, + messages: Sequence | None = None, start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index a39ea342f0a..9805c7d7fae 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -4,7 +4,7 @@ OpenAI Passthrough Logging Handler Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime from typing import Final from urllib.parse import urlparse @@ -516,6 +516,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): all_chunks: Sequence[str], litellm_logging_obj: LiteLLMLoggingObj, model: str, + messages: Sequence[Mapping[str, object]] | None = None, ) -> ModelResponse | TextCompletionResponse | None: """ Builds complete response from raw chunks for OpenAI streaming responses. @@ -559,7 +560,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return None # Build complete response from chunks - complete_streaming_response: Final = litellm.stream_chunk_builder(chunks=all_openai_chunks) + complete_streaming_response: Final = litellm.stream_chunk_builder( + chunks=all_openai_chunks, messages=messages + ) return complete_streaming_response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 3354610bfcf..cfe737867f1 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2914 + "limit": 2912 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 115 + "limit": 114 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 68fbc99b5d5..d0657679983 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -3,6 +3,8 @@ from unittest.mock import MagicMock import httpx +import litellm + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig from litellm.types.utils import ModelResponse @@ -101,8 +103,15 @@ def _sse_line(payload: dict) -> str: def _azure_chat_completion_chunks() -> list[str]: head = {"id": "chatcmpl-abc123", "object": "chat.completion.chunk", "created": 1700000000, "model": "gpt-4.1-mini"} return [ - _sse_line({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}]}), - _sse_line({**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]}), + _sse_line( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}], + } + ), + _sse_line( + {**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]} + ), _sse_line({**head, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), _sse_line({**head, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}}), "data: [DONE]", @@ -124,6 +133,26 @@ def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response(): assert response.usage.completion_tokens == 8 +def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request(): + messages = [{"role": "user", "content": "Say hi in three words"}] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens > 0 + assert response.usage.prompt_tokens == litellm.token_counter(model="gpt-4.1-mini", messages=messages) + assert response.usage.completion_tokens > 0 + + def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): response = AzurePassthroughConfig().handle_logging_collected_chunks( all_chunks=_azure_chat_completion_chunks(), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 30bd580c4e2..74b54148418 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22178 + "limit": 22172 }, "LIT002": { "limit": 26745 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16458 + "limit": 16452 }, "LIT011": { "limit": 5506 From bbbdccb82d88f3489a1a311109f86d87bb1782e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:30:13 -0700 Subject: [PATCH 05/24] fix(azure_ai): count relayed image prompt tokens without fetching the image --- basedpyright-code-budget.json | 8 +++--- .../streaming_chunk_builder_utils.py | 5 +++- litellm/main.py | 3 ++ .../openai_passthrough_logging_handler.py | 2 +- ruff-strict-budget.json | 4 +-- .../test_azure_passthrough_transformation.py | 28 +++++++++++++++++++ type-discipline-budget.json | 4 +-- 7 files changed, 44 insertions(+), 10 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 4dc7a0d8671..f63fb2a497a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15277 + "limit": 15276 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38266 }, "reportUnknownParameterType": { - "limit": 19580 + "limit": 19579 }, "reportUnknownVariableType": { - "limit": 29821 + "limit": 29817 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1beb06f52e7..1dca64f1a80 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -994,6 +994,7 @@ class ChunkProcessor: completion_output: str, messages: Sequence | None = None, reasoning_tokens: int | None = None, + use_default_image_token_count: bool = False, ) -> Usage: """ Calculate usage for the given chunks. @@ -1018,7 +1019,9 @@ class ChunkProcessor: cost: Final[float | None] = calculated_usage_per_chunk["cost"] try: - returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages) + returned_usage.prompt_tokens = prompt_tokens or token_counter( + model=model, messages=messages, use_default_image_token_count=use_default_image_token_count + ) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") returned_usage.prompt_tokens = 0 diff --git a/litellm/main.py b/litellm/main.py index a7fa5560d30..d125f1c1d7f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8674,6 +8674,7 @@ def stream_chunk_builder( start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, + use_default_image_token_count: bool = False, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -8747,6 +8748,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=0, + use_default_image_token_count=use_default_image_token_count, ) setattr(response, "usage", usage) @@ -8924,6 +8926,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=reasoning_tokens, + use_default_image_token_count=use_default_image_token_count, ) setattr(response, "usage", usage) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 9805c7d7fae..e21105f760c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -561,7 +561,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Build complete response from chunks complete_streaming_response: Final = litellm.stream_chunk_builder( - chunks=all_openai_chunks, messages=messages + chunks=all_openai_chunks, messages=messages, use_default_image_token_count=True ) return complete_streaming_response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index cfe737867f1..907317e81b1 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2912 + "limit": 2910 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 114 + "limit": 113 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index d0657679983..b07b7915277 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -153,6 +153,34 @@ def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_fr assert response.usage.completion_tokens > 0 +def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_without_fetching_the_image(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png"}}, + ], + } + ] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.usage.prompt_tokens > 0 + assert response.usage.prompt_tokens == litellm.token_counter( + model="gpt-4.1-mini", messages=messages, use_default_image_token_count=True + ) + + def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): response = AzurePassthroughConfig().handle_logging_collected_chunks( all_chunks=_azure_chat_completion_chunks(), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 74b54148418..78cd84f5b35 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22172 + "limit": 22166 }, "LIT002": { "limit": 26745 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16452 + "limit": 16446 }, "LIT011": { "limit": 5506 From 3f695846b3cbb688d0d581bc64c4da5258546240 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:09:25 -0700 Subject: [PATCH 06/24] fix(router): rewrite the passthrough model group as a whole path segment --- .../base_llm/passthrough/transformation.py | 20 ++++++++++++++-- litellm/router.py | 3 ++- tests/test_litellm/test_router.py | 23 +++++++++++++++++++ 3 files changed, 43 insertions(+), 3 deletions(-) diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 34533e4c603..d97ecddc75d 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -8,11 +8,27 @@ if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.utils import CostResponseTypes + from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException +def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: + path: Final = endpoint.lstrip("/") + for model_name in model_names: + if not model_name: + continue + if path == model_name: + return "" + if path.startswith(f"{model_name}/"): + return path[len(model_name) + 1 :] + return path + + +def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: + return "/".join(replacement if part == segment else part for part in endpoint.split("/")) + + class BasePassthroughConfig(BaseLLMModelInfo): @abstractmethod def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: @@ -104,7 +120,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): request_data: dict, logging_obj: "LiteLLMLoggingObj", endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> Optional["CostResponseTypes | StandardPassThroughResponseObject"]: pass def handle_logging_collected_chunks( diff --git a/litellm/router.py b/litellm/router.py index 9b7a7ee7ce8..273c1f9d99c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -86,6 +86,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_credentials_in_payload, mask_sensitive_structure, ) +from litellm.llms.base_llm.passthrough.transformation import replace_path_segment from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, @@ -5010,7 +5011,7 @@ class Router: # If get_llm_provider fails, fall back to using model_name as-is replacement_model_name = model_name - kwargs["endpoint"] = kwargs["endpoint"].replace(model, replacement_model_name) + kwargs["endpoint"] = replace_path_segment(kwargs["endpoint"], model, replacement_model_name) return kwargs async def _ageneric_api_call_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b220b23c338..0f4e199694d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4469,6 +4469,29 @@ def test_get_deployment_model_info_base_model_merge_priority(): print("✓ Base model merge priority test passed!") +def test_add_deployment_model_to_endpoint_rewrites_whole_path_segments_only(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.services.ai.azure.com", + "api_key": "key", + }, + } + ], + ) + + result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( + kwargs={"endpoint": "gpt/openai/deployments/gpt-4o/chat/completions", "custom_llm_provider": "azure_ai"}, + model="gpt", + model_name="azure_ai/gpt-5.4-mini", + ) + + assert result["endpoint"] == "gpt-5.4-mini/openai/deployments/gpt-4o/chat/completions" + + def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): """ Test that _add_deployment_model_to_endpoint_for_llm_passthrough_route correctly strips bedrock provider prefix From df6fb9e5d95e97b7b207cd278d107e82b134fe45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:09:26 -0700 Subject: [PATCH 07/24] fix(azure_ai): relay from the Foundry root and log non-chat relays for spend tracking --- .../llms/azure/passthrough/transformation.py | 14 +-- .../azure_ai/passthrough/transformation.py | 69 +++++++++++---- ...est_azure_ai_passthrough_transformation.py | 87 +++++++++++++++++-- 3 files changed, 140 insertions(+), 30 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 0b8b216faae..983033a8b14 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -7,7 +7,11 @@ from pydantic import BaseModel, ValidationError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + replace_path_segment, + strip_leading_model_segment, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams @@ -36,7 +40,7 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: - return "stream" in request_data + return bool(request_data.get("stream")) def get_complete_url( self, @@ -54,13 +58,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): litellm_metadata: Final = litellm_params.get("litellm_metadata") or {} model_group: Final = litellm_metadata.get("model_group") - if model_group and model_group in endpoint: - endpoint = endpoint.replace(model_group, model) + routed_endpoint: Final = replace_path_segment(endpoint, model_group, model) if model_group else endpoint + native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,)) complete_url: Final = BaseAzureLLM._get_base_azure_url( api_base=base_target_url, litellm_params=litellm_params, - route=endpoint, + route=native_endpoint, default_api_version=request_query_params.get("api-version") if request_query_params else None, ) return ( diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 5948c865bad..4776bbe0755 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -1,17 +1,20 @@ from __future__ import annotations from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final -from pydantic import BaseModel, ConfigDict, ValidationError +import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, api_key_header_for_base, get_azure_ai_auth_headers, ) -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, strip_leading_model_segment from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import StandardPassThroughResponseObject if TYPE_CHECKING: from httpx import URL, Response @@ -20,16 +23,7 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes -def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: - path: Final = endpoint.lstrip("/") - for model_name in model_names: - if not model_name: - continue - if path == model_name: - return "" - if path.startswith(f"{model_name}/"): - return path[len(model_name) + 1 :] - return path +EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({}) class PassthroughMetadata(BaseModel): @@ -45,9 +39,44 @@ def model_group_from(litellm_params: Mapping[str, object]) -> str: return "" +def api_version_from(litellm_params: Mapping[str, object]) -> str | None: + try: + return TypeAdapter(str | None).validate_python(litellm_params.get("api_version")) + except ValidationError: + return None + + +def foundry_root(api_base: str) -> str: + url: Final = httpx.URL(api_base) + segments: Final = tuple(segment for segment in url.path.split("/") if segment) + root_segments: Final = segments[: segments.index("models")] if "models" in segments else segments + return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/") + + +def relay_query_params( + request_query_params: Mapping[str, object] | None, + deployment_api_version: str | None, + api_base: str, +) -> Mapping[str, object] | None: + if request_query_params and "api-version" in request_query_params: + return request_query_params + api_version: Final = deployment_api_version or httpx.URL(api_base).params.get("api-version") + if api_version is None: + return request_query_params + return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version}) + + +def relayed_body(httpx_response: Response) -> str | dict: + try: + body: Final[object] = httpx_response.json() + except ValueError: + return httpx_response.text + return body if isinstance(body, dict) else httpx_response.text + + class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: - return request_data.get("stream") is True + return bool(request_data.get("stream")) def get_complete_url( self, @@ -62,11 +91,12 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): if base_target_url is None: raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") + root: Final = foundry_root(base_target_url) native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) - return ( - self.format_url(native_endpoint, base_target_url, request_query_params), - base_target_url, + query_params: Final = relay_query_params( + request_query_params, api_version_from(litellm_params), base_target_url ) + return (self.format_url(native_endpoint, root, query_params), root) def validate_environment( self, @@ -93,10 +123,10 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> CostResponseTypes | None: + ) -> CostResponseTypes | StandardPassThroughResponseObject | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig - return AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict + chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict model=model, custom_llm_provider=custom_llm_provider, httpx_response=httpx_response, @@ -104,6 +134,9 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): logging_obj=logging_obj, endpoint=endpoint, ) + if chat_result is not None: + return chat_result + return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) def handle_logging_collected_chunks( self, diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 67851d2d58e..e343d05c5df 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -66,6 +66,72 @@ def test_model_inside_the_path_stays_and_query_params_are_forwarded(): assert str(url) == f"{FOUNDRY_BASE}/openai/deployments/gpt-5.4-mini/chat/completions?api-version=2024-10-21" +def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models", + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params={"api-version": "2024-05-01-preview"}, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + assert base == FOUNDRY_BASE + + +def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models", + api_key="key", + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + + +def test_deployment_api_version_fills_in_when_the_caller_sends_none(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params=None, + litellm_params={"api_version": "2024-05-01-preview"}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + + +def test_callers_api_version_beats_the_deployments(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params={"api-version": "2025-04-01-preview"}, + litellm_params={"api_version": "2024-05-01-preview"}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2025-04-01-preview" + + +def test_api_version_on_the_configured_api_base_is_the_last_fallback(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview", + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + + def test_missing_api_base_raises_instead_of_building_a_relative_url(): with pytest.raises(ValueError, match="AZURE_AI_API_BASE"): AzureAIPassthroughConfig().get_complete_url( @@ -116,7 +182,7 @@ def test_no_credentials_at_all_raises(): @pytest.mark.parametrize( "request_data, expected", - [({"stream": True}, True), ({"stream": False}, False), ({}, False)], + [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], ) def test_is_streaming_request_reads_the_stream_flag(request_data, expected): assert AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) is expected @@ -155,15 +221,14 @@ def test_chat_completions_relay_yields_a_model_response_for_cost_tracking(): assert result.usage.completion_tokens == 8 -def test_non_chat_relay_yields_no_cost_response(): +def _non_chat_logging_result(content: bytes, content_type: str): parse_response = httpx.Response( status_code=200, - headers={"content-type": "application/json"}, - content=b'{"id":"parse-1","pages":[]}', + headers={"content-type": content_type}, + content=content, request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), ) - - result = AzureAIPassthroughConfig().logging_non_streaming_response( + return AzureAIPassthroughConfig().logging_non_streaming_response( model="Cohere-parse-v5", custom_llm_provider="azure_ai", httpx_response=parse_response, @@ -172,7 +237,15 @@ def test_non_chat_relay_yields_no_cost_response(): endpoint="providers/cohere/v2/parse", ) - assert result is None + +def test_non_chat_relay_logs_the_parsed_body_so_spend_tracking_sees_the_call(): + result = _non_chat_logging_result(b'{"id":"parse-1","pages":[],"meta":{"billed_units":{"pages":1}}}', "application/json") + + assert result == {"response": {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 1}}}} + + +def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): + assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} def test_streaming_chat_completion_chunks_are_costed_like_azure(): From 2f981d14e4262db3e573891ba94f0e66209b216d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:09:26 -0700 Subject: [PATCH 08/24] refactor(passthrough): inject the streaming prompt-token counter instead of a default-image flag --- .../streaming_chunk_builder_utils.py | 8 ++-- litellm/main.py | 8 ++-- .../openai_passthrough_logging_handler.py | 44 +++++++++++++++++- ruff-strict-budget.json | 4 +- .../test_azure_passthrough_transformation.py | 44 ++++++++++++++++-- ...test_openai_passthrough_logging_handler.py | 46 +++++++++++++++++++ type-discipline-budget.json | 6 +-- 7 files changed, 142 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1dca64f1a80..f30a8e5f8bd 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,6 +1,6 @@ import base64 import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast @@ -994,7 +994,7 @@ class ChunkProcessor: completion_output: str, messages: Sequence | None = None, reasoning_tokens: int | None = None, - use_default_image_token_count: bool = False, + count_prompt_tokens: Callable[[], int] | None = None, ) -> Usage: """ Calculate usage for the given chunks. @@ -1019,8 +1019,8 @@ class ChunkProcessor: cost: Final[float | None] = calculated_usage_per_chunk["cost"] try: - returned_usage.prompt_tokens = prompt_tokens or token_counter( - model=model, messages=messages, use_default_image_token_count=use_default_image_token_count + returned_usage.prompt_tokens = prompt_tokens or ( + count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages) ) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") diff --git a/litellm/main.py b/litellm/main.py index d125f1c1d7f..a3830192e7a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -19,7 +19,7 @@ import random import sys import time import traceback -from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy @@ -8674,7 +8674,7 @@ def stream_chunk_builder( start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, - use_default_image_token_count: bool = False, + count_prompt_tokens: Callable[[], int] | None = None, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -8748,7 +8748,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=0, - use_default_image_token_count=use_default_image_token_count, + count_prompt_tokens=count_prompt_tokens, ) setattr(response, "usage", usage) @@ -8926,7 +8926,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=reasoning_tokens, - use_default_image_token_count=use_default_image_token_count, + count_prompt_tokens=count_prompt_tokens, ) setattr(response, "usage", usage) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index e21105f760c..6a142d08be5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -13,6 +13,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, @@ -97,6 +98,45 @@ def _is_openai_compatible_url(url_route: str | None) -> bool: return False +def _is_remote_high_detail_image(part: object) -> bool: + if not isinstance(part, Mapping) or part.get("type") != "image_url": + return False + image_url: Final = part.get("image_url") + if not isinstance(image_url, Mapping): + return False + url: Final = image_url.get("url") + return isinstance(url, str) and url.startswith(("http://", "https://")) and image_url.get("detail") == "high" + + +def _content_parts(message: Mapping[str, object]) -> Sequence[object]: + content: Final = message.get("content") + return content if isinstance(content, list) else () + + +def _without_remote_high_detail_images(message: Mapping[str, object]) -> Mapping[str, object]: + if not isinstance(message.get("content"), list): + return message + kept_parts: Final = [ # mutable-ok: token_counter reads message content only when it is a list + part for part in _content_parts(message) if not _is_remote_high_detail_image(part) + ] + return {**message, "content": kept_parts} # mutable-ok: token_counter rejects any message that is not a dict + + +def count_relayed_prompt_tokens(model: str, messages: Sequence[Mapping[str, object]] | None) -> int: + if messages is None: + return 0 + remote_high_detail_images: Final = sum( + 1 for message in messages for part in _content_parts(message) if _is_remote_high_detail_image(part) + ) + local_messages: Final = [ # mutable-ok: token_counter takes a list of messages + _without_remote_high_detail_images(message) for message in messages + ] + return ( + litellm.token_counter(model=model, messages=local_messages) + + DEFAULT_IMAGE_TOKEN_COUNT * remote_high_detail_images + ) + + class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): """ OpenAI-specific passthrough logging handler that provides cost tracking for /chat/completions endpoints. @@ -561,7 +601,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Build complete response from chunks complete_streaming_response: Final = litellm.stream_chunk_builder( - chunks=all_openai_chunks, messages=messages, use_default_image_token_count=True + chunks=all_openai_chunks, + messages=messages, + count_prompt_tokens=lambda: count_relayed_prompt_tokens(model, messages), ) return complete_streaming_response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 907317e81b1..574bf01cc34 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2910 + "limit": 2904 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 113 + "limit": 110 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index b07b7915277..5d9e4312b5e 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -2,8 +2,10 @@ import json from unittest.mock import MagicMock import httpx +import pytest import litellm +from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig @@ -159,7 +161,7 @@ def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_wit "role": "user", "content": [ {"type": "text", "text": "Describe this"}, - {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png"}}, + {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png", "detail": "high"}}, ], } ] @@ -174,10 +176,10 @@ def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_wit endpoint="openai/deployments/gpt-4.1-mini/chat/completions", ) + text_only_messages = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] assert isinstance(response, ModelResponse) - assert response.usage.prompt_tokens > 0 - assert response.usage.prompt_tokens == litellm.token_counter( - model="gpt-4.1-mini", messages=messages, use_default_image_token_count=True + assert response.usage.prompt_tokens == ( + litellm.token_counter(model="gpt-4.1-mini", messages=text_only_messages) + DEFAULT_IMAGE_TOKEN_COUNT ) @@ -218,3 +220,37 @@ def test_azure_passthrough_url_prefers_the_deployments_api_version(): ) assert url.params["api-version"] == "2024-10-21" + + +def test_azure_passthrough_url_strips_the_leading_router_model_segment(): + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="gpt-4.1-mini/openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={}, + ) + + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + + +def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment(): + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="gpt/openai/deployments/gpt-4o/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "gpt"}}, + ) + + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21" + + +@pytest.mark.parametrize( + "request_data, expected", + [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], +) +def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_data, expected): + assert AzurePassthroughConfig().is_streaming_request(endpoint="openai/deployments/x/chat/completions", request_data=request_data) is expected diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 6f9142c85df..23c26a0b8a3 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -8,9 +8,11 @@ import pytest import litellm +from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, + count_relayed_prompt_tokens, ) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, @@ -2037,3 +2039,47 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: if __name__ == "__main__": pytest.main([__file__]) + + +ONE_PIXEL_PNG_DATA_URL = ( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" +) +UNREACHABLE_IMAGE_URL = "http://127.0.0.1:9/doc.png" +TEXT_ONLY_MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] + + +def _image_messages(url: str, detail: str) -> list[dict]: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": url, "detail": detail}}, + ], + } + ] + + +def test_count_relayed_prompt_tokens_counts_a_data_url_image_exactly(): + messages = _image_messages(ONE_PIXEL_PNG_DATA_URL, "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( + model="gpt-4.1-mini", messages=messages + ) + + +def test_count_relayed_prompt_tokens_keeps_a_low_detail_remote_image_at_the_base_count(): + messages = _image_messages(UNREACHABLE_IMAGE_URL, "low") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( + model="gpt-4.1-mini", messages=messages + ) + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) < DEFAULT_IMAGE_TOKEN_COUNT + + +def test_count_relayed_prompt_tokens_estimates_only_the_remote_high_detail_image(): + messages = _image_messages(UNREACHABLE_IMAGE_URL, "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + DEFAULT_IMAGE_TOKEN_COUNT + ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 78cd84f5b35..3a4abfbc20a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22166 + "limit": 22156 }, "LIT002": { "limit": 26745 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16446 + "limit": 16434 }, "LIT011": { - "limit": 5506 + "limit": 5504 }, "LIT012": { "limit": 4486 From 4706acef958a726f1407ad5413ea2db896d6c825 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 03:48:06 -0700 Subject: [PATCH 09/24] fix(passthrough): charge remote high-detail images at the high-detail upper bound --- litellm/litellm_core_utils/token_counter.py | 7 ++++++ .../openai_passthrough_logging_handler.py | 4 ++-- ruff-strict-budget.json | 4 ++-- .../litellm_core_utils/test_token_counter.py | 22 ++++++++++++++++++- .../test_azure_passthrough_transformation.py | 4 ++-- ...test_openai_passthrough_logging_handler.py | 8 +++---- type-discipline-budget.json | 6 ++--- 7 files changed, 41 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3732ffd734c..4c66d878f14 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -172,6 +172,13 @@ def calculate_tiles_needed( return total_tiles +def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int: + largest_tile_count: Final = calculate_tiles_needed( + MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES + ) + return base_tokens + (base_tokens * 2) * largest_tile_count + + def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]: return struct.unpack(fmt, buffer) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 6a142d08be5..206002d13f7 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -13,11 +13,11 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, ) +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig @@ -133,7 +133,7 @@ def count_relayed_prompt_tokens(model: str, messages: Sequence[Mapping[str, obje ] return ( litellm.token_counter(model=model, messages=local_messages) - + DEFAULT_IMAGE_TOKEN_COUNT * remote_high_detail_images + + high_detail_image_token_upper_bound() * remote_high_detail_images ) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 574bf01cc34..134eec578bc 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2904 + "limit": 2900 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 110 + "limit": 108 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 4694fa8fbed..1898fd57220 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1,5 +1,6 @@ #### What this tests #### # This tests litellm.token_counter.token_counter() function +import base64 import importlib import time import traceback @@ -14,7 +15,11 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function +from litellm.litellm_core_utils.token_counter import ( + _get_tiktoken_count_function, + calculate_img_tokens, + high_detail_image_token_upper_bound, +) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text from tests.test_litellm.litellm_core_utils.messages_with_counts import ( @@ -1412,3 +1417,18 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): assert _count_user_content([prompt, named]) == _count_user_content( [prompt, {"type": "text", "text": "report.pdf"}] ) + + +def _png_data_url(width: int, height: int) -> str: + ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big") + return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() + + +@pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)]) +def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None: + assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound() + + +def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None: + assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound() + assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 5d9e4312b5e..8ffd632fe85 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -5,9 +5,9 @@ import httpx import pytest import litellm -from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig from litellm.types.utils import ModelResponse @@ -179,7 +179,7 @@ def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_wit text_only_messages = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] assert isinstance(response, ModelResponse) assert response.usage.prompt_tokens == ( - litellm.token_counter(model="gpt-4.1-mini", messages=text_only_messages) + DEFAULT_IMAGE_TOKEN_COUNT + litellm.token_counter(model="gpt-4.1-mini", messages=text_only_messages) + high_detail_image_token_upper_bound() ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 23c26a0b8a3..0ddafd20e68 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -8,8 +8,8 @@ import pytest import litellm -from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, count_relayed_prompt_tokens, @@ -2074,12 +2074,12 @@ def test_count_relayed_prompt_tokens_keeps_a_low_detail_remote_image_at_the_base assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( model="gpt-4.1-mini", messages=messages ) - assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) < DEFAULT_IMAGE_TOKEN_COUNT + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) < high_detail_image_token_upper_bound() -def test_count_relayed_prompt_tokens_estimates_only_the_remote_high_detail_image(): +def test_count_relayed_prompt_tokens_charges_only_the_remote_high_detail_image_at_the_upper_bound(): messages = _image_messages(UNREACHABLE_IMAGE_URL, "high") assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( - litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + DEFAULT_IMAGE_TOKEN_COUNT + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3a4abfbc20a..b4c17ffd47b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22156 + "limit": 22146 }, "LIT002": { "limit": 26745 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16434 + "limit": 16422 }, "LIT011": { - "limit": 5504 + "limit": 5502 }, "LIT012": { "limit": 4486 From 5c23d296d31c3eba4a362fd0fa0e8cca96f51426 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 16:51:56 -0700 Subject: [PATCH 10/24] fix(azure_ai): cost OCR relays per page so non-chat relays debit budgets --- .../azure_ai/passthrough/transformation.py | 45 +++++- .../base_llm/passthrough/transformation.py | 3 +- ...est_azure_ai_passthrough_transformation.py | 134 +++++++++++++++++- 3 files changed, 172 insertions(+), 10 deletions(-) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 4776bbe0755..4cb93b4f8f9 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -1,25 +1,28 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from litellm._logging import verbose_logger from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, api_key_header_for_base, get_azure_ai_auth_headers, ) +from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, strip_leading_model_segment from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import StandardPassThroughResponseObject +from litellm.types.utils import CallTypes, StandardPassThroughResponseObject if TYPE_CHECKING: from httpx import URL, Response from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.types.utils import CostResponseTypes @@ -75,6 +78,10 @@ def relayed_body(httpx_response: Response) -> str | dict: class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): + def __init__(self, ocr_config_for: Callable[[str], BaseOCRConfig | None] = get_azure_ai_ocr_config) -> None: + super().__init__() + self.ocr_config_for: Final = ocr_config_for + def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: return bool(request_data.get("stream")) @@ -123,7 +130,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> CostResponseTypes | StandardPassThroughResponseObject | None: + ) -> CostResponseTypes | OCRResponse | StandardPassThroughResponseObject | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict @@ -136,8 +143,40 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): ) if chat_result is not None: return chat_result + ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) + if ocr_result is not None: + return ocr_result return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) + def logged_ocr_response( + self, model: str, httpx_response: Response, logging_obj: Logging, endpoint: str + ) -> OCRResponse | None: + ocr_config: Final = self.ocr_config_for(model) + if ocr_config is None or httpx_response.status_code != 200: + return None + relayed_url: Final = httpx_response.request.url + relayed_origin: Final = str(relayed_url.copy_with(path="/", query=None, fragment=None)).rstrip("/") + ocr_url: Final = httpx.URL( + ocr_config.get_complete_url( + api_base=relayed_origin, + model=model, + optional_params={}, # mutable-ok: BaseOCRConfig wants a dict + ) + ) + known_prefixes: Final = (model, model_group_from(logging_obj.litellm_params)) + native_endpoint: Final = strip_leading_model_segment(endpoint, known_prefixes) + if f"/{native_endpoint.strip('/')}" != ocr_url.path: + return None + try: + ocr_response: Final = ocr_config.transform_ocr_response( + model=model, raw_response=httpx_response, logging_obj=logging_obj + ) + except (ValueError, AttributeError) as error: + verbose_logger.warning("azure_ai passthrough: OCR body from %s is not costable: %s", ocr_url, error) + return None + logging_obj.call_type = CallTypes.aocr.value # rebind-ok: routes cost calculation to the per-page OCR path + return ocr_response + def handle_logging_collected_chunks( self, all_chunks: Sequence[str], diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index d97ecddc75d..464c121ecc1 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException + from ..ocr.transformation import OCRResponse def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: @@ -120,7 +121,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): request_data: dict, logging_obj: "LiteLLMLoggingObj", endpoint: str, - ) -> Optional["CostResponseTypes | StandardPassThroughResponseObject"]: + ) -> Optional["CostResponseTypes | OCRResponse | StandardPassThroughResponseObject"]: pass def handle_logging_collected_chunks( diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index e343d05c5df..644316c2e37 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -1,11 +1,14 @@ import json +from datetime import datetime from unittest.mock import MagicMock import httpx import pytest import litellm +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig +from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager @@ -238,16 +241,135 @@ def _non_chat_logging_result(content: bytes, content_type: str): ) -def test_non_chat_relay_logs_the_parsed_body_so_spend_tracking_sees_the_call(): - result = _non_chat_logging_result(b'{"id":"parse-1","pages":[],"meta":{"billed_units":{"pages":1}}}', "application/json") - - assert result == {"response": {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 1}}}} - - def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} +def _relay_logging_obj(model: str, api_base: str) -> Logging: + logging_obj = Logging( + model=model, + messages=[], + stream=False, + call_type="allm_passthrough_route", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + model=model, + litellm_params={"api_base": api_base, "custom_llm_provider": "azure_ai"}, + optional_params={}, + custom_llm_provider="azure_ai", + ) + return logging_obj + + +def _relay_logging_result( + config: AzureAIPassthroughConfig, model: str, native_path: str, body, api_base: str = FOUNDRY_BASE, status_code: int = 200 +): + relayed_url = f"{FOUNDRY_BASE}/{native_path}?api-version=2024-05-01-preview" + logging_obj = _relay_logging_obj(model, api_base) + response = httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", relayed_url), + ) + result = config.logging_non_streaming_response( + model=model, + custom_llm_provider="azure_ai", + httpx_response=response, + request_data={"model": model}, + logging_obj=logging_obj, + endpoint=f"{model}/{native_path}", + ) + return result, logging_obj + + +MISTRAL_OCR_BODY = { + "pages": [{"index": 0, "markdown": "page one"}, {"index": 1, "markdown": "page two"}], + "model": "mistral-document-ai-2512", + "usage_info": {"pages_processed": 2, "doc_size_bytes": 4321}, +} + + +def test_mistral_document_ai_relay_is_costed_per_page(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY + ) + per_page = litellm.get_model_info("azure_ai/mistral-document-ai-2512")["ocr_cost_per_page"] + + assert isinstance(result, OCRResponse) + assert result.usage_info.pages_processed == 2 + assert per_page > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_page) + + +def test_ocr_route_under_a_models_api_base_is_still_recognised(): + result, _ = _relay_logging_result( + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + MISTRAL_OCR_BODY, + api_base=f"{FOUNDRY_BASE}/models", + ) + + assert isinstance(result, OCRResponse) + + +def test_relay_to_a_non_ocr_route_keeps_the_passthrough_object_and_call_type(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "models/info", {"name": "mistral-document-ai-2512"} + ) + + assert result == {"response": {"name": "mistral-document-ai-2512"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +COHERE_PARSE_BODY = {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 3}}} + + +def test_cohere_parse_relay_is_costed_per_billed_page(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "Cohere-parse-v5", "providers/cohere/v2/parse", COHERE_PARSE_BODY + ) + per_page = litellm.get_model_info("azure_ai/Cohere-parse-v5")["ocr_cost_per_page"] + + assert isinstance(result, OCRResponse) + assert result.usage_info.pages_processed == 3 + assert logging_obj.call_type == "aocr" + assert per_page > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(3 * per_page) + + +def test_deployment_without_an_ocr_config_is_never_costed_as_ocr(): + config = AzureAIPassthroughConfig(ocr_config_for=lambda model: None) + result, logging_obj = _relay_logging_result( + config, "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY + ) + + assert result == {"response": MISTRAL_OCR_BODY} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_accepted_ocr_job_without_a_result_body_is_not_costed(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", {"status": "running"}, status_code=202 + ) + + assert result == {"response": {"status": "running"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", ["not", "an", "ocr", "body"] + ) + + assert result == {"response": '["not", "an", "ocr", "body"]'} + assert logging_obj.call_type == "allm_passthrough_route" + + def test_streaming_chat_completion_chunks_are_costed_like_azure(): head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} chunks = [ From 6f99917b332a1038ac83f5f5a5f8614fa573b1e6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:57:54 -0700 Subject: [PATCH 11/24] fix(router): rewrite multi-segment model groups as whole passthrough path segments --- .../base_llm/passthrough/transformation.py | 4 +- tests/test_litellm/test_router.py | 46 ++++++++++++------- 2 files changed, 32 insertions(+), 18 deletions(-) diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 464c121ecc1..2c21df01b5b 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,3 +1,4 @@ +import re from abc import abstractmethod from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Union @@ -27,7 +28,8 @@ def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: - return "/".join(replacement if part == segment else part for part in endpoint.split("/")) + bounded_segment: Final = re.compile(rf"(? Date: Mon, 7 Sep 2026 18:00:48 -0700 Subject: [PATCH 12/24] fix(azure): let the caller's api-version win over the deployment's on passthrough relays --- litellm/llms/azure/passthrough/transformation.py | 4 ++-- .../test_azure_passthrough_transformation.py | 10 ++++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 983033a8b14..062c60d7231 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -61,11 +61,11 @@ class AzurePassthroughConfig(BasePassthroughConfig): routed_endpoint: Final = replace_path_segment(endpoint, model_group, model) if model_group else endpoint native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,)) + caller_api_version: Final = request_query_params.get("api-version") if request_query_params else None complete_url: Final = BaseAzureLLM._get_base_azure_url( api_base=base_target_url, - litellm_params=litellm_params, + litellm_params={**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")}, route=native_endpoint, - default_api_version=request_query_params.get("api-version") if request_query_params else None, ) return ( httpx.URL(complete_url), diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 8ffd632fe85..262a281aeb3 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -207,18 +207,24 @@ def _complete_url(request_query_params: dict, litellm_params: dict) -> httpx.URL return url -def test_azure_passthrough_url_falls_back_to_the_callers_api_version(): +def test_azure_passthrough_url_forwards_the_callers_api_version(): url = _complete_url(request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={}) assert url.path == "/openai/deployments/gpt-4.1-mini/chat/completions" assert url.params["api-version"] == "2025-04-01-preview" -def test_azure_passthrough_url_prefers_the_deployments_api_version(): +def test_azure_passthrough_url_prefers_the_callers_api_version_over_the_deployments(): url = _complete_url( request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={"api_version": "2024-10-21"} ) + assert url.params["api-version"] == "2025-04-01-preview" + + +def test_azure_passthrough_url_fills_in_the_deployments_api_version_when_the_caller_sends_none(): + url = _complete_url(request_query_params={}, litellm_params={"api_version": "2024-10-21"}) + assert url.params["api-version"] == "2024-10-21" From 17b70035924b7bb2e97065a7f3512c540e105d66 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:11:17 -0700 Subject: [PATCH 13/24] fix(azure_ai): cost embeddings, responses, images, and rerank relays instead of logging zero --- .../llms/azure/passthrough/transformation.py | 45 ++++++- .../azure_ai/passthrough/transformation.py | 32 ++++- .../base_llm/passthrough/transformation.py | 20 ++- .../test_azure_passthrough_transformation.py | 120 +++++++++++++++--- ...est_azure_ai_passthrough_transformation.py | 68 +++++++++- 5 files changed, 261 insertions(+), 24 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 062c60d7231..d7078d94a0b 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Optional import httpx @@ -9,12 +10,14 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, + relayed_json_object, replace_path_segment, strip_leading_model_segment, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse if TYPE_CHECKING: from httpx import URL @@ -38,6 +41,40 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj return details.request_data.messages if details.request_data else None +@dataclass(frozen=True, slots=True) +class OpenAIRelayShape: + path_suffix: str + call_type: CallTypes + parse: Callable[[Mapping[str, object]], EmbeddingResponse | ImageResponse | ResponsesAPIResponse] + + +OPENAI_RELAY_SHAPES: Final = ( + OpenAIRelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), + OpenAIRelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), + OpenAIRelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), +) + + +def logged_openai_response( + httpx_response: Response, logging_obj: Logging, endpoint: str +) -> EmbeddingResponse | ImageResponse | ResponsesAPIResponse | None: + relayed_path: Final = f"/{endpoint.strip('/')}" + shape: Final = next( + (candidate for candidate in OPENAI_RELAY_SHAPES if relayed_path.endswith(candidate.path_suffix)), None + ) + body: Final = relayed_json_object(httpx_response) if shape else None + if shape is None or body is None: + return None + try: + parsed: Final = shape.parse(body) + except ValidationError: + return None + logging_obj.call_type = ( + shape.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return parsed + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -114,13 +151,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): request_data: dict, logging_obj: Logging, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> Optional["CostResponseTypes | ResponsesAPIResponse"]: from litellm import encoding from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.utils import ModelResponse if "chat/completions" not in endpoint: - return None + return logged_openai_response(httpx_response, logging_obj, endpoint) openai_chat_config: Final = OpenAIGPTConfig() diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 4cb93b4f8f9..c49e10566ed 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -14,8 +14,13 @@ from litellm.llms.azure_ai.common_utils import ( get_azure_ai_auth_headers, ) from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, strip_leading_model_segment +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + relayed_json_object, + strip_leading_model_segment, +) from litellm.types.llms.openai import AllMessageValues +from litellm.types.rerank import RerankResponse from litellm.types.utils import CallTypes, StandardPassThroughResponseObject if TYPE_CHECKING: @@ -23,6 +28,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse + from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import CostResponseTypes @@ -77,6 +83,18 @@ def relayed_body(httpx_response: Response) -> str | dict: return body if isinstance(body, dict) else httpx_response.text +def logged_rerank_response(httpx_response: Response, logging_obj: Logging, endpoint: str) -> RerankResponse | None: + body: Final = relayed_json_object(httpx_response) if f"/{endpoint.strip('/')}".endswith("/rerank") else None + if body is None: + return None + try: + rerank_response: Final = RerankResponse.model_validate(body) + except ValidationError: + return None + logging_obj.call_type = CallTypes.arerank.value # rebind-ok: routes cost calculation to the per-query rerank path + return rerank_response + + class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): def __init__(self, ocr_config_for: Callable[[str], BaseOCRConfig | None] = get_azure_ai_ocr_config) -> None: super().__init__() @@ -130,7 +148,14 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> CostResponseTypes | OCRResponse | StandardPassThroughResponseObject | None: + ) -> ( + CostResponseTypes + | OCRResponse + | RerankResponse + | ResponsesAPIResponse + | StandardPassThroughResponseObject + | None + ): from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict @@ -146,6 +171,9 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) if ocr_result is not None: return ocr_result + rerank_result: Final = logged_rerank_response(httpx_response, logging_obj, endpoint) + if rerank_result is not None: + return rerank_result return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) def logged_ocr_response( diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 2c21df01b5b..975b74137db 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -3,18 +3,25 @@ from abc import abstractmethod from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Union +from pydantic import TypeAdapter, ValidationError + from ..base_utils import BaseLLMModelInfo if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.rerank import RerankResponse from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException from ..ocr.transformation import OCRResponse +RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) + + def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: path: Final = endpoint.lstrip("/") for model_name in model_names: @@ -32,6 +39,15 @@ def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: return bounded_segment.sub(lambda _: replacement, endpoint) +def relayed_json_object(httpx_response: "Response") -> Mapping[str, object] | None: + if httpx_response.status_code != 200: + return None + try: + return RELAYED_JSON_OBJECT.validate_python(httpx_response.json()) + except (ValueError, ValidationError): + return None + + class BasePassthroughConfig(BaseLLMModelInfo): @abstractmethod def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: @@ -123,7 +139,9 @@ class BasePassthroughConfig(BaseLLMModelInfo): request_data: dict, logging_obj: "LiteLLMLoggingObj", endpoint: str, - ) -> Optional["CostResponseTypes | OCRResponse | StandardPassThroughResponseObject"]: + ) -> Optional[ + "CostResponseTypes | OCRResponse | RerankResponse | ResponsesAPIResponse | StandardPassThroughResponseObject" + ]: pass def handle_logging_collected_chunks( diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 262a281aeb3..2723a13c3fc 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -1,15 +1,16 @@ import json +from datetime import datetime from unittest.mock import MagicMock import httpx import pytest import litellm - - +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig -from litellm.types.utils import ModelResponse +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import EmbeddingResponse, ModelResponse def _azure_chat_completion_body(): @@ -77,25 +78,112 @@ def test_azure_passthrough_logging_non_streaming_response_chat_completions(): assert result.usage.total_tokens == 18 -def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_returns_none(): - """ - Endpoints other than chat/completions (responses, messages, images) fall - through to None — matches base-class behavior and Bedrock's "unknown - endpoint" handling. Not a regression; just scoping. - """ - config = AzurePassthroughConfig() - logging_obj = MagicMock() - - result = config.logging_non_streaming_response( - model="gpt-4.1-mini", +def _relay_logging_obj(model: str) -> Logging: + logging_obj = Logging( + model=model, + messages=[], + stream=False, + call_type="allm_passthrough_route", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + model=model, + litellm_params={"api_base": "https://my-resource.openai.azure.com", "custom_llm_provider": "azure"}, + optional_params={}, custom_llm_provider="azure", - httpx_response=_make_httpx_response(_azure_chat_completion_body()), + ) + return logging_obj + + +def _relay_logging_result(model: str, endpoint: str, body, status_code: int = 200): + logging_obj = _relay_logging_obj(model) + response = httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview"), + ) + result = AzurePassthroughConfig().logging_non_streaming_response( + model=model, + custom_llm_provider="azure", + httpx_response=response, request_data={}, logging_obj=logging_obj, - endpoint="openai/responses", + endpoint=endpoint, + ) + return result, logging_obj + + +EMBEDDINGS_BODY = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1000, "total_tokens": 1000}, +} + +RESPONSES_BODY = { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4.1-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, +} + + +def test_azure_passthrough_embeddings_relay_is_costed_per_input_token(): + result, logging_obj = _relay_logging_result( + "text-embedding-3-small", "openai/deployments/text-embedding-3-small/embeddings", EMBEDDINGS_BODY + ) + per_token = litellm.get_model_info("azure/text-embedding-3-small")["input_cost_per_token"] + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(1000 * per_token) + + +def test_azure_passthrough_responses_relay_is_costed_per_token(): + result, logging_obj = _relay_logging_result("gpt-4.1-mini", "openai/responses", RESPONSES_BODY) + info = litellm.get_model_info("azure/gpt-4.1-mini") + + assert isinstance(result, ResponsesAPIResponse) + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +def test_azure_passthrough_failed_embeddings_relay_is_not_costed(): + result, logging_obj = _relay_logging_result( + "text-embedding-3-small", + "openai/deployments/text-embedding-3-small/embeddings", + {"error": {"code": "429", "message": "rate limited"}}, + status_code=429, ) assert result is None + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_returns_none(): + result, logging_obj = _relay_logging_result( + "gpt-4o-mini-tts", "openai/deployments/gpt-4o-mini-tts/audio/speech", {"audio": "..."} + ) + + assert result is None + assert logging_obj.call_type == "allm_passthrough_route" def _sse_line(payload: dict) -> str: diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 644316c2e37..788aced50f8 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -9,7 +9,8 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.types.utils import LlmProviders, ModelResponse +from litellm.types.rerank import RerankResponse +from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" @@ -370,6 +371,71 @@ def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): assert logging_obj.call_type == "allm_passthrough_route" +EMBEDDINGS_BODY = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "embed-v-4-0", + "usage": {"prompt_tokens": 1200, "total_tokens": 1200}, +} + +RERANK_BODY = { + "id": "rerank-1", + "results": [{"index": 1, "relevance_score": 0.9}, {"index": 0, "relevance_score": 0.2}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 2}}, +} + +IMAGE_BODY = {"created": 1, "data": [{"b64_json": "AAAA"}]} + + +def test_foundry_embeddings_relay_is_costed_per_input_token(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "embed-v-4-0", "models/embeddings", EMBEDDINGS_BODY + ) + per_token = litellm.get_model_info("azure_ai/embed-v-4-0")["input_cost_per_token"] + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(1200 * per_token) + + +def test_cohere_rerank_relay_is_costed_per_search_unit(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "cohere-rerank-v4.0-fast", "providers/cohere/v2/rerank", RERANK_BODY + ) + per_query = litellm.get_model_info("azure_ai/cohere-rerank-v4.0-fast")["input_cost_per_query"] + + assert isinstance(result, RerankResponse) + assert logging_obj.call_type == "arerank" + assert per_query > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_query) + + +def test_image_generation_relay_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "openai/deployments/FLUX.2-pro/images/generations", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert per_image > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + +def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), + "cohere-rerank-v4.0-fast", + "providers/cohere/v2/rerank", + {"message": "invalid request"}, + status_code=400, + ) + + assert result == {"response": {"message": "invalid request"}} + assert logging_obj.call_type == "allm_passthrough_route" + + def test_streaming_chat_completion_chunks_are_costed_like_azure(): head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} chunks = [ From 12d9413860e96f1cf2b0d138466b34fcae53f925 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:37:58 -0700 Subject: [PATCH 14/24] refactor(passthrough): share the relay shape table and price FLUX 2 provider relays --- .../llms/azure/passthrough/transformation.py | 44 +++----------- .../azure_ai/passthrough/transformation.py | 36 ++++-------- .../base_llm/passthrough/transformation.py | 58 ++++++++++++++----- ...est_azure_ai_passthrough_transformation.py | 49 +++++++++++++--- 4 files changed, 106 insertions(+), 81 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index d7078d94a0b..bdb854bb87a 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,5 +1,4 @@ -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional import httpx @@ -10,7 +9,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, - relayed_json_object, + RelayShape, + logged_relay_shape, replace_path_segment, strip_leading_model_segment, ) @@ -22,6 +22,7 @@ from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse if TYPE_CHECKING: from httpx import URL + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse from litellm.types.utils import CostResponseTypes @@ -41,40 +42,13 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj return details.request_data.messages if details.request_data else None -@dataclass(frozen=True, slots=True) -class OpenAIRelayShape: - path_suffix: str - call_type: CallTypes - parse: Callable[[Mapping[str, object]], EmbeddingResponse | ImageResponse | ResponsesAPIResponse] - - OPENAI_RELAY_SHAPES: Final = ( - OpenAIRelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), - OpenAIRelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), - OpenAIRelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), + RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), + RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), + RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), ) -def logged_openai_response( - httpx_response: Response, logging_obj: Logging, endpoint: str -) -> EmbeddingResponse | ImageResponse | ResponsesAPIResponse | None: - relayed_path: Final = f"/{endpoint.strip('/')}" - shape: Final = next( - (candidate for candidate in OPENAI_RELAY_SHAPES if relayed_path.endswith(candidate.path_suffix)), None - ) - body: Final = relayed_json_object(httpx_response) if shape else None - if shape is None or body is None: - return None - try: - parsed: Final = shape.parse(body) - except ValidationError: - return None - logging_obj.call_type = ( - shape.call_type.value - ) # rebind-ok: routes cost calculation to the relayed shape's pricing path - return parsed - - class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -151,13 +125,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): request_data: dict, logging_obj: Logging, endpoint: str, - ) -> Optional["CostResponseTypes | ResponsesAPIResponse"]: + ) -> Optional["LoggedRelayResponse"]: from litellm import encoding from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.utils import ModelResponse if "chat/completions" not in endpoint: - return logged_openai_response(httpx_response, logging_obj, endpoint) + return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint) openai_chat_config: Final = OpenAIGPTConfig() diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index c49e10566ed..e854874eaee 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -16,19 +16,20 @@ from litellm.llms.azure_ai.common_utils import ( from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, - relayed_json_object, + RelayShape, + logged_relay_shape, strip_leading_model_segment, ) from litellm.types.llms.openai import AllMessageValues from litellm.types.rerank import RerankResponse -from litellm.types.utils import CallTypes, StandardPassThroughResponseObject +from litellm.types.utils import CallTypes, ImageResponse, StandardPassThroughResponseObject if TYPE_CHECKING: from httpx import URL, Response from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse - from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse from litellm.types.utils import CostResponseTypes @@ -83,16 +84,10 @@ def relayed_body(httpx_response: Response) -> str | dict: return body if isinstance(body, dict) else httpx_response.text -def logged_rerank_response(httpx_response: Response, logging_obj: Logging, endpoint: str) -> RerankResponse | None: - body: Final = relayed_json_object(httpx_response) if f"/{endpoint.strip('/')}".endswith("/rerank") else None - if body is None: - return None - try: - rerank_response: Final = RerankResponse.model_validate(body) - except ValidationError: - return None - logging_obj.call_type = CallTypes.arerank.value # rebind-ok: routes cost calculation to the per-query rerank path - return rerank_response +FOUNDRY_RELAY_SHAPES: Final = ( + RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate), + RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate), +) class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): @@ -148,14 +143,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> ( - CostResponseTypes - | OCRResponse - | RerankResponse - | ResponsesAPIResponse - | StandardPassThroughResponseObject - | None - ): + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict @@ -171,9 +159,9 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) if ocr_result is not None: return ocr_result - rerank_result: Final = logged_rerank_response(httpx_response, logging_obj, endpoint) - if rerank_result is not None: - return rerank_result + foundry_result: Final = logged_relay_shape(FOUNDRY_RELAY_SHAPES, httpx_response, logging_obj, endpoint) + if foundry_result is not None: + return foundry_result return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) def logged_ocr_response( diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 975b74137db..93ec5a09e62 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,10 +1,15 @@ +from __future__ import annotations + import re from abc import abstractmethod -from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Optional, Union +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias from pydantic import TypeAdapter, ValidationError +from litellm.types.utils import CallTypes + from ..base_utils import BaseLLMModelInfo if TYPE_CHECKING: @@ -18,6 +23,8 @@ if TYPE_CHECKING: from ..chat.transformation import BaseLLMException from ..ocr.transformation import OCRResponse + LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse + RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) @@ -39,7 +46,7 @@ def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: return bounded_segment.sub(lambda _: replacement, endpoint) -def relayed_json_object(httpx_response: "Response") -> Mapping[str, object] | None: +def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None: if httpx_response.status_code != 200: return None try: @@ -48,6 +55,31 @@ def relayed_json_object(httpx_response: "Response") -> Mapping[str, object] | No return None +@dataclass(frozen=True, slots=True) +class RelayShape: + path_suffix: str + call_type: CallTypes + parse: Callable[[Mapping[str, object]], LoggedRelayResponse] + + +def logged_relay_shape( + shapes: Sequence[RelayShape], httpx_response: Response, logging_obj: LiteLLMLoggingObj, endpoint: str +) -> LoggedRelayResponse | None: + relayed_path: Final = f"/{endpoint.strip('/')}" + shape: Final = next((candidate for candidate in shapes if relayed_path.endswith(candidate.path_suffix)), None) + body: Final = relayed_json_object(httpx_response) if shape else None + if shape is None or body is None: + return None + try: + parsed: Final = shape.parse(body) + except ValidationError: + return None + logging_obj.call_type = ( + shape.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return parsed + + class BasePassthroughConfig(BaseLLMModelInfo): @abstractmethod def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: @@ -60,7 +92,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): endpoint: str, base_target_url: str, request_query_params: Mapping[str, object] | None, - ) -> "URL": + ) -> URL: """ Helper function to add query params to the url Args: @@ -94,7 +126,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): endpoint: str, request_query_params: dict | None, litellm_params: dict, - ) -> tuple["URL", str]: + ) -> tuple[URL, str]: """ Get the complete url for the request Returns: @@ -124,9 +156,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): """ return headers, None - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, "Headers"] - ) -> "BaseLLMException": + def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: from litellm.llms.base_llm.chat.transformation import BaseLLMException return BaseLLMException(status_code=status_code, message=error_message, headers=headers) @@ -135,23 +165,21 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, model: str, custom_llm_provider: str, - httpx_response: "Response", + httpx_response: Response, request_data: dict, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, endpoint: str, - ) -> Optional[ - "CostResponseTypes | OCRResponse | RerankResponse | ResponsesAPIResponse | StandardPassThroughResponseObject" - ]: + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: pass def handle_logging_collected_chunks( self, all_chunks: list[str], - litellm_logging_obj: "LiteLLMLoggingObj", + litellm_logging_obj: LiteLLMLoggingObj, model: str, custom_llm_provider: str, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> CostResponseTypes | None: return None def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 788aced50f8..8b6c76952d6 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -25,7 +25,9 @@ def clear_azure_ai_env(monkeypatch): def test_provider_config_manager_resolves_azure_ai_passthrough_config(): - config = ProviderConfigManager.get_provider_passthrough_config(model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI) + config = ProviderConfigManager.get_provider_passthrough_config( + model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI + ) assert isinstance(config, AzureAIPassthroughConfig) @@ -189,7 +191,10 @@ def test_no_credentials_at_all_raises(): [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], ) def test_is_streaming_request_reads_the_stream_flag(request_data, expected): - assert AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) is expected + assert ( + AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) + is expected + ) def _chat_completion_response() -> httpx.Response: @@ -266,7 +271,12 @@ def _relay_logging_obj(model: str, api_base: str) -> Logging: def _relay_logging_result( - config: AzureAIPassthroughConfig, model: str, native_path: str, body, api_base: str = FOUNDRY_BASE, status_code: int = 200 + config: AzureAIPassthroughConfig, + model: str, + native_path: str, + body, + api_base: str = FOUNDRY_BASE, + status_code: int = 200, ): relayed_url = f"{FOUNDRY_BASE}/{native_path}?api-version=2024-05-01-preview" logging_obj = _relay_logging_obj(model, api_base) @@ -355,7 +365,11 @@ def test_deployment_without_an_ocr_config_is_never_costed_as_ocr(): def test_accepted_ocr_job_without_a_result_body_is_not_costed(): result, logging_obj = _relay_logging_result( - AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", {"status": "running"}, status_code=202 + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + {"status": "running"}, + status_code=202, ) assert result == {"response": {"status": "running"}} @@ -364,7 +378,10 @@ def test_accepted_ocr_job_without_a_result_body_is_not_costed(): def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): result, logging_obj = _relay_logging_result( - AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", ["not", "an", "ocr", "body"] + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + ["not", "an", "ocr", "body"], ) assert result == {"response": '["not", "an", "ocr", "body"]'} @@ -423,6 +440,17 @@ def test_image_generation_relay_is_costed_per_image(): assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) +def test_flux_2_relay_through_the_provider_route_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "providers/blackforestlabs/v1/flux-2-pro", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): result, logging_obj = _relay_logging_result( AzureAIPassthroughConfig(), @@ -439,8 +467,15 @@ def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): def test_streaming_chat_completion_chunks_are_costed_like_azure(): head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} chunks = [ - "data: " + json.dumps({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]}), - "data: " + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), + "data: " + + json.dumps( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + } + ), + "data: " + + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), "data: [DONE]", ] From 02b18b4fc638dfba8fd77e70db796b3ea38c8757 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 19:10:06 -0700 Subject: [PATCH 15/24] fix(passthrough): match image URL schemes case-insensitively when counting relayed prompt tokens --- .../openai_passthrough_logging_handler.py | 4 +++- .../test_openai_passthrough_logging_handler.py | 9 +++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 206002d13f7..93fe3c5b31b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -105,7 +105,9 @@ def _is_remote_high_detail_image(part: object) -> bool: if not isinstance(image_url, Mapping): return False url: Final = image_url.get("url") - return isinstance(url, str) and url.startswith(("http://", "https://")) and image_url.get("detail") == "high" + return ( + isinstance(url, str) and url.lower().startswith(("http://", "https://")) and image_url.get("detail") == "high" + ) def _content_parts(message: Mapping[str, object]) -> Sequence[object]: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 0ddafd20e68..a3d3ae32169 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -2083,3 +2083,12 @@ def test_count_relayed_prompt_tokens_charges_only_the_remote_high_detail_image_a assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() ) + + +@pytest.mark.parametrize("scheme", ["HTTPS://", "Http://"]) +def test_count_relayed_prompt_tokens_charges_an_uppercase_scheme_remote_high_detail_image_at_the_upper_bound(scheme): + messages = _image_messages(scheme + UNREACHABLE_IMAGE_URL.split("://", 1)[1], "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() + ) From 1de369a1a446b7ecc1ed60a0c04d26bed1601c3a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:54:37 -0700 Subject: [PATCH 16/24] fix(passthrough): bound azure relays to the key's model group and reject foreign deployment segments --- litellm/proxy/auth/auth_utils.py | 10 ++ .../llm_passthrough_endpoints.py | 41 ++++++ .../test_azure_passthrough_transformation.py | 4 +- .../proxy/auth/test_auth_utils.py | 32 +++++ .../test_llm_pass_through_endpoints.py | 134 +++++++++++++++++- tests/test_litellm/test_router.py | 4 +- 6 files changed, 217 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 1e4836654a1..dd7df2b3fd9 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1985,9 +1985,19 @@ def get_model_from_request( bedrock_model: Final = _model_from_bedrock_route(route) return model if bedrock_model is None else bedrock_model + if route.lower().startswith(("/azure/", "/azure_ai/")): + azure_model: Final = _router_model_from_azure_route(route, llm_router) + return model if azure_model is None else azure_model + return model +def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import azure_router_model_in_endpoint + + return azure_router_model_in_endpoint(re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE), llm_router) + + def _model_from_bedrock_route(route: str) -> str | None: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _extract_model_from_bedrock_endpoint, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 839c9d47c6f..47b3c56e66a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -78,6 +78,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials +from litellm.types.router import LiteLLMParamsTypedDict from litellm.types.utils import LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager @@ -120,6 +121,37 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li return False +def azure_router_model_in_endpoint(endpoint: str, llm_router: litellm.Router | None) -> str | None: + parts: Final = endpoint.split("/") + if len(parts) < 2: + return None + return next((part for part in parts if is_known_model(part, llm_router)), None) + + +AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? str: + model: Final = litellm_params.get("model", "") + try: + return get_llm_provider(model=model, custom_llm_provider=litellm_params.get("custom_llm_provider"))[0] + except litellm.BadRequestError: + return model + + +def foreign_azure_deployment(endpoint: str, model_group: str, llm_router: litellm.Router) -> str | None: + match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint) + if match is None: + return None + deployment: Final = match.group(1) + if deployment == model_group: + return None + served: Final = frozenset( + _deployment_model_name(row["litellm_params"]) for row in llm_router.get_model_list(model_name=model_group) or () + ) + return None if deployment in served else deployment + + def is_passthrough_request_streaming(request_body: object) -> bool: """ Returns True if the request is streaming. @@ -1523,6 +1555,15 @@ async def _relay_azure_router_model( is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, ) -> Response: + foreign_deployment: Final = foreign_azure_deployment(endpoint, model, llm_router) + if foreign_deployment is not None: + raise HTTPException( + status_code=400, + detail={ + "error": f"deployment '{foreign_deployment}' in the path is not served by model group '{model}'; " + "put the model group name in the deployments segment" + }, + ) try: result: Final = await llm_router.allm_passthrough_route( model=model, diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 2723a13c3fc..54e27b34a59 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -334,12 +334,12 @@ def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment( api_base="https://my-resource.openai.azure.com", api_key="key", model="gpt-4.1-mini", - endpoint="gpt/openai/deployments/gpt-4o/chat/completions", + endpoint="gpt/openai/deployments/gpt-4.1-mini/chat/completions", request_query_params={"api-version": "2024-10-21"}, litellm_params={"litellm_metadata": {"model_group": "gpt"}}, ) - assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21" + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index a996de4d40c..ccb707a8e52 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -528,6 +528,38 @@ def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): ) +def _azure_relay_router(): + from litellm.router import Router + + return Router( + model_list=[ + { + "model_name": "gpt", + "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://a.services.ai.azure.com", "api_key": "k"}, + }, + { + "model_name": "other-group", + "litellm_params": {"model": "azure/gpt-5.4", "api_base": "https://b.openai.azure.com", "api_key": "k"}, + }, + ] + ) + + +@pytest.mark.parametrize( + "route, request_data, expected", + [ + ("/azure_ai/other-group/openai/deployments/other-group/chat/completions", {"model": "gpt"}, "other-group"), + ("/azure_ai/other-group/models/chat/completions", {}, "other-group"), + ("/azure/openai/deployments/gpt/chat/completions", {"model": "other-group"}, "gpt"), + ("/azure/openai/deployments/gpt/chat/completions", {}, "gpt"), + ("/azure/openai/deployments/my-azure-deployment/chat/completions", {"model": "gpt"}, "gpt"), + ("/azure_ai/gpt", {"model": "other-group"}, "other-group"), + ], +) +def test_get_model_from_request_azure_relay_routes_use_the_model_group_in_the_path(route, request_data, expected): + assert get_model_from_request(request_data=request_data, route=route, llm_router=_azure_relay_router()) == expected + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 6ff14157468..fe4400df704 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -4999,7 +4999,11 @@ class TestPassthroughRouterModelBudgetReservation: monkeypatch.setattr(proxy_server, "llm_router", RecordingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) return captured def _assert_metadata_carries_attribution(self, captured: list[dict], user_api_key_dict: UserAPIKeyAuth) -> None: @@ -5098,7 +5102,11 @@ class TestAzureRouterModelStreamingDispatch: monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) request = MagicMock(spec=Request) request.method = "POST" @@ -5158,7 +5166,11 @@ class TestAzureRouterModelStreamingKeepalive: monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) request = MagicMock(spec=Request) request.method = "POST" @@ -5225,7 +5237,11 @@ class TestRouterModelRelayUpstreamContract: monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) def _recording_router(self, captured: list[dict]): class RecordingRouter: @@ -5327,3 +5343,113 @@ async def test_bedrock_count_tokens_error_forwards_provider_headers(): assert exc_info.value.status_code == 500 assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" + + +class _AzureGroupRouter: + def __init__(self, captured: list[dict]) -> None: + self.captured = captured + + def get_model_names(self, team_id=None): + return ["gpt", "other-group"] + + def get_model_list(self, model_name=None, team_id=None): + rows = [ + {"model_name": "gpt", "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_key": "k"}}, + {"model_name": "other-group", "litellm_params": {"model": "azure/gpt-5.4", "api_key": "k"}}, + ] + return [row for row in rows if model_name is None or row["model_name"] == model_name] + + async def allm_passthrough_route(self, **kwargs): + self.captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + +class TestAzureRelayDeploymentSegment: + """A key allowed one model group must not reach another deployment by naming it in the + ``openai/deployments/`` segment while the group segment picks the credential.""" + + @pytest.mark.parametrize( + "endpoint, expected", + [ + ("gpt/openai/deployments/gpt/chat/completions", None), + ("openai/deployments/gpt/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), + ("gpt/models/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), + ("gpt/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/victim/gpt/chat/completions", "victim"), + ], + ) + def test_foreign_azure_deployment_names_a_segment_outside_the_group(self, endpoint, expected): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import foreign_azure_deployment + + assert foreign_azure_deployment(endpoint, "gpt", _AzureGroupRouter([])) == expected + + @pytest.mark.parametrize( + "endpoint, expected", + [ + ("other-group/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/gpt/chat/completions", "gpt"), + ("openai/deployments/my-azure-deployment/chat/completions", None), + ("gpt", None), + ], + ) + def test_azure_router_model_in_endpoint_matches_the_relay_decision(self, endpoint, expected): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import azure_router_model_in_endpoint + + assert azure_router_model_in_endpoint(endpoint, _AzureGroupRouter([])) == expected + + def _install(self, monkeypatch, body: dict) -> list[dict]: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + captured: list[dict] = [] + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", _AzureGroupRouter(captured)) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + return captured + + def _request(self) -> Request: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": "application/json"} + request.query_params = {} + return request + + @pytest.mark.asyncio + async def test_azure_relay_rejects_a_deployment_the_group_does_not_serve(self, monkeypatch): + from fastapi import HTTPException + + captured = self._install(monkeypatch, {"model": "gpt", "messages": []}) + + with pytest.raises(HTTPException) as exc_info: + await azure_proxy_route( + endpoint="gpt/openai/deployments/gpt-5.4/chat/completions", + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token", models=["gpt"]), + ) + + assert exc_info.value.status_code == 400 + assert "gpt-5.4" in exc_info.value.detail["error"] + assert captured == [] + + @pytest.mark.asyncio + async def test_azure_relay_dispatches_the_group_and_its_own_deployment_name(self, monkeypatch): + captured = self._install(monkeypatch, {"model": "gpt", "messages": []}) + + for endpoint in ( + "gpt/openai/deployments/gpt/chat/completions", + "gpt/openai/deployments/gpt-5.4-mini/chat/completions", + ): + await azure_proxy_route( + endpoint=endpoint, + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token", models=["gpt"]), + ) + + assert [call["model"] for call in captured] == ["gpt", "gpt"] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 204c6f109ff..a766f0c9db1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5053,8 +5053,8 @@ def test_get_deployment_model_info_base_model_merge_priority(): ( "gpt", {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, - "gpt/openai/deployments/gpt-4o/chat/completions", - "gpt-5.4-mini/openai/deployments/gpt-4o/chat/completions", + "gpt/openai/deployments/gpt-5.4-mini/chat/completions", + "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", ), ( "aws/anthropic/bedrock-claude", From 52e9cfe254355b2ffdb75fe314c2b62b467f5018 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:57:00 -0700 Subject: [PATCH 17/24] fix(azure): price streamed Responses relays and keep a full-URL api_base from doubling the native path --- litellm/litellm_core_utils/litellm_logging.py | 5 +- .../llms/azure/passthrough/transformation.py | 23 +++++++- .../azure_ai/passthrough/transformation.py | 5 +- .../base_llm/passthrough/transformation.py | 2 +- .../test_azure_passthrough_transformation.py | 47 +++++++++++++++ ...est_azure_ai_passthrough_transformation.py | 57 +++++++++++++++++++ 6 files changed, 129 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ca2cca5360f..ec988b200d0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -119,7 +119,6 @@ from litellm.types.utils import ( CachingDetails, CallTypes, CostBreakdown, - CostResponseTypes, CustomPricingLiteLLMParams, DynamicPromptManagementParamLiteral, EmbeddingResponse, @@ -201,7 +200,7 @@ if TYPE_CHECKING: from mcp.types import EmbeddedResource, ImageContent, TextContent from litellm.integrations.otel.logger import OpenTelemetryV2 - from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, LoggedRelayResponse try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -2363,7 +2362,7 @@ class Logging(LiteLLMLoggingBaseClass): self, raw_bytes: list[bytes], provider_config: "BasePassthroughConfig", - ) -> Optional["CostResponseTypes"]: + ) -> Optional["LoggedRelayResponse"]: all_chunks: Final = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) complete_streaming_response: Final = provider_config.handle_logging_collected_chunks( all_chunks=all_chunks, diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index bdb854bb87a..ab7fe412628 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -23,7 +23,6 @@ if TYPE_CHECKING: from httpx import URL from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse - from litellm.types.utils import CostResponseTypes class RelayedChatRequest(BaseModel): @@ -42,13 +41,29 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj return details.request_data.messages if details.request_data else None +RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate) + OPENAI_RELAY_SHAPES: Final = ( RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), - RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), + RESPONSES_RELAY_SHAPE, RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), ) +def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesAPIResponse | None: + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + terminal_response: Final = OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks( + all_chunks=list(all_chunks) + ) + if terminal_response is None: + return None + logging_obj.call_type = ( + RESPONSES_RELAY_SHAPE.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return terminal_response + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -157,11 +172,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): model: str, custom_llm_provider: str, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> Optional["LoggedRelayResponse"]: from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, ) + if f"/{endpoint.strip('/')}".endswith(RESPONSES_RELAY_SHAPE.path_suffix): + return logged_responses_stream(all_chunks, litellm_logging_obj) if "chat/completions" not in endpoint: return None diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index e854874eaee..1689755b042 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -30,7 +30,6 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse - from litellm.types.utils import CostResponseTypes EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({}) @@ -111,8 +110,8 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): if base_target_url is None: raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") - root: Final = foundry_root(base_target_url) native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) + root: Final = foundry_root(base_target_url).removesuffix(f"/{native_endpoint.strip('/')}") query_params: Final = relay_query_params( request_query_params, api_version_from(litellm_params), base_target_url ) @@ -200,7 +199,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): model: str, custom_llm_provider: str, endpoint: str, - ) -> CostResponseTypes | None: + ) -> LoggedRelayResponse | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig return AzurePassthroughConfig().handle_logging_collected_chunks( diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 93ec5a09e62..dc857dfc808 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -179,7 +179,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): model: str, custom_llm_provider: str, endpoint: str, - ) -> CostResponseTypes | None: + ) -> LoggedRelayResponse | None: return None def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 54e27b34a59..42dd65ea661 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -283,6 +283,53 @@ def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): assert response is None +def _azure_responses_stream_chunks(terminal_event: str | None = "response.completed") -> list[str]: + in_progress = {**RESPONSES_BODY, "status": "in_progress", "output": [], "usage": None} + events = [ + ("response.created", {"type": "response.created", "sequence_number": 0, "response": in_progress}), + ( + "response.output_text.delta", + {"type": "response.output_text.delta", "sequence_number": 1, "item_id": "msg_1", "delta": "hi"}, + ), + ] + ([(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] if terminal_event else []) + return [line for name, payload in events for line in (f"event: {name}", _sse_line(payload))] + + +def test_azure_passthrough_streaming_responses_chunks_are_costed_per_token(): + logging_obj = _relay_logging_obj("gpt-4.1-mini") + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_responses_stream_chunks(), + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/responses", + ) + info = litellm.get_model_info("azure/gpt-4.1-mini") + + assert isinstance(response, ResponsesAPIResponse) + assert response.usage.input_tokens == 1000 + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +def test_azure_passthrough_streaming_responses_without_a_terminal_event_are_not_costed(): + logging_obj = _relay_logging_obj("gpt-4.1-mini") + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_responses_stream_chunks(terminal_event=None), + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/responses", + ) + + assert response is None + assert logging_obj.call_type == "allm_passthrough_route" + + def _complete_url(request_query_params: dict, litellm_params: dict) -> httpx.URL: url, _ = AzurePassthroughConfig().get_complete_url( api_base="https://my-resource.openai.azure.com", diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 8b6c76952d6..0b23da984e8 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -86,6 +86,22 @@ def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): assert base == FOUNDRY_BASE +def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled(): + model_router_url = "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{model_router_url}?api-version=2025-01-01-preview", + api_key="key", + model="model_router/model-router", + endpoint="model-router/chat/completions", + request_query_params=None, + litellm_params={"litellm_metadata": {"model_group": "model-router"}}, + ) + + assert str(url) == f"{model_router_url}?api-version=2025-01-01-preview" + assert base == "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router" + + def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): url, _ = AzureAIPassthroughConfig().get_complete_url( api_base=f"{FOUNDRY_BASE}/models", @@ -490,3 +506,44 @@ def test_streaming_chat_completion_chunks_are_costed_like_azure(): assert isinstance(response, ModelResponse) assert response.choices[0].message.content == "hi" assert response.usage.total_tokens == 4 + + +def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure(): + completed = { + "type": "response.completed", + "sequence_number": 2, + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + }, + } + logging_obj = _relay_logging_obj("gpt-5.4-mini", FOUNDRY_BASE) + + response = AzureAIPassthroughConfig().handle_logging_collected_chunks( + all_chunks=["event: response.completed", "data: " + json.dumps(completed)], + litellm_logging_obj=logging_obj, + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + endpoint="gpt/openai/responses", + ) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert response is not None + assert response.usage.output_tokens == 100 + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) From 13db79da46f41b9775e781ac0df9c1aa14c29f9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:38:41 -0700 Subject: [PATCH 18/24] fix(azure_ai): price streamed Responses relays from their terminal event A streamed Responses API relay handed the success handler a bare ResponsesAPIResponse, which the streaming assembly step drops, so the relay never reached the spend callbacks. Hand it the terminal response.completed event instead, which the assembly step already converts, and cover the whole flush path with a regression test that fails on the previous tip. --- .../llms/azure/passthrough/transformation.py | 11 +-- .../base_llm/passthrough/transformation.py | 4 +- .../llms/openai/responses/transformation.py | 9 +- litellm/types/llms/openai.py | 3 + .../test_azure_passthrough_transformation.py | 35 ++++++-- ...est_azure_ai_passthrough_transformation.py | 90 +++++++++++++------ 6 files changed, 107 insertions(+), 45 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index ab7fe412628..85647ae02c9 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -15,7 +15,7 @@ from litellm.llms.base_llm.passthrough.transformation import ( strip_leading_model_segment, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse, ResponsesTerminalEvent from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse @@ -50,18 +50,19 @@ OPENAI_RELAY_SHAPES: Final = ( ) -def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesAPIResponse | None: +def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None: + """A streaming logging object assembles the logged response from the terminal event, not from its body.""" from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - terminal_response: Final = OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks( + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks( all_chunks=list(all_chunks) ) - if terminal_response is None: + if terminal_event is None: return None logging_obj.call_type = ( RESPONSES_RELAY_SHAPE.call_type.value ) # rebind-ok: routes cost calculation to the relayed shape's pricing path - return terminal_response + return terminal_event class AzurePassthroughConfig(BasePassthroughConfig): diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index dc857dfc808..20180c5cfa2 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -16,14 +16,14 @@ if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesTerminalEvent from litellm.types.rerank import RerankResponse from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException from ..ocr.transformation import OCRResponse - LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse + LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse | ResponsesTerminalEvent RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 926de3e8854..cde399065fd 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -620,15 +620,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return event_pydantic_model.model_construct(**parsed_chunk) @staticmethod - def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + def parse_terminal_event_from_stream_chunks(all_chunks: list[str]) -> ResponsesTerminalEvent | None: for chunk_str in reversed(all_chunks): for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent): try: - return event_model.model_validate_json(chunk_str.removeprefix("data: ")).response + return event_model.model_validate_json(chunk_str.removeprefix("data: ")) except ValueError: continue return None + @staticmethod + def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks) + return None if terminal_event is None else terminal_event.response + @staticmethod def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]: """ diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b6da9490e01..b7c4371f32f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1564,6 +1564,9 @@ class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): response: ResponsesAPIResponse +ResponsesTerminalEvent: TypeAlias = ResponseCompletedEvent | ResponseIncompleteEvent | ResponseFailedEvent + + class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.RESPONSE_PART_ADDED] item_id: str diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 42dd65ea661..6e85b0cdca6 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -9,7 +9,7 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -103,7 +103,9 @@ def _relay_logging_result(model: str, endpoint: str, body, status_code: int = 20 status_code=status_code, headers={"content-type": "application/json"}, content=json.dumps(body).encode("utf-8"), - request=httpx.Request("POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview"), + request=httpx.Request( + "POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview" + ), ) result = AzurePassthroughConfig().logging_non_streaming_response( model=model, @@ -291,7 +293,11 @@ def _azure_responses_stream_chunks(terminal_event: str | None = "response.comple "response.output_text.delta", {"type": "response.output_text.delta", "sequence_number": 1, "item_id": "msg_1", "delta": "hi"}, ), - ] + ([(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] if terminal_event else []) + ] + ( + [(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] + if terminal_event + else [] + ) return [line for name, payload in events for line in (f"event: {name}", _sse_line(payload))] @@ -307,10 +313,10 @@ def test_azure_passthrough_streaming_responses_chunks_are_costed_per_token(): ) info = litellm.get_model_info("azure/gpt-4.1-mini") - assert isinstance(response, ResponsesAPIResponse) - assert response.usage.input_tokens == 1000 + assert isinstance(response, ResponseCompletedEvent) + assert response.response.usage.input_tokens == 1000 assert logging_obj.call_type == "aresponses" - assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] ) @@ -373,7 +379,10 @@ def test_azure_passthrough_url_strips_the_leading_router_model_segment(): litellm_params={}, ) - assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment(): @@ -386,7 +395,10 @@ def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment( litellm_params={"litellm_metadata": {"model_group": "gpt"}}, ) - assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) @pytest.mark.parametrize( @@ -394,4 +406,9 @@ def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment( [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], ) def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_data, expected): - assert AzurePassthroughConfig().is_streaming_request(endpoint="openai/deployments/x/chat/completions", request_data=request_data) is expected + assert ( + AzurePassthroughConfig().is_streaming_request( + endpoint="openai/deployments/x/chat/completions", request_data=request_data + ) + is expected + ) diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 0b23da984e8..064be954518 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -6,6 +6,7 @@ import httpx import pytest import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -14,6 +15,36 @@ from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, from litellm.utils import ProviderConfigManager FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" +RESPONSES_COMPLETED_EVENT = { + "type": "response.completed", + "sequence_number": 2, + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + }, +} + + +class _SpendProbe(CustomLogger): + logged_call_type: str | None = None + logged_cost: float | None = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_call_type = kwargs["call_type"] + self.logged_cost = kwargs["response_cost"] @pytest.fixture(autouse=True) @@ -87,7 +118,9 @@ def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled(): - model_router_url = "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + model_router_url = ( + "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + ) url, base = AzureAIPassthroughConfig().get_complete_url( api_base=f"{model_router_url}?api-version=2025-01-01-preview", @@ -267,21 +300,29 @@ def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} -def _relay_logging_obj(model: str, api_base: str) -> Logging: +def _relay_logging_obj( + model: str, + api_base: str, + stream: bool = False, + callbacks: list[CustomLogger] | None = None, + endpoint: str = "", +) -> Logging: logging_obj = Logging( model=model, messages=[], - stream=False, + stream=stream, call_type="allm_passthrough_route", start_time=datetime.now(), litellm_call_id="call-1", function_id="fn-1", + dynamic_async_success_callbacks=callbacks, ) logging_obj.update_environment_variables( model=model, litellm_params={"api_base": api_base, "custom_llm_provider": "azure_ai"}, optional_params={}, custom_llm_provider="azure_ai", + endpoint=endpoint, ) return logging_obj @@ -509,31 +550,10 @@ def test_streaming_chat_completion_chunks_are_costed_like_azure(): def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure(): - completed = { - "type": "response.completed", - "sequence_number": 2, - "response": { - "id": "resp_1", - "object": "response", - "created_at": 1, - "status": "completed", - "model": "gpt-5.4-mini", - "output": [ - { - "type": "message", - "id": "msg_1", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "hi", "annotations": []}], - } - ], - "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, - }, - } logging_obj = _relay_logging_obj("gpt-5.4-mini", FOUNDRY_BASE) response = AzureAIPassthroughConfig().handle_logging_collected_chunks( - all_chunks=["event: response.completed", "data: " + json.dumps(completed)], + all_chunks=["event: response.completed", "data: " + json.dumps(RESPONSES_COMPLETED_EVENT)], litellm_logging_obj=logging_obj, model="gpt-5.4-mini", custom_llm_provider="azure_ai", @@ -542,8 +562,24 @@ def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure info = litellm.get_model_info("azure_ai/gpt-5.4-mini") assert response is not None - assert response.usage.output_tokens == 100 + assert response.response.usage.output_tokens == 100 assert logging_obj.call_type == "aresponses" - assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] ) + + +async def test_streaming_responses_relay_flush_reaches_the_success_callbacks_with_a_price(): + probe = _SpendProbe() + logging_obj = _relay_logging_obj( + "gpt-5.4-mini", FOUNDRY_BASE, stream=True, callbacks=[probe], endpoint="gpt/openai/responses" + ) + stream = "event: response.completed\ndata: " + json.dumps(RESPONSES_COMPLETED_EVENT) + "\n\n" + + await logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=[stream.encode()], provider_config=AzureAIPassthroughConfig() + ) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert probe.logged_call_type == "allm_passthrough_route" + assert probe.logged_cost == pytest.approx(1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"]) From db59e99932f9438ac364c731cb1584092a76c94d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:21:34 -0700 Subject: [PATCH 19/24] fix(azure): let the caller's api-version override a full-URL api_base on relays --- .../llms/azure/passthrough/transformation.py | 9 +++++- .../test_azure_passthrough_transformation.py | 32 +++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 85647ae02c9..dc71c53e126 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -65,6 +65,12 @@ def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> return terminal_event +def without_api_version(api_base: str) -> str: + url: Final = httpx.URL(api_base) + kept_params: Final = tuple((key, value) for key, value in url.params.multi_items() if key != "api-version") + return str(url.copy_with(params=httpx.QueryParams(kept_params))) + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -89,8 +95,9 @@ class AzurePassthroughConfig(BasePassthroughConfig): native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,)) caller_api_version: Final = request_query_params.get("api-version") if request_query_params else None + relay_base: Final = without_api_version(base_target_url) if caller_api_version else base_target_url complete_url: Final = BaseAzureLLM._get_base_azure_url( - api_base=base_target_url, + api_base=relay_base, litellm_params={**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")}, route=native_endpoint, ) diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 6e85b0cdca6..a3f5c415417 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -369,6 +369,38 @@ def test_azure_passthrough_url_fills_in_the_deployments_api_version_when_the_cal assert url.params["api-version"] == "2024-10-21" +FULL_URL_API_BASE = ( + "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" +) + + +def _full_url_complete_url(request_query_params: dict) -> httpx.URL: + url, _ = AzurePassthroughConfig().get_complete_url( + api_base=FULL_URL_API_BASE, + api_key="key", + model="gpt-4.1-mini", + endpoint="chat/completions", + request_query_params=request_query_params, + litellm_params={}, + ) + return url + + +def test_azure_passthrough_url_prefers_the_callers_api_version_over_a_full_url_api_bases(): + url = _full_url_complete_url(request_query_params={"api-version": "2025-04-01-preview"}) + + assert str(url) == ( + "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions" + "?api-version=2025-04-01-preview" + ) + + +def test_azure_passthrough_url_keeps_a_full_url_api_bases_api_version_when_the_caller_sends_none(): + url = _full_url_complete_url(request_query_params={}) + + assert url.params["api-version"] == "2024-10-21" + + def test_azure_passthrough_url_strips_the_leading_router_model_segment(): url, _ = AzurePassthroughConfig().get_complete_url( api_base="https://my-resource.openai.azure.com", From 5e056a264ecff102a72ee36fed333e5e2ecd9c37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:45:05 -0700 Subject: [PATCH 20/24] refactor(azure): move the passthrough deployment-segment helpers under llms --- .../llms/azure/passthrough/transformation.py | 23 +++++++++++- litellm/proxy/auth/auth_utils.py | 8 +++-- .../llm_passthrough_endpoints.py | 26 ++++---------- .../test_azure_passthrough_transformation.py | 35 ++++++++++++++++++- .../test_llm_pass_through_endpoints.py | 33 +++-------------- 5 files changed, 71 insertions(+), 54 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index dc71c53e126..51a409c61da 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping, Sequence +import re +from collections.abc import Callable, Collection, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional import httpx @@ -65,6 +66,26 @@ def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> return terminal_event +AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? str | None: + parts: Final = endpoint.split("/") + if len(parts) < 2: + return None + return next((part for part in parts if part in router_models), None) + + +def foreign_azure_deployment( + endpoint: str, model_group: str, served_models: Callable[[], Collection[str]] +) -> str | None: + match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint) + if match is None: + return None + deployment: Final = match.group(1) + return None if deployment == model_group or deployment in served_models() else deployment + + def without_api_version(api_base: str) -> str: url: Final = httpx.URL(api_base) kept_params: Final = tuple((key, value) for key, value in url.params.multi_items() if key != "api-version") diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index bb6ebddf77a..aa77835e858 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.url_utils import ( provider_url_destination_candidates, validate_url, ) +from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint from litellm.proxy._types import * from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -2011,9 +2012,10 @@ def get_model_from_request( def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None: - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import azure_router_model_in_endpoint - - return azure_router_model_in_endpoint(re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE), llm_router) + if llm_router is None: + return None + endpoint: Final = re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE) + return azure_router_model_in_endpoint(endpoint, frozenset(llm_router.get_model_names())) def _model_from_bedrock_route(route: str) -> str | None: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 47b3c56e66a..66987042962 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -33,6 +33,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse @@ -121,16 +122,6 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li return False -def azure_router_model_in_endpoint(endpoint: str, llm_router: litellm.Router | None) -> str | None: - parts: Final = endpoint.split("/") - if len(parts) < 2: - return None - return next((part for part in parts if is_known_model(part, llm_router)), None) - - -AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? str: model: Final = litellm_params.get("model", "") try: @@ -139,17 +130,10 @@ def _deployment_model_name(litellm_params: LiteLLMParamsTypedDict) -> str: return model -def foreign_azure_deployment(endpoint: str, model_group: str, llm_router: litellm.Router) -> str | None: - match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint) - if match is None: - return None - deployment: Final = match.group(1) - if deployment == model_group: - return None - served: Final = frozenset( +def _models_served_by_group(llm_router: litellm.Router, model_group: str) -> frozenset[str]: + return frozenset( _deployment_model_name(row["litellm_params"]) for row in llm_router.get_model_list(model_name=model_group) or () ) - return None if deployment in served else deployment def is_passthrough_request_streaming(request_body: object) -> bool: @@ -1555,7 +1539,9 @@ async def _relay_azure_router_model( is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, ) -> Response: - foreign_deployment: Final = foreign_azure_deployment(endpoint, model, llm_router) + foreign_deployment: Final = foreign_azure_deployment( + endpoint, model, lambda: _models_served_by_group(llm_router, model) + ) if foreign_deployment is not None: raise HTTPException( status_code=400, diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index a3f5c415417..69fdd59a748 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -8,7 +8,11 @@ import pytest import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound -from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig +from litellm.llms.azure.passthrough.transformation import ( + AzurePassthroughConfig, + azure_router_model_in_endpoint, + foreign_azure_deployment, +) from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -444,3 +448,32 @@ def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_da ) is expected ) + + +@pytest.mark.parametrize( + "endpoint, expected", + [ + ("gpt/openai/deployments/gpt/chat/completions", None), + ("openai/deployments/gpt/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), + ("gpt/models/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), + ("gpt/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/victim/gpt/chat/completions", "victim"), + ], +) +def test_foreign_azure_deployment_names_a_segment_outside_the_group(endpoint, expected): + assert foreign_azure_deployment(endpoint, "gpt", lambda: frozenset({"gpt-5.4-mini"})) == expected + + +@pytest.mark.parametrize( + "endpoint, expected", + [ + ("other-group/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/gpt/chat/completions", "gpt"), + ("openai/deployments/my-azure-deployment/chat/completions", None), + ("gpt", None), + ], +) +def test_azure_router_model_in_endpoint_picks_the_first_router_model_segment(endpoint, expected): + assert azure_router_model_in_endpoint(endpoint, frozenset({"gpt", "other-group"})) == expected diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index fe4400df704..d1f9e5c4c1d 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -5368,36 +5368,11 @@ class TestAzureRelayDeploymentSegment: """A key allowed one model group must not reach another deployment by naming it in the ``openai/deployments/`` segment while the group segment picks the credential.""" - @pytest.mark.parametrize( - "endpoint, expected", - [ - ("gpt/openai/deployments/gpt/chat/completions", None), - ("openai/deployments/gpt/chat/completions", None), - ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), - ("gpt/models/chat/completions", None), - ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), - ("gpt/openai/deployments/other-group/chat/completions", "other-group"), - ("openai/deployments/victim/gpt/chat/completions", "victim"), - ], - ) - def test_foreign_azure_deployment_names_a_segment_outside_the_group(self, endpoint, expected): - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import foreign_azure_deployment + def test_models_served_by_group_resolves_each_deployment_to_its_model_name(self): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _models_served_by_group - assert foreign_azure_deployment(endpoint, "gpt", _AzureGroupRouter([])) == expected - - @pytest.mark.parametrize( - "endpoint, expected", - [ - ("other-group/openai/deployments/other-group/chat/completions", "other-group"), - ("openai/deployments/gpt/chat/completions", "gpt"), - ("openai/deployments/my-azure-deployment/chat/completions", None), - ("gpt", None), - ], - ) - def test_azure_router_model_in_endpoint_matches_the_relay_decision(self, endpoint, expected): - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import azure_router_model_in_endpoint - - assert azure_router_model_in_endpoint(endpoint, _AzureGroupRouter([])) == expected + assert _models_served_by_group(_AzureGroupRouter([]), "gpt") == frozenset({"gpt-5.4-mini"}) + assert _models_served_by_group(_AzureGroupRouter([]), "missing-group") == frozenset() def _install(self, monkeypatch, body: dict) -> list[dict]: import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep From a892e67c40c3d84064f4c9f3d83c74bb1e4295a5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:05:06 -0700 Subject: [PATCH 21/24] fix(azure): match deployment segments case-insensitively and keep relay helpers immutable --- litellm/llms/azure/passthrough/transformation.py | 15 ++++++++++----- litellm/llms/openai/responses/transformation.py | 2 +- .../llm_passthrough_endpoints.py | 16 +++++++++------- .../test_azure_passthrough_transformation.py | 10 ++++++++++ 4 files changed, 30 insertions(+), 13 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 51a409c61da..69767bb29a3 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,5 +1,6 @@ import re from collections.abc import Callable, Collection, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional import httpx @@ -55,9 +56,7 @@ def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> """A streaming logging object assembles the logged response from the terminal event, not from its body.""" from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks( - all_chunks=list(all_chunks) - ) + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks) if terminal_event is None: return None logging_obj.call_type = ( @@ -83,7 +82,11 @@ def foreign_azure_deployment( if match is None: return None deployment: Final = match.group(1) - return None if deployment == model_group or deployment in served_models() else deployment + folded: Final = deployment.casefold() + if folded == model_group.casefold(): + return None + served: Final = frozenset(name.casefold() for name in served_models()) + return None if folded in served else deployment def without_api_version(api_base: str) -> str: @@ -119,7 +122,9 @@ class AzurePassthroughConfig(BasePassthroughConfig): relay_base: Final = without_api_version(base_target_url) if caller_api_version else base_target_url complete_url: Final = BaseAzureLLM._get_base_azure_url( api_base=relay_base, - litellm_params={**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")}, + litellm_params=MappingProxyType( + {**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")} + ), route=native_endpoint, ) return ( diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index cde399065fd..97e7bcc60b7 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -620,7 +620,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return event_pydantic_model.model_construct(**parsed_chunk) @staticmethod - def parse_terminal_event_from_stream_chunks(all_chunks: list[str]) -> ResponsesTerminalEvent | None: + def parse_terminal_event_from_stream_chunks(all_chunks: Sequence[str]) -> ResponsesTerminalEvent | None: for chunk_str in reversed(all_chunks): for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent): try: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 66987042962..086e1654f70 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -122,6 +122,10 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li return False +class RelayRejection(TypedDict): + error: ReadOnly[str] + + def _deployment_model_name(litellm_params: LiteLLMParamsTypedDict) -> str: model: Final = litellm_params.get("model", "") try: @@ -1543,13 +1547,11 @@ async def _relay_azure_router_model( endpoint, model, lambda: _models_served_by_group(llm_router, model) ) if foreign_deployment is not None: - raise HTTPException( - status_code=400, - detail={ - "error": f"deployment '{foreign_deployment}' in the path is not served by model group '{model}'; " - "put the model group name in the deployments segment" - }, - ) + rejection: Final[RelayRejection] = { + "error": f"deployment '{foreign_deployment}' in the path is not served by model group '{model}'; " + "put the model group name in the deployments segment" + } + raise HTTPException(status_code=400, detail=rejection) try: result: Final = await llm_router.allm_passthrough_route( model=model, diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 69fdd59a748..798ef1f5a4f 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -456,16 +456,26 @@ def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_da ("gpt/openai/deployments/gpt/chat/completions", None), ("openai/deployments/gpt/chat/completions", None), ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), + ("gpt/openai/deployments/GPT-5.4-MINI/chat/completions", None), + ("gpt/openai/deployments/Gpt/chat/completions", None), ("gpt/models/chat/completions", None), ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), ("gpt/openai/deployments/other-group/chat/completions", "other-group"), ("openai/deployments/victim/gpt/chat/completions", "victim"), + ("gpt/openai/deployments/GPT-5.4/chat/completions", "GPT-5.4"), ], ) def test_foreign_azure_deployment_names_a_segment_outside_the_group(endpoint, expected): assert foreign_azure_deployment(endpoint, "gpt", lambda: frozenset({"gpt-5.4-mini"})) == expected +def test_foreign_azure_deployment_skips_the_router_when_the_segment_is_the_group_itself(): + def served_models(): + raise AssertionError("the router must not be consulted for the group's own name") + + assert foreign_azure_deployment("gpt/openai/deployments/Gpt/chat/completions", "gpt", served_models) is None + + @pytest.mark.parametrize( "endpoint, expected", [ From 5ea447e295109b4079d1b41a70f82f1d88024e6a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:19:23 -0700 Subject: [PATCH 22/24] fix(azure): keep the model group segment exact and casefold only deployment names in the relay guard --- litellm/llms/azure/passthrough/transformation.py | 5 ++--- .../passthrough/test_azure_passthrough_transformation.py | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 69767bb29a3..1b5a4083ebe 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -82,11 +82,10 @@ def foreign_azure_deployment( if match is None: return None deployment: Final = match.group(1) - folded: Final = deployment.casefold() - if folded == model_group.casefold(): + if deployment == model_group: return None served: Final = frozenset(name.casefold() for name in served_models()) - return None if folded in served else deployment + return None if deployment.casefold() in served else deployment def without_api_version(api_base: str) -> str: diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 798ef1f5a4f..c7e86616ee2 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -457,7 +457,7 @@ def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_da ("openai/deployments/gpt/chat/completions", None), ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), ("gpt/openai/deployments/GPT-5.4-MINI/chat/completions", None), - ("gpt/openai/deployments/Gpt/chat/completions", None), + ("gpt/openai/deployments/Gpt/chat/completions", "Gpt"), ("gpt/models/chat/completions", None), ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), ("gpt/openai/deployments/other-group/chat/completions", "other-group"), @@ -473,7 +473,7 @@ def test_foreign_azure_deployment_skips_the_router_when_the_segment_is_the_group def served_models(): raise AssertionError("the router must not be consulted for the group's own name") - assert foreign_azure_deployment("gpt/openai/deployments/Gpt/chat/completions", "gpt", served_models) is None + assert foreign_azure_deployment("gpt/openai/deployments/gpt/chat/completions", "gpt", served_models) is None @pytest.mark.parametrize( From a25eccafc952aa654c9d826cfff18e568c32d817 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:32:13 -0700 Subject: [PATCH 23/24] fix(azure_ai): drop the api_base path segments a relay already repeats --- .../azure_ai/passthrough/transformation.py | 18 +++++++++++++++++- ...test_azure_ai_passthrough_transformation.py | 18 ++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 1689755b042..ed630bc6f2b 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -62,6 +62,22 @@ def foundry_root(api_base: str) -> str: return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/") +def without_repeated_native_prefix(root: str, native_endpoint: str) -> str: + url: Final = httpx.URL(root) + root_segments: Final = tuple(segment for segment in url.path.split("/") if segment) + native_segments: Final = tuple(segment.casefold() for segment in native_endpoint.split("/") if segment) + overlap: Final = next( + ( + length + for length in range(min(len(root_segments), len(native_segments)), 0, -1) + if tuple(segment.casefold() for segment in root_segments[-length:]) == native_segments[:length] + ), + 0, + ) + kept_segments: Final = root_segments[: len(root_segments) - overlap] + return str(url.copy_with(path="/" + "/".join(kept_segments), query=None)).rstrip("/") + + def relay_query_params( request_query_params: Mapping[str, object] | None, deployment_api_version: str | None, @@ -111,7 +127,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) - root: Final = foundry_root(base_target_url).removesuffix(f"/{native_endpoint.strip('/')}") + root: Final = without_repeated_native_prefix(foundry_root(base_target_url), native_endpoint) query_params: Final = relay_query_params( request_query_params, api_version_from(litellm_params), base_target_url ) diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 064be954518..1b9036c50c5 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -135,6 +135,24 @@ def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled assert base == "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router" +@pytest.mark.parametrize("relayed_deployment", ["gpt-4o", "GPT-4o"]) +def test_deployment_root_api_base_is_not_repeated_when_the_relay_carries_the_deployment_path(relayed_deployment): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com/openai/deployments/gpt-4o", + api_key="key", + model="gpt-4o", + endpoint=f"aoai-gpt-4o/openai/deployments/{relayed_deployment}/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "aoai-gpt-4o"}}, + ) + + assert str(url) == ( + f"https://my-resource.openai.azure.com/openai/deployments/{relayed_deployment}/chat/completions" + "?api-version=2024-10-21" + ) + assert base == "https://my-resource.openai.azure.com" + + def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): url, _ = AzureAIPassthroughConfig().get_complete_url( api_base=f"{FOUNDRY_BASE}/models", From 8627576c9c37bc475ac35d98ab6f44d1ea8978d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:11:04 -0700 Subject: [PATCH 24/24] fix(azure_ai): keep a deployment named like the first native path segment in the relayed URL --- .../llms/azure_ai/passthrough/transformation.py | 5 +++++ .../test_azure_ai_passthrough_transformation.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index ed630bc6f2b..f2be1d95593 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -62,6 +62,10 @@ def foundry_root(api_base: str) -> str: return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/") +def is_repeated_native_prefix(native_segments: tuple[str, ...], overlap: int) -> bool: + return overlap == len(native_segments) or native_segments[0] == "openai" + + def without_repeated_native_prefix(root: str, native_endpoint: str) -> str: url: Final = httpx.URL(root) root_segments: Final = tuple(segment for segment in url.path.split("/") if segment) @@ -71,6 +75,7 @@ def without_repeated_native_prefix(root: str, native_endpoint: str) -> str: length for length in range(min(len(root_segments), len(native_segments)), 0, -1) if tuple(segment.casefold() for segment in root_segments[-length:]) == native_segments[:length] + and is_repeated_native_prefix(native_segments, length) ), 0, ) diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 1b9036c50c5..c8007acf70f 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -153,6 +153,20 @@ def test_deployment_root_api_base_is_not_repeated_when_the_relay_carries_the_dep assert base == "https://my-resource.openai.azure.com" +def test_deployment_named_like_the_first_native_segment_keeps_its_deployment_root(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com/openai/deployments/chat", + api_key="key", + model="chat", + endpoint="aoai-chat/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "aoai-chat"}}, + ) + + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/chat/chat/completions?api-version=2024-10-21" + assert base == "https://my-resource.openai.azure.com/openai/deployments/chat" + + def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): url, _ = AzureAIPassthroughConfig().get_complete_url( api_base=f"{FOUNDRY_BASE}/models",