From c62e6f3fb92299d88ffc16e3b76954c43202ffc2 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Tue, 25 Aug 2026 12:43:52 -0700 Subject: [PATCH] chore(streaming): give the per-index helpers real types The strict gate flagged the new helpers for ANN401 (bare `Any` on `choice`, `chunk` and the `_replace_chunk_choices` return) and ANN001 (`start_time` and `end_time` unannotated). Each one has a concrete type: the chunk and choice helpers already branch on `isinstance(..., dict)`, so the union of the pydantic model and a dict is what they actually take, and the two times are the `datetime.datetime | None` pair threaded into stream_chunk_builder. That also puts the already-imported StreamingChoices to use, clearing an F401. Net ruff difference on the file is 7 findings removed and none added. --- litellm/main.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 2ddded81a23..a333a2bafd5 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8439,12 +8439,12 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N return TextCompletionResponse(**response) -def _streaming_choice_index(choice: Any) -> int | None: +def _streaming_choice_index(choice: StreamingChoices | dict[str, Any]) -> int | None: index = choice.get("index") if isinstance(choice, dict) else getattr(choice, "index", None) return index if isinstance(index, int) else None -def _streaming_chunk_choices(chunk: Any) -> list[Any]: +def _streaming_chunk_choices(chunk: ModelResponseStream | dict[str, Any]) -> list[Any]: choices = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None) return choices or [] @@ -8456,7 +8456,9 @@ def _distinct_streaming_choice_indices(chunks: list[Any]) -> list[int]: return sorted(index for index in indices if index is not None) -def _replace_chunk_choices(chunk: Any, choices: list[Any]) -> Any: +def _replace_chunk_choices( + chunk: ModelResponseStream | dict[str, Any], choices: list[Any] +) -> ModelResponseStream | dict[str, Any]: if not choices: return chunk if isinstance(chunk, dict): @@ -8480,8 +8482,8 @@ def _build_streaming_choices_per_index( chunks: list[Any], indices: list[int], messages: list | None, - start_time, - end_time, + start_time: datetime.datetime | None, + end_time: datetime.datetime | None, ) -> list[Choices] | None: built: Final = [ stream_chunk_builder(