Implementation:BerriAI Litellm Custom LLM Handler

Attribute	Value
Sources	litellm/llms/custom_llm.py
Domains	LLM Providers, Custom LLM, Extensibility, Streaming
Last Updated	2026-02-15 16:00 GMT

Overview

CustomLLM is a handler class for implementing custom chat LLM providers that supports completion, async completion, streaming, async streaming, image generation, image editing, and embeddings via a pluggable interface.

Description

The CustomLLM class extends BaseLLM and provides a comprehensive set of method stubs for building a custom LLM provider. Each method raises CustomLLMError (with status code 500) by default, requiring the user to override the methods they need. The class supports: synchronous completion (completion), async completion (acompletion), synchronous streaming (streaming), async streaming (astreaming), image generation (image_generation, aimage_generation), image editing (image_edit, aimage_edit), and embeddings (embedding, aembedding). The module also includes a custom_chat_llm_router function that routes calls to the appropriate method based on whether the call is async and/or streaming. The CustomLLMError exception class provides structured error reporting with status codes.

Usage

Import and subclass CustomLLM when you want to implement a custom LLM provider that integrates with LiteLLM's routing, streaming, and async infrastructure. Override the methods relevant to your use case (e.g., completion for basic chat, streaming for streaming responses).

Code Reference

Source Location

litellm/llms/custom_llm.py

Signature

class CustomLLMError(Exception):
    def __init__(self, status_code, message)

class CustomLLM(BaseLLM):
    def __init__(self) -> None
    def completion(self, model, messages, api_base, custom_prompt_dict, model_response, print_verbose,
                   encoding, api_key, logging_obj, optional_params, acompletion=None, litellm_params=None,
                   logger_fn=None, headers={}, timeout=None, client=None) -> Union[ModelResponse, "CustomStreamWrapper"]
    def streaming(self, model, messages, api_base, ..., client=None) -> Iterator[GenericStreamingChunk]
    async def acompletion(self, model, messages, api_base, ..., client=None) -> Union[ModelResponse, "CustomStreamWrapper"]
    async def astreaming(self, model, messages, api_base, ..., client=None) -> AsyncIterator[GenericStreamingChunk]
    def image_generation(self, model, prompt, api_key, api_base, model_response, optional_params, logging_obj, timeout=None, client=None) -> ImageResponse
    async def aimage_generation(self, model, prompt, model_response, api_key, api_base, optional_params, logging_obj, timeout=None, client=None) -> ImageResponse
    def embedding(self, model, input, model_response, print_verbose, logging_obj, optional_params, api_key=None, api_base=None, timeout=None, litellm_params=None) -> EmbeddingResponse
    async def aembedding(self, model, input, model_response, ...) -> EmbeddingResponse
    def image_edit(self, model, image, prompt, model_response, api_key, api_base, optional_params, logging_obj, timeout=None, client=None) -> ImageResponse
    async def aimage_edit(self, model, image, prompt, model_response, ...) -> ImageResponse

def custom_chat_llm_router(async_fn: bool, stream: Optional[bool], custom_llm: CustomLLM) -> Callable

Import

from litellm.llms.custom_llm import CustomLLM, CustomLLMError, custom_chat_llm_router

I/O Contract

Inputs

Parameter	Type	Description
model	str	The model identifier.
messages	list	The chat messages list.
api_base	str	The base URL for the API.
custom_prompt_dict	dict	Custom prompt formatting dictionary.
model_response	ModelResponse	The response object to populate.
optional_params	dict	Optional LLM parameters (temperature, max_tokens, etc.).
timeout	Optional[Union[float, httpx.Timeout]]	Request timeout.
client	Optional[HTTPHandler/AsyncHTTPHandler]	HTTP client to use for requests.
prompt	str	Text prompt for image generation.
input	list	Input texts for embedding.

Outputs

Method	Return Type	Description
completion	Union[ModelResponse, CustomStreamWrapper]	Chat completion response.
streaming	Iterator[GenericStreamingChunk]	Synchronous streaming iterator.
acompletion	Union[ModelResponse, CustomStreamWrapper]	Async chat completion response.
astreaming	AsyncIterator[GenericStreamingChunk]	Async streaming iterator.
image_generation / aimage_generation	ImageResponse	Image generation response.
embedding / aembedding	EmbeddingResponse	Embedding response.
custom_chat_llm_router	Callable	The appropriate method reference based on async/stream flags.

Usage Examples

from litellm.llms.custom_llm import CustomLLM
from litellm.utils import ModelResponse

class MyCustomLLM(CustomLLM):
    def completion(self, model, messages, api_base, custom_prompt_dict,
                   model_response, print_verbose, encoding, api_key,
                   logging_obj, optional_params, **kwargs):
        # Implement your custom completion logic
        model_response.choices[0].message.content = "Hello from my custom LLM!"
        return model_response

# Register the custom LLM
import litellm
litellm.custom_provider_map = [
    {"provider": "my-provider", "custom_handler": MyCustomLLM()}
]

# Use it through litellm
response = litellm.completion(
    model="my-provider/my-model",
    messages=[{"role": "user", "content": "Hello!"}],
)

Related Pages

• BerriAI_Litellm_LLM_Base_Class - The base class that CustomLLM extends