From 1bb8263c922bdb133c13db2ece74b2d783900fbe Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 2 Mar 2024 14:43:01 -0800 Subject: [PATCH] (feat) set soft_budget with /key/generate --- litellm/proxy/_types.py | 1 + litellm/proxy/proxy_server.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e4b28001d5..6196f18a2e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -151,6 +151,7 @@ class GenerateRequestBase(LiteLLMBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None allowed_cache_controls: Optional[list] = [] + soft_budget: Optional[float] = None class GenerateKeyRequest(GenerateRequestBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dcd4283ba0..482397b86e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1995,6 +1995,9 @@ async def generate_key_helper_fn( except Exception as e: traceback.print_exc() raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) + + # Add budget related info in key_data - this ensures it's returned + key_data["soft_budget"] = key_soft_budget return key_data