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 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(
KeyManagementEventHooks.async_key_generated_hook(
data=data, data=data,
response=response, response=response,
user_api_key_dict=user_api_key_dict, user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by, litellm_changed_by=litellm_changed_by,
) )
)
return GenerateKeyResponse(**response) return GenerateKeyResponse(**response)
except Exception as e: except Exception as e:
@ -374,13 +376,15 @@ 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(
KeyManagementEventHooks.async_key_updated_hook(
data=data, data=data,
existing_key_row=existing_key_row, existing_key_row=existing_key_row,
response=response, response=response,
user_api_key_dict=user_api_key_dict, user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by, litellm_changed_by=litellm_changed_by,
) )
)
if response is None: if response is None:
raise ValueError("Failed to update key got response = 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}" 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, data=data,
user_api_key_dict=user_api_key_dict, user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by, litellm_changed_by=litellm_changed_by,
response=number_deleted_keys, response=number_deleted_keys,
) )
)
return {"deleted_keys": keys} return {"deleted_keys": keys}
except Exception as e: except Exception as e: