Add "Circuit Breaker" to Fallback model

[ChuckJonas]

Description

The Fallback model allows for requests to be sent to different model on the event of a request failure.

However, in the event the model is temporarily down (service outage or rate-limiting), pydantic-ai will continue to retry the failing model first on every request.

It would be nice if there were a built in way to "time out" a model (potentially using a circuit breaker pattern or just something simple):

def handle_fallback(e: Exception) -> bool:
    fallback = isinstance(e, (ModelHTTPError, ThrottlingException))
    return fallback

def curcuit_breaker(model, e: Exception):
   if isinstance(e, ThrottlingException)
     return {"state": "open", "expires": e.retry_after}

default_model = FallbackModel(
    primary_model, fallback_model, fallback_on=handle_fallback, circuit_breaker=curcuit_breaker
)

If a model state was set to "open" it would effectively be removed from the FallbackModel sequence until the expires.

A less opinioned interface might just be to allow for a enabled(model) -> boolean hook check before selecting the model. That would give the user flexibility to implement with whatever strategy they choose.

References

No response

[DouweM]

@ChuckJonas For rate limit handling, we have https://ai.pydantic.dev/retries/#rate-limit-handling-with-retry-after-support, but that will just retry the same model after waiting, while what you'd want is for it to use a fallback model until the retry-after time has expired.

You could implement a CircuitBreakerModel manually by subclassing WrapperModel, keeping track of the number/type of exceptions, and raising one immediately once the limit has been reached. FallbackModel would then immediately move on to the next one.

It'd be nice to let that retries integration raise an error like your ThrottlingException that can be intelligently handled by the FallbackModel, and if at the very least, it had a max_retries argument and tracked the number of failed calls to each model before skipping over it from then on -- I'd welcome a PR for that!

More complete circuit breaker functionality could also be useful, but I feel like we'd cover most cases with those 2 ideas.

[ChuckJonas]

@DouweM curious, does the approach outline in rate-limit-handling-with-retry-after-support actually work for bedrock models since they use the BOTA3 client behind the scenes? I know you can do something similar with BotoConfig, but AFIAK, it's not quite as flexible.

The WrapperModel approach sounds promising... Do you happen to have an code example using the WrapperModel in a similar fashion? I assume I'd have to override request & request_stream to be able to catch the exceptions?

I don't think the max_retries really solves my problem as I only want the model to be disabled until the Retry-After expiration.

On a semi-related note, it might be nice if the FallbackModel.fallback_on handler provided the failing model along with the exception...

Then if you were to add a enabled hook on the FallbackModel, it would allow you to do something like this, which seems simple yet powerful:

timed_out_models: dict[Model, datetime] = {
}

def handle_fallback(e: Exception, model: Model) -> bool:
    if isinstance(e, ThrottlingException):
        retry_after = e.response.get('ResponseMetadata', {}).get("retry-after", 60)
        retry_after = datetime.now() + timedelta(seconds=retry_after)
        timed_out_models[model] = retry_after
        return true

    # other fallback conditions
    return isinstance(
        e, (ModelHTTPError, ConnectionClosedError)
    )

def handle_enabled_models(model: Model) -> bool:
    if model in timed_out_models:
        if timed_out_models[model] < datetime.now():
            return False
    return True

default_model = FallbackModel(
    primary_model,
    secondary_model,
    enabled_models=handle_enabled_models,
    fallback_on=handle_fallback,
)

[DouweM]

curious, does the approach outline in rate-limit-handling-with-retry-after-support actually work for bedrock models since they use the BOTA3 client behind the scenes?

@ChuckJonas That's a good point, it doesn't as it requires the provider to take a custom http_client.

The WrapperModel approach sounds promising... Do you happen to have an code example using the WrapperModel in a similar fashion? I assume I'd have to override request & request_stream to be able to catch the exceptions?

Yep, exactly. Let me know if an example would be useful, but it's really just a subclass with a failure counter and those 2 methods with a check for the failure count and try/except/else to increment or reset the counter.

I don't think the max_retries really solves my problem as I only want the model to be disabled until the Retry-After expiration.

Right, that'd require some more intelligence in FallbackModel, or the custom WrapperModel that knows what errors the wrapped model will raise in case of a rate limit (like ThrottlingException for BedrockModel). We could implement our own ThrottlingException and wrap different models' exceptions as appropriate.

On a semi-related note, it might be nice if the FallbackModel.fallback_on handler provided the failing model along with the exception...

Then if you were to add a enabled hook on the FallbackModel, it would allow you to do something like this, which seems simple yet powerful:

What do you think of a new hook that takes the model and the most recent exception for that model (or None), and lets you return a boolean on whether to use it or not? That way you wouldn't need to store timed_out_models anymore.

[ChuckJonas]

What do you think of a new hook that takes the model and the most recent exception for that model (or None), and lets you return a boolean on whether to use it or not? That way you wouldn't need to store timed_out_models anymore.

That would be nice.

I think that would work for a lot of simple use cases (IE: I want to wait until Retry-After has expired to call this model again), but wouldn't support more complex/traditional circuit breaker scenarios (open breaker when failures exceed x% or X Failures per Y Time period).

However, I think adding the Model to fallback_on would actually be more useful/flexible. Would support more complicated patterns without the need to subclass the ModelWrapper.

Even just for the use case where I want do add some logging inside of the fallback_on handle... right now it's hard to figure out which model is actually failing (especially if I have more than 2).

[DouweM]

@ChuckJonas Yeah I think adding the model to fallback_on (with backward compatibility for callables that only take exception) and a new model_enabled that takes the model and the most recent exception, would give us all the flexibility and ease of use we need. Would you be up for filing a PR?

[ChuckJonas]

Would you be up for filing a PR?

Yup! Already working on it :)

[ChuckJonas]

Just wanted to provide an update since this issue is getting stale...

I implemented the approach outlined above, but just wasn't happy with how it came together in practice. The way the "last exceptions" needed to be tracked at the model just felt... off? I also came to the realization that it wasn't really going to support the flow I really needed, which is that if all models are "timed-out", then we might just want to wait for one of those timeouts to expire.

I still think adding the Model to fallback_on callback (backwards compatible) is still a useful addition.

[DouweM]

if all models are "timed-out", then we might just want to wait for one of those timeouts to expire.

@ChuckJonas Do you have an idea on how that could be implemented on FallbackModel? Because we'd want to move to the model that has the lowest wait time, the entire model selection would need to be customizable... Maybe a select_next_model function that takes the models and their most recent exceptions, so you can determine that, do an await asyncio.sleep() if necessary, and then return the first model?

[ChuckJonas]

Do you have an idea on how that could be implemented on FallbackModel?

My initial reaction is that keeping true the original purpose of the FallbackModel, maybe a "last chance" callback would be more aligned (on_models_exhausted or all_none_enable)?

That said, I think select_next_model would be a more flexible and powerful interface that maybe just warrants an entirely new class (maybe something like AdaptiveModel)?

I think with select_next_model(List[Model] -> Model) you could use this to do things like "load balancing", "extended-context upgrade", "circuit breaking", "tier-selection", etc.

I went ahead an drafted a "proposal/poc" PR for what an AdaptiveModel might look like like and how some interesting use-cases might be implemented

#3335