Errors: one set of exceptions for every provider

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

from shop_llm import shop


class FlakyLLM(CustomLLM):
    """Fails the first `failures` calls with a rate limit, then answers like ShopLLM."""

    def __init__(self, failures):
        super().__init__()
        self.failures = failures
        self.calls = 0

    def completion(self, model, messages, *args, **kwargs):
        self.calls += 1
        if self.calls <= self.failures:
            raise litellm.RateLimitError(message="slow down", llm_provider="flaky", model=model)
        return shop.completion(model, messages)

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


class DownLLM(CustomLLM):
    """A provider that is having an outage."""

    def completion(self, model, messages, *args, **kwargs):
        raise litellm.ServiceUnavailableError(message="provider is down", llm_provider="down", model=model)

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


flaky = FlakyLLM(failures=2)
down = DownLLM()
litellm.custom_provider_map = [
    {"provider": "shop", "custom_handler": shop},
    {"provider": "flaky", "custom_handler": flaky},
    {"provider": "down", "custom_handler": down},
]
custom_llm_setup()

import litellm
from litellm import completion

import flaky_llm

litellm.suppress_debug_info = True

for model in ["flaky/small", "down/large"]:
    try:
        completion(model=model, messages=[{"role": "user", "content": "I was charged twice for one order"}])
    except litellm.RateLimitError as error:
        print("rate limited:", error.status_code, error.llm_provider)
    except litellm.ServiceUnavailableError as error:
        print("unavailable:", error.status_code, error.llm_provider)

import openai

try:
    completion(model="down/large", messages=[{"role": "user", "content": "I was charged twice for one order"}])
except openai.APIStatusError as error:
    print(type(error).__name__, isinstance(error, litellm.ServiceUnavailableError))