diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 1f739a60b4..6b83ded9d2 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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 diff --git a/ruff.toml b/ruff.toml index a310446671..43ff802a68 100644 --- a/ruff.toml +++ b/ruff.toml @@ -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"] \ No newline at end of file +"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"]