From 679c47d196110268068ccc014001591ab8a42de9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 2 Mar 2024 19:11:37 -0800 Subject: [PATCH] fix(proxy_server.py): fix budget creation --- litellm/proxy/proxy_server.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ddd702c22..ab8e7b2d6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1881,16 +1881,20 @@ async def generate_key_helper_fn( allowed_cache_controls = allowed_cache_controls # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable - if prisma_client is not None: + if prisma_client is not None and key_soft_budget is not None: # create the Budget Row for the LiteLLM Verification Token budget_row = LiteLLM_BudgetTable( - soft_budget=key_soft_budget or litellm.default_soft_budget, + soft_budget=key_soft_budget, model_max_budget=model_max_budget or {}, - created_by=user_id, - updated_by=user_id, ) new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - _budget = await prisma_client.db.litellm_budgettable.create(data={**new_budget}) # type: ignore + _budget = await prisma_client.db.litellm_budgettable.create( + data={ + **new_budget, # type: ignore + "created_by": user_id, + "updated_by": user_id, + } + ) _budget_id = getattr(_budget, "id", None) try: