phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

use tasks for key management hooks

Browse source
This commit is contained in:
Ishaan Jaff 2024-11-13 11:38:48 -08:00
parent 5db942a39e
commit 7788747bfe
1 changed files with 22 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

38
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -235,11 +235,13 @@ async def generate_key_fn( # noqa: PLR0915
data.soft_budget data.soft_budget
) # include the user-input soft budget in the response ) # include the user-input soft budget in the response
await KeyManagementEventHooks.async_key_generated_hook( asyncio.create_task(
data=data, KeyManagementEventHooks.async_key_generated_hook(
response=response, data=data,
user_api_key_dict=user_api_key_dict, response=response,
litellm_changed_by=litellm_changed_by, user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
) )
return GenerateKeyResponse(**response) return GenerateKeyResponse(**response)
@ -374,12 +376,14 @@ async def update_key_fn(
proxy_logging_obj=proxy_logging_obj, proxy_logging_obj=proxy_logging_obj,
) )
await KeyManagementEventHooks.async_key_updated_hook( asyncio.create_task(
data=data, KeyManagementEventHooks.async_key_updated_hook(
existing_key_row=existing_key_row, data=data,
response=response, existing_key_row=existing_key_row,
user_api_key_dict=user_api_key_dict, response=response,
litellm_changed_by=litellm_changed_by, user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
) )
if response is None: if response is None:
@ -499,11 +503,13 @@ async def delete_key_fn(
f"/keys/delete - cache after delete: {user_api_key_cache.in_memory_cache.cache_dict}" f"/keys/delete - cache after delete: {user_api_key_cache.in_memory_cache.cache_dict}"
) )
await KeyManagementEventHooks.async_key_deleted_hook( asyncio.create_task(
data=data, KeyManagementEventHooks.async_key_deleted_hook(
user_api_key_dict=user_api_key_dict, data=data,
litellm_changed_by=litellm_changed_by, user_api_key_dict=user_api_key_dict,
response=number_deleted_keys, litellm_changed_by=litellm_changed_by,
response=number_deleted_keys,
)
) )
return {"deleted_keys": keys} return {"deleted_keys": keys}