phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy_server.py): speed up proxy startup time

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-30 16:17:23 -08:00
parent 37de964da4
commit 2019347f0d
1 changed files with 39 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/proxy_server.py
View file

@ -1267,7 +1267,7 @@ async def generate_key_helper_fn(
update_key_values: Optional[dict] = None, update_key_values: Optional[dict] = None,
key_alias: Optional[str] = None, key_alias: Optional[str] = None,
): ):
global prisma_client, custom_db_client global prisma_client, custom_db_client, user_api_key_cache
if prisma_client is None and custom_db_client is None: if prisma_client is None and custom_db_client is None:
raise Exception( raise Exception(
@ -1357,6 +1357,18 @@ async def generate_key_helper_fn(
} }
if general_settings.get("allow_user_auth", False) == True: if general_settings.get("allow_user_auth", False) == True:
key_data["key_name"] = f"sk-...{token[-4:]}" key_data["key_name"] = f"sk-...{token[-4:]}"
saved_token = copy.deepcopy(key_data)
if isinstance(saved_token["aliases"], str):
saved_token["aliases"] = json.loads(saved_token["aliases"])
if isinstance(saved_token["config"], str):
saved_token["config"] = json.loads(saved_token["config"])
if isinstance(saved_token["metadata"], str):
saved_token["metadata"] = json.loads(saved_token["metadata"])
user_api_key_cache.set_cache(
key=key_data["token"],
value=LiteLLM_VerificationToken(**saved_token), # type: ignore
ttl=60,
)
if prisma_client is not None: if prisma_client is not None:
## CREATE USER (If necessary) ## CREATE USER (If necessary)
verbose_proxy_logger.debug(f"prisma_client: Creating User={user_data}") verbose_proxy_logger.debug(f"prisma_client: Creating User={user_data}")
@ -1671,7 +1683,8 @@ async def startup_event():
if prisma_client is not None and master_key is not None: if prisma_client is not None and master_key is not None:
# add master key to db # add master key to db
await generate_key_helper_fn( asyncio.create_task(
generate_key_helper_fn(
duration=None, duration=None,
models=[], models=[],
aliases={}, aliases={},
@ -1680,6 +1693,7 @@ async def startup_event():
token=master_key, token=master_key,
user_id="default_user_id", user_id="default_user_id",
) )
)
if prisma_client is not None and litellm.max_budget > 0: if prisma_client is not None and litellm.max_budget > 0:
if litellm.budget_duration is None: if litellm.budget_duration is None:
@ -1688,7 +1702,8 @@ async def startup_event():
) )
# add proxy budget to db in the user table # add proxy budget to db in the user table
await generate_key_helper_fn( asyncio.create_task(
generate_key_helper_fn(
user_id=litellm_proxy_budget_name, user_id=litellm_proxy_budget_name,
duration=None, duration=None,
models=[], models=[],
@ -1703,6 +1718,7 @@ async def startup_event():
"budget_duration": litellm.budget_duration, "budget_duration": litellm.budget_duration,
}, },
) )
)
verbose_proxy_logger.debug( verbose_proxy_logger.debug(
f"custom_db_client client {custom_db_client}. Master_key: {master_key}" f"custom_db_client client {custom_db_client}. Master_key: {master_key}"