Custom Hooks and Callbacks
Relevant source files
The following files were used as context for generating this wiki page:
Custom hooks and callbacks in LiteLLM provide a powerful extensibility framework that allows developers to intercept, modify, and augment the request-response lifecycle of the AI Gateway and SDK. These components enable advanced functionalities such as custom authentication, logging, spend tracking, and automated agentic loops without modifying the core codebase.
The system is built on a plugin architecture where specialized classes, typically inheriting from a base CustomLogger or similar interface, are registered to receive events at various stages of an LLM call. This includes pre-call interceptions for input validation or prompt injection and post-call handlers for success or failure events.
Sources: ARCHITECTURE.md, litellm/proxy/example_config_yaml/custom_callbacks.py
Architecture and Event Flow
The LiteLLM AI Gateway (Proxy) integrates hooks into its main request flow to enforce management features like rate limiting and budget controls. Hooks are executed both before the request is forwarded to the LLM provider and after the response is received.
Proxy Request Lifecycle with Hooks
The following sequence diagram illustrates how proxy-level hooks interact with the internal SDK and external providers.
This diagram shows the interception points where custom logic can be applied to manage traffic and data.
Sources: ARCHITECTURE.md:16-56
Internal Hook Execution
Proxy hooks are registered in proxy/hooks/__init__.py and are managed by the litellm_logging.py module. When an event occurs, such as a successful LLM completion, the logging object triggers the async_success_handler(), which in turn iterates through all registered hooks.
Sources: ARCHITECTURE.md:120-135
Key Hook Components and Features
LiteLLM provides several specialized hook types categorized by their purpose, ranging from infrastructure management to agentic behavior.
Core Proxy Hooks
| Hook Name | Description | Purpose |
|---|---|---|
max_budget_limiter | Enforces budget limits for keys/users. | Cost Management |
parallel_request_limiter | Limits concurrent requests per key. | Rate Limiting |
cache_control_check | Validates cache headers and state. | Performance |
responses_id_security | Validates Response IDs for security. | Security |
litellm_skills | Injects skills into the prompt. | Prompt Engineering |
Sources: ARCHITECTURE.md:120-130
SSO and Authentication Hooks
Custom hooks can be used to extend the LiteLLM Dashboard's authentication capabilities. For example, a CustomUISSO hook can be implemented to handle user redirection and token exchange for Single Sign-On (SSO).
Key methods for SSO hooks include:
get_ui_redirect_url: Generates the URL for the SSO provider login.get_user_info: Exchanges a code for user details and tokens.
Sources: litellm/proxy/custom_hooks/custom_ui_sso_hook.py:8-33
Agentic and Skills Hooks
The SkillsInjectionHook and CodeExecutionHandler facilitate complex workflows where LiteLLM manages an automated loop. The CodeExecutionHandler specifically looks for tool calls in the LLM response that require Python execution.
The async_post_call_success_deployment_hook implements this "Agentic Loop," repeatedly calling the LLM until no further code execution is required.
Sources: litellm/llms/litellm_proxy/skills/code_execution.py:19-75
Implementation Example: Custom Callbacks
Developers can define custom callbacks by creating a class that inherits from litellm.integrations.custom_logger.CustomLogger. This class can then be registered in the LiteLLM configuration.
import litellm
from litellm.integrations.custom_logger import CustomLogger
class MyCustomHandler(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
print(f"Pre-call for model: {model}")
def log_success_event(self, kwargs, response_obj, start_time, end_time):
print("Call succeeded")
# Registration
custom_handler = MyCustomHandler()
litellm.callbacks = [custom_handler]
Sources: litellm/proxy/example_config_yaml/custom_callbacks.py:1-25
Logging and Observability Integrations
LiteLLM's callback system is the primary mechanism for supporting third-party logging and guardrail providers. These are organized within the AI Integrations section of the project's release and management structure.
Supported Logging Categories
- Providers: DataDog, Langfuse, Prometheus, PostHog, SQS, Opik, Arize Phoenix.
- General Logging Features: Callback controls and sensitive data masking.
- Guardrails: LakeraAI, Presidio, IBM Guardrails, etc.
Sources: cookbook/misc/RELEASE_NOTES_GENERATION_INSTRUCTIONS.md:121-155
Conclusion
Custom hooks and callbacks are fundamental to LiteLLM's flexibility, allowing it to function as a robust AI Gateway. By utilizing pre-call and post-call interception points, developers can implement sophisticated logic for security, cost control, and agentic task completion while maintaining a consistent interface for the end user.