From 961290fc46429d48de4664d49fb3d64dfa62467a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:28:31 +0000 Subject: [PATCH] refactor(redis): move search free-tier allocation into a typed store --- .../proxy/hooks/proxy_track_cost_callback.py | 7 +- litellm/proxy/hooks/search_free_tier.py | 33 ++- .../proxy/spend_tracking/cost_finalization.py | 12 +- litellm/redis/stores/search_free_tier.py | 41 +++ tests/_support/fake_redis.py | 30 ++- .../hooks/test_proxy_track_cost_callback.py | 59 +++-- .../unit/proxy/hooks/test_search_free_tier.py | 192 +++++++++++---- .../test_success_handler.py | 45 ++-- .../spend_tracking/test_cost_finalization.py | 194 ++++++++++----- .../redis/stores/test_search_free_tier.py | 233 ++++++++++++++++++ 10 files changed, 670 insertions(+), 176 deletions(-) create mode 100644 litellm/redis/stores/search_free_tier.py create mode 100644 tests/unit/redis/stores/test_search_free_tier.py diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0d333d4fd08..bca9f8fb843 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -68,6 +68,7 @@ from litellm.proxy.utils import ( BilledFailedAttempt, ProxyUpdateSpend, ) +from litellm.redis.stores.search_free_tier import TransportSearchFreeTierStore from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, @@ -711,7 +712,11 @@ class _ProxyDBLogger(CustomLogger): sl_object=sl_object, # pyright: ignore[reportUnknownArgumentType] # payload carries untyped lists response_cost=priced_cost, completion_response=completion_response, - counter=proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, + counter=( + TransportSearchFreeTierStore(search_redis.transport()) + if (search_redis := proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache) is not None + else None + ), now=datetime.now(timezone.utc), ) if litellm.apply_search_free_tier diff --git a/litellm/proxy/hooks/search_free_tier.py b/litellm/proxy/hooks/search_free_tier.py index 57123efbcec..76edc8f0d2d 100644 --- a/litellm/proxy/hooks/search_free_tier.py +++ b/litellm/proxy/hooks/search_free_tier.py @@ -1,8 +1,8 @@ from collections.abc import MutableMapping from dataclasses import dataclass -from datetime import datetime +from datetime import datetime, timedelta from types import MappingProxyType -from typing import Final, Protocol +from typing import Final from pydantic import BaseModel, ConfigDict, ValidationError @@ -12,6 +12,8 @@ from litellm.litellm_core_utils.duration_parser import get_next_standardized_res from litellm.litellm_core_utils.llm_cost_calc.finalization import finalize_cost from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests_from_usage from litellm.llms.gemini.cost_calculator import cost_per_web_search_request +from litellm.redis.stores.search_free_tier import SearchFreeTierStore, allocation_key +from litellm.redis.transport.contract import RedisFailure from litellm.types.utils import ModelInfo, SearchFreeRequests, StandardLoggingPayload, Usage _PERIOD_DURATION: Final = MappingProxyType({"day": "1d", "month": "1mo"}) @@ -20,10 +22,6 @@ _TTL_GRACE_SECONDS: Final = 3600 _COUNTED_KWARG: Final = "search_free_tier_counted" -class SharedCounter(Protocol): - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: ... - - class _UsageCarrier(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) @@ -37,7 +35,7 @@ def freed_units(allowance: int, units: int, total_after: int) -> int: def search_free_tier_key(provider: str, allowance: SearchFreeRequests, now: datetime) -> str: period_stamp: Final = now.strftime(_PERIOD_STAMP[allowance["period"]]) - return f"litellm:search_free_tier:{provider}:{allowance['shared_group']}:{period_stamp}" + return allocation_key(provider, allowance["shared_group"], period_stamp) def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int: @@ -91,7 +89,7 @@ async def apply_search_free_tier( sl_object: StandardLoggingPayload, response_cost: float, completion_response: object, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, ) -> float: """Zero the grounding queries that landed under the model's shared free allowance, everything else is unchanged.""" @@ -109,7 +107,7 @@ async def allocate_search_free_tier( sl_object: StandardLoggingPayload, response_cost: float, completion_response: object, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, ) -> SearchAllowanceAllocation | None: model_info: Final = sl_object["model_map_information"]["model_map_value"] @@ -126,14 +124,15 @@ async def allocate_search_free_tier( ) return None provider: Final = sl_object["custom_llm_provider"] or model_info["litellm_provider"] - key: Final = search_free_tier_key(provider, allowance, now) - try: - total_after: Final = await counter.async_increment(key, units, ttl=search_free_tier_ttl(allowance, now)) - except Exception: # noqa: BLE001 # any counter failure means full price, never a lost spend row - verbose_proxy_logger.warning( - "search free tier counter %s is unavailable, grounding is billed at full price", key - ) + total_after: Final = await counter.allocate_units( + provider, + allowance["shared_group"], + now.strftime(_PERIOD_STAMP[allowance["period"]]), + units, + ttl=timedelta(seconds=search_free_tier_ttl(allowance, now)), + ) + if isinstance(total_after, RedisFailure): return None freed: Final = freed_units(allowance["count"], units, int(total_after)) freed_cost: Final = min(response_cost, cost_per_web_search_request(usage, model_info) * freed / units) - return SearchAllowanceAllocation(key, units, freed, freed_cost) + return SearchAllowanceAllocation(search_free_tier_key(provider, allowance, now), units, freed, freed_cost) diff --git a/litellm/proxy/spend_tracking/cost_finalization.py b/litellm/proxy/spend_tracking/cost_finalization.py index b260412331e..7dbf1e42689 100644 --- a/litellm/proxy/spend_tracking/cost_finalization.py +++ b/litellm/proxy/spend_tracking/cost_finalization.py @@ -7,7 +7,8 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost, PreparedCost, finalize_cost -from litellm.proxy.hooks.search_free_tier import SharedCounter, allocate_search_free_tier +from litellm.proxy.hooks.search_free_tier import allocate_search_free_tier +from litellm.redis.stores.search_free_tier import SearchFreeTierStore, TransportSearchFreeTierStore from litellm.types.utils import CostBreakdown, StandardLoggingPayload @@ -22,7 +23,7 @@ def prepare_response_cost(logging_obj: Logging, response: object, now: datetime) async def finalize_response_cost( logging_obj: Logging, response: object, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, guardrail_cost: float | None = None, ) -> FinalizedCost | None: @@ -40,7 +41,7 @@ async def finalize_response_cost( async def _finalize_response_cost( logging_obj: Logging, response: object, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, guardrail_cost: float | None, ) -> FinalizedCost | None: @@ -80,7 +81,7 @@ async def _finish_finalization( payload: StandardLoggingPayload, amount: float, breakdown: CostBreakdown | None, - counter: SharedCounter | None, + counter: SearchFreeTierStore | None, now: datetime, ) -> FinalizedCost: if logging_obj.cost_allowance_discount is None: @@ -102,10 +103,11 @@ async def finalize_proxy_response_cost( ) -> FinalizedCost | None: from litellm.proxy.proxy_server import proxy_logging_obj + redis_cache: Final = proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache return await finalize_response_cost( logging_obj, response, - proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, + TransportSearchFreeTierStore(redis_cache.transport()) if redis_cache is not None else None, datetime.now(timezone.utc), guardrail_cost, ) diff --git a/litellm/redis/stores/search_free_tier.py b/litellm/redis/stores/search_free_tier.py new file mode 100644 index 00000000000..8fd0cf40a62 --- /dev/null +++ b/litellm/redis/stores/search_free_tier.py @@ -0,0 +1,41 @@ +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +def allocation_key(provider: str, shared_group: str, period_stamp: str) -> str: + return f"litellm:search_free_tier:{provider}:{shared_group}:{period_stamp}" + + +class SearchFreeTierStore(Protocol): + async def allocate_units( + self, provider: str, shared_group: str, period_stamp: str, units: int, *, ttl: timedelta | None + ) -> float | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportSearchFreeTierStore: + transport: RedisTransport + + async def allocate_units( + self, provider: str, shared_group: str, period_stamp: str, units: int, *, ttl: timedelta | None + ) -> float | RedisFailure: + key: Final = allocation_key(provider, shared_group, period_stamp) + if ttl is None: + return await self.transport.incrbyfloat(key, units) + batch: Final = self.transport.post_call_batch() or self.transport.batch() + increment: Final = batch.incrbyfloat(key, units) + expiry: Final = batch.ttl(key) + total: Final = await increment + current_ttl: Final = await expiry + if isinstance(total, RedisFailure): + return total + if isinstance(current_ttl, RedisFailure): + return current_ttl + if current_ttl == -1: + expiration: Final = await self.transport.expire_seconds(key, ttl) + if isinstance(expiration, RedisFailure): + return expiration + return total diff --git a/tests/_support/fake_redis.py b/tests/_support/fake_redis.py index cc5fd18a305..79a6ef3eb20 100644 --- a/tests/_support/fake_redis.py +++ b/tests/_support/fake_redis.py @@ -1,7 +1,8 @@ from collections.abc import Awaitable, Callable -from typing import TypeVar +from typing import Final, TypeVar import fakeredis +from pydantic import TypeAdapter import litellm from litellm._service_logger import ServiceLogging @@ -55,3 +56,30 @@ def fake_redis_transport( request_engine=lambda: None, post_call_engine=lambda: None, ) + + +_FLOAT: Final = TypeAdapter(float) +_INT: Final = TypeAdapter(int) +_KEYS: Final = TypeAdapter(list[bytes]) + + +def set_reply_callback(cache: RedisCache, command: str, callback: Callable[..., object]) -> None: + cache.init_async_client().set_response_callback(command, callback) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # redis-py types the callback for decoding only, tests also raise or await in it + + +async def legacy_increment(cache: RedisCache, key: str, value: float, *, ttl: int | None) -> float: + return _FLOAT.validate_python(await cache.async_increment(key, value, ttl=ttl)) # pyright: ignore[reportUnknownMemberType] # legacy RedisCache API is untyped + + +def legacy_default_ttl(cache: RedisCache) -> int: + return _INT.validate_python(cache.get_ttl(ttl=None)) # pyright: ignore[reportUnknownMemberType] # legacy RedisCache API is untyped + + +async def keys_matching(client: fakeredis.FakeAsyncRedis, pattern: str) -> tuple[bytes, ...]: + return tuple(_KEYS.validate_python(await client.keys(pattern))) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # fakeredis is untyped + + +def trip_breaker(cache: RedisCache, failures: int | None = None) -> None: + breaker: Final = cache._circuit_breaker # pyright: ignore[reportPrivateUsage] # tests open the breaker the way RedisCache does after hard failures + for _ in range(breaker.failure_threshold if failures is None else failures): + breaker.record_failure() diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 8d4f8cd87ad..2d754373f4a 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, create_autospec, patch +import fakeredis import httpx import pytest from fastapi import HTTPException @@ -40,6 +41,10 @@ from litellm.proxy.utils import ( ) from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, +) @pytest.mark.asyncio @@ -536,7 +541,10 @@ async def test_async_post_call_failure_hook_carries_used_client_oauth_token_from @pytest.mark.parametrize( "metadata_buckets, expected", [ - ({"metadata": {"used_client_oauth_token": True}, "litellm_metadata": {"used_client_oauth_token": False}}, False), + ( + {"metadata": {"used_client_oauth_token": True}, "litellm_metadata": {"used_client_oauth_token": False}}, + False, + ), ({"metadata": {"used_client_oauth_token": True}, "litellm_metadata": {"user_id": "caller"}}, None), ({"metadata": {"used_client_oauth_token": "yes"}}, None), ], @@ -623,7 +631,6 @@ async def test_async_post_call_failure_hook_reads_used_client_oauth_token_from_t assert json.loads(payload["metadata"])["used_client_oauth_token"] is expected - @pytest.mark.asyncio @pytest.mark.parametrize( "request_route, body", @@ -743,7 +750,10 @@ async def test_async_post_call_failure_hook_adds_guardrail_cost_to_recovered_str [ (None, ["detect-only-scanner"]), ([], ["detect-only-scanner"]), - ([{"guardrail_name": "logged-before-the-failure", "guardrail_status": "success"}], ["logged-before-the-failure"]), + ( + [{"guardrail_name": "logged-before-the-failure", "guardrail_status": "success"}], + ["logged-before-the-failure"], + ), ], ) async def test_async_post_call_failure_hook_keeps_a_guardrail_verdict_recorded_after_the_failure_was_logged( @@ -3368,7 +3378,9 @@ async def test_track_cost_callback_failure_alert_never_carries_request_metadata_ @pytest.mark.asyncio @pytest.mark.parametrize("identity_field", ["agent_id", "billing_agent_id"]) -async def test_autonomous_llm_callback_persists_without_human_or_key(identity_field: str) -> None: # test-quality-ok: verifies anonymous-agent charges reach the persistence boundary; no injection seam +async def test_autonomous_llm_callback_persists_without_human_or_key( # test-quality-ok: verifies anonymous-agent charges reach the persistence boundary; no injection seam + identity_field: str, +) -> None: kwargs: Final = { "call_type": "acompletion", "model": "test-model", @@ -3392,9 +3404,12 @@ async def test_autonomous_llm_callback_persists_without_human_or_key(identity_fi @pytest.mark.parametrize("agent_id,expected", [(None, False), ("autonomous-agent", True)]) def test_autonomous_agent_cost_tracking_needs_no_human_or_virtual_key(agent_id: str | None, expected: bool) -> None: - assert _should_track_cost_callback( - user_api_key=None, user_id=None, team_id=None, end_user_id=None, call_type="acompletion", agent_id=agent_id - ) is expected + assert ( + _should_track_cost_callback( + user_api_key=None, user_id=None, team_id=None, end_user_id=None, call_type="acompletion", agent_id=agent_id + ) + is expected + ) _CALL_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) @@ -3435,14 +3450,6 @@ async def test_track_cost_callback_enqueue_emits_no_service_span(): # test-qual assert update_database.await_count == 1, "the spend enqueue itself must still run" assert success_hook.await_count == 0, [call.kwargs for call in success_hook.await_args_list] -class _SharedCounter: - def __init__(self) -> None: - self.totals: dict[str, float] = {} - - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.totals[key] = self.totals.get(key, 0) + value - return self.totals[key] - def _grounded_gemini_kwargs(response_cost: float) -> dict: return { @@ -3478,9 +3485,13 @@ def _grounded_gemini_kwargs(response_cost: float) -> dict: @pytest.mark.asyncio @pytest.mark.parametrize(("opted_in", "finalized"), [(True, False), (False, False), (True, True)]) -async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when_opted_in(monkeypatch, opted_in, finalized): +async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when_opted_in( + monkeypatch, opted_in, finalized +): monkeypatch.setattr(litellm, "apply_search_free_tier", opted_in, raising=False) - counter = _SharedCounter() + server: Final = fakeredis.FakeServer() + counter: Final = fake_redis_cache(server, namespace="search-callback") + raw: Final = fakeredis.FakeAsyncRedis(server=server) kwargs = _grounded_gemini_kwargs(response_cost=1.0) kwargs["standard_logging_object"]["cost_finalized"] = finalized response = ModelResponse( @@ -3520,10 +3531,12 @@ async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when if opted_in and not finalized: assert kwargs["response_cost"] == pytest.approx(1.0 - 5 * 0.1) assert breakdown["original_cost"] == 1.0 and breakdown["discount_amount"] == pytest.approx(5 * 0.1) - assert sum(counter.totals.values()) == 5 + keys: Final = await keys_matching(raw, "search-callback:litellm:search_free_tier:*") + assert len(keys) == 1 + assert await raw.get(keys[0]) == b"5" else: assert kwargs["response_cost"] == 1.0 and "discount_amount" not in breakdown - assert counter.totals == {} + assert await keys_matching(raw, "*") == () written_cost = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs["response_cost"] assert written_cost == pytest.approx(kwargs["standard_logging_object"]["response_cost"]) @@ -3603,7 +3616,9 @@ async def test_batch_accounting_does_not_retry_after_a_search_allowance_update(m from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting monkeypatch.setattr(litellm, "apply_search_free_tier", True) - counter: Final = _SharedCounter() + server: Final = fakeredis.FakeServer() + counter: Final = fake_redis_cache(server, namespace="search-callback") + raw: Final = fakeredis.FakeAsyncRedis(server=server) dependency: Final = SimpleNamespace( failed_tracking_alert=AsyncMock(), internal_usage_cache=SimpleNamespace(dual_cache=SimpleNamespace(redis_cache=counter)), @@ -3621,7 +3636,9 @@ async def test_batch_accounting_does_not_retry_after_a_search_allowance_update(m with track_batch_accounting("grounded-batch") as outcome: await ProxyDBLogger().async_log_success_event(kwargs, response, _CALL_START, _CALL_START) - assert sum(counter.totals.values()) == 5 + keys: Final = await keys_matching(raw, "search-callback:litellm:search_free_tier:*") + assert len(keys) == 1 + assert await raw.get(keys[0]) == b"5" assert outcome.failed is True assert outcome.retryable_failure is False diff --git a/tests/unit/proxy/hooks/test_search_free_tier.py b/tests/unit/proxy/hooks/test_search_free_tier.py index b57908378e5..a6f9f9ea836 100644 --- a/tests/unit/proxy/hooks/test_search_free_tier.py +++ b/tests/unit/proxy/hooks/test_search_free_tier.py @@ -1,7 +1,11 @@ from datetime import datetime, timedelta, timezone from typing import Final +from unittest.mock import patch +import fakeredis import pytest +import pytest_asyncio +from redis.exceptions import ConnectionError as RedisConnectionError from litellm.proxy.hooks.search_free_tier import ( apply_search_free_tier, @@ -9,29 +13,34 @@ from litellm.proxy.hooks.search_free_tier import ( search_free_tier_key, search_free_tier_ttl, ) -from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +from litellm.redis.stores.search_free_tier import TransportSearchFreeTierStore +from litellm.types.utils import CostBreakdown, ModelResponse, PromptTokensDetailsWrapper, StandardLoggingPayload, Usage +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + set_reply_callback, +) NOW: Final = datetime(2026, 10, 20, 12, 30, tzinfo=timezone.utc) RATE: Final = 0.01 -class _Counter: - def __init__(self) -> None: - self.totals: dict[str, float] = {} - self.ttls: dict[str, int | None] = {} - - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.totals[key] = self.totals.get(key, 0) + value - self.ttls[key] = ttl - return self.totals[key] +@pytest.fixture +def server() -> fakeredis.FakeServer: + return fakeredis.FakeServer() -class _BrokenCounter: - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - raise ConnectionError("redis is down") +@pytest.fixture +def raw(server: fakeredis.FakeServer) -> fakeredis.FakeAsyncRedis: + return fakeredis.FakeAsyncRedis(server=server) -def _model_info(billing_unit: str, count: int, period: str = "month", provider: str = "vertex_ai") -> dict: +@pytest_asyncio.fixture(loop_scope="function") +async def counter(server: fakeredis.FakeServer) -> TransportSearchFreeTierStore: + return TransportSearchFreeTierStore(fake_redis_cache(server, namespace="search-test").transport()) + + +def _model_info(billing_unit: str, count: int, period: str = "month", provider: str = "vertex_ai") -> dict[str, object]: return { "litellm_provider": provider, "mode": "chat", @@ -48,8 +57,16 @@ def _response(queries: int | None) -> ModelResponse: ) -def _call(model_info: dict, response_cost: float, cache_hit: bool | None = None) -> tuple[dict, dict]: - sl_object = { +def _breakdown(sl_object: StandardLoggingPayload) -> CostBreakdown: + breakdown: Final = sl_object["cost_breakdown"] + assert breakdown is not None + return breakdown + + +def _call( + model_info: dict[str, object], response_cost: float, cache_hit: bool | None = None +) -> tuple[dict[str, object], StandardLoggingPayload]: + sl_object: StandardLoggingPayload = { # pyright: ignore[reportAssignmentType] # the discount path reads only these keys "response_cost": response_cost, "cache_hit": cache_hit, "custom_llm_provider": model_info["litellm_provider"], @@ -93,8 +110,31 @@ def test_counter_ttl_outlives_the_period_boundary_by_under_two_hours(period, bou @pytest.mark.asyncio -async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_crossing_call(): - counter = _Counter() +async def test_search_discount_works_without_lua_millisecond_or_transaction_permissions( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +) -> None: + await raw.acl_setuser( + "default", + enabled=True, + nopass=True, + reset=True, + keys=["*"], + commands=["+incrbyfloat", "+ttl", "+expire", "+get", "+pttl", "+ping", "+info", "+client"], + ) + kwargs, sl_object = _call(_model_info("per_query", count=7), response_cost=1.0) + key: Final = "search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10" + + with patch("time.time", return_value=NOW.timestamp()): + cost: Final = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) + assert cost == kwargs["response_cost"] == sl_object["response_cost"] == 0.95 + assert await raw.get(key) == b"5" + assert await raw.pttl(key) == 995_400_000 + + +@pytest.mark.asyncio +async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_crossing_call( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): model_info = _model_info("per_query", count=7) reduced = [] for _ in range(3): @@ -108,13 +148,12 @@ async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_cross assert c2 == k2 == s2 == pytest.approx(1.0 - 2 * RATE) assert b2["discount_amount"] == pytest.approx(2 * RATE) assert c3 == k3 == s3 == 1.0 and "discount_amount" not in b3 - assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 15} - assert counter.ttls["litellm:search_free_tier:vertex_ai:test-group:2026-10"] > 0 + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"15" + assert await raw.pttl("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") > 0 @pytest.mark.asyncio -async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances(): - counter = _Counter() +async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances(counter: TransportSearchFreeTierStore): gemini_api = _model_info("per_query", count=3, provider="gemini") vertex_ai = _model_info("per_query", count=3, provider="vertex_ai") kwargs, sl_object = _call(gemini_api, response_cost=1.0) @@ -128,8 +167,9 @@ async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances(): @pytest.mark.asyncio -async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_query_count(): - counter = _Counter() +async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_query_count( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): model_info = _model_info("per_prompt", count=1, period="day") kwargs, sl_object = _call(model_info, response_cost=1.0) first = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) @@ -137,12 +177,11 @@ async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_q second = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) assert first == pytest.approx(1.0 - RATE) assert second == 1.0 - assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10-20": 2} + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10-20") == b"2" @pytest.mark.asyncio -async def test_sidecar_shaped_response_mapping_is_counted_like_a_model_response(): - counter = _Counter() +async def test_sidecar_shaped_response_mapping_is_counted_like_a_model_response(counter: TransportSearchFreeTierStore): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) response = { "usage": { @@ -156,55 +195,112 @@ async def test_sidecar_shaped_response_mapping_is_counted_like_a_model_response( @pytest.mark.asyncio -@pytest.mark.parametrize("counter", [None, _BrokenCounter()]) -async def test_full_price_when_no_counter_is_available(counter): +@pytest.mark.parametrize("connected", [True, False]) +async def test_full_price_when_no_counter_is_available( + counter: TransportSearchFreeTierStore, server: fakeredis.FakeServer, connected: bool +): + server.connected = connected kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) - assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) == 1.0 - assert kwargs["response_cost"] == 1.0 and "discount_amount" not in sl_object["cost_breakdown"] + assert ( + await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), None if connected else counter, NOW) == 1.0 + ) + assert kwargs["response_cost"] == 1.0 and "discount_amount" not in _breakdown(sl_object) @pytest.mark.asyncio @pytest.mark.parametrize("response", [_response(0), _response(None), ModelResponse()]) -async def test_requests_without_grounding_never_touch_the_counter(response): - counter = _Counter() +async def test_requests_without_grounding_never_touch_the_counter( + response, counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) assert await apply_search_free_tier(kwargs, sl_object, 1.0, response, counter, NOW) == 1.0 - assert counter.totals == {} + assert await keys_matching(raw, "*") == () @pytest.mark.asyncio -async def test_models_without_an_allowance_never_touch_the_counter(): - counter = _Counter() +async def test_models_without_an_allowance_never_touch_the_counter( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): model_info = {k: v for k, v in _model_info("per_query", count=100).items() if k != "search_free_requests"} kwargs, sl_object = _call(model_info, response_cost=1.0) assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) == 1.0 - assert counter.totals == {} + assert await keys_matching(raw, "*") == () @pytest.mark.asyncio -async def test_a_second_callback_run_on_the_same_request_counts_it_once(): - counter = _Counter() +async def test_a_second_callback_run_on_the_same_request_counts_it_once( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) first = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) second = await apply_search_free_tier(kwargs, sl_object, first, _response(5), counter, NOW) assert first == second == pytest.approx(1.0 - 5 * RATE) - assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5} + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"5" @pytest.mark.asyncio -async def test_a_cache_hit_with_grounded_usage_is_not_counted_and_stays_at_zero(): - counter = _Counter() +async def test_lost_increment_reply_keeps_full_price_and_only_recounts_on_an_explicit_callback_retry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="search-test") + store: Final = TransportSearchFreeTierStore(cache.transport()) + kwargs, sl_object = _call(_model_info("per_query", count=7), response_cost=1.0) + + def lost_reply(_: object) -> None: + raise RedisConnectionError("reply lost after allocation") + + set_reply_callback(cache, "INCRBYFLOAT", lost_reply) + assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), store, NOW) == 1.0 + assert kwargs["response_cost"] == sl_object["response_cost"] == 1.0 + assert not kwargs.get("search_free_tier_counted", False) + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"5" + + set_reply_callback(cache, "INCRBYFLOAT", float) + assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), store, NOW) == pytest.approx(0.98) + assert kwargs["search_free_tier_counted"] is True + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"10" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("period", "stamp", "boundary"), + [ + ("month", "2026-10", datetime(2026, 11, 1, tzinfo=timezone.utc)), + ("day", "2026-10-20", datetime(2026, 10, 21, tzinfo=timezone.utc)), + ], +) +async def test_caller_places_expiry_at_period_boundary_plus_one_hour( + server: fakeredis.FakeServer, + raw: fakeredis.FakeAsyncRedis, + period: str, + stamp: str, + boundary: datetime, +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="search-test").transport()) + kwargs, sl_object = _call(_model_info("per_query", count=7, period=period), response_cost=1.0) + with patch("time.time", return_value=NOW.timestamp()): + assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), store, NOW) == pytest.approx(0.95) # pyright: ignore[reportUnknownMemberType] # pytest.approx is untyped + assert await raw.pttl(f"search-test:litellm:search_free_tier:vertex_ai:test-group:{stamp}") == ( + (boundary - NOW + timedelta(hours=1)) // timedelta(milliseconds=1) + ) + + +@pytest.mark.asyncio +async def test_a_cache_hit_with_grounded_usage_is_not_counted_and_stays_at_zero( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=0.0, cache_hit=True) assert await apply_search_free_tier(kwargs, sl_object, 0.0, _response(2), counter, NOW) == 0.0 - assert counter.totals == {} - assert "discount_amount" not in sl_object["cost_breakdown"] + assert await keys_matching(raw, "*") == () + assert "discount_amount" not in _breakdown(sl_object) @pytest.mark.asyncio -async def test_freed_cost_never_exceeds_the_billed_cost(): - counter = _Counter() +async def test_freed_cost_never_exceeds_the_billed_cost( + counter: TransportSearchFreeTierStore, raw: fakeredis.FakeAsyncRedis +): billed = 3 * RATE kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=billed) assert await apply_search_free_tier(kwargs, sl_object, billed, _response(5), counter, NOW) == 0.0 - assert sl_object["cost_breakdown"]["discount_amount"] == pytest.approx(billed) - assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5} + assert _breakdown(sl_object).get("discount_amount") == pytest.approx(billed) # pyright: ignore[reportUnknownMemberType] # pytest.approx is untyped + assert await raw.get("search-test:litellm:search_free_tier:vertex_ai:test-group:2026-10") == b"5" diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index c721e808ed5..ea881d5cac2 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -659,6 +659,8 @@ async def test_completed_passthrough_cancellation_preserves_accounting( from unittest.mock import create_autospec import anyio + import fakeredis + from redis.asyncio.client import Pipeline from litellm.caching.dual_cache import DualCache from litellm.proxy import proxy_server @@ -667,20 +669,29 @@ async def test_completed_passthrough_cancellation_preserves_accounting( from litellm.proxy.spend_tracking.budget_reservation import release_unbound_budget_reservation from litellm.proxy.utils import ProxyLogging from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage + from tests._support.fake_redis import fake_redis_cache - class AllowanceCounter: - def __init__(self) -> None: - self.started = asyncio.Event() - self.release = asyncio.Event() - self.calls = 0 - self.units = 0.0 + allowance_started: Final = asyncio.Event() + allowance_release: Final = asyncio.Event() + executions: Final[list[int]] = [] # mutable-ok: counts pipeline flushes observed by the gated fake client - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.started.set() - await self.release.wait() - self.calls += 1 - self.units += value - return self.units + class GatedClient(fakeredis.FakeAsyncRedis): + def pipeline(self, transaction: bool = True, shard_hint: str | None = None) -> "Pipeline[bytes]": + pipe: Final = super().pipeline(transaction=transaction, shard_hint=shard_hint) # pyright: ignore[reportUnknownVariableType] # fakeredis pipeline is untyped + native_execute: Final = pipe.execute # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis-py pipeline.execute reply is untyped + + async def gated_execute(raise_on_error: bool = True) -> list[object]: + allowance_started.set() + await allowance_release.wait() + executions.append(1) + return await native_execute(raise_on_error=raise_on_error) # pyright: ignore[reportUnknownVariableType] # redis-py erases the reply type + + monkeypatch.setattr(pipe, "execute", gated_execute) # pyright: ignore[reportUnknownArgumentType] # fakeredis pipeline is untyped + return pipe # pyright: ignore[reportUnknownVariableType] # fakeredis pipeline is untyped + + server: Final = fakeredis.FakeServer() + allowance: Final = fake_redis_cache(server, client=GatedClient) # pyright: ignore[reportUnknownArgumentType] # fakeredis client class is partially untyped + verifier: Final = fakeredis.FakeAsyncRedis(server=server) now: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) model: Final = "passthrough-cancellation-test" @@ -714,7 +725,6 @@ async def test_completed_passthrough_cancellation_preserves_accounting( writer: Final = create_autospec(DBSpendUpdateWriter, instance=True) proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.db_spend_update_writer = writer - allowance: Final = AllowanceCounter() monkeypatch.setattr(proxy_logging.internal_usage_cache.dual_cache, "redis_cache", allowance) monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging) logging_obj: Final = LiteLLMLoggingObj( @@ -771,7 +781,7 @@ async def test_completed_passthrough_cancellation_preserves_accounting( pending: Final = asyncio.create_task(dispatch()) await entered.wait() if cancel_phase == "allocation": - await allowance.started.wait() + await allowance_started.wait() if cancel_mode == "scope": scope.cancel() elif cancel_mode != "none": @@ -781,7 +791,7 @@ async def test_completed_passthrough_cancellation_preserves_accounting( pending.cancel() if cancel_phase == "lock": logging_obj.cost_finalization_lock.release() - allowance.release.set() + allowance_release.set() if cancel_mode in ("once", "twice"): with pytest.raises(asyncio.CancelledError): await pending @@ -791,8 +801,9 @@ async def test_completed_passthrough_cancellation_preserves_accounting( assert reservation["callback_bound"] is True await release_unbound_budget_reservation(reservation) await GLOBAL_LOGGING_WORKER.flush() - assert allowance.calls == 1 - assert allowance.units == 3 + assert executions == [1] + (allocation_key,) = await verifier.keys("*") # pyright: ignore[reportUnknownVariableType] # fakeredis replies are untyped + assert await verifier.get(allocation_key) == b"3" # pyright: ignore[reportUnknownArgumentType] # fakeredis replies are untyped assert writer.update_database.await_count == 1 recorded: Final = writer.update_database.call_args.kwargs assert recorded["response_cost"] == pytest.approx(0.48) diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index aefc4a96a28..47eb37c414f 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -1,44 +1,61 @@ import asyncio +from dataclasses import dataclass from datetime import datetime, timezone from typing import Final +import fakeredis import pytest +import pytest_asyncio +from redis.exceptions import ConnectionError as RedisConnectionError import litellm +from litellm.caching.redis_cache import RedisCache from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.spend_tracking.cost_finalization import finalize_response_cost +from litellm.redis.stores.search_free_tier import TransportSearchFreeTierStore from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +from tests._support.fake_redis import ( + fake_redis_cache, + set_reply_callback, +) NOW: Final = datetime(2026, 10, 20, 12, 30, tzinfo=timezone.utc) MODEL: Final = "gemini-finalization-test" -class Counter: - def __init__(self, total: float = 0.0, fail: bool = False) -> None: - self.total = total - self.calls = 0 - self.fail = fail - - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.calls += 1 - if self.fail: - raise ConnectionError("counter unavailable") - self.total += value - return self.total +_ALLOCATION_KEY: Final = "ns:litellm:search_free_tier:vertex_ai:finalization-test:2026-10" -class BlockingCounter(Counter): - def __init__(self) -> None: - super().__init__() - self.started = asyncio.Event() - self.release = asyncio.Event() +@dataclass(frozen=True, slots=True) +class AllocationRedis: + server: fakeredis.FakeServer + raw: fakeredis.FakeAsyncRedis + cache: RedisCache + store: TransportSearchFreeTierStore - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.started.set() - await self.release.wait() - return await super().async_increment(key, value, ttl) + +@pytest_asyncio.fixture(loop_scope="function") +async def allocation_redis() -> AllocationRedis: + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server, namespace="ns") + return AllocationRedis( + server, fakeredis.FakeAsyncRedis(server=server), cache, TransportSearchFreeTierStore(cache.transport()) + ) + + +def block_increment(cache: RedisCache) -> tuple[asyncio.Event, asyncio.Event]: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def receive(value: bytes) -> float: + started.set() + await release.wait() + return float(value) + + set_reply_callback(cache, "INCRBYFLOAT", receive) + return started, release @pytest.fixture @@ -89,15 +106,16 @@ def response(cost: float = 0.5) -> ModelResponse: @pytest.mark.asyncio @pytest.mark.parametrize(("used", "expected"), [(0, 0.48), (1, 0.49), (2, 0.5)]) async def test_header_and_logging_share_one_allowance_allocation( - logging_obj: Logging, used: int, expected: float + logging_obj: Logging, allocation_redis: AllocationRedis, used: int, expected: float ) -> None: - counter: Final = Counter(used) + await allocation_redis.raw.set(_ALLOCATION_KEY, used) + counter: Final = allocation_redis.store result: Final = response() costs: Final = await asyncio.gather(*(finalize_response_cost(logging_obj, result, counter, NOW) for _ in range(3))) assert costs[0] is costs[1] is costs[2] assert costs[0] is not None assert costs[0].amount == pytest.approx(expected) - assert counter.calls == 1 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == str(used + 3).encode() logging_obj._success_handler_helper_fn(result=result, start_time=NOW, end_time=NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -109,19 +127,47 @@ async def test_header_and_logging_share_one_allowance_allocation( @pytest.mark.asyncio -async def test_failed_counter_freezes_full_price_without_a_retry(logging_obj: Logging) -> None: - counter: Final = Counter(fail=True) +async def test_failed_counter_freezes_full_price_without_a_retry( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: + allocation_redis.server.connected = False + counter: Final = allocation_redis.store result: Final = response() finalized: Final = await finalize_response_cost(logging_obj, result, counter, NOW) assert finalized is not None and finalized.amount == 0.5 + allocation_redis.server.connected = True assert await finalize_response_cost(logging_obj, result, counter, NOW) is finalized - assert counter.calls == 1 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None + assert logging_obj.cost_allowance_discount == 0.0 @pytest.mark.asyncio -async def test_guardrail_cost_is_added_once_after_allowance(logging_obj: Logging) -> None: +async def test_lost_allocation_reply_freezes_full_price_after_allowance_was_consumed( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: + def lost_reply(_: object) -> None: + raise RedisConnectionError("allocation response lost") + + set_reply_callback(allocation_redis.cache, "INCRBYFLOAT", lost_reply) result: Final = response() - finalized: Final = await finalize_response_cost(logging_obj, result, Counter(), NOW, guardrail_cost=0.2) + finalized: Final = await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) + assert finalized is not None and finalized.amount == 0.5 + assert logging_obj.cost_allowance_discount == 0.0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" + + set_reply_callback(allocation_redis.cache, "INCRBYFLOAT", float) + assert await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) is finalized + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" + + +@pytest.mark.asyncio +async def test_guardrail_cost_is_added_once_after_allowance( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: + result: Final = response() + finalized: Final = await finalize_response_cost( + logging_obj, result, allocation_redis.store, NOW, guardrail_cost=0.2 + ) assert finalized is not None and finalized.amount == pytest.approx(0.68) logging_obj._success_handler_helper_fn(result=result, start_time=NOW, end_time=NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] @@ -130,18 +176,22 @@ async def test_guardrail_cost_is_added_once_after_allowance(logging_obj: Logging @pytest.mark.asyncio -async def test_cache_hit_finalizes_zero_without_allocating(logging_obj: Logging) -> None: +async def test_cache_hit_finalizes_zero_without_allocating( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: result: Final = response() result._hidden_params["cache_hit"] = True result._hidden_params["response_cost"] = 0.0 - counter: Final = Counter() + counter: Final = allocation_redis.store finalized: Final = await finalize_response_cost(logging_obj, result, counter, NOW) assert finalized is not None and finalized.amount == 0.0 - assert counter.calls == 0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None @pytest.mark.asyncio -async def test_both_callback_paths_wait_for_finalization(logging_obj: Logging) -> None: +async def test_both_callback_paths_wait_for_finalization( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.utils import dispatch_success_logging @@ -169,18 +219,19 @@ async def test_both_callback_paths_wait_for_finalization(logging_obj: Logging) - assert not logging_obj._async_success_enqueued assert logging_obj.deferred_success_result is not None assert logging_obj.deferred_success_result[0] is inner - counter: Final = BlockingCounter() + counter: Final = allocation_redis.store + started, release = block_increment(allocation_redis.cache) pending: Final = asyncio.create_task(finalize_response_cost(logging_obj, outer, counter, NOW)) - await counter.started.wait() + await started.wait() dispatch_success_logging(logging_obj, outer, NOW, NOW, False, False) assert not logging_obj._sync_success_scheduled assert not logging_obj._async_success_enqueued concurrent: Final = asyncio.create_task(finalize_response_cost(logging_obj, outer, counter, NOW)) - counter.release.set() + release.set() finalized: Final = await pending assert finalized is not None assert await concurrent is finalized - assert counter.calls == 1 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logging_obj, exception_raised=False) observed: Final = await asyncio.wait_for(asyncio.gather(sync_cost, async_cost), timeout=10) await GLOBAL_LOGGING_WORKER.flush() @@ -201,12 +252,14 @@ async def test_rejected_response_does_not_release_success_callbacks(logging_obj: @pytest.mark.asyncio -async def test_prebuilt_payload_is_reconciled_with_final_cost(logging_obj: Logging) -> None: +async def test_prebuilt_payload_is_reconciled_with_final_cost( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: result: Final = response() prepared: Final = logging_obj.prepare_success_logging(result, NOW, NOW, None, None, True, False) original: Final = prepared[3] assert original is not None and original["response_cost"] == 0.5 - finalized: Final = await finalize_response_cost(logging_obj, result, Counter(), NOW) + finalized: Final = await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) assert finalized is not None logging_obj._success_handler_helper_fn(result, NOW, NOW, standard_logging_object=original) payload: Final = logging_obj.model_call_details["standard_logging_object"] @@ -215,7 +268,9 @@ async def test_prebuilt_payload_is_reconciled_with_final_cost(logging_obj: Loggi @pytest.mark.asyncio -async def test_cached_response_waits_for_guardrail_cost(logging_obj: Logging) -> None: +async def test_cached_response_waits_for_guardrail_cost( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: from litellm.caching.caching_handler import LLMCachingHandler result: Final = response(0.0) @@ -225,17 +280,19 @@ async def test_cached_response_waits_for_guardrail_cost(logging_obj: Logging) -> handler._async_log_cache_hit_on_callbacks(logging_obj, result, NOW, NOW, True) assert not logging_obj._sync_success_scheduled assert not logging_obj._async_success_enqueued - counter: Final = Counter() + counter: Final = allocation_redis.store finalized: Final = await finalize_response_cost(logging_obj, result, counter, NOW, guardrail_cost=0.2) assert finalized is not None and finalized.amount == 0.2 - assert counter.calls == 0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None logging_obj._success_handler_helper_fn(result, NOW, NOW) assert logging_obj.model_call_details["standard_logging_object"]["response_cost"] == 0.2 ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logging_obj, exception_raised=True) @pytest.mark.asyncio -async def test_collector_event_preserves_finalized_cost(logging_obj: Logging) -> None: +async def test_collector_event_preserves_finalized_cost( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: from litellm.proxy.spend_tracking.spend_event import ( SpendEvent, build_spend_event, @@ -244,7 +301,7 @@ async def test_collector_event_preserves_finalized_cost(logging_obj: Logging) -> ) result: Final = response() - await finalize_response_cost(logging_obj, result, Counter(), NOW) + await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) logging_obj._success_handler_helper_fn(result, NOW, NOW) event: Final = build_spend_event(logging_obj.model_call_details, result, NOW, NOW, store_bodies=False) assert isinstance(event, bytes) @@ -259,30 +316,33 @@ async def test_collector_event_preserves_finalized_cost(logging_obj: Logging) -> @pytest.mark.asyncio @pytest.mark.parametrize("cancel_count", [1, 2]) async def test_cancellation_finishes_an_inflight_allocation_before_releasing_callbacks( - logging_obj: Logging, cancel_count: int + logging_obj: Logging, allocation_redis: AllocationRedis, cancel_count: int ) -> None: - counter: Final = BlockingCounter() + counter: Final = allocation_redis.store + started, release = block_increment(allocation_redis.cache) result: Final = response() pending: Final = asyncio.create_task(finalize_response_cost(logging_obj, result, counter, NOW)) - await counter.started.wait() + await started.wait() for _ in range(cancel_count): pending.cancel() await asyncio.sleep(0) assert not pending.done() - counter.release.set() + release.set() with pytest.raises(asyncio.CancelledError): await pending assert logging_obj.finalized_cost is not None assert logging_obj.finalized_cost.amount == pytest.approx(0.48) assert await finalize_response_cost(logging_obj, result, counter, NOW) is logging_obj.finalized_cost - assert counter.calls == 1 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" @pytest.mark.asyncio -async def test_existing_calculator_prices_usage_before_finalizing(logging_obj: Logging) -> None: +async def test_existing_calculator_prices_usage_before_finalizing( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: result: Final = response() result._hidden_params.pop("response_cost") - finalized: Final = await finalize_response_cost(logging_obj, result, Counter(), NOW) + finalized: Final = await finalize_response_cost(logging_obj, result, allocation_redis.store, NOW) assert finalized is not None and finalized.amount == pytest.approx(0.21) logging_obj._success_handler_helper_fn(result, NOW, NOW) assert logging_obj.model_call_details["standard_logging_object"]["response_cost"] == pytest.approx(0.21) @@ -371,20 +431,22 @@ async def test_malformed_passthrough_logging_preserves_response_dispatch_contrac @pytest.mark.asyncio async def test_finalization_failure_leaves_the_request_and_callback_billing_intact( - logging_obj: Logging, monkeypatch: pytest.MonkeyPatch + logging_obj: Logging, allocation_redis: AllocationRedis, monkeypatch: pytest.MonkeyPatch ) -> None: def explode(**kwargs: object) -> object: raise RuntimeError("payload build failed") monkeypatch.setattr(logging_obj, "prepare_success_logging", explode) - counter: Final = Counter() + counter: Final = allocation_redis.store assert await finalize_response_cost(logging_obj, response(), counter, NOW) is None assert logging_obj.finalized_cost is None - assert counter.calls == 0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None @pytest.mark.asyncio -async def test_deferred_cache_hit_finalizes_zero_for_results_without_hidden_params(logging_obj: Logging) -> None: +async def test_deferred_cache_hit_finalizes_zero_for_results_without_hidden_params( + logging_obj: Logging, allocation_redis: AllocationRedis +) -> None: from litellm import completion from litellm.caching.caching_handler import LLMCachingHandler @@ -395,10 +457,10 @@ async def test_deferred_cache_hit_finalizes_zero_for_results_without_hidden_para )._async_log_cache_hit_on_callbacks( logging_obj=logging_obj, cached_result=cached_result, start_time=NOW, end_time=NOW, cache_hit=True ) - counter: Final = Counter() + counter: Final = allocation_redis.store finalized: Final = await finalize_response_cost(logging_obj, cached_result, counter, NOW) assert finalized is not None and finalized.amount == 0.0 - assert counter.calls == 0 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) is None logging_obj._success_handler_helper_fn(cached_result, NOW, NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] assert payload["cache_hit"] is True @@ -423,18 +485,17 @@ async def test_completed_call_keeps_its_charge_when_post_call_processing_fails( disable_spend_updates: bool, ) -> None: from unittest.mock import create_autospec + import anyio - from fastapi import Request, Response - from litellm.proxy._types import ProxyException - from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.proxy_server import ProxyConfig - from litellm.proxy import proxy_server from litellm.proxy.utils import ProxyLogging scope: Final = anyio.CancelScope() @@ -545,7 +606,7 @@ async def test_completed_call_keeps_its_charge_when_post_call_processing_fails( @pytest.mark.asyncio async def test_redacted_completed_call_preserves_usage_and_search_allowance( - logging_obj: Logging, monkeypatch: pytest.MonkeyPatch + logging_obj: Logging, allocation_redis: AllocationRedis, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.proxy.spend_tracking.cost_finalization import prepare_response_cost @@ -554,11 +615,11 @@ async def test_redacted_completed_call_preserves_usage_and_search_allowance( prepared: Final = prepare_response_cost(logging_obj, result, NOW) assert prepared is not None result.usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) - counter: Final = Counter() + counter: Final = allocation_redis.store finalized: Final = await finalize_response_cost(logging_obj, result, counter, NOW) assert finalized is not None assert finalized.amount == pytest.approx(0.48) - assert counter.total == 3 + assert await allocation_redis.raw.get(_ALLOCATION_KEY) == b"3" logging_obj._failure_handler_helper_fn(RuntimeError("post-call failed"), "", NOW, NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] assert (payload["prompt_tokens"], payload["completion_tokens"]) == (10, 5) @@ -568,9 +629,10 @@ async def test_redacted_completed_call_preserves_usage_and_search_allowance( @pytest.mark.asyncio async def test_guardrail_json_rejection_retains_completed_provider_usage( - logging_obj: Logging, monkeypatch: pytest.MonkeyPatch + logging_obj: Logging, allocation_redis: AllocationRedis, monkeypatch: pytest.MonkeyPatch ) -> None: import httpx + from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.utils import ProxyLogging @@ -599,7 +661,7 @@ async def test_guardrail_json_rejection_retains_completed_provider_usage( ) assert logging_obj.prepared_cost is not None normalized.usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) - finalized: Final = await finalize_response_cost(logging_obj, upstream, Counter(), NOW) + finalized: Final = await finalize_response_cost(logging_obj, upstream, allocation_redis.store, NOW) assert finalized is not None and finalized.amount == 0.5 logging_obj._failure_handler_helper_fn(RuntimeError("blocked JSON response"), "", NOW, NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] diff --git a/tests/unit/redis/stores/test_search_free_tier.py b/tests/unit/redis/stores/test_search_free_tier.py new file mode 100644 index 00000000000..993e6d45fd1 --- /dev/null +++ b/tests/unit/redis/stores/test_search_free_tier.py @@ -0,0 +1,233 @@ +from collections.abc import Iterator +from datetime import timedelta +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import ResponseError + +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.redis.stores.search_free_tier import TransportSearchFreeTierStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + legacy_default_ttl, + legacy_increment, + set_reply_callback, +) + +_KEY: Final = "ns:litellm:search_free_tier:vertex_ai:shared:2026-10" +_LOGICAL_KEY: Final = "litellm:search_free_tier:vertex_ai:shared:2026-10" + + +def _not_a_number(_: object) -> bytes: + return b"not-a-number" + + +@pytest.fixture(autouse=True) +def frozen_clock() -> Iterator[None]: + with patch("time.time", return_value=1792497600.0): + yield + + +@pytest.mark.asyncio +async def test_old_and_new_writers_share_numeric_bytes_and_do_not_refresh_expiry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + assert await legacy_increment(cache, _LOGICAL_KEY, 2.5, ttl=600) == 2.5 + assert await raw.get(_KEY) == b"2.5" + assert await raw.pttl(_KEY) == 600_000 + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=900)) == 5.5 + assert await raw.get(_KEY) == b"5.5" + assert await raw.pttl(_KEY) == 600_000 + assert await legacy_increment(cache, _LOGICAL_KEY, 1.5, ttl=1200) == 7.0 + assert await raw.get(_KEY) == b"7" + assert await raw.pttl(_KEY) == 600_000 + assert await keys_matching(raw, "*") == (_KEY.encode(),) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [None, timedelta(seconds=60)]) +async def test_new_writer_is_readable_by_legacy_increment( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: timedelta | None +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 4, ttl=ttl) == 4.0 + assert await raw.get(_KEY) == b"4" + assert await raw.pttl(_KEY) == (-1 if ttl is None else 60_000) + assert await legacy_increment(cache, _LOGICAL_KEY, 2.5, ttl=None) == 6.5 + assert await raw.get(_KEY) == b"6.5" + assert await raw.pttl(_KEY) == (legacy_default_ttl(cache) * 1000 if ttl is None else 60_000) + + +@pytest.mark.asyncio +async def test_provider_group_and_period_have_independent_allowance_counters( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="ns").transport()) + identities: Final = ( + ("vertex_ai", "shared", "2026-10"), + ("gemini", "shared", "2026-10"), + ("vertex_ai", "other", "2026-10"), + ("vertex_ai", "shared", "2026-11"), + ("vertex_ai", "shared", "2026-10-20"), + ("vertex_ai", "shared 雪", "2026-10"), + ) + for provider, group, period in identities: + assert await store.allocate_units(provider, group, period, 2, ttl=None) == 2.0 + assert sorted(await keys_matching(raw, "*")) == sorted( + f"ns:litellm:search_free_tier:{provider}:{group}:{period}".encode() for provider, group, period in identities + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [-1, 0, 300]) +async def test_unexpiring_legacy_counter_gets_seconds_based_expiry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, seconds: int +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="ns").transport()) + legacy: Final = fake_redis_cache(server, namespace="legacy") + legacy_key: Final = f"legacy:{_LOGICAL_KEY}" + await raw.set(_KEY, b"5") + await raw.set(legacy_key, b"5") + assert await legacy_increment(legacy, _LOGICAL_KEY, 2, ttl=seconds) == 7.0 + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 2, ttl=timedelta(seconds=seconds)) == 7.0 + assert await raw.get(_KEY) == await raw.get(legacy_key) + assert await raw.pttl(_KEY) == await raw.pttl(legacy_key) + + +@pytest.mark.asyncio +async def test_expiry_remains_a_separate_write_after_increment( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + async def ttl_reply(value: int) -> int: + assert value == -1 + assert await raw.get(_KEY) == b"3" + assert await raw.pttl(_KEY) == -1 + await raw.pexpire(_KEY, 900_000) + return value + + set_reply_callback(cache, "TTL", ttl_reply) + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) == 3.0 + assert await raw.pttl(_KEY) == 60_000 + + +@pytest.mark.asyncio +async def test_counter_disappearing_before_expire_keeps_successful_allocation_result( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + async def ttl_reply(value: int) -> int: + await raw.delete(_KEY) + return value + + set_reply_callback(cache, "TTL", ttl_reply) + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) == 3.0 + assert await raw.get(_KEY) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("command,remaining_ttl", [("INCRBYFLOAT", -1), ("TTL", -1), ("EXPIRE", 60_000)]) +async def test_malformed_replies_are_failures_after_the_counter_was_consumed( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, command: str, remaining_ttl: int +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + set_reply_callback(cache, command, _not_a_number) + + outcome: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) + assert isinstance(outcome, RedisFailure) and outcome.kind == "decode" + assert await raw.get(_KEY) == b"3" + assert await raw.pttl(_KEY) == remaining_ttl + + +@pytest.mark.asyncio +@pytest.mark.parametrize("command", ["INCRBYFLOAT", "TTL", "EXPIRE"]) +async def test_lost_reply_reports_unknown_applied_write_without_retry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, command: str +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + def lose_reply(_: object) -> None: + raise RedisConnectionError("allocation reply lost") + + set_reply_callback(cache, command, lose_reply) + outcome: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) + assert isinstance(outcome, RedisFailure) and outcome.kind == "transport" + assert await raw.get(_KEY) == b"3" + assert await raw.pttl(_KEY) == (60_000 if command == "EXPIRE" else -1) + + +@pytest.mark.asyncio +async def test_corrupt_counter_is_not_reset_or_expired( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="ns").transport()) + await raw.set(_KEY, b"corrupt") + outcome: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) + assert isinstance(outcome, RedisFailure) and outcome.kind == "server" + assert await raw.get(_KEY) == b"corrupt" + assert await raw.pttl(_KEY) == -1 + + +@pytest.mark.asyncio +async def test_outage_opens_existing_cache_breaker_and_rejects_writes_after_reconnection( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + store: Final = TransportSearchFreeTierStore(fake_redis_cache(server, namespace="ns").transport()) + server.connected = False + for _ in range(10): + outcome = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=None) + assert isinstance(outcome, RedisFailure) and outcome.kind in ("transport", "breaker_open") + server.connected = True + blocked: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=None) + assert isinstance(blocked, RedisFailure) and blocked.kind == "breaker_open" + assert await raw.get(_KEY) is None + + +@pytest.mark.asyncio +async def test_expiry_server_failure_reports_consumed_allocation_without_retry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + store: Final = TransportSearchFreeTierStore(cache.transport()) + + def denied(_: object) -> None: + raise ResponseError("NOPERM expiry denied") + + set_reply_callback(cache, "EXPIRE", denied) + outcome: Final = await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) + assert isinstance(outcome, RedisFailure) and outcome.kind == "server" + assert await raw.get(_KEY) == b"3" + + +@pytest.mark.asyncio +async def test_allocation_flushes_pending_post_call_work_on_the_same_transport( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + transport: Final = cache.transport() + store: Final = TransportSearchFreeTierStore(transport) + with request_redis_batch_scope(): + batch: Final = transport.post_call_batch() + assert batch is not None + pending: Final = batch.set("other-keyspace", b"queued", ttl=None) + assert await store.allocate_units("vertex_ai", "shared", "2026-10", 3, ttl=timedelta(seconds=60)) == 3.0 + assert pending.done and await pending is True + assert await raw.get("ns:other-keyspace") == b"queued" + assert await raw.get(_KEY) == b"3" + assert await raw.pttl(_KEY) == 60_000