A stand-in provider: writing your own CustomLLM

mock_response always gives the same answer. A custom provider decides, which retries, fallbacks and a gateway need to be worth watching.

Exampleshop_llm.py

import litellm
from litellm import CustomLLM
from litellm.utils import custom_llm_setup


class ShopLLM(CustomLLM):
    def completion(self, model, messages, *args, **kwargs):
        text = messages[-1]["content"].lower()
        if "refund" in text or "charged" in text:
            answer = '{"category": "billing", "priority": 4}'
        elif "parcel" in text or "arrived" in text:
            answer = '{"category": "shipping", "priority": 3}'
        else:
            answer = "I am not sure how to sort this one."
        return litellm.ModelResponse(
            model=f"shop/{model}",
            choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": answer}}],
            usage=litellm.Usage(prompt_tokens=len(text.split()), completion_tokens=len(answer.split()),
                                total_tokens=len(text.split()) + len(answer.split())),
        )

    async def acompletion(self, model, messages, *args, **kwargs):
        return self.completion(model, messages)


shop = ShopLLM()
litellm.custom_provider_map = [{"provider": "shop", "custom_handler": shop}]
custom_llm_setup()

A provider is a subclass of CustomLLM. completion receives the model id, the part after the slash, and the messages, and returns a ModelResponse, the same object LiteLLM returns for every provider. This one sorts tickets by keywords, like the stand-in from Python for AI, and counts words as tokens. acompletion is the async version the gateway will call.

The last lines register it. litellm.custom_provider_map maps the prefix shop to the handler, as the custom provider docs show for connecting an internal model server.

custom_llm_setup() adds the prefixes in that map to LiteLLM's list of known providers. LiteLLM does this itself at the start of every completion call, which is enough for a single call. The Router in part 3 and register_model in lesson 6 check model strings before any call has happened, and without this line they reject shop/ as an unknown provider. It lives in litellm.utils, an internal module, so check it still exists when you upgrade.

Example

from litellm import completion

import shop_llm

for text in ["I was charged twice for one order", "My parcel has not arrived", "How do I change my password?"]:
    response = completion(model="shop/small", messages=[{"role": "user", "content": text}])
    print(response.choices[0].message.content, "|", response.model, "|", response.usage.total_tokens)

Same completion call as lesson 2, now answered by your class, and the answer differs per ticket. From here on, shop/small stands in for a cheap hosted model and shop/large for an expensive one; the class does not care which name it gets.

Try it yourself

Add an "account" branch for passwords.
Print response.choices[0].finish_reason.
Call completion(model="shop/large", ...) and print response.model.

PreviousProviders and keys: the model string

This is what real progress feels like.