(feat) set soft_budget with /key/generate

This commit is contained in:
ishaan-jaff 2024-03-02 14:43:01 -08:00
parent fd9f8b7010
commit 1bb8263c92
2 changed files with 4 additions and 0 deletions

View file

@ -151,6 +151,7 @@ class GenerateRequestBase(LiteLLMBase):
rpm_limit: Optional[int] = None rpm_limit: Optional[int] = None
budget_duration: Optional[str] = None budget_duration: Optional[str] = None
allowed_cache_controls: Optional[list] = [] allowed_cache_controls: Optional[list] = []
soft_budget: Optional[float] = None
class GenerateKeyRequest(GenerateRequestBase): class GenerateKeyRequest(GenerateRequestBase):

View file

@ -1995,6 +1995,9 @@ async def generate_key_helper_fn(
except Exception as e: except Exception as e:
traceback.print_exc() traceback.print_exc()
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) 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 return key_data