mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-24 12:26:58 +00:00
Merge pull request #22526 from BerriAI/fix/router-plr0915-noqa
Fix PLR0915 lint error in _completion_streaming_iterator
This commit is contained in:
+24
-57
@@ -1551,7 +1551,28 @@ class Router:
|
||||
)
|
||||
raise e
|
||||
|
||||
async def _acompletion_streaming_iterator( # noqa: PLR0915
|
||||
@staticmethod
|
||||
def _combine_fallback_usage(
|
||||
fallback_item: ModelResponseStream,
|
||||
complete_response_object_usage: Optional[Usage],
|
||||
) -> None:
|
||||
"""Merge partial-stream usage with fallback-stream usage on the chunk."""
|
||||
from litellm.cost_calculator import BaseTokenUsageProcessor
|
||||
|
||||
usage = cast(Optional[Usage], getattr(fallback_item, "usage", None))
|
||||
usage_objects = [usage] if usage is not None else []
|
||||
if (
|
||||
complete_response_object_usage is not None
|
||||
and hasattr(complete_response_object_usage, "usage")
|
||||
and complete_response_object_usage.usage is not None # type: ignore
|
||||
):
|
||||
usage_objects.append(complete_response_object_usage)
|
||||
combined_usage = BaseTokenUsageProcessor.combine_usage_objects(
|
||||
usage_objects=usage_objects
|
||||
)
|
||||
setattr(fallback_item, "usage", combined_usage)
|
||||
|
||||
async def _acompletion_streaming_iterator(
|
||||
self,
|
||||
model_response: CustomStreamWrapper,
|
||||
messages: List[Dict[str, str]],
|
||||
@@ -1654,32 +1675,7 @@ class Router:
|
||||
and isinstance(fallback_item, ModelResponseStream)
|
||||
and hasattr(fallback_item, "usage")
|
||||
):
|
||||
from litellm.cost_calculator import (
|
||||
BaseTokenUsageProcessor,
|
||||
)
|
||||
|
||||
usage = cast(
|
||||
Optional[Usage],
|
||||
getattr(fallback_item, "usage", None),
|
||||
)
|
||||
if usage is not None:
|
||||
usage_objects = [usage]
|
||||
else:
|
||||
usage_objects = []
|
||||
|
||||
if (
|
||||
complete_response_object_usage is not None
|
||||
and hasattr(complete_response_object_usage, "usage")
|
||||
and complete_response_object_usage.usage is not None # type: ignore
|
||||
):
|
||||
usage_objects.append(complete_response_object_usage)
|
||||
|
||||
combined_usage = (
|
||||
BaseTokenUsageProcessor.combine_usage_objects(
|
||||
usage_objects=usage_objects
|
||||
)
|
||||
)
|
||||
setattr(fallback_item, "usage", combined_usage)
|
||||
self._combine_fallback_usage(fallback_item, complete_response_object_usage)
|
||||
yield fallback_item
|
||||
else:
|
||||
# If fallback returns a non-streaming response, yield None
|
||||
@@ -1815,36 +1811,7 @@ class Router:
|
||||
and isinstance(fallback_item, ModelResponseStream)
|
||||
and hasattr(fallback_item, "usage")
|
||||
):
|
||||
from litellm.cost_calculator import (
|
||||
BaseTokenUsageProcessor,
|
||||
)
|
||||
|
||||
usage = cast(
|
||||
Optional[Usage],
|
||||
getattr(fallback_item, "usage", None),
|
||||
)
|
||||
if usage is not None:
|
||||
usage_objects = [usage]
|
||||
else:
|
||||
usage_objects = []
|
||||
|
||||
if (
|
||||
complete_response_object_usage is not None
|
||||
and hasattr(
|
||||
complete_response_object_usage, "usage"
|
||||
)
|
||||
and complete_response_object_usage.usage is not None # type: ignore
|
||||
):
|
||||
usage_objects.append(
|
||||
complete_response_object_usage
|
||||
)
|
||||
|
||||
combined_usage = (
|
||||
BaseTokenUsageProcessor.combine_usage_objects(
|
||||
usage_objects=usage_objects
|
||||
)
|
||||
)
|
||||
setattr(fallback_item, "usage", combined_usage)
|
||||
router_self._combine_fallback_usage(fallback_item, complete_response_object_usage)
|
||||
yield fallback_item
|
||||
else:
|
||||
yield None
|
||||
|
||||
Reference in New Issue
Block a user