From 80104f9785b39b284df620af25656e3796794cb6 Mon Sep 17 00:00:00 2001 From: Jason Kidd Date: Wed, 8 Apr 2026 08:42:38 -0700 Subject: [PATCH] fix: use model_group instead of model for per-model budget cache keys on write path The write path in async_log_success_event used standard_logging_payload["model"] (the deployment name, e.g. "deepseek/deepseek-chat") to build cache keys. The read path in is_key_within_model_budget uses request_data["model"] (the public model name, e.g. "deepseek-chat") to build cache keys. These never match when the deployment name includes a provider prefix. For example, a model group "deepseek-chat" backed by deployments "deepseek/deepseek-chat" and "openrouter/deepseek/deepseek-chat": Write path (after deepseek/deepseek-chat call): cache key = virtual_key_spend:{hash}:deepseek/deepseek-chat:1d Write path (after openrouter/deepseek/deepseek-chat call): cache key = virtual_key_spend:{hash}:openrouter/deepseek/deepseek-chat:1d Read path (budget check): cache key = virtual_key_spend:{hash}:deepseek-chat:1d -> miss Spend is written to keys the read path never looks up. The budget is never enforced, and spend from different deployments in the same group is tracked in separate counters instead of being aggregated. The fix: prefer standard_logging_payload["model_group"] (the public router name that matches how budgets are configured) over ["model"], falling back to ["model"] when model_group is not set. This aligns the write path cache keys with the read path and correctly aggregates spend across all deployments in a model group. Includes 9 new tests covering model_group usage with single-slash and multi-slash deployment names, fallback behavior, cache key alignment between write and read paths, cross-deployment spend aggregation, and an end-to-end write-then-read round trip. --- .../proxy/hooks/model_max_budget_limiter.py | 8 +- ...st_model_budget_write_path_key_mismatch.py | 350 ++++++++++++++++++ 2 files changed, 357 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/hooks/test_model_budget_write_path_key_mismatch.py diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 5e48ef2879e..0fb985a7485 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -255,7 +255,13 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return response_cost: float = standard_logging_payload.get("response_cost", 0) - model = standard_logging_payload.get("model") + # Use model_group (public/router name, e.g. "gpt-4") rather than + # model (deployment name, e.g. "azure/gpt-4-deployment"). Budget + # configs and the read path both use the public model name, so the + # write path must match to produce cache hits. + model = standard_logging_payload.get( + "model_group" + ) or standard_logging_payload.get("model") virtual_key = standard_logging_payload.get("metadata", {}).get( "user_api_key_hash" ) diff --git a/tests/test_litellm/proxy/hooks/test_model_budget_write_path_key_mismatch.py b/tests/test_litellm/proxy/hooks/test_model_budget_write_path_key_mismatch.py new file mode 100644 index 00000000000..854e7ccb8e9 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_model_budget_write_path_key_mismatch.py @@ -0,0 +1,350 @@ +""" +Tests for the model vs model_group key mismatch fix on the write path. + +The write path in async_log_success_event previously used +``standard_logging_payload["model"]`` (the deployment name, e.g. +``"deepseek/deepseek-chat"``) for cache keys. Budget configs and the +read path use the public model name (``model_group``, e.g. +``"deepseek-chat"``), so the cache keys never matched when the +deployment name included a provider prefix. + +The fix: prefer ``model_group`` over ``model`` on the write path. + +These tests verify: +- Write path uses model_group when available +- Write path falls back to model when model_group is None +- Cache keys from write path match what the read path expects +- Spend from multiple deployments in the same group is aggregated +- End-to-end: spend written via model_group is visible on read +""" + +import os +import sys +from unittest.mock import AsyncMock, patch + +sys.path.insert( + 0, os.path.abspath("../../..") +) + +import pytest + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.model_max_budget_limiter import ( + END_USER_SPEND_CACHE_KEY_PREFIX, + VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, + _PROXY_VirtualKeyModelMaxBudgetLimiter, +) + + +@pytest.fixture +def budget_limiter(): + dual_cache = DualCache() + return _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + + +def _make_kwargs( + *, + model: str, + model_group: str = None, + response_cost: float = 0.05, + virtual_key: str = "test-key-hash", + end_user_id: str = None, + key_budget: dict = None, + end_user_budget: dict = None, +): + """Build kwargs dict for async_log_success_event.""" + metadata = {"user_api_key_hash": virtual_key} + if end_user_id: + metadata["user_api_key_end_user_id"] = end_user_id + + slp = { + "response_cost": response_cost, + "model": model, + "metadata": metadata, + } + if model_group is not None: + slp["model_group"] = model_group + if end_user_id: + slp["end_user"] = end_user_id + + litellm_metadata = {} + if key_budget is not None: + litellm_metadata["user_api_key_model_max_budget"] = key_budget + if end_user_budget is not None: + litellm_metadata["user_api_key_end_user_model_max_budget"] = end_user_budget + + return { + "standard_logging_object": slp, + "litellm_params": {"metadata": litellm_metadata}, + } + + +# Real model names from model_prices_and_context_window.json: +# deepseek/deepseek-chat: input=$0.28/M, output=$0.42/M +# openrouter/deepseek/deepseek-chat: input=$0.14/M, output=$0.28/M +# Both are deployments a user might group under model_name "deepseek-chat". + +_MODEL_GROUP = "deepseek-chat" +_DEPLOYMENT_A = "deepseek/deepseek-chat" # single-slash prefix +_DEPLOYMENT_B = "openrouter/deepseek/deepseek-chat" # multi-slash prefix +_BUDGET = {_MODEL_GROUP: {"budget_limit": 100.0, "time_period": "1d"}} + + +class TestWritePathUsesModelGroup: + """The write path should use model_group (public name) for cache keys.""" + + @pytest.mark.asyncio + async def test_should_use_model_group_for_virtual_key_spend_key( + self, budget_limiter + ): + """When model_group is set, the spend_key passed to + _increment_spend_for_key should use model_group, not model.""" + kwargs = _make_kwargs( + model=_DEPLOYMENT_A, + model_group=_MODEL_GROUP, + key_budget=_BUDGET, + ) + + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + spend_key = mock_increment.call_args.kwargs["spend_key"] + assert f":{_MODEL_GROUP}:" in spend_key + assert _DEPLOYMENT_A not in spend_key + + @pytest.mark.asyncio + async def test_should_use_model_group_with_multi_slash_deployment( + self, budget_limiter + ): + """Deployment names with multiple slashes (e.g. + openrouter/deepseek/deepseek-chat) should still use model_group.""" + kwargs = _make_kwargs( + model=_DEPLOYMENT_B, + model_group=_MODEL_GROUP, + key_budget=_BUDGET, + ) + + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + spend_key = mock_increment.call_args.kwargs["spend_key"] + assert f":{_MODEL_GROUP}:" in spend_key + assert _DEPLOYMENT_B not in spend_key + + @pytest.mark.asyncio + async def test_should_use_model_group_for_end_user_spend_key( + self, budget_limiter + ): + """End-user spend key should also use model_group.""" + kwargs = _make_kwargs( + model=_DEPLOYMENT_A, + model_group=_MODEL_GROUP, + end_user_id="user-1", + end_user_budget=_BUDGET, + ) + + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + spend_key = mock_increment.call_args.kwargs["spend_key"] + assert f":{_MODEL_GROUP}:" in spend_key + assert _DEPLOYMENT_A not in spend_key + + @pytest.mark.asyncio + async def test_should_fall_back_to_model_when_model_group_is_none( + self, budget_limiter + ): + """When model_group is not in the payload, fall back to model.""" + kwargs = _make_kwargs( + model=_MODEL_GROUP, + model_group=None, + key_budget=_BUDGET, + ) + + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + spend_key = mock_increment.call_args.kwargs["spend_key"] + assert f":{_MODEL_GROUP}:" in spend_key + + +class TestWriteReadCacheKeyAlignment: + """Verify that what the write path writes can be found by the read path.""" + + @pytest.mark.asyncio + async def test_should_produce_cache_key_readable_by_virtual_key_lookup( + self, budget_limiter + ): + """The cache key from the write path should exactly match + what _get_virtual_key_spend_for_model reads.""" + virtual_key = "sk-hash-123" + + expected_cache_key = ( + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}" + f":{_MODEL_GROUP}:1d" + ) + + kwargs = _make_kwargs( + model=_DEPLOYMENT_A, + model_group=_MODEL_GROUP, + virtual_key=virtual_key, + key_budget=_BUDGET, + ) + + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + actual_spend_key = mock_increment.call_args.kwargs["spend_key"] + assert actual_spend_key == expected_cache_key + + @pytest.mark.asyncio + async def test_should_produce_cache_key_readable_by_end_user_lookup( + self, budget_limiter + ): + """The end-user cache key from the write path should exactly match + what _get_end_user_spend_for_model reads.""" + end_user_id = "user-abc" + budget_duration = "7d" + budget = {_MODEL_GROUP: {"budget_limit": 50.0, "time_period": budget_duration}} + + expected_cache_key = ( + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}" + f":{_MODEL_GROUP}:{budget_duration}" + ) + + kwargs = _make_kwargs( + model=_DEPLOYMENT_B, + model_group=_MODEL_GROUP, + end_user_id=end_user_id, + end_user_budget=budget, + ) + + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + actual_spend_key = mock_increment.call_args.kwargs["spend_key"] + assert actual_spend_key == expected_cache_key + + @pytest.mark.asyncio + async def test_should_aggregate_spend_across_deployments( + self, budget_limiter + ): + """Two requests to different deployments (deepseek/deepseek-chat + and openrouter/deepseek/deepseek-chat) in the same model group + should write to the same cache key, aggregating spend.""" + virtual_key = "sk-hash-agg" + + # Request 1: deepseek/deepseek-chat, costs $2.00 + kwargs1 = _make_kwargs( + model=_DEPLOYMENT_A, + model_group=_MODEL_GROUP, + response_cost=2.0, + virtual_key=virtual_key, + key_budget=_BUDGET, + ) + await budget_limiter.async_log_success_event( + kwargs1, response_obj=None, start_time=None, end_time=None + ) + + # Request 2: openrouter/deepseek/deepseek-chat, costs $0.70 + kwargs2 = _make_kwargs( + model=_DEPLOYMENT_B, + model_group=_MODEL_GROUP, + response_cost=0.7, + virtual_key=virtual_key, + key_budget=_BUDGET, + ) + await budget_limiter.async_log_success_event( + kwargs2, response_obj=None, start_time=None, end_time=None + ) + + # Read back: should see aggregated spend of $2.70 + from litellm.types.utils import BudgetConfig + + spend = await budget_limiter._get_virtual_key_spend_for_model( + user_api_key_hash=virtual_key, + model=_MODEL_GROUP, + key_budget_config=BudgetConfig( + budget_limit=100.0, time_period="1d" + ), + ) + assert spend == pytest.approx(2.7) + + @pytest.mark.asyncio + async def test_end_to_end_write_then_read_with_model_group( + self, budget_limiter + ): + """Write spend via async_log_success_event with model_group, then + read it back via is_key_within_model_budget. The values should + be visible because the cache keys now align.""" + virtual_key = "sk-hash-e2e" + + user_api_key = UserAPIKeyAuth( + token=virtual_key, + key_alias="test-alias", + model_max_budget=_BUDGET, + ) + + kwargs = _make_kwargs( + model=_DEPLOYMENT_A, + model_group=_MODEL_GROUP, + response_cost=42.0, + virtual_key=virtual_key, + key_budget=_BUDGET, + ) + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + + # Read: check the budget — should see the 42.0 spend + result = await budget_limiter.is_key_within_model_budget( + user_api_key, _MODEL_GROUP + ) + assert result is True # 42.0 < 100.0 + + # Verify the spend is actually there in cache + from litellm.types.utils import BudgetConfig + + spend = await budget_limiter._get_virtual_key_spend_for_model( + user_api_key_hash=virtual_key, + model=_MODEL_GROUP, + key_budget_config=BudgetConfig( + budget_limit=100.0, time_period="1d" + ), + ) + assert spend == 42.0