diff --git a/litellm/main.py b/litellm/main.py index 122de9a02c0..617ce65be9d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -13,6 +13,7 @@ import asyncio import contextvars import datetime import inspect +import itertools import json import os import random @@ -8855,36 +8856,46 @@ def config_completion(**kwargs): ) -def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse: - id: Final = chunks[0]["id"] - object: Final = chunks[0]["object"] - created: Final = chunks[0]["created"] - model: Final = chunks[0]["model"] +def stream_chunk_builder_text_completion( + chunks: Sequence[TextCompletionResponse], messages: Sequence | None = None +) -> TextCompletionResponse: + id: Final = chunks[0].id + object: Final = chunks[0].object + created: Final = chunks[0].created + model: Final = chunks[0].model system_fingerprint: Final = chunks[0].get("system_fingerprint", None) - finish_reason: Final = chunks[-1]["choices"][0]["finish_reason"] - logprobs: Final = chunks[-1]["choices"][0]["logprobs"] - - content_list: Final = [] - for chunk in chunks: - choices = chunk["choices"] - for choice in choices: - if choice is not None and hasattr(choice, "text") and choice.get("text") is not None: - _choice = choice.get("text") - content_list.append(_choice) - - # Combine the "content" strings into a single string || combine the 'function' strings into a single string - combined_content: Final = "".join(content_list) - - try: - prompt_tokens = token_counter(model=model, messages=messages) - except Exception: # don't allow this failing to block a complete streaming response from being returned - print_verbose("token_counter failed, assuming prompt tokens is 0") - prompt_tokens = 0 - completion_tokens: Final = token_counter( - model=model, - text=combined_content, - count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages + # With stream_options.include_usage the last chunk is a usage-only trailer, and some providers + # (e.g. vLLM) send finish_reason on a text-less chunk before it, so scan rather than read chunks[-1]. + chunks_with_choices: Final = [chunk for chunk in chunks if chunk.choices] + finish_reason: Final = next( + (c.choices[0].finish_reason for c in reversed(chunks_with_choices) if c.choices[0].finish_reason), + None, ) + logprobs: Final = chunks_with_choices[-1].choices[0].logprobs if chunks_with_choices else None + + all_choices: Final = itertools.chain.from_iterable(chunk.choices for chunk in chunks) + combined_content: Final = "".join(choice.text for choice in all_choices if choice.text) + + # Prefer the usage the provider reported (the include_usage trailer) over a local recount, which + # cannot see a text-completion prompt (it is not in `messages`) and so reports 0 prompt tokens. + provider_usage: Final = next( + (c.usage for c in reversed(chunks) if c.usage and c.usage.total_tokens), + None, + ) + if provider_usage is not None: + prompt_tokens = provider_usage.prompt_tokens + completion_tokens = provider_usage.completion_tokens + else: + try: + prompt_tokens = token_counter(model=model or "", messages=messages) + except Exception: # don't allow this failing to block a complete streaming response from being returned + print_verbose("token_counter failed, assuming prompt tokens is 0") + prompt_tokens = 0 + completion_tokens = token_counter( + model=model or "", + text=combined_content, + count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages + ) response: Final = { "id": id, diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index ffb17a17e3e..801a4b0d1a3 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -3294,6 +3294,33 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens +@pytest.mark.parametrize("trailer_choices", [[], [{"text": None, "index": 0, "logprobs": None, "finish_reason": None}]]) +def test_stream_chunk_builder_text_completion_keeps_finish_reason_and_provider_usage(trailer_choices): + """vLLM-style include_usage stream: finish_reason arrives on a text-less chunk, then a + usage-only trailer. The rebuilt response must keep both instead of reading chunks[-1] + and recounting the prompt from `messages` (which a text completion doesn't have).""" + from litellm.main import stream_chunk_builder_text_completion + from litellm.types.utils import TextCompletionResponse + + def chunk(choices, **extra): + return TextCompletionResponse( + id="cmpl-1", object="text_completion", created=1, model="my-model", choices=choices, **extra + ) + + chunks = [ + chunk([{"text": "Hello", "index": 0, "logprobs": None, "finish_reason": None}]), + chunk([{"text": " world", "index": 0, "logprobs": None, "finish_reason": None}]), + chunk([{"text": "", "index": 0, "logprobs": None, "finish_reason": "length"}]), + chunk(trailer_choices, usage={"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}), + ] + + response = stream_chunk_builder_text_completion(chunks=chunks, messages=None) + + assert response.choices[0].text == "Hello world" + assert response.choices[0].finish_reason == "length" + assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (7, 2, 9) + + def test_completion_forwards_store_and_prompt_cache_key_to_openai(): """ Regression test for https://github.com/BerriAI/litellm/issues/33184