diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2276efc4e34..76af64c8ccd 100644 --- a/.github/workflows/test-unit.yml +++ b/.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" diff --git a/litellm/main.py b/litellm/main.py index 27700cf9def..32a7c290bb8 100644 --- a/litellm/main.py +++ b/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, diff --git a/litellm/text_completions/__init__.py b/litellm/text_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py new file mode 100644 index 00000000000..ec1b96fd432 --- /dev/null +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index 7de7ce6188d..73057f140a8 100644 --- a/litellm/utils.py +++ b/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, diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index cdff03af3b4..ad069aa987b 100644 --- a/tests/unit/test_utils.py +++ b/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 ( diff --git a/tests/unit/text_completions/__init__.py b/tests/unit/text_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py new file mode 100644 index 00000000000..b7edffc0753 --- /dev/null +++ b/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")