Merge 9eb748effd into 97fc72e5bb

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,

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