From 9eb748effdc3eb42ad17010bcab8c47fa59f818c Mon Sep 17 00:00:00 2001 From: Lai Wei Date: Sat, 3 Oct 2026 00:34:12 -0700 Subject: [PATCH] fix(streaming): type the text completion chunk builder so its scans add no basedpyright debt Read chunk fields as typed attributes instead of untyped subscripts, which clears the reportUnknownArgumentType breach the lint gate flagged --- litellm/main.py | 41 ++++++++++++++++++----------------------- 1 file changed, 18 insertions(+), 23 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 373b7f7df59..99d560ac70e 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 @@ -8837,49 +8838,43 @@ 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) # 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"]] + 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"]), + (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 + logprobs: Final = chunks_with_choices[-1].choices[0].logprobs if chunks_with_choices else None - 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) + 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.get("usage") for c in reversed(chunks) if c.get("usage") and c["usage"].get("total_tokens")), + (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.get("prompt_tokens") or 0 - completion_tokens = provider_usage.get("completion_tokens") or 0 + prompt_tokens = provider_usage.prompt_tokens + completion_tokens = provider_usage.completion_tokens else: try: - prompt_tokens = token_counter(model=model, messages=messages) + 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, + 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 )