fix(proxy_server.py): enable default new user params

This commit is contained in:
Krrish Dholakia 2024-02-23 16:39:50 -08:00
parent 1e769bcb77
commit 0806aa8da1
7 changed files with 462 additions and 389 deletions

View file

@ -730,6 +730,8 @@ async def user_api_key_auth(
"/user",
"/model/info",
"/v2/model/info",
"/models",
"/v1/models",
]
# check if the current route startswith any of the allowed routes
if (
@ -1758,6 +1760,7 @@ async def generate_key_helper_fn(
allowed_cache_controls: Optional[list] = [],
permissions: Optional[dict] = {},
model_max_budget: Optional[dict] = {},
table_name: Optional[Literal["key", "user"]] = None,
):
global prisma_client, custom_db_client, user_api_key_cache
@ -1884,8 +1887,10 @@ async def generate_key_helper_fn(
table_name="user",
update_key_values=update_key_values,
)
if user_id == litellm_proxy_budget_name:
# do not create a key for litellm_proxy_budget_name
if user_id == litellm_proxy_budget_name or (
table_name is not None and table_name == "user"
):
# do not create a key for litellm_proxy_budget_name or if table name is set to just 'user'
# we only need to ensure this exists in the user table
# the LiteLLM_VerificationToken table will increase in size if we don't do this check
return key_data
@ -5461,27 +5466,50 @@ async def auth_callback(request: Request):
user_id_models: List = []
# User might not be already created on first generation of key
# But if it is, we want its models preferences
# But if it is, we want their models preferences
default_ui_key_values = {
"duration": "1hr",
"key_max_budget": 0.01,
"aliases": {},
"config": {},
"spend": 0,
"team_id": "litellm-dashboard",
}
user_defined_values = {
"models": user_id_models,
"user_id": user_id,
"user_email": user_email,
}
try:
if prisma_client is not None:
user_info = await prisma_client.get_data(user_id=user_id, table_name="user")
verbose_proxy_logger.debug(
f"user_info: {user_info}; litellm.default_user_params: {litellm.default_user_params}"
)
if user_info is not None:
user_id_models = getattr(user_info, "models", [])
user_defined_values = {
"models": getattr(user_info, "models", []),
"user_id": getattr(user_info, "user_id", user_id),
"user_email": getattr(user_info, "user_id", user_email),
}
elif litellm.default_user_params is not None and isinstance(
litellm.default_user_params, dict
):
user_defined_values = {
"models": litellm.default_user_params.get("models", user_id_models),
"user_id": litellm.default_user_params.get("user_id", user_id),
"user_email": litellm.default_user_params.get(
"user_email", user_email
),
}
except Exception as e:
pass
verbose_proxy_logger.info(
f"user_defined_values for creating ui key: {user_defined_values}"
)
response = await generate_key_helper_fn(
**{
"duration": "1hr",
"key_max_budget": 0.01,
"models": user_id_models,
"aliases": {},
"config": {},
"spend": 0,
"user_id": user_id,
"team_id": "litellm-dashboard",
"user_email": user_email,
} # type: ignore
**default_ui_key_values, **user_defined_values # type: ignore
)
key = response["token"] # type: ignore
user_id = response["user_id"] # type: ignore