scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

refactor: build budget window usage without in-place mutation

Browse source 
Replace _attach_budget_limits_usage, which rewrote the caller's key_info dict, with _budget_limits_with_usage returning a new list. Callers assign the result once. Keeps the response shape and spend-counter read path identical while following the repo's no-mutation rule.
This commit is contained in:
Thijmen Stavenuiter 2026-08-17 16:21:31 +02:00
parent db67bafd59
commit 31130036c0
2 changed files with 121 additions and 81 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

119
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -3517,23 +3517,44 @@ async def _build_model_max_budget_usage(
return result
def _budget_window_to_dict(window: object) -> dict | None:
def _budget_window_to_dict(window: object) -> Mapping[str, object] | None:
"""Coerce a budget_limits entry to a dict; None when the entry is unusable."""
if isinstance(window, dict):
return window
model_dump = getattr(window, "model_dump", None)
if callable(model_dump):
model_dump: Final = getattr(window, "model_dump", None)
if not callable(model_dump):
return None
try:
dumped: Final = model_dump()
except Exception: # noqa: BLE001 # model_dump implementations can raise arbitrary errors
return None
return dumped if isinstance(dumped, dict) else None
def _coerce_budget_limits(budget_limits: object) -> Sequence[object] | None:
"""Coerce budget_limits to a sequence of windows, parsing JSON strings; None when unusable."""
if isinstance(budget_limits, str):
try:
dumped: Any = model_dump()
except Exception: # noqa: BLE001
parsed: Final = json.loads(budget_limits)
except (TypeError, ValueError):
return None
return parsed if isinstance(parsed, list) else None
return budget_limits if isinstance(budget_limits, list) else None
def _parse_window_max_budget(value: object) -> float | None:
"""Coerce a window's max_budget to float; None when absent or unparseable."""
if isinstance(value, (int, float, str)):
try:
return float(value)
except (TypeError, ValueError):
return None
return dumped if isinstance(dumped, dict) else None
return None
async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None:
async def _budget_window_with_usage(window: Mapping[str, object], api_key_hash: str) -> Mapping[str, object]:
"""
Attach current-window spend to each entry in key_info["budget_limits"], in place.
Return a copy of a budget window with current-window spend attached.
Per-window spend is not persisted in the DB; it lives in the cross-pod spend
counters (spend:key:{hashed_token}:window:{budget_duration}) that
@ -3544,39 +3565,43 @@ async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None
"""
from litellm.proxy.proxy_server import get_current_spend
budget_limits: Any = key_info.get("budget_limits")
if isinstance(budget_limits, str):
try:
budget_limits = json.loads(budget_limits)
except (TypeError, ValueError):
return
key_info["budget_limits"] = budget_limits
if not isinstance(budget_limits, list):
return
duration: Final = window.get("budget_duration")
if not duration:
return dict(window) # mutable-ok: per-window response copy, built once per window
spend: Final = await get_current_spend(
counter_key=f"spend:key:{api_key_hash}:window:{duration}",
fallback_spend=0.0,
max_budget=_parse_window_max_budget(window.get("max_budget")),
window_entity_type="Key",
window_entity_id=api_key_hash,
window_start=get_budget_window_start(window),
)
return {**window, "current_spend": round(spend, 4)} # mutable-ok: per-window response copy, built once per window
for idx, window in enumerate(budget_limits):
w = _budget_window_to_dict(window)
if not w:
continue
budget_limits[idx] = w
duration: Any = w.get("budget_duration")
max_budget: Any = w.get("max_budget")
if not duration:
continue
try:
max_budget = float(max_budget) if max_budget is not None else None
except (TypeError, ValueError):
max_budget = None
counter_key = f"spend:key:{api_key_hash}:window:{duration}"
spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
max_budget=max_budget,
window_entity_type="Key",
window_entity_id=api_key_hash,
window_start=get_budget_window_start(w),
)
w["current_spend"] = round(spend, 4)
async def _budget_limits_entry_with_usage(window: object, api_key_hash: str) -> object:
"""Return the window as an enriched dict when dict-coercible; the original entry otherwise."""
coerced: Final = _budget_window_to_dict(window)
if not coerced:
return window
return await _budget_window_with_usage(window=coerced, api_key_hash=api_key_hash)
async def _budget_limits_with_usage(budget_limits: object, api_key_hash: str) -> Sequence[object] | None:
"""
Return budget_limits as window dicts with current-window spend attached.
None when budget_limits is not a usable (possibly JSON-encoded) list; the
caller keeps the original value then. Entries that are not dict-coercible
are preserved as-is.
"""
windows: Final = _coerce_budget_limits(budget_limits)
if windows is None:
return None
return [ # mutable-ok: entries are awaited, so they cannot be built inside a frozen wrapper
await _budget_limits_entry_with_usage(window=window, api_key_hash=api_key_hash)
for window in windows
]
@router.post(
@ -3661,7 +3686,12 @@ async def info_key_fn_v2(
user_api_key_cache=user_api_key_cache,
)
if k_token_hash:
await _attach_budget_limits_usage(key_info=k_dict, api_key_hash=k_token_hash)
budget_limits_usage = await _budget_limits_with_usage(
budget_limits=k_dict.get("budget_limits"),
api_key_hash=k_token_hash,
)
if budget_limits_usage is not None:
k_dict["budget_limits"] = budget_limits_usage
filtered_key_info.append(k_dict)
return {"key": data.keys, "info": filtered_key_info}
@ -3777,7 +3807,12 @@ async def info_key_fn(
model_max_budget=model_max_budget,
user_api_key_cache=user_api_key_cache,
)
await _attach_budget_limits_usage(key_info=key_info, api_key_hash=key_token_hash)
budget_limits_usage: Final = await _budget_limits_with_usage(
budget_limits=key_info.get("budget_limits"),
api_key_hash=key_token_hash,
)
if budget_limits_usage is not None:
key_info["budget_limits"] = budget_limits_usage
# Attach object_permission if object_permission_id is set
key_info = await attach_object_permission_to_dict(key_info, prisma_client)

83
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -13996,13 +13996,13 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch):
@pytest.mark.asyncio
async def test_attach_budget_limits_usage_json_string_input(monkeypatch):
async def test_budget_limits_with_usage_json_string_input(monkeypatch):
"""budget_limits stored as a JSON string should be parsed and annotated."""
import json as json_module
from unittest.mock import AsyncMock
from litellm.proxy.management_endpoints.key_management_endpoints import (
_attach_budget_limits_usage,
_budget_limits_with_usage,
)
mock_get_current_spend = AsyncMock(return_value=0.5)
@ -14010,24 +14010,29 @@ async def test_attach_budget_limits_usage_json_string_input(monkeypatch):
"litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
)
key_info = {
"budget_limits": json_module.dumps(
[{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}]
)
}
await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1")
raw = json_module.dumps(
[{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}]
)
result = await _budget_limits_with_usage(budget_limits=raw, api_key_hash="hash-1")
assert isinstance(key_info["budget_limits"], list)
assert key_info["budget_limits"][0]["current_spend"] == 0.5
assert result == [
{
"budget_duration": "1h",
"max_budget": 2.0,
"reset_at": None,
"current_spend": 0.5,
}
]
mock_get_current_spend.assert_awaited_once()
@pytest.mark.asyncio
async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch):
async def test_budget_limits_with_usage_skips_unusable_inputs(monkeypatch):
"""Invalid JSON strings, non-list values, and malformed windows are skipped."""
from unittest.mock import AsyncMock
from litellm.proxy.management_endpoints.key_management_endpoints import (
_attach_budget_limits_usage,
_budget_limits_with_usage,
)
mock_get_current_spend = AsyncMock(return_value=0.0)
@ -14035,25 +14040,18 @@ async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch):
"litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
)
# invalid JSON string
key_info = {"budget_limits": "{not json"}
await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1")
assert key_info["budget_limits"] == "{not json"
# non-list value
key_info = {"budget_limits": {"budget_duration": "1h"}}
await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1")
# invalid JSON string and non-list values return None: callers keep the original
assert await _budget_limits_with_usage(budget_limits="{not json", api_key_hash="hash-1") is None
assert await _budget_limits_with_usage(budget_limits={"budget_duration": "1h"}, api_key_hash="hash-1") is None
# windows that are falsy, missing budget_duration, or not dict-like
key_info = {
"budget_limits": [
{},
{"max_budget": 2.0},
{"budget_duration": "1h", "max_budget": "not-a-number"},
42,
]
}
await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1")
windows = [
{},
{"max_budget": 2.0},
{"budget_duration": "1h", "max_budget": "not-a-number"},
42,
]
result = await _budget_limits_with_usage(budget_limits=windows, api_key_hash="hash-1")
# only the well-formed window (with unparseable max_budget coerced to None)
# triggers a spend lookup
@ -14061,17 +14059,22 @@ async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch):
call_kwargs = mock_get_current_spend.await_args.kwargs
assert call_kwargs["counter_key"] == "spend:key:hash-1:window:1h"
assert call_kwargs["max_budget"] is None
assert key_info["budget_limits"][2]["current_spend"] == 0.0
assert key_info["budget_limits"][3] == 42
assert result is not None
assert result[0] == {}
assert result[1] == {"max_budget": 2.0}
assert result[2] == {"budget_duration": "1h", "max_budget": "not-a-number", "current_spend": 0.0}
assert result[3] == 42
# input is not mutated
assert windows[2] == {"budget_duration": "1h", "max_budget": "not-a-number"}
@pytest.mark.asyncio
async def test_attach_budget_limits_usage_pydantic_windows(monkeypatch):
"""Window objects with model_dump() are converted to dicts in place."""
async def test_budget_limits_with_usage_pydantic_windows(monkeypatch):
"""Window objects with model_dump() are converted to dicts; failing windows pass through."""
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.management_endpoints.key_management_endpoints import (
_attach_budget_limits_usage,
_budget_limits_with_usage,
)
mock_get_current_spend = AsyncMock(return_value=1.0)
@ -14088,13 +14091,15 @@ async def test_attach_budget_limits_usage_pydantic_windows(monkeypatch):
bad_window = MagicMock()
bad_window.model_dump.side_effect = ValueError("boom")
key_info = {"budget_limits": [good_window, bad_window]}
await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-2")
result = await _budget_limits_with_usage(
budget_limits=[good_window, bad_window], api_key_hash="hash-2"
)
# good window converted to dict and annotated; failing window left as-is
assert isinstance(key_info["budget_limits"][0], dict)
assert key_info["budget_limits"][0]["current_spend"] == 1.0
assert key_info["budget_limits"][1] is bad_window
assert result is not None
assert isinstance(result[0], dict)
assert result[0]["current_spend"] == 1.0
assert result[1] is bad_window
mock_get_current_spend.assert_awaited_once()