Merge pull request #39614 from BerriAI/litellm_fix_stream_usage_default_openai_hosts

fix(openai): default stream usage on PrivateLink and regional api.openai.com hosts

litellm/llms/openai/common_utils.py

@@ -11,6 +11,7 @@ import time
 import uuid
 from collections.abc import AsyncIterator, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Optional
+from urllib.parse import urlsplit
 
 import httpx
 import openai
@@ -43,6 +44,14 @@ _OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(OpenAI)
 _AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(AzureOpenAI)
 
 
+_OPENAI_API_HOST: Final[str] = "api.openai.com"
+
+
+def is_openai_backed_api_base(api_base: str) -> bool:
+    hostname: Final = urlsplit(api_base).hostname
+    return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}"))
+
+
 class OpenAIError(BaseLLMException):
     def __init__(
         self,

litellm/llms/openai/openai.py

@@ -2,7 +2,6 @@ import time
 import types
 from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
-from urllib.parse import urlparse
 
 import httpx
 
@@ -55,6 +54,7 @@ from .common_utils import (
     OpenAIError,
     build_output_token_limit_response,
     drop_params_from_unprocessable_entity_error,
+    is_openai_backed_api_base,
     is_output_token_limit_error,
 )
 from .workload_identity import resolve_openai_workload_identity_config
@@ -1190,10 +1190,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         """
         if stream_options is not None:
             return {"stream_options": stream_options}
-        else:
-            # by default litellm will include usage for openai endpoints
-            if api_base is None or urlparse(api_base).hostname == "api.openai.com":
-                return {"stream_options": {"include_usage": True}}
+        if api_base is None or is_openai_backed_api_base(api_base):
+            return {"stream_options": {"include_usage": True}}
         return {}
 
     # Embedding

tests/test_litellm/llms/openai/test_openai.py

@@ -0,0 +1,52 @@
+import pytest
+
+from litellm.llms.openai.openai import OpenAIChatCompletion
+
+
+@pytest.mark.parametrize(
+    "api_base",
+    [
+        None,
+        "https://api.openai.com/v1",
+        "https://api.openai.com:443/v1",
+        "https://southcentralus.privatelink.api.openai.com/v1",
+        "https://eu.api.openai.com/v1",
+        "https://us.api.openai.com/v1",
+        "HTTPS://API.OPENAI.COM/v1/",
+    ],
+)
+def test_get_stream_options_defaults_include_usage_on_every_openai_backed_host(api_base):
+    """
+    PrivateLink and regional hostnames reach the real OpenAI backend, so a stream with no caller
+    stream_options must ask for the usage chunk exactly as the default base does. Regression guard
+    for LIT-6875: spend for those deployments fell back to local token counting.
+    """
+    assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {
+        "stream_options": {"include_usage": True}
+    }
+
+
+@pytest.mark.parametrize(
+    "api_base",
+    [
+        "https://my-gateway.example/v1",
+        "https://api.openai.com.evil.example/v1",
+        "https://notapi.openai.com/v1",
+        "https://gateway.example/v1?upstream=api.openai.com",
+        "https://openai.internal.example/api.openai.com/v1",
+    ],
+)
+def test_get_stream_options_leaves_foreign_hosts_without_a_usage_default(api_base):
+    """Only the host decides: an OpenAI-compatible backend elsewhere may not support stream_options at all."""
+    assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {}
+
+
+@pytest.mark.parametrize(
+    "api_base",
+    ["https://southcentralus.privatelink.api.openai.com/v1", "https://my-gateway.example/v1"],
+)
+def test_get_stream_options_passes_caller_stream_options_through_on_any_host(api_base):
+    caller_options = {"include_usage": False}
+    assert OpenAIChatCompletion().get_stream_options(stream_options=caller_options, api_base=api_base) == {
+        "stream_options": caller_options
+    }

tests/test_litellm/llms/openai/test_openai_common_utils.py

@@ -7,7 +7,7 @@ import pytest
 
 import litellm
 from litellm.litellm_core_utils.token_counter import token_counter
-from litellm.llms.openai.common_utils import BaseOpenAILLM
+from litellm.llms.openai.common_utils import BaseOpenAILLM, is_openai_backed_api_base
 
 # Test parameters for different API functions
 API_FUNCTION_PARAMS = [
@@ -392,3 +392,22 @@ async def test_async_genuine_bad_request_still_raises(provider, stream):
 
     with pytest.raises(litellm.BadRequestError):
         await _call_and_drain()
+
+
+@pytest.mark.parametrize(
+    ("api_base", "expected"),
+    [
+        ("https://api.openai.com/v1", True),
+        ("https://api.openai.com:443/v1/", True),
+        ("https://southcentralus.privatelink.api.openai.com/v1", True),
+        ("https://eu.api.openai.com/v1", True),
+        ("HTTPS://API.OPENAI.COM/v1", True),
+        ("https://my-gateway.example/v1", False),
+        ("https://api.openai.com.evil.example/v1", False),
+        ("https://notapi.openai.com/v1", False),
+        ("https://gateway.example/v1?upstream=api.openai.com", False),
+        ("not a url", False),
+    ],
+)
+def test_is_openai_backed_api_base_decides_by_hostname_only(api_base, expected):
+    assert is_openai_backed_api_base(api_base) is expected