diff --git a/litellm/litellm_core_utils/stream_usage_chunks.py b/litellm/litellm_core_utils/stream_usage_chunks.py new file mode 100644 index 00000000000..eda57c7ea12 --- /dev/null +++ b/litellm/litellm_core_utils/stream_usage_chunks.py @@ -0,0 +1,27 @@ +from collections.abc import Iterable +from typing import Final + +from litellm.types.utils import ModelResponseStream, StreamingChoices + + +def is_empty_streaming_choice(choice: StreamingChoices) -> bool: + if choice.finish_reason is not None: + return False + if getattr(choice, "logprobs", None) is not None: + return False + delta: Final = choice.delta if hasattr(choice, "delta") else None + if delta is None: + return True + return _all_values_none(delta.model_dump().values()) + + +def _all_values_none(values: Iterable[object]) -> bool: + return all(value is None for value in values) + + +def is_injected_stream_usage_artifact(chunk: object) -> bool: + if not isinstance(chunk, ModelResponseStream): + return False + if chunk.provider_specific_fields is not None: + return False + return all(is_empty_streaming_choice(choice) for choice in chunk.choices or []) diff --git a/litellm/main.py b/litellm/main.py index 50ab274a060..0162742de18 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7975,6 +7975,8 @@ async def atext_completion(*args, **kwargs) -> TextCompletionResponse | TextComp model=model, custom_llm_provider=custom_llm_provider, stream_options=kwargs.get("stream_options"), + usage_injected_by_proxy="_litellm_strip_stream_usage" in kwargs + and kwargs["_litellm_strip_stream_usage"] is True, ) else: ## OpenAI / Azure Text Completion Returns here diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e810a5ab329..666269be2eb 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -588,6 +588,7 @@ def model_deployments_support_stream_options( _STREAM_USAGE_FLAGS: Final = ("_litellm_strip_stream_usage", STREAM_USAGE_INJECTED_FIELD) +_STREAM_USAGE_TRACKED_ROUTES: Final = frozenset({"acompletion", "atext_completion"}) def stream_usage_tracking_updates( @@ -613,7 +614,7 @@ def stream_usage_tracking_updates( return scrub if always_include is False and route_type == "agenerate_content_stream" and "stream_options" not in data: return {**scrub, "stream_options": {"include_usage": False}} - if always_include is False or route_type != "acompletion": + if always_include is False or route_type not in _STREAM_USAGE_TRACKED_ROUTES: return scrub if isinstance(stream_options, dict) and stream_options.get("include_usage") is True: return scrub diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8f60311e491..a87fc815402 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2505,6 +2505,7 @@ def refresh_proxy_server_request_body_snapshot( "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG, STREAM_USAGE_INJECTED_FIELD, + "_litellm_strip_stream_usage", } ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d4ec26614c0..e6d157a7f81 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -201,7 +201,6 @@ from litellm.types.utils import ( CountContextManagement, ModelResponse, ModelResponseStream, - StreamingChoices, TextCompletionResponse, TokenCountMessagesFormat, TokenCountResponse, @@ -384,6 +383,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, mask_sensitive_keys, ) +from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.llms.anthropic.messages.convertible_blocks import ( deployment_drops_params, drop_non_string_tool_result_texts, @@ -10214,29 +10214,6 @@ def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes: return chunk.model_dump_json(exclude_none=True, exclude_unset=True) -def _is_injected_stream_usage_artifact(chunk: object) -> bool: - if not isinstance(chunk, ModelResponseStream): - return False - if chunk.provider_specific_fields is not None: - return False - return all(_is_empty_streaming_choice(choice) for choice in chunk.choices or []) - - -def _is_empty_streaming_choice(choice: StreamingChoices) -> bool: - if choice.finish_reason is not None: - return False - if getattr(choice, "logprobs", None) is not None: - return False - delta: Final = choice.delta if hasattr(choice, "delta") else None - if delta is None: - return True - return _all_values_none(delta.model_dump().values()) - - -def _all_values_none(values: Iterable[object]) -> bool: - return all(value is None for value in values) - - async def _apply_streaming_chunk_hooks( *, chunk: Any, @@ -10616,7 +10593,7 @@ async def async_data_generator( fallback_model_from_metadata=fallback_model_from_metadata, ) - if strip_stream_usage and _is_injected_stream_usage_artifact(chunk): + if strip_stream_usage and is_injected_stream_usage_artifact(chunk): if pending_fallback_event: yield _format_fallback_metadata_sse_event( fallback_model=fallback_model_from_metadata, diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index ec1b96fd432..61e47e9d7af 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,6 +1,7 @@ from typing import Final -from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse +from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact +from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse class TextCompletionStreamWrapper: @@ -10,11 +11,18 @@ class TextCompletionStreamWrapper: model, stream_options: dict | None = None, custom_llm_provider: str | None = None, + usage_injected_by_proxy: bool = False, ): self.completion_stream = completion_stream self.model = model self.stream_options = stream_options self.custom_llm_provider = custom_llm_provider + self.usage_injected_by_proxy: Final = usage_injected_by_proxy + self.include_usage: Final = ( + not usage_injected_by_proxy + and stream_options is not None + and stream_options.get("include_usage", False) is True + ) def __iter__(self): return self @@ -22,7 +30,7 @@ class TextCompletionStreamWrapper: def __aiter__(self): return self - def convert_to_text_completion_object(self, chunk: ModelResponse): + def convert_to_text_completion_object(self, chunk: ModelResponse | ModelResponseStream): try: response: Final = TextCompletionResponse() response["id"] = chunk.get("id", None) @@ -43,8 +51,7 @@ class TextCompletionStreamWrapper: 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: + if self.include_usage: response.usage = chunk.get("usage", None) else: response.usage = None @@ -53,6 +60,9 @@ class TextCompletionStreamWrapper: except Exception as e: raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {e}") + def _is_injected_usage_artifact(self, chunk: ModelResponseStream) -> bool: + return self.usage_injected_by_proxy and is_injected_stream_usage_artifact(chunk) + def __next__(self): # model_response = ModelResponse(stream=True, model=self.model) TextCompletionResponse() @@ -60,8 +70,9 @@ class TextCompletionStreamWrapper: 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 + if isinstance(chunk, ModelResponseStream) and self._is_injected_usage_artifact(chunk): + continue + return self.convert_to_text_completion_object(chunk=chunk) raise StopIteration except StopIteration: raise StopIteration @@ -81,8 +92,9 @@ class TextCompletionStreamWrapper: 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 + if isinstance(chunk, ModelResponseStream) and self._is_injected_usage_artifact(chunk): + continue + return self.convert_to_text_completion_object(chunk=chunk) raise StopIteration except StopIteration: raise StopAsyncIteration diff --git a/tests/integration/spend/test_text_completion_stream_usage_cost.py b/tests/integration/spend/test_text_completion_stream_usage_cost.py new file mode 100644 index 00000000000..4d66e60f3e7 --- /dev/null +++ b/tests/integration/spend/test_text_completion_stream_usage_cost.py @@ -0,0 +1,185 @@ +import json +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +PROMPT_TOKENS: Final = 19 +COMPLETION_TOKENS: Final = 5 +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +EXPECTED_SPEND: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE +USAGE: Final[dict[str, JsonValue]] = { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, +} +PROMPT: Final = "Say hi" + + +def _sse_frame(payload: dict[str, JsonValue]) -> bytes: + return f"data: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def _chunk(request_id: str, model: str, delta: dict[str, JsonValue], finish_reason: str | None) -> dict[str, JsonValue]: + return { + "id": request_id, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + + +def _respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=json.dumps({"object": "list", "data": [{"id": "gpt-5.4-mini", "object": "model"}]}).encode()) + assert request.method == "POST", request.method + assert request.target == "/v1/chat/completions", request.target + body: Final = json.loads(request.body) + assert body["messages"] == [{"role": "user", "content": PROMPT}], body + request_id: Final = f"chatcmpl-{uuid4().hex}" + model: Final = str(body["model"]) + if body.get("stream") is not True: + return Reply( + body=json.dumps( + { + "id": request_id, + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "Hi there"}, "finish_reason": "length"} + ], + "usage": USAGE, + } + ).encode() + ) + stream_options: Final = body.get("stream_options") or {} + usage_frames: Final = ( + ( + _sse_frame( + { + "id": request_id, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [], + "usage": USAGE, + } + ), + ) + if stream_options.get("include_usage") is True + else () + ) + return Reply( + content_type="text/event-stream", + chunks=( + _sse_frame(_chunk(request_id, model, {"role": "assistant", "content": "Hi"}, None)), + _sse_frame(_chunk(request_id, model, {"content": " there"}, None)), + _sse_frame(_chunk(request_id, model, {}, "length")), + *usage_frames, + b"data: [DONE]\n\n", + ), + ) + + +def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path: + config: Final = directory / "text_completion_stream_usage_cost_config.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": model, + "litellm_params": { + "model": "openai/gpt-5.4-mini", + "api_base": f"{upstream_url}/v1", + "api_key": "scripted-upstream", + "input_cost_per_token": INPUT_RATE, + "output_cost_per_token": OUTPUT_RATE, + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + "disable_spend_logs": False, + "proxy_batch_write_at": 1, + "proxy_batch_polling_interval": 1, + }, + "router_settings": {"disable_cooldowns": True}, + } + ) + ) + return config + + +def _data_events(body: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + json.loads(line.removeprefix("data:")) + for line in body.splitlines() + if line.startswith("data:") and not line.endswith("[DONE]") + ) + + +def _carries_text_or_finish(frame: dict[str, JsonValue]) -> bool: + choices: Final = frame["choices"] + assert isinstance(choices, list) and len(choices) == 1, frame + choice: Final = choices[0] + assert isinstance(choice, dict), frame + return choice.get("text") is not None or choice.get("finish_reason") is not None + + +def _billed(row: dict[str, object]) -> tuple[object, object, float]: + return (row["prompt_tokens"], row["completion_tokens"], float(str(row["spend"]))) + + +def test_text_completion_stream_without_stream_options_bills_like_its_twins(gateway: Gateway, tmp_path: Path) -> None: + model: Final = f"gpt-5.4-mini-{uuid4().hex}" + body: Final[dict[str, JsonValue]] = {"model": model, "prompt": PROMPT, "max_tokens": COMPLETION_TOKENS} + + with wire_server(_respond) as wire: + config: Final = _proxy_config(tmp_path, model, wire.url) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + silent: Final = candidate.request("POST", "/v1/completions", {**body, "stream": True}) + asked: Final = candidate.request( + "POST", "/v1/completions", {**body, "stream": True, "stream_options": {"include_usage": True}} + ) + whole: Final = candidate.request("POST", "/v1/completions", body) + assert (silent.status_code, asked.status_code, whole.status_code) == (200, 200, 200), ( + silent.text, + asked.text, + whole.text, + ) + silent_frames: Final = _data_events(silent.text) + asked_frames: Final = _data_events(asked.text) + request_ids: Final = (silent_frames[0]["id"], asked_frames[0]["id"], whole.json()["id"]) + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" ' + "WHERE request_id = ANY(%s)", + (list(request_ids),), + database_url=candidate.database_url, + ), + lambda values: len(values) == 3, + seconds=70, + ) + upstream_stream_options: Final = [ + json.loads(sent.body).get("stream_options") for sent in wire.drain() if sent.method == "POST" + ] + + assert upstream_stream_options == [{"include_usage": True}, {"include_usage": True}, None] + assert all(_carries_text_or_finish(frame) for frame in silent_frames), silent.text + assert all("usage" not in frame for frame in silent_frames), silent.text + assert asked_frames[-1]["usage"]["total_tokens"] == PROMPT_TOKENS + COMPLETION_TOKENS, asked.text + billed: Final = {str(row["request_id"]): _billed(row) for row in rows} + assert billed[str(request_ids[0])] == billed[str(request_ids[1])] == billed[str(request_ids[2])], billed + assert billed[str(request_ids[0])] == (PROMPT_TOKENS, COMPLETION_TOKENS, pytest.approx(EXPECTED_SPEND)), billed diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 55b913aa7e4..cc1fc5cd6c3 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9400,6 +9400,14 @@ class TestApplyStreamUsageTracking: assert data["_litellm_strip_stream_usage"] is True assert data["_litellm_stream_usage_injected"] is True + def test_default_injects_usage_and_marks_strip_for_text_completions(self): + data: Final = {"stream": True, "model": "gpt-5.4-nano"} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="atext_completion") + + assert data["stream_options"] == {"include_usage": True} + assert data["_litellm_strip_stream_usage"] is True + def test_default_preserves_other_client_stream_options_keys(self): data = {"stream": True, "stream_options": {"include_obfuscation": True}} diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 00051462b3c..1ce61a8e9b3 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1203,24 +1203,35 @@ def test_body_snapshot_excludes_team_callback_credentials() -> None: }, proxy_request -def test_body_snapshot_keeps_the_stream_usage_injection_marker_out_of_the_stored_request() -> None: +@pytest.mark.parametrize( + ("route_type", "request_input"), + [ + ("acompletion", {"messages": [{"role": "user", "content": "hi"}]}), + ("atext_completion", {"prompt": "hi"}), + ], +) +def test_body_snapshot_keeps_the_stream_usage_injection_marker_out_of_the_stored_request( + route_type: str, request_input: dict[str, object] +) -> None: from litellm.proxy.common_request_processing import stream_usage_tracking_updates from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD - request: Final = {"model": "gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True} + request: Final = {"model": "gpt-5-mini", **request_input, "stream": True} proxy_request: Final = {"body": dict(request)} data: Final = { **request, - **stream_usage_tracking_updates(request, {}, "acompletion", lambda: True), + **stream_usage_tracking_updates(request, {}, route_type, lambda: True), "proxy_server_request": proxy_request, } refresh_proxy_server_request_body_snapshot(data) assert data[STREAM_USAGE_INJECTED_FIELD] is True + assert data["_litellm_strip_stream_usage"] is True assert proxy_request["body"]["stream_options"] == {"include_usage": True} assert STREAM_USAGE_INJECTED_FIELD not in proxy_request["body"] + assert "_litellm_strip_stream_usage" not in proxy_request["body"] @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 02f88d444e4..023ef0b4606 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13583,29 +13583,29 @@ def _stream_usage_generator_chunks(): def test_is_injected_stream_usage_artifact(): - from litellm.proxy.proxy_server import _is_injected_stream_usage_artifact + from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.types.utils import ModelResponseStream, Usage content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks() - assert _is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True + assert is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True synthetic_final_chunk = ModelResponseStream(model="gpt-5.4-nano") synthetic_final_chunk.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238) - assert _is_injected_stream_usage_artifact(synthetic_final_chunk) is True + assert is_injected_stream_usage_artifact(synthetic_final_chunk) is True azure_prompt_filter_chunk = ModelResponseStream(model="gpt-5.4-nano", choices=[]) - assert _is_injected_stream_usage_artifact(azure_prompt_filter_chunk) is True + assert is_injected_stream_usage_artifact(azure_prompt_filter_chunk) is True - assert _is_injected_stream_usage_artifact(content_chunk) is False - assert _is_injected_stream_usage_artifact(finish_chunk) is False + assert is_injected_stream_usage_artifact(content_chunk) is False + assert is_injected_stream_usage_artifact(finish_chunk) is False content_chunk_with_usage, finish_chunk_with_usage, _ = _stream_usage_test_chunks() content_chunk_with_usage.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238) finish_chunk_with_usage.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238) - assert _is_injected_stream_usage_artifact(content_chunk_with_usage) is False - assert _is_injected_stream_usage_artifact(finish_chunk_with_usage) is False + assert is_injected_stream_usage_artifact(content_chunk_with_usage) is False + assert is_injected_stream_usage_artifact(finish_chunk_with_usage) is False - assert _is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False + assert is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False async def _collect_async_data_generator_frames(request_data: dict) -> list: diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index e10352e0b3d..11104cf7bca 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -48,7 +48,14 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.litellm_params import ControlOptions from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent from litellm.types.prompts.init_prompts import PromptSpec -from litellm.types.utils import Delta, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, Usage +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StandardCallbackDynamicParams, + StreamingChoices, + TextCompletionResponse, + Usage, +) from tests.unit.llms.bedrock.slow_upstream import ( CONVERSE_ANSWER, STREAM_TIMEOUT_SECONDS, @@ -4136,6 +4143,95 @@ def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admissi assert non_stream.usage.prompt_tokens == 0 +def _carries_text_or_finish(frame: TextCompletionResponse) -> bool: + return frame.choices[0].text is not None or frame.choices[0].finish_reason is not None + + +@pytest.mark.asyncio +async def test_mock_atext_completion_stream_hides_the_usage_chunk_the_proxy_asked_for_on_the_clients_behalf(): + response: Final = await litellm.atext_completion( + model="openai/gpt-5.4-mini", + prompt="hi", + mock_response="ok then", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + _litellm_strip_stream_usage=True, + ) + + frames: Final[list[TextCompletionResponse]] = [frame async for frame in response] + + assert len(frames) == 4 + assert all(_carries_text_or_finish(frame) for frame in frames) + assert all(frame.usage is None for frame in frames) + + +def _native_completions_reply(outbound: list[dict[str, object]]) -> Callable[[httpx.Request], httpx.Response]: + def respond(request: httpx.Request) -> httpx.Response: + outbound.append(json.loads(request.content)) + base: Final = {"id": "cmpl-native", "object": "text_completion", "created": 1, "model": "code-model"} + frames: Final = ( + {**base, "choices": [{"index": 0, "text": "def ", "finish_reason": None, "logprobs": None}]}, + {**base, "choices": [{"index": 0, "text": "main", "finish_reason": "length", "logprobs": None}]}, + {**base, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}, + ) + content: Final = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) + "data: [DONE]\n\n" + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=content.encode()) + + return respond + + +def _native_completions_upstream(outbound: list[dict[str, object]]) -> OpenAI: + transport: Final = httpx.MockTransport(_native_completions_reply(outbound)) + return OpenAI(api_key="transport-only", http_client=httpx.Client(transport=transport)) + + +def _async_native_completions_upstream(outbound: list[dict[str, object]]) -> AsyncOpenAI: + transport: Final = httpx.MockTransport(_native_completions_reply(outbound)) + return AsyncOpenAI(api_key="transport-only", http_client=httpx.AsyncClient(transport=transport)) + + +@pytest.mark.asyncio +async def test_native_atext_completion_stream_hides_the_usage_chunk_the_proxy_asked_for_on_the_clients_behalf(): + outbound: Final[list[dict[str, object]]] = [] + response: Final = await litellm.atext_completion( + model="text-completion-openai/code-model", + prompt="hi", + client=_async_native_completions_upstream(outbound), + api_key="transport-only", + stream=True, + stream_options={"include_usage": True}, + _litellm_strip_stream_usage=True, + ) + + frames: Final[list[TextCompletionResponse]] = [frame async for frame in response] + + assert [sent.get("stream_options") for sent in outbound] == [{"include_usage": True}] + assert [(frame.choices[0].text, frame.choices[0].finish_reason) for frame in frames] == [ + ("def ", None), + ("main", None), + (None, "length"), + ] + assert all(frame.usage is None for frame in frames) + + +def test_native_text_completion_stream_keeps_the_usage_chunk_the_client_asked_for(): + frames: Final = list( + litellm.text_completion( + model="text-completion-openai/code-model", + prompt="hi", + client=_native_completions_upstream([]), + api_key="transport-only", + stream=True, + stream_options={"include_usage": True}, + ) + ) + + assert len(frames) == 4 + assert all(_carries_text_or_finish(frame) for frame in frames[:3]) + assert (frames[-1].usage.prompt_tokens, frames[-1].usage.completion_tokens) == (10, 5) + + def test_mock_completion_stream_with_model_response(): """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" from litellm import completion @@ -4293,7 +4389,6 @@ class TestCallTypesOCR: def test_stream_chunk_builder_text_completion_combines_text_and_usage(): from litellm.main import stream_chunk_builder_text_completion - from litellm.types.utils import TextCompletionResponse chunks = [ TextCompletionResponse( diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index b7edffc0753..e5459a5a0cd 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -142,3 +142,90 @@ def test_sync_text_completion_stream_maps_upstream_failure_after_content() -> No next(wrapper) assert (error.value.model, error.value.llm_provider) == ("test-model", "openai") + + +def _stream_with_trailing_usage_chunk() -> tuple[ModelResponseStream, ...]: + return ( + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hi"))]), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="length")]), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta())], + usage=Usage(prompt_tokens=19, completion_tokens=5, total_tokens=24), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_stream", (False, True)) +@pytest.mark.parametrize( + ("usage_injected_by_proxy", "expected_frames", "expected_usage"), + ( + (True, [("Hi", None), (None, "length")], [None, None]), + (False, [("Hi", None), (None, "length"), (None, None)], [None, None, (19, 5)]), + ), +) +async def test_text_completion_stream_hides_only_the_usage_chunk_the_proxy_asked_for( + async_stream: bool, + usage_injected_by_proxy: bool, + expected_frames: list[tuple[str | None, str | None]], + expected_usage: list[tuple[int, int] | None], +) -> None: + async def chunks() -> AsyncIterator[ModelResponseStream]: + for chunk in _stream_with_trailing_usage_chunk(): + yield chunk + + wrapper: Final = TextCompletionStreamWrapper( + completion_stream=chunks() if async_stream else iter(_stream_with_trailing_usage_chunk()), + model="gpt-5.4-mini", + stream_options={"include_usage": True}, + usage_injected_by_proxy=usage_injected_by_proxy, + ) + frames: Final = [frame async for frame in wrapper] if async_stream else list(wrapper) + + assert [(frame.choices[0].text, frame.choices[0].finish_reason) for frame in frames] == expected_frames + assert [ + None if frame.usage is None else (frame.usage.prompt_tokens, frame.usage.completion_tokens) for frame in frames + ] == expected_usage + + +def test_text_completion_stream_keeps_a_finish_frame_but_hides_the_usage_the_proxy_asked_for() -> None: + finish_with_usage: Final = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + usage=Usage(prompt_tokens=27, completion_tokens=5, total_tokens=32), + ) + wrapper: Final = TextCompletionStreamWrapper( + completion_stream=iter((finish_with_usage,)), + model="claude-haiku-4-5", + stream_options={"include_usage": True}, + usage_injected_by_proxy=True, + ) + + frames: Final = list(wrapper) + + assert [(frame.choices[0].finish_reason, frame.usage) for frame in frames] == [("stop", None)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_stream", (False, True)) +async def test_text_completion_stream_hides_the_empty_prompt_filter_frame_when_the_proxy_asked_for_usage( + async_stream: bool, +) -> None: + prompt_filter_frame: Final = ModelResponseStream(id="", choices=[]) + upstream: Final = (prompt_filter_frame, *_stream_with_trailing_usage_chunk()) + + async def chunks() -> AsyncIterator[ModelResponseStream]: + for chunk in upstream: + yield chunk + + wrapper: Final = TextCompletionStreamWrapper( + completion_stream=chunks() if async_stream else iter(upstream), + model="gpt-5.4-mini", + stream_options={"include_usage": True}, + usage_injected_by_proxy=True, + ) + frames: Final = [frame async for frame in wrapper] if async_stream else list(wrapper) + + assert [(frame.choices[0].text, frame.choices[0].finish_reason) for frame in frames] == [ + ("Hi", None), + (None, "length"), + ]