fix(azure_ai): add passthrough config so router-model relays reach the deployment's own endpoint

basedpyright-code-budget.json

@@ -57,7 +57,7 @@
     "limit": 5570
   },
   "reportMissingTypeArgument": {
-    "limit": 15281
+    "limit": 15276
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -105,13 +105,13 @@
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 38283
+    "limit": 38266
   },
   "reportUnknownParameterType": {
-    "limit": 19584
+    "limit": 19579
   },
   "reportUnknownVariableType": {
-    "limit": 29829
+    "limit": 29817
   },
   "reportUnnecessaryCast": {
     "limit": 110

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
@@ -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,8 +992,9 @@ def calculate_usage(
         chunks: Sequence["_UsageBearingChunk | ModelResponse"],
         model: str,
         completion_output: str,
-        messages: list | None = None,
+        messages: Sequence | None = None,
         reasoning_tokens: int | None = None,
+        count_prompt_tokens: Callable[[], int] | None = None,
     ) -> Usage:
         """
         Calculate usage for the given chunks.
@@ -1018,7 +1019,9 @@ def calculate_usage(
         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 (
+                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")
             returned_usage.prompt_tokens = 0

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)
 

litellm/llms/azure/passthrough/transformation.py

@@ -1,11 +1,17 @@
+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
-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
@@ -16,9 +22,25 @@
     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
+        return bool(request_data.get("stream"))
 
     def get_complete_url(
         self,
@@ -36,14 +58,14 @@ def get_complete_url(
 
         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,
-            default_api_version=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),
@@ -116,3 +138,25 @@ def logging_non_streaming_response(
         )
 
         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,
+            messages=_relayed_messages(litellm_logging_obj),
+        )

litellm/llms/azure_ai/chat/transformation.py

@@ -2,7 +2,6 @@
 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 @@
     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 @@ def _should_use_api_key_header(self, api_base: str) -> bool:
         """
         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,

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.

litellm/llms/azure_ai/passthrough/transformation.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final
+
+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, 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
+
+    from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm.types.utils import CostResponseTypes
+
+
+EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({})
+
+
+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 ""
+
+
+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 bool(request_data.get("stream"))
+
+    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")
+
+        root: Final = foundry_root(base_target_url)
+        native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params)))
+        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,
+        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 | 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
+            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,
+        )
+        if chat_result is not None:
+            return chat_result
+        return StandardPassThroughResponseObject(response=relayed_body(httpx_response))
+
+    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,
+        )