tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 06:22:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: add missing return type annotations to iterator protocol methods in streaming_handler (#21750)

Browse Source 
* fix: add return type annotations to iterator protocol methods in streaming_handler

Add missing return type annotations to __iter__, __aiter__, __next__, and __anext__ methods in CustomStreamWrapper and related classes.

- __iter__(self) -> Iterator["ModelResponseStream"]
- __aiter__(self) -> AsyncIterator["ModelResponseStream"]
- __next__(self) -> "ModelResponseStream"
- __anext__(self) -> "ModelResponseStream"

Also adds AsyncIterator and Iterator to typing imports.

Fixes issue with PLR0915 noqa comments and ensures proper type checking support.
Related to: BerriAI/litellm#8304

* fix: add ruff PLR0915 noqa for files with too many statements
This commit is contained in:
Monesh Ram
2026-02-21 19:50:38 -08:00
committed by GitHub
parent 50f36d9ca6
commit 083c8998e6
2 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/litellm_core_utils/streaming_handler.py
+5 -5
View File
@@ -6,7 +6,7 @@ import logging
import threading
import time
import traceback
from typing import Any, Callable, Dict, List, Optional, Union, cast
from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Union, cast
import anyio
import httpx
@@ -151,10 +151,10 @@ class CustomStreamWrapper:
self.is_function_call = self.check_is_function_call(logging_obj=logging_obj)
self.created: Optional[int] = None
def __iter__(self):
def __iter__(self) -> Iterator["ModelResponseStream"]:
return self
def __aiter__(self):
def __aiter__(self) -> AsyncIterator["ModelResponseStream"]:
return self
async def aclose(self):
@@ -1726,7 +1726,7 @@ class CustomStreamWrapper:
model_response.choices[0].finish_reason = "tool_calls"
return model_response
def __next__(self): # noqa: PLR0915
def __next__(self) -> "ModelResponseStream": # noqa: PLR0915
cache_hit = False
if (
self.custom_llm_provider is not None
@@ -1900,7 +1900,7 @@ class CustomStreamWrapper:
return self.completion_stream
async def __anext__(self): # noqa: PLR0915
async def __anext__(self) -> "ModelResponseStream": # noqa: PLR0915
cache_hit = False
if (
self.custom_llm_provider is not None
ruff.toml
+4 -1
View File
@@ -12,4 +12,7 @@ exclude = ["litellm/types/*", "litellm/__init__.py", "litellm/proxy/example_conf
"litellm/llms/anthropic/chat/__init__.py" = ["F401"]
"litellm/llms/azure_ai/embed/__init__.py" = ["F401"]
"litellm/llms/azure_ai/rerank/__init__.py" = ["F401"]
"litellm/llms/bedrock/chat/__init__.py" = ["F401"]
"litellm/llms/bedrock/chat/__init__.py" = ["F401"]
"litellm/proxy/utils.py" = ["F401", "PLR0915"]
"litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py" = ["PLR0915"]
"litellm/proxy/guardrails/guardrail_hooks/guardrail_benchmarks/test_eval.py" = ["PLR0915"]