Support budget/rate limit tiers for keys (#7429)

* feat(proxy/utils.py): get associated litellm budget from db in combined_view for key allows user to create rate limit tiers and associate those to keys * feat(proxy/_types.py): update the value of key-level tpm/rpm/model max budget metrics with the associated budget table values if set allows rate limit tiers to be easily applied to keys * docs(rate_limit_tiers.md): add doc on setting rate limit / budget tiers make feature discoverable * feat(key_management_endpoints.py): return litellm_budget_table value in key generate make it easy for user to know associated budget on key creation * fix(key_management_endpoints.py): document 'budget_id' param in `/key/generate` * docs(key_management_endpoints.py): document budget_id usage * refactor(budget_management_endpoints.py): refactor budget endpoints into separate file - makes it easier to run documentation testing against it * docs(test_api_docs.py): add budget endpoints to ci/cd doc test + add missing param info to docs * fix(customer_endpoints.py): use new pydantic obj name * docs(user_management_heirarchy.md): add simple doc explaining teams/keys/org/users on litellm * Litellm dev 12 26 2024 p2 (#7432) * (Feat) Add logging for `POST v1/fine_tuning/jobs` (#7426) * init commit ft jobs logging * add ft logging * add logging for FineTuningJob * simple FT Job create test * (docs) - show all supported Azure OpenAI endpoints in overview (#7428) * azure batches * update doc * docs azure endpoints * docs endpoints on azure * docs azure batches api * docs azure batches api * fix(key_management_endpoints.py): fix key update to actually work * test(test_key_management.py): add e2e test asserting ui key update call works * fix: proxy/_types - fix linting erros * test: update test --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> * fix: test * fix(parallel_request_limiter.py): enforce tpm/rpm limits on key from tiers * fix: fix linting errors * test: fix test * fix: remove unused import * test: update test * docs(customer_endpoints.py): document new model_max_budget param * test: specify unique key alias * docs(budget_management_endpoints.py): document new model_max_budget param * test: fix test * test: fix tests --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
2026-07-13 23:07:47 +00:00 · 2024-12-26 19:05:27 -08:00
parent 12c4e7e695
commit 539f166166
25 changed files with 764 additions and 376 deletions
@@ -1,6 +1,7 @@
 import asyncio
 import os
 import sys
+from typing import Any, Dict
 from unittest.mock import Mock
 from litellm.proxy.utils import _get_redoc_url, _get_docs_url
 import json
@@ -1104,3 +1105,89 @@ def test_proxy_config_state_post_init_callback_call():

    config = pc.get_config_state()
    assert config["litellm_settings"]["default_team_settings"][0]["team_id"] == "test"
+
+
+@pytest.mark.parametrize(
+    "associated_budget_table, expected_user_api_key_auth_key, expected_user_api_key_auth_value",
+    [
+        (
+            {
+                "litellm_budget_table_max_budget": None,
+                "litellm_budget_table_tpm_limit": None,
+                "litellm_budget_table_rpm_limit": 1,
+                "litellm_budget_table_model_max_budget": None,
+            },
+            "rpm_limit",
+            1,
+        ),
+        (
+            {},
+            None,
+            None,
+        ),
+        (
+            {
+                "litellm_budget_table_max_budget": None,
+                "litellm_budget_table_tpm_limit": None,
+                "litellm_budget_table_rpm_limit": None,
+                "litellm_budget_table_model_max_budget": {"gpt-4o": 100},
+            },
+            "model_max_budget",
+            {"gpt-4o": 100},
+        ),
+    ],
+)
+def test_litellm_verification_token_view_response_with_budget_table(
+    associated_budget_table,
+    expected_user_api_key_auth_key,
+    expected_user_api_key_auth_value,
+):
+    from litellm.proxy._types import LiteLLM_VerificationTokenView
+
+    args: Dict[str, Any] = {
+        "token": "78b627d4d14bc3acf5571ae9cb6834e661bc8794d1209318677387add7621ce1",
+        "key_name": "sk-...if_g",
+        "key_alias": None,
+        "soft_budget_cooldown": False,
+        "spend": 0.011441999999999997,
+        "expires": None,
+        "models": [],
+        "aliases": {},
+        "config": {},
+        "user_id": None,
+        "team_id": "test",
+        "permissions": {},
+        "max_parallel_requests": None,
+        "metadata": {},
+        "blocked": None,
+        "tpm_limit": None,
+        "rpm_limit": None,
+        "max_budget": None,
+        "budget_duration": None,
+        "budget_reset_at": None,
+        "allowed_cache_controls": [],
+        "model_spend": {},
+        "model_max_budget": {},
+        "budget_id": "my-test-tier",
+        "created_at": "2024-12-26T02:28:52.615+00:00",
+        "updated_at": "2024-12-26T03:01:51.159+00:00",
+        "team_spend": 0.012134999999999998,
+        "team_max_budget": None,
+        "team_tpm_limit": None,
+        "team_rpm_limit": None,
+        "team_models": [],
+        "team_metadata": {},
+        "team_blocked": False,
+        "team_alias": None,
+        "team_members_with_roles": [{"role": "admin", "user_id": "default_user_id"}],
+        "team_member_spend": None,
+        "team_model_aliases": None,
+        "team_member": None,
+        **associated_budget_table,
+    }
+    resp = LiteLLM_VerificationTokenView(**args)
+    if expected_user_api_key_auth_key is not None:
+        assert (
+            getattr(resp, expected_user_api_key_auth_key)
+            == expected_user_api_key_auth_value
+        )