fix(azure_ai): relay from the Foundry root and log non-chat relays for spend tracking

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 (

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,

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():