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
2026-09-15 23:31:29 +00:00 · 2026-09-04 23:32:27 -07:00 · 2026-09-04 23:32:27 -07:00 · cd25eb9189
commit cd25eb9189
parent a276690ce2
11 changed files with 260 additions and 35 deletions
--- 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
--- 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,
+        )
--- 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,
+        )
--- 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)
--- 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"}
--- 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:
--- 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
--- 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"
--- 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
--- 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"
--- 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