tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 22:22:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2724 from BerriAI/litellm_auto_create_user_fix

Browse Source 
fix(proxy_server.py): don't auto-create user when creating key
This commit is contained in:
Krish Dholakia
2024-03-28 14:57:44 -07:00
committed by GitHub
parent 934a9ac2b4 7d3244c012
commit 72e1f776f4
2 changed files with 67 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+1 -1
View File
@@ -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):
litellm/proxy/proxy_server.py
+66 -54
View File
@@ -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"],