From a408c46a67fa1f5448b33e8ea815c9ac97ba1d85 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 27 Mar 2024 16:48:57 -0700 Subject: [PATCH 1/3] fix(proxy_server.py): don't auto-create user when creating key --- litellm/proxy/_new_secret_config.yaml | 15 +++---- litellm/proxy/_types.py | 2 +- litellm/proxy/proxy_server.py | 58 +++++++++++++++------------ 3 files changed, 41 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index bd277bbdfc..ed82946f5a 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,21 +1,22 @@ model_list: -- model_name: fake_openai +- model_name: fake-openai-endpoint litellm_params: model: openai/my-fake-model api_key: my-fake-key - api_base: http://0.0.0.0:8080 + api_base: https://exampleopenaiendpoint-production.up.railway.app/ - model_name: gpt-3.5-turbo litellm_params: model: gpt-3.5-turbo-1106 api_key: os.environ/OPENAI_API_KEY -litellm_settings: - cache: true - cache_params: - type: redis - callbacks: ["batch_redis_requests"] +# litellm_settings: +# cache: true +# cache_params: +# type: redis +# callbacks: ["batch_redis_requests"] # success_callbacks: ["langfuse"] general_settings: master_key: sk-1234 + # disable_spend_logs: true database_url: "postgresql://neondb_owner:hz8tyUlJ5ivV@ep-cool-sunset-a5ywubeh.us-east-2.aws.neon.tech/neondb?sslmode=require" \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9704a2f199..3da645da43 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 8fa2862f27..69f75115ab 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2320,8 +2320,6 @@ 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 @@ -2409,20 +2407,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" ): @@ -2435,16 +2436,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 @@ -4051,7 +4055,7 @@ async def generate_key_fn( 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") return GenerateKeyResponse(**response) except Exception as e: traceback.print_exc() @@ -5042,6 +5046,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"], From ad43138f282a2fae9ef8ad7f2c27850fdfbd024c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 27 Mar 2024 19:43:15 -0700 Subject: [PATCH 2/3] fix(proxy_server.py): fix budget add logic to accurately log who created it --- litellm/proxy/proxy_server.py | 62 +++++++++++++++++++---------------- 1 file changed, 34 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 69f75115ab..04f35a02be 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1008,7 +1008,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 ) @@ -2262,9 +2264,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 @@ -2287,7 +2287,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( @@ -2325,25 +2325,6 @@ async def generate_key_helper_fn( 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 = { @@ -2381,7 +2362,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 @@ -2461,7 +2442,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 @@ -3941,6 +3922,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), ): """ @@ -4044,18 +4026,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, + "updated_by": user_api_key_dict.user_id, + } + ) + _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, 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() From 34790c7e12c63eaca576aefd8c7183f7059ac858 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 27 Mar 2024 19:54:25 -0700 Subject: [PATCH 3/3] fix(proxy_server.py): fix linting error --- litellm/proxy/proxy_server.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04f35a02be..e8abe6fe4d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4041,8 +4041,8 @@ async def generate_key_fn( _budget = await prisma_client.db.litellm_budgettable.create( data={ **new_budget, # type: ignore - "created_by": user_api_key_dict.user_id, - "updated_by": user_api_key_dict.user_id, + "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)