fix(completions): omit stream usage unless requested (#47953) (#48052)

* fix(completions): omit stream usage unless requested

Co-authored-by: HootRock <wtr1274970944@163.com>

* refactor(completions): move stream wrapper into text completions

* refactor(completions): preserve typed compatibility exports

* fix(completions): preserve lazy export typing

* ci(completions): run streaming tests in endpoint shard

---------

Co-authored-by: HootRock <wtr1274970944@163.com>

.github/workflows/test-unit.yml

@@ -425,6 +425,7 @@ jobs:
               tests/unit/realtime_api
               tests/unit/rerank_api
               tests/unit/search
+              tests/unit/text_completions
               tests/unit/vector_stores
               tests/unit/videos
             python-version: "3.12"

litellm/main.py

@@ -143,6 +143,7 @@ from litellm.router_utils.reasoning_effort_capability import (
     resolve_chat_reasoning,
 )
 from litellm.secret_managers.main import get_secret_bool, get_secret_str
+from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
 from litellm.types.completion import (
     CompletionDispatchContext,
     _CompletionDispatchResult,
@@ -169,7 +170,6 @@ from litellm.utils import (
     ProviderConfigManager,
     TextChoices,
     TextCompletionResponse,
-    TextCompletionStreamWrapper,
     TranscriptionResponse,
     Usage,
     add_provider_specific_params_to_optional_params,

litellm/text_completions/streaming_iterator.py

@@ -0,0 +1,88 @@
+from typing import Final
+
+from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse
+
+
+class TextCompletionStreamWrapper:
+    def __init__(
+        self,
+        completion_stream,
+        model,
+        stream_options: dict | None = None,
+        custom_llm_provider: str | None = None,
+    ):
+        self.completion_stream = completion_stream
+        self.model = model
+        self.stream_options = stream_options
+        self.custom_llm_provider = custom_llm_provider
+
+    def __iter__(self):
+        return self
+
+    def __aiter__(self):
+        return self
+
+    def convert_to_text_completion_object(self, chunk: ModelResponse):
+        try:
+            response: Final = TextCompletionResponse()
+            response["id"] = chunk.get("id", None)
+            response["object"] = "text_completion"
+            response["created"] = chunk.get("created", None)
+            response["model"] = chunk.get("model", None)
+            text_choices: Final = TextChoices()
+            if isinstance(chunk, Choices):  # chunk should always be of type StreamingChoices
+                raise Exception
+
+            if chunk["choices"]:
+                delta: Final = chunk["choices"][0]["delta"]
+                text_choices["text"] = delta["content"]
+                text_choices["reasoning_content"] = delta.get("reasoning_content")
+                text_choices["index"] = chunk["choices"][0]["index"]
+                text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"]
+                response["choices"] = [text_choices]
+            else:
+                response["choices"] = []
+
+            # only pass usage when stream_options["include_usage"] is True
+            if self.stream_options and self.stream_options.get("include_usage", False) is True:
+                response.usage = chunk.get("usage", None)
+            else:
+                response.usage = None
+
+            return response
+        except Exception as e:
+            raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {e}")
+
+    def __next__(self):
+        # model_response = ModelResponse(stream=True, model=self.model)
+        TextCompletionResponse()
+        try:
+            for chunk in self.completion_stream:
+                if chunk == "None" or chunk is None:
+                    raise Exception
+                processed_chunk = self.convert_to_text_completion_object(chunk=chunk)
+                return processed_chunk
+            raise StopIteration
+        except StopIteration:
+            raise StopIteration
+        except Exception as e:
+            from litellm.litellm_core_utils.exception_mapping_utils import exception_type
+
+            raise exception_type(
+                model=self.model,
+                custom_llm_provider=self.custom_llm_provider or "",
+                original_exception=e,
+                completion_kwargs={},
+                extra_kwargs={},
+            )
+
+    async def __anext__(self):
+        try:
+            async for chunk in self.completion_stream:
+                if chunk == "None" or chunk is None:
+                    raise Exception
+                processed_chunk = self.convert_to_text_completion_object(chunk=chunk)
+                return processed_chunk
+            raise StopIteration
+        except StopIteration:
+            raise StopAsyncIteration

litellm/utils.py

@@ -218,6 +218,7 @@ def _get_cached_audio_utils():
     return _audio_utils_module
 
 
+from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper as _TextCompletionStreamWrapper
 from litellm.types.llms.openai import (
     AllMessageValues,
     AllPromptValues,
@@ -262,13 +263,17 @@ from litellm.types.utils import (
     SearchProviders,
     SelectTokenizerResponse,
     StreamingChoices,
-    TextChoices,
-    TextCompletionResponse,
     TranscriptionResponse,
     Usage,
     credential_api_key_secret_reference,
     is_litellm_owned_kwarg,
 )
+from litellm.types.utils import (
+    TextChoices as _TextChoices,
+)
+from litellm.types.utils import (
+    TextCompletionResponse as _TextCompletionResponse,
+)
 
 _CALL_TYPE_ENUM_MAP: Final[dict] = {ct.value: ct for ct in CallTypes}
 _BACKFILL_MODES: Final = frozenset({"chat", "responses"})
@@ -299,10 +304,14 @@ except (ImportError, AttributeError, TypeError):
 claude_json_str = json.dumps(json_data)
 import importlib.metadata
 from collections.abc import AsyncIterator, Callable, Collection, Iterable, Iterator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable
 
 from typing_extensions import assert_never
 
+TextChoices: TypeAlias = _TextChoices
+TextCompletionResponse: TypeAlias = _TextCompletionResponse
+TextCompletionStreamWrapper: TypeAlias = _TextCompletionStreamWrapper
+
 from litellm import utils as litellm_utils
 
 # These are lazy loaded via __getattr__
@@ -340,7 +349,8 @@ if TYPE_CHECKING:
 
     # Type stubs for lazy-loaded functions to help mypy understand their types
     # These imports allow mypy to understand the types when these are accessed via __getattr__
-    from litellm.litellm_core_utils.exception_mapping_utils import exception_type
+    exception_type: Callable[..., Exception]
+
     from litellm.litellm_core_utils.get_litellm_params import (
         get_base_model_from_litellm_call_metadata,
         get_litellm_params,
@@ -7945,88 +7955,6 @@ def register_prompt_template(
     return litellm.custom_prompt_dict
 
 
-class TextCompletionStreamWrapper:
-    def __init__(
-        self,
-        completion_stream,
-        model,
-        stream_options: dict | None = None,
-        custom_llm_provider: str | None = None,
-    ):
-        self.completion_stream = completion_stream
-        self.model = model
-        self.stream_options = stream_options
-        self.custom_llm_provider = custom_llm_provider
-
-    def __iter__(self):
-        return self
-
-    def __aiter__(self):
-        return self
-
-    def convert_to_text_completion_object(self, chunk: ModelResponse):
-        try:
-            response: Final = TextCompletionResponse()
-            response["id"] = chunk.get("id", None)
-            response["object"] = "text_completion"
-            response["created"] = chunk.get("created", None)
-            response["model"] = chunk.get("model", None)
-            text_choices: Final = TextChoices()
-            if isinstance(chunk, Choices):  # chunk should always be of type StreamingChoices
-                raise Exception
-
-            if chunk["choices"]:
-                delta: Final = chunk["choices"][0]["delta"]
-                text_choices["text"] = delta["content"]
-                text_choices["reasoning_content"] = delta.get("reasoning_content")
-                text_choices["index"] = chunk["choices"][0]["index"]
-                text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"]
-                response["choices"] = [text_choices]
-            else:
-                response["choices"] = []
-
-            # only pass usage when stream_options["include_usage"] is True
-            if self.stream_options and self.stream_options.get("include_usage", False) is True:
-                response["usage"] = chunk.get("usage", None)
-
-            return response
-        except Exception as e:
-            raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {e}")
-
-    def __next__(self):
-        # model_response = ModelResponse(stream=True, model=self.model)
-        TextCompletionResponse()
-        try:
-            for chunk in self.completion_stream:
-                if chunk == "None" or chunk is None:
-                    raise Exception
-                processed_chunk = self.convert_to_text_completion_object(chunk=chunk)
-                return processed_chunk
-            raise StopIteration
-        except StopIteration:
-            raise StopIteration
-        except Exception as e:
-            exception_type: Final = litellm_utils.exception_type
-            raise exception_type(
-                model=self.model,
-                custom_llm_provider=self.custom_llm_provider or "",
-                original_exception=e,
-                completion_kwargs={},
-                extra_kwargs={},
-            )
-
-    async def __anext__(self):
-        try:
-            async for chunk in self.completion_stream:
-                if chunk == "None" or chunk is None:
-                    raise Exception
-                processed_chunk = self.convert_to_text_completion_object(chunk=chunk)
-                return processed_chunk
-            raise StopIteration
-        except StopIteration:
-            raise StopAsyncIteration
-
-
 def mock_stream_usage_chunk(model_response: ModelResponseStream, model: str, prompt_tokens: int) -> ModelResponseStream:
     return ModelResponseStream(
         id=model_response.id,

tests/unit/test_utils.py

@@ -90,7 +90,6 @@ from litellm.utils import (
     get_llm_provider,
     get_optional_params_embeddings,
     ProviderConfigManager,
-    TextCompletionStreamWrapper,
     _check_provider_match,
     get_potential_model_names,
     _is_litellm_router_call,
@@ -2528,41 +2527,6 @@ def test_priced_router_deployment_id_stays_out_of_wildcard_model_listing():
     assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.000001
 
 
-def test_reasoning_content_preserved_in_text_completion_wrapper():
-    """Ensure reasoning_content is copied from delta to text_choices."""
-    chunk = ModelResponseStream(
-        id="test-id",
-        created=1234567890,
-        model="test-model",
-        object="chat.completion.chunk",
-        choices=[
-            StreamingChoices(
-                finish_reason=None,
-                index=0,
-                delta=Delta(
-                    content="Some answer text",
-                    role="assistant",
-                    reasoning_content="Here's my chain of thought...",
-                ),
-            )
-        ],
-    )
-
-    wrapper = TextCompletionStreamWrapper(
-        completion_stream=None,  # Not used in convert_to_text_completion_object
-        model="test-model",
-        stream_options=None,
-    )
-
-    transformed = wrapper.convert_to_text_completion_object(chunk)
-
-    assert "choices" in transformed
-    assert len(transformed["choices"]) == 1
-    choice = transformed["choices"][0]
-    assert choice["text"] == "Some answer text"
-    assert choice["reasoning_content"] == "Here's my chain of thought..."
-
-
 def test_anthropic_claude_4_invoke_chat_provider_config():
     """Test that the Anthropic Claude 4 Invoke chat provider config is correct."""
     from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (

tests/unit/text_completions/test_streaming_iterator.py

@@ -0,0 +1,144 @@
+from collections.abc import AsyncIterator, Iterator
+from typing import Final
+
+import pytest
+
+from litellm.exceptions import APIConnectionError
+from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
+from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, Usage
+
+
+def test_reasoning_content_preserved_in_text_completion_wrapper():
+    """Ensure reasoning_content is copied from delta to text_choices."""
+    chunk = ModelResponseStream(
+        id="test-id",
+        created=1234567890,
+        model="test-model",
+        object="chat.completion.chunk",
+        choices=[
+            StreamingChoices(
+                finish_reason=None,
+                index=0,
+                delta=Delta(
+                    content="Some answer text",
+                    role="assistant",
+                    reasoning_content="Here's my chain of thought...",
+                ),
+            )
+        ],
+    )
+
+    wrapper = TextCompletionStreamWrapper(
+        completion_stream=None,  # Not used in convert_to_text_completion_object
+        model="test-model",
+        stream_options=None,
+    )
+
+    transformed = wrapper.convert_to_text_completion_object(chunk)
+
+    assert "choices" in transformed
+    assert len(transformed["choices"]) == 1
+    choice = transformed["choices"][0]
+    assert choice["text"] == "Some answer text"
+    assert choice["reasoning_content"] == "Here's my chain of thought..."
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("async_stream", (False, True))
+@pytest.mark.parametrize("empty_choices", (False, True))
+@pytest.mark.parametrize("stream_options", (None, {}, {"include_usage": False}, {"include_usage": True}))
+@pytest.mark.parametrize("usage", (None, Usage(prompt_tokens=7, completion_tokens=2, total_tokens=9), Usage()))
+async def test_text_completion_stream_usage_is_opt_in(
+    async_stream: bool, empty_choices: bool, stream_options: dict[str, bool] | None, usage: Usage | None
+) -> None:
+    chunk: Final = ModelResponseStream(
+        id="test-id",
+        created=1234567890,
+        model="test-model",
+        choices=[] if empty_choices else [StreamingChoices(index=0, delta=Delta(content="Hello"))],
+        usage=usage,
+    )
+
+    async def chunks() -> AsyncIterator[ModelResponseStream]:
+        yield chunk
+
+    wrapper: Final = TextCompletionStreamWrapper(
+        completion_stream=chunks() if async_stream else iter((chunk,)),
+        model="test-model",
+        stream_options=stream_options,
+    )
+    responses: Final = [response async for response in wrapper] if async_stream else list(wrapper)
+    expected_usage: Final = (
+        {"usage": usage.model_dump(exclude_none=True)}
+        if stream_options and stream_options.get("include_usage") is True and usage is not None
+        else {}
+    )
+    assert [response.model_dump(exclude_none=True) for response in responses] == [
+        {
+            "id": "test-id",
+            "object": "text_completion",
+            "created": 1234567890,
+            "model": "test-model",
+            "choices": [] if empty_choices else [{"index": 0, "text": "Hello"}],
+            **expected_usage,
+        }
+    ]
+
+
+def test_text_completion_conversion_rejects_non_streaming_choices() -> None:
+    chunk: Final = Choices()
+    wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model")
+
+    with pytest.raises(Exception, match="Error occurred converting to text completion object") as error:
+        wrapper.convert_to_text_completion_object(chunk)
+
+    assert str(chunk) in str(error.value)
+
+
+@pytest.mark.parametrize("chunk", (None, "None"))
+def test_sync_text_completion_stream_rejects_invalid_chunks(chunk: str | None) -> None:
+    wrapper: Final = TextCompletionStreamWrapper(
+        completion_stream=iter((chunk,)), model="test-model", custom_llm_provider="openai"
+    )
+
+    with pytest.raises(APIConnectionError, match="OpenAIException") as error:
+        next(wrapper)
+
+    assert (error.value.model, error.value.llm_provider) == ("test-model", "openai")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("chunk", (None, "None"))
+async def test_async_text_completion_stream_rejects_invalid_chunks(chunk: str | None) -> None:
+    async def chunks() -> AsyncIterator[str | None]:
+        yield chunk
+
+    wrapper: Final = TextCompletionStreamWrapper(completion_stream=chunks(), model="test-model")
+
+    with pytest.raises(Exception, match=r"^$") as error:
+        await anext(wrapper)
+
+    assert not isinstance(error.value, StopAsyncIteration)
+
+
+def test_sync_text_completion_stream_maps_upstream_failure_after_content() -> None:
+    chunk: Final = ModelResponseStream(
+        id="test-id",
+        created=1234567890,
+        model="test-model",
+        choices=[StreamingChoices(index=0, delta=Delta(content="Hello"))],
+    )
+
+    def chunks() -> Iterator[ModelResponseStream]:
+        yield chunk
+        raise RuntimeError("upstream disconnected")
+
+    wrapper: Final = TextCompletionStreamWrapper(
+        completion_stream=chunks(), model="test-model", custom_llm_provider="openai"
+    )
+    assert next(wrapper).choices[0].text == "Hello"
+
+    with pytest.raises(APIConnectionError, match="upstream disconnected") as error:
+        next(wrapper)
+
+    assert (error.value.model, error.value.llm_provider) == ("test-model", "openai")