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

feat(dynamic_rate_limiter.py): working e2e

Browse source
This commit is contained in:
Krrish Dholakia 2024-06-22 14:41:22 -07:00
parent 532f24bfb7
commit a31a05d45d
7 changed files with 420 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

49
litellm/proxy/utils.py
View file

@ -584,8 +584,15 @@ class ProxyLogging:
for callback in litellm.callbacks:
try:
if isinstance(callback, CustomLogger):
await callback.async_post_call_failure_hook(
_callback: Optional[CustomLogger] = None
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
callback
)
else:
_callback = callback # type: ignore
if _callback is not None and isinstance(_callback, CustomLogger):
await _callback.async_post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=original_exception,
)
@ -606,8 +613,15 @@ class ProxyLogging:
"""
for callback in litellm.callbacks:
try:
if isinstance(callback, CustomLogger):
await callback.async_post_call_success_hook(
_callback: Optional[CustomLogger] = None
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
callback
)
else:
_callback = callback # type: ignore
if _callback is not None and isinstance(_callback, CustomLogger):
await _callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict, response=response
)
except Exception as e:
@ -625,14 +639,25 @@ class ProxyLogging:
Covers:
1. /chat/completions
"""
for callback in litellm.callbacks:
try:
if isinstance(callback, CustomLogger):
await callback.async_post_call_streaming_hook(
user_api_key_dict=user_api_key_dict, response=response
)
except Exception as e:
raise e
response_str: Optional[str] = None
if isinstance(response, ModelResponse):
response_str = litellm.get_response_string(response_obj=response)
if response_str is not None:
for callback in litellm.callbacks:
try:
_callback: Optional[CustomLogger] = None
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
callback
)
else:
_callback = callback # type: ignore
if _callback is not None and isinstance(_callback, CustomLogger):
await _callback.async_post_call_streaming_hook(
user_api_key_dict=user_api_key_dict, response=response_str
)
except Exception as e:
raise e
return response
async def post_call_streaming_hook(