Spaces:
Paused
Paused
| from typing import Literal, Optional | |
| import litellm | |
| from litellm.integrations.custom_logger import CustomLogger | |
| from litellm.proxy.proxy_server import DualCache, UserAPIKeyAuth | |
| # This file includes the custom callbacks for LiteLLM Proxy | |
| # Once defined, these can be passed in proxy_config.yaml | |
| class MyCustomHandler( | |
| CustomLogger | |
| ): # https://docs.litellm.ai/docs/observability/custom_callback#callback-class | |
| # Class variables or attributes | |
| def __init__(self): | |
| pass | |
| #### CALL HOOKS - proxy only #### | |
| async def async_pre_call_hook( | |
| self, | |
| user_api_key_dict: UserAPIKeyAuth, | |
| cache: DualCache, | |
| data: dict, | |
| call_type: Literal[ | |
| "completion", | |
| "text_completion", | |
| "embeddings", | |
| "image_generation", | |
| "moderation", | |
| "audio_transcription", | |
| "pass_through_endpoint", | |
| "rerank", | |
| ], | |
| ): | |
| return data | |
| async def async_post_call_failure_hook( | |
| self, | |
| request_data: dict, | |
| original_exception: Exception, | |
| user_api_key_dict: UserAPIKeyAuth, | |
| ): | |
| pass | |
| async def async_post_call_success_hook( | |
| self, | |
| data: dict, | |
| user_api_key_dict: UserAPIKeyAuth, | |
| response, | |
| ): | |
| # print("in async_post_call_success_hook") | |
| pass | |
| async def async_moderation_hook( # call made in parallel to llm api call | |
| self, | |
| data: dict, | |
| user_api_key_dict: UserAPIKeyAuth, | |
| call_type: Literal[ | |
| "completion", | |
| "embeddings", | |
| "image_generation", | |
| "moderation", | |
| "audio_transcription", | |
| "responses", | |
| ], | |
| ): | |
| pass | |
| async def async_post_call_streaming_hook( | |
| self, | |
| user_api_key_dict: UserAPIKeyAuth, | |
| response: str, | |
| ): | |
| # print("in async_post_call_streaming_hook") | |
| pass | |
| proxy_handler_instance = MyCustomHandler() | |