mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
fix(proxy): declare the file usage counters as rate limit calls
Main's redis producer gate requires every module that writes a shared cache to name its key family, and the file usage counters wrote theirs without one.
This commit is contained in:
parent
0e61b32516
commit
c74eb7d1c9
2 changed files with 21 additions and 0 deletions
|
|
@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, NoReturn, TypeAlias
|
|||
from pydantic import Field, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
EMPTY_MAPPING,
|
||||
|
|
@ -148,6 +149,7 @@ async def _increment_all_or_none(
|
|||
return exceeded
|
||||
|
||||
|
||||
@with_service_target("rate_limits")
|
||||
async def consume_file_usage(
|
||||
cache: "InternalUsageCache",
|
||||
limits: tuple[ScopedFileUsageLimit, ...],
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
|
|
@ -141,6 +142,24 @@ async def test_a_rejection_does_not_use_up_the_scope_that_rejected_it():
|
|||
assert await consume_file_usage(cache, raised_limit, DAY, "", MIDDAY) is None
|
||||
|
||||
|
||||
class _ServiceTargetSeen(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _TargetReportingCache(InternalUsageCache):
|
||||
async def async_increment_cache(self, key, value, litellm_parent_otel_span, local_only=False, **kwargs):
|
||||
raise _ServiceTargetSeen(current_service_target())
|
||||
|
||||
|
||||
async def test_counter_writes_are_declared_as_rate_limit_calls_so_their_redis_spans_are_named():
|
||||
cache: Final = _TargetReportingCache(dual_cache=DualCache())
|
||||
|
||||
with pytest.raises(_ServiceTargetSeen) as seen:
|
||||
await consume_file_usage(cache, (_scoped("key", "hashed", 1),), DAY, "", MIDDAY)
|
||||
|
||||
assert seen.value.args == ("rate_limits",)
|
||||
|
||||
|
||||
async def test_download_counters_are_per_file():
|
||||
cache: Final = _cache()
|
||||
limits: Final = (_scoped("key", "hashed", 1, setting="max_file_downloads_per_minute"),)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue