From 640f2d865bfa954d37bbbffc447f4e504e1dad21 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 15:01:53 -0700 Subject: [PATCH] fix(proxy): scope the search free-tier counter to the request's provider (#47646) * fix(proxy): scope the search free-tier counter to the request's provider * fix(proxy): read the provider with get so partial spend payloads still get the free tier * test(proxy): give the grounded spend payload fixture its provider --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/hooks/search_free_tier.py | 8 ++-- .../hooks/test_proxy_track_cost_callback.py | 1 + .../unit/proxy/hooks/test_search_free_tier.py | 40 +++++++++++++------ 3 files changed, 34 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/hooks/search_free_tier.py b/litellm/proxy/hooks/search_free_tier.py index 4ef300dfbac..34f4d50550a 100644 --- a/litellm/proxy/hooks/search_free_tier.py +++ b/litellm/proxy/hooks/search_free_tier.py @@ -34,8 +34,9 @@ def freed_units(allowance: int, units: int, total_after: int) -> int: return max(0, min(units, allowance - (total_after - units))) -def search_free_tier_key(allowance: SearchFreeRequests, now: datetime) -> str: - return f"litellm:search_free_tier:{allowance['shared_group']}:{now.strftime(_PERIOD_STAMP[allowance['period']])}" +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}" def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int: @@ -107,7 +108,8 @@ async def apply_search_free_tier( "apply_search_free_tier is on but the proxy has no Redis, grounding is billed at full price" ) return response_cost - key: Final = search_free_tier_key(allowance, now) + 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 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 289f8cd4a1a..ba1b5536a07 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -3202,6 +3202,7 @@ def _grounded_gemini_kwargs(response_cost: float) -> dict: "response_cost": response_cost, "request_tags": None, "cache_hit": None, + "custom_llm_provider": "vertex_ai", "cost_breakdown": { "input_cost": 0.0, "output_cost": 0.0, diff --git a/tests/unit/proxy/hooks/test_search_free_tier.py b/tests/unit/proxy/hooks/test_search_free_tier.py index 3dc89c0dc72..b57908378e5 100644 --- a/tests/unit/proxy/hooks/test_search_free_tier.py +++ b/tests/unit/proxy/hooks/test_search_free_tier.py @@ -31,9 +31,9 @@ class _BrokenCounter: raise ConnectionError("redis is down") -def _model_info(billing_unit: str, count: int, period: str = "month") -> dict: +def _model_info(billing_unit: str, count: int, period: str = "month", provider: str = "vertex_ai") -> dict: return { - "litellm_provider": "vertex_ai", + "litellm_provider": provider, "mode": "chat", "web_search_billing_unit": billing_unit, "search_context_cost_per_query": {"search_context_size_medium": RATE}, @@ -52,6 +52,7 @@ def _call(model_info: dict, response_cost: float, cache_hit: bool | None = None) sl_object = { "response_cost": response_cost, "cache_hit": cache_hit, + "custom_llm_provider": model_info["litellm_provider"], "cost_breakdown": {"input_cost": 0.1, "output_cost": 0.2, "tool_usage_cost": 0.5, "total_cost": response_cost}, "model_map_information": {"model_map_key": "gemini-test", "model_map_value": model_info}, } @@ -73,12 +74,12 @@ def test_freed_units_covers_only_what_landed_under_the_allowance(allowance, unit assert freed_units(allowance, units, total_after) == expected -def test_counter_key_is_per_group_and_utc_period(): - assert search_free_tier_key({"count": 1, "period": "month", "shared_group": "g3"}, NOW) == ( - "litellm:search_free_tier:g3:2026-10" +def test_counter_key_is_per_provider_group_and_utc_period(): + assert search_free_tier_key("gemini", {"count": 1, "period": "month", "shared_group": "g3"}, NOW) == ( + "litellm:search_free_tier:gemini:g3:2026-10" ) - assert search_free_tier_key({"count": 1, "period": "day", "shared_group": "g2"}, NOW) == ( - "litellm:search_free_tier:g2:2026-10-20" + assert search_free_tier_key("vertex_ai", {"count": 1, "period": "day", "shared_group": "g2"}, NOW) == ( + "litellm:search_free_tier:vertex_ai:g2:2026-10-20" ) @@ -107,8 +108,23 @@ 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:test-group:2026-10": 15} - assert counter.ttls["litellm:search_free_tier:test-group:2026-10"] > 0 + 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 + + +@pytest.mark.asyncio +async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances(): + counter = _Counter() + 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) + await apply_search_free_tier(kwargs, sl_object, 1.0, _response(2), counter, NOW) + kwargs, sl_object = _call(vertex_ai, response_cost=1.0) + vertex_cost = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(4), counter, NOW) + kwargs, sl_object = _call(gemini_api, response_cost=1.0) + gemini_cost = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(1), counter, NOW) + assert vertex_cost == pytest.approx(1.0 - 3 * RATE), "vertex ai had its whole allowance left" + assert gemini_cost == pytest.approx(1.0 - RATE), "gemini api had one free query left" @pytest.mark.asyncio @@ -121,7 +137,7 @@ 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:test-group:2026-10-20": 2} + assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10-20": 2} @pytest.mark.asyncio @@ -172,7 +188,7 @@ async def test_a_second_callback_run_on_the_same_request_counts_it_once(): 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:test-group:2026-10": 5} + assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5} @pytest.mark.asyncio @@ -191,4 +207,4 @@ async def test_freed_cost_never_exceeds_the_billed_cost(): 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:test-group:2026-10": 5} + assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5}