Can custom middleware be developed and registered in Hosted LiteLLM?

[jenoOvchi]

Hi!

While exploring the LiteLLM source code, I found the middleware directory (e.g., https://github.com/BerriAI/litellm/blob/main/litellm/proxy/middleware/prometheus_auth_middleware.py), which handles authentication for specific types of requests, such as Prometheus queries.

Is it possible to develop and register custom middleware instances in the Hosted version of LiteLLM? If so, what is the process for doing this?

Specifically, I’m interested in:

Whether the hosted version supports custom middleware.

Steps to develop and integrate custom middleware (e.g., for auth, logging, or request/response modification).

Any limitations or best practices to consider.

Thanks in advance for your guidance!

[KeepALifeUS]

Custom middleware in hosted LiteLLM is valuable. Here is a pattern for building middleware:

Middleware Architecture

from typing import Callable, Dict, Any, Optional
from dataclasses import dataclass
import time

@dataclass
class MiddlewareContext:
    request: Dict[str, Any]
    response: Optional[Dict[str, Any]] = None
    metadata: Dict[str, Any] = None
    start_time: float = 0
    
    def __post_init__(self):
        self.metadata = self.metadata or {}
        self.start_time = time.time()

class MiddlewareChain:
    def __init__(self):
        self.pre_request: list[Callable] = []
        self.post_response: list[Callable] = []
    
    def add_pre_request(self, middleware: Callable):
        self.pre_request.append(middleware)
    
    def add_post_response(self, middleware: Callable):
        self.post_response.append(middleware)
    
    async def process_request(self, ctx: MiddlewareContext) -> MiddlewareContext:
        for middleware in self.pre_request:
            ctx = await middleware(ctx)
            if ctx.response:  # Short circuit if middleware returns response
                return ctx
        return ctx
    
    async def process_response(self, ctx: MiddlewareContext) -> MiddlewareContext:
        for middleware in self.post_response:
            ctx = await middleware(ctx)
        return ctx

Example Middleware

# Rate limiting middleware
async def rate_limit_middleware(ctx: MiddlewareContext) -> MiddlewareContext:
    user_id = ctx.request.get("user")
    if not check_rate_limit(user_id):
        ctx.response = {"error": "Rate limit exceeded"}
    return ctx

# Cost tracking middleware
async def cost_tracking_middleware(ctx: MiddlewareContext) -> MiddlewareContext:
    if ctx.response:
        tokens = ctx.response.get("usage", {})
        cost = calculate_cost(ctx.request["model"], tokens)
        ctx.metadata["cost"] = cost
        await log_cost(ctx.request.get("user"), cost)
    return ctx

# PII redaction middleware
async def pii_redaction_middleware(ctx: MiddlewareContext) -> MiddlewareContext:
    messages = ctx.request.get("messages", [])
    for msg in messages:
        msg["content"] = redact_pii(msg["content"])
    return ctx

# Register middleware
chain = MiddlewareChain()
chain.add_pre_request(rate_limit_middleware)
chain.add_pre_request(pii_redaction_middleware)
chain.add_post_response(cost_tracking_middleware)

For Hosted LiteLLM

If self-hosting is not an option, consider:

1. Webhook callbacks for post-processing
2. Custom headers to pass metadata
3. Proxy layer in front of hosted LiteLLM

More on agent patterns: https://github.com/KeepALifeUS/autonomous-agents

[jenoOvchi]

Thank you for the explanation regarding the middleware architecture in LiteLLM.

To proceed with implementing custom middleware in the Hosted LiteLLM environment, I need clarification on the deployment process. Please address the following questions:

1. Where should custom middleware files (e.g., my_custom_middleware.py) be located in the hosted setup? Is there a designated directory or upload mechanism for middleware code?
2. Which configuration file or interface is used to register custom middleware?
3. How can I verify that middleware has been successfully loaded? Where can I access logs related to middleware execution?