fix(vertex): use module-level import for ModelResponseStream instead of type: ignore

Move ModelResponseStream to module-level import so MyPy can properly
narrow the Union type after isinstance checks, removing the need for
type: ignore suppressions.
This commit is contained in:
Chesars
2026-02-22 09:37:26 -03:00
parent eaf38ed3a2
commit 7e4f07c45d
@@ -80,6 +80,7 @@ from litellm.types.utils import (
ChatCompletionTokenLogprob,
ChoiceLogprobs,
CompletionTokensDetailsWrapper,
ModelResponseStream,
PromptTokensDetailsWrapper,
TopLogprob,
Usage,
@@ -1931,7 +1932,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
@staticmethod
def _process_candidates( # noqa: PLR0915
_candidates: List[Candidates],
model_response: Union[ModelResponse, "ModelResponseStream"],
model_response: Union[ModelResponse, ModelResponseStream],
standard_optional_params: dict,
cumulative_tool_call_index: int = 0,
) -> Tuple[List[dict], List[dict], List, List, int]:
@@ -2093,7 +2094,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
chat_completion_logprobs=chat_completion_logprobs,
image_response=image_response,
)
model_response.choices.append(choice) # type: ignore[arg-type]
model_response.choices.append(choice)
elif isinstance(model_response, ModelResponse):
choice = litellm.Choices(
finish_reason=VertexGeminiConfig._check_finish_reason(
@@ -2104,7 +2105,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
logprobs=chat_completion_logprobs,
enhancements=None,
)
model_response.choices.append(choice) # type: ignore[arg-type]
model_response.choices.append(choice)
return (
grounding_metadata,