diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 701af27e9..d962c9de3 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -278,6 +278,18 @@ Request Params: } ``` +## Default /key/generate params +Use this, if you need to control the default `max_budget` or any `key/generate` param per key + +Set `litellm_settings:default_key_generate_params`: +litellm_settings: + default_key_generate_params: + max_budget: 1.5000 + models: ["azure-gpt-3.5"] + duration: # blank means `null` + metadata: {"setting":"default"} + team_id: "core-infra" + ## Set Budgets - Per Key Set `max_budget` in (USD $) param in the `key/generate` request. By default the `max_budget` is set to `null` and is not checked for keys diff --git a/litellm/__init__.py b/litellm/__init__.py index 2e23191a1..226078ed4 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -143,6 +143,7 @@ model_cost_map_url: str = "https://raw.githubusercontent.com/BerriAI/litellm/mai suppress_debug_info = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None +default_key_generate_params: Optional[Dict] = None #### RELIABILITY #### request_timeout: Optional[float] = 6000 num_retries: Optional[int] = None # per model endpoint diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index aa950c035..31ff4ef57 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -69,11 +69,16 @@ litellm_settings: success_callback: ['langfuse'] max_budget: 10 # global budget for proxy budget_duration: 30d # global budget duration, will reset after 30d + default_key_generate_params: + max_budget: 1.5000 + models: ["azure-gpt-3.5"] + duration: None # cache: True # setting callback class # callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance] general_settings: + allow_user_auth: True master_key: sk-1234 alerting: ["slack"] alerting_threshold: 10 # sends alerts if requests hang for 2 seconds diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7b82d4d1c..d9944e6e1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1117,6 +1117,9 @@ class ProxyConfig: # see usage here: https://docs.litellm.ai/docs/proxy/caching pass else: + verbose_proxy_logger.debug( + f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}" + ) setattr(litellm, key, value) ## GENERAL SERVER SETTINGS (e.g. master key,..) # do this after initializing litellm, to ensure sentry logging works for proxylogging @@ -2385,6 +2388,26 @@ async def generate_key_fn( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=message ) + # check if user set default key/generate params on config.yaml + if litellm.default_key_generate_params is not None: + for elem in data: + key, value = elem + if value is None and key in [ + "max_budget", + "user_id", + "team_id", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "budget_duration", + ]: + setattr( + data, key, litellm.default_key_generate_params.get(key, None) + ) + elif key == "models" and value == []: + setattr(data, key, litellm.default_key_generate_params.get(key, [])) + elif key == "metadata" and value == {}: + setattr(data, key, litellm.default_key_generate_params.get(key, {})) data_json = data.json() # type: ignore @@ -2959,6 +2982,11 @@ async def google_callback(code: str, request: Request): user_id = response["user_id"] # type: ignore litellm_dashboard_ui = "https://litellm-dashboard.vercel.app/" + # if user set LITELLM_UI_LINK in .env, use that + litellm_ui_link_in_env = os.getenv("LITELLM_UI_LINK", None) + if litellm_ui_link_in_env is not None: + litellm_dashboard_ui = litellm_ui_link_in_env + litellm_dashboard_ui += ( "?userID=" + user_id diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 9f6bfd1be..728514342 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -1188,3 +1188,27 @@ async def test_key_name_set(prisma_client): except Exception as e: print("Got Exception", e) pytest.fail(f"Got exception {e}") + + +@pytest.mark.asyncio() +async def test_default_key_params(prisma_client): + """ + - create key + - get key info + - assert key_name is not null + """ + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + setattr(litellm.proxy.proxy_server, "general_settings", {"allow_user_auth": True}) + litellm.default_key_generate_params = {"max_budget": 0.000122} + await litellm.proxy.proxy_server.prisma_client.connect() + try: + request = GenerateKeyRequest() + key = await generate_key_fn(request) + generated_key = key.key + result = await info_key_fn(key=generated_key) + print("result from info_key_fn", result) + assert result["info"]["max_budget"] == 0.000122 + except Exception as e: + print("Got Exception", e) + pytest.fail(f"Got exception {e}") diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index edd601570..b0a0ede62 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -15,8 +15,6 @@ export const keyCreateCall = async ( "Content-Type": "application/json", }, body: JSON.stringify({ - team_id: "core-infra-4", - max_budget: 10, user_id: userID, }), }); diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx index 5471ad05e..e500ffe72 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/view_key_table.tsx @@ -58,7 +58,8 @@ const ViewKeyTable: React.FC = ({ Secret Key - Spend + Spend (USD) + Key Budget (USD) Expires @@ -68,11 +69,24 @@ const ViewKeyTable: React.FC = ({ return ( - {item.key_name} + {item.key_name != null ? ( + {item.key_name} + ) : ( + {item.token} + ) + } {item.spend} + + {item.max_budget != null ? ( + {item.max_budget} + ) : ( + Unlimited Budget + ) + } + {item.expires != null ? ( {item.expires}