diff --git a/litellm/proxy/openai_files_endpoints/file_usage_caps.py b/litellm/proxy/openai_files_endpoints/file_usage_caps.py index 770db11b2ec..5de5bbce27e 100644 --- a/litellm/proxy/openai_files_endpoints/file_usage_caps.py +++ b/litellm/proxy/openai_files_endpoints/file_usage_caps.py @@ -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, ...], diff --git a/tests/unit/proxy/openai_files_endpoint/test_file_usage_caps.py b/tests/unit/proxy/openai_files_endpoint/test_file_usage_caps.py index 1facf69bab7..7cf2505dd38 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_file_usage_caps.py +++ b/tests/unit/proxy/openai_files_endpoint/test_file_usage_caps.py @@ -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"),)