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

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

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