scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

change logic to match Kriish's input

Browse source
This commit is contained in:
shivam 2026-02-14 10:53:58 -08:00
parent ba1b466480
commit cf14f0c821
2 changed files with 18 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -592,15 +592,12 @@ async def _common_key_generation_helper( # noqa: PLR0915
if _budget_id is not None:
data_json["budget_id"] = _budget_id
# Only set budget_duration on key when explicitly provided. Keys with budget_id
# but no explicit budget_duration follow their linked budget tier's schedule;
# reset_budget_for_keys_linked_to_budgets() resets them when the tier resets.
# This avoids duplicating budget_duration on keys so tier updates apply automatically.
if "budget_duration" in data_json:
data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
elif _budget_id is not None and prisma_client is not None:
# Inherit budget_duration from linked budget tier if not explicitly set on the key
budget_row = await prisma_client.db.litellm_budgettable.find_unique(
where={"budget_id": _budget_id}
)
if budget_row is not None and budget_row.budget_duration is not None:
data_json["key_budget_duration"] = budget_row.budget_duration
if user_api_key_dict.user_id is not None:
data_json["created_by"] = user_api_key_dict.user_id

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

@ -5562,26 +5562,19 @@ async def test_validate_key_list_check_key_hash_not_found():
@pytest.mark.asyncio
async def test_key_inherits_budget_duration_from_budget_tier():
async def test_key_with_budget_id_does_not_store_budget_duration():
"""
Test that when a key is created with budget_id pointing to a budget tier
that has budget_duration, the key inherits budget_duration from the tier
even when budget_duration is not explicitly set on the key request.
Test that when a key is created with budget_id but without explicit
budget_duration, the key does NOT get budget_duration stored on it.
This verifies the fix for the bug where keys created with budget_id
would have null budget_duration and budget_reset_at, causing the
budget reset job to never reset their spend.
Keys with budget_id follow their linked budget tier's reset schedule;
reset_budget_for_keys_linked_to_budgets() resets them when the tier resets.
This avoids duplicating budget_duration on keys so tier updates apply
automatically to all linked keys.
"""
from unittest.mock import AsyncMock, MagicMock, patch
# Mock the budget tier lookup to return a budget with budget_duration="7d"
mock_budget_row = MagicMock()
mock_budget_row.budget_duration = "7d"
mock_prisma = MagicMock()
mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(
return_value=mock_budget_row
)
mock_generate_key = AsyncMock(
return_value={
@ -5619,20 +5612,17 @@ async def test_key_inherits_budget_duration_from_budget_tier():
team_table=None,
)
# Verify generate_key_helper_fn was called
mock_generate_key.assert_awaited_once()
call_kwargs = mock_generate_key.call_args.kwargs
# The key should have inherited key_budget_duration from the budget tier
assert call_kwargs.get("key_budget_duration") == "7d", (
"key_budget_duration should be inherited from the linked budget tier "
f"but got: {call_kwargs.get('key_budget_duration')}"
# Key should NOT have key_budget_duration - it follows the budget tier's schedule
assert call_kwargs.get("key_budget_duration") is None, (
"key_budget_duration should be None for budget-linked keys without explicit "
f"budget_duration; got: {call_kwargs.get('key_budget_duration')}"
)
# Verify the budget tier was looked up with the correct budget_id
mock_prisma.db.litellm_budgettable.find_unique.assert_awaited_once_with(
where={"budget_id": "7d-budget-tier"}
)
# No budget tier lookup - we don't copy budget_duration onto the key
mock_prisma.db.litellm_budgettable.find_unique.assert_not_called()
@pytest.mark.asyncio
@ -5698,4 +5688,4 @@ async def test_key_does_not_override_explicit_budget_duration():
)
# The budget tier should NOT have been looked up since budget_duration was explicit
mock_prisma.db.litellm_budgettable.find_unique.assert_not_awaited()
mock_prisma.db.litellm_budgettable.find_unique.assert_not_called()