fix(streaming): narrow the built choices with isinstance instead of cast

The per-index builder asserted its result with two nested cast() calls, which the
repo's LIT006 budget counts as unchecked assertions. Move the narrowing into a
small helper that isinstance-checks the response and its first choice and returns
None when either is not what the merge needs, so the caller's fallback to the
single-response path is reached by a real check rather than by faith. The two
one-shot locals in the index/choices readers pick up the Final annotation LIT010
asks for.

No behaviour change: the same inputs produce the same choices, and the same
inputs fall back.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
This commit is contained in:
Vineeth Sai 2026-08-27 09:47:55 -07:00
parent 8473655cea
commit e77568dca1

View file

@ -8702,12 +8702,12 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o
def _streaming_choice_index(choice: StreamingChoices | Mapping[str, Any]) -> int | None:
index = choice.get("index") if isinstance(choice, dict) else getattr(choice, "index", None)
index: Final = 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: ModelResponseStream | Mapping[str, Any]) -> Sequence[Any]:
choices = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None)
choices: Final = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None)
return choices or [] # mutable-ok: the caller iterates it; an empty list is the natural absent value
@ -8743,6 +8743,21 @@ def _chunks_for_streaming_choice_index(
]
def _renumbered_first_choice(response: Any, index: int) -> Choices | None:
"""The single choice a per-index build produced, renumbered to that index.
None when the build did not yield a usable non-streaming choice: `stream_chunk_builder`
is typed to admit a streaming response and a streaming choice, and either one leaves
nothing to merge.
"""
if not isinstance(response, ModelResponse) or not response.choices:
return None
first: Final = response.choices[0]
if not isinstance(first, Choices):
return None
return first.model_copy(update={"index": index}) # mutable-ok: model_copy stores what it is given
def _build_streaming_choices_per_index(
chunks: Sequence[Any],
indices: Sequence[int],
@ -8750,21 +8765,23 @@ def _build_streaming_choices_per_index(
start_time: datetime.datetime | None,
end_time: datetime.datetime | None,
) -> list[Choices] | None: # mutable-ok: assigned to ModelResponse.choices, which is a list
built: Final = [ # mutable-ok: one built response per choice index, zipped below
stream_chunk_builder(
chunks=_chunks_for_streaming_choice_index(chunks, index),
messages=messages,
start_time=start_time,
end_time=end_time,
built: Final = [ # mutable-ok: one choice per index, returned as ModelResponse.choices
_renumbered_first_choice(
stream_chunk_builder(
chunks=_chunks_for_streaming_choice_index(chunks, index),
messages=messages,
start_time=start_time,
end_time=end_time,
),
index,
)
for index in indices
]
if any(not isinstance(response, ModelResponse) or not response.choices for response in built):
# One unusable build means there is nothing to merge, so the caller falls back to
# the single-response path rather than being handed a partial list.
if any(choice is None for choice in built):
return None
return [ # mutable-ok: assigned to ModelResponse.choices, which is a list
cast(Choices, cast(ModelResponse, response).choices[0]).model_copy(update={"index": index}) # mutable-ok: model_copy stores what it is given
for index, response in zip(indices, built)
]
return [choice for choice in built if choice is not None] # mutable-ok: assigned to ModelResponse.choices, which is a list
def stream_chunk_builder(