mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
refactor(redis): move search free-tier allocation into a typed store
This commit is contained in:
Devin AI
parent
0cba26b520
commit
961290fc46
10 changed files with 670 additions and 176 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
41
litellm/redis/stores/search_free_tier.py
Normal file
41
litellm/redis/stores/search_free_tier.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
233
tests/unit/redis/stores/test_search_free_tier.py
Normal file
233
tests/unit/redis/stores/test_search_free_tier.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue