diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 75396233ef..577231c26f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -305,7 +305,7 @@ class GenerateKeyResponse(GenerateKeyRequest): key: str key_name: Optional[str] = None expires: Optional[datetime] - user_id: str + user_id: Optional[str] = None @root_validator(pre=True) def set_model_info(cls, values): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4e91afb843..2ca33554a3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1028,7 +1028,9 @@ async def user_api_key_auth( # Do something if the current route starts with any of the allowed routes pass else: - if _is_user_proxy_admin(user_id_information): + if user_id_information is not None and _is_user_proxy_admin( + user_id_information + ): return UserAPIKeyAuth( api_key=api_key, user_role="proxy_admin", **valid_token_dict ) @@ -2292,9 +2294,7 @@ async def generate_key_helper_fn( spend: float, key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key key_budget_duration: Optional[str] = None, - key_soft_budget: Optional[ - float - ] = None, # key_soft_budget is used to Budget Per key + budget_id: Optional[float] = None, # budget id <-> LiteLLM_BudgetTable soft_budget: Optional[ float ] = None, # soft_budget is used to set soft Budgets Per user @@ -2317,7 +2317,7 @@ async def generate_key_helper_fn( model_max_budget: Optional[dict] = {}, table_name: Optional[Literal["key", "user"]] = None, ): - global prisma_client, custom_db_client, user_api_key_cache + global prisma_client, custom_db_client, user_api_key_cache, litellm_proxy_admin_name if prisma_client is None and custom_db_client is None: raise Exception( @@ -2350,32 +2350,11 @@ async def generate_key_helper_fn( permissions_json = json.dumps(permissions) metadata_json = json.dumps(metadata) model_max_budget_json = json.dumps(model_max_budget) - - user_id = user_id or str(uuid.uuid4()) user_role = user_role or "app_user" tpm_limit = tpm_limit rpm_limit = rpm_limit allowed_cache_controls = allowed_cache_controls - # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable - _budget_id = 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, - model_max_budget=model_max_budget or {}, - ) - 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 - "created_by": user_id, - "updated_by": user_id, - } - ) - _budget_id = getattr(_budget, "budget_id", None) - try: # Create a new verification token (you may want to enhance this logic based on your needs) user_data = { @@ -2413,7 +2392,7 @@ async def generate_key_helper_fn( "allowed_cache_controls": allowed_cache_controls, "permissions": permissions_json, "model_max_budget": model_max_budget_json, - "budget_id": _budget_id, + "budget_id": budget_id, } if ( general_settings.get("allow_user_auth", False) == True @@ -2439,20 +2418,23 @@ async def generate_key_helper_fn( ): saved_token["expires"] = saved_token["expires"].isoformat() if prisma_client is not None: - ## CREATE USER (If necessary) - if query_type == "insert_data": - user_row = await prisma_client.insert_data( - data=user_data, table_name="user" - ) - ## use default user model list if no key-specific model list provided - if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore - key_data["models"] = user_row.models - elif query_type == "update_data": - user_row = await prisma_client.update_data( - data=user_data, - table_name="user", - update_key_values=update_key_values, - ) + if ( + table_name is None or table_name == "user" + ): # do not auto-create users for `/key/generate` + ## CREATE USER (If necessary) + if query_type == "insert_data": + user_row = await prisma_client.insert_data( + data=user_data, table_name="user" + ) + ## use default user model list if no key-specific model list provided + if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore + key_data["models"] = user_row.models + elif query_type == "update_data": + user_row = await prisma_client.update_data( + data=user_data, + table_name="user", + update_key_values=update_key_values, + ) if user_id == litellm_proxy_budget_name or ( table_name is not None and table_name == "user" ): @@ -2465,16 +2447,19 @@ async def generate_key_helper_fn( verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data) await prisma_client.insert_data(data=key_data, table_name="key") elif custom_db_client is not None: - ## CREATE USER (If necessary) - verbose_proxy_logger.debug("CustomDBClient: Creating User= %s", user_data) - user_row = await custom_db_client.insert_data( - value=user_data, table_name="user" - ) - if user_row is None: - # GET USER ROW - user_row = await custom_db_client.get_data( - key=user_id, table_name="user" + if table_name is None or table_name == "user": + ## CREATE USER (If necessary) + verbose_proxy_logger.debug( + "CustomDBClient: Creating User= %s", user_data ) + user_row = await custom_db_client.insert_data( + value=user_data, table_name="user" + ) + if user_row is None: + # GET USER ROW + user_row = await custom_db_client.get_data( + key=user_id, table_name="user" # type: ignore + ) ## use default user model list if no key-specific model list provided if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore @@ -2487,7 +2472,7 @@ async def generate_key_helper_fn( 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 + key_data["budget_id"] = budget_id return key_data @@ -3934,6 +3919,7 @@ async def moderations( ) async def generate_key_fn( data: GenerateKeyRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), Authorization: Optional[str] = Header(None), ): """ @@ -4037,18 +4023,42 @@ async def generate_key_fn( data, key, litellm.upperbound_key_generate_params[key] ) + # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable + _budget_id = None + if prisma_client is not None and data.soft_budget is not None: + # create the Budget Row for the LiteLLM Verification Token + budget_row = LiteLLM_BudgetTable( + soft_budget=data.soft_budget, + model_max_budget=data.model_max_budget or {}, + ) + 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 + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + ) + _budget_id = getattr(_budget, "budget_id", None) data_json = data.json() # type: ignore # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users if "max_budget" in data_json: data_json["key_max_budget"] = data_json.pop("max_budget", None) - if "soft_budget" in data_json: - data_json["key_soft_budget"] = data_json.pop("soft_budget", None) + if _budget_id is not None: + data_json["budget_id"] = _budget_id if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) - response = await generate_key_helper_fn(**data_json) + response = await generate_key_helper_fn(**data_json, table_name="key") + + response["soft_budget"] = ( + data.soft_budget + ) # include the user-input soft budget in the response return GenerateKeyResponse(**response) except Exception as e: traceback.print_exc() @@ -5110,6 +5120,8 @@ async def new_user(data: NewUserRequest): param="user_role", code=status.HTTP_400_BAD_REQUEST, ) + if "user_id" in data_json and data_json["user_id"] is None: + data_json["user_id"] = str(uuid.uuid4()) response = await generate_key_helper_fn(**data_json) return NewUserResponse( key=response["token"],