Class ModelFallback

tests/test_middleware.py:133–154

class ModelFallback(Middleware):
    def __init__(self, fallback_model: str) -> None:
        self.fallback_model = fallback_model

    def _fallback(self, request: APIRequest) -> APIRequest:
        body = request.json
        assert isinstance(body, dict)
        return request.copy(body={**cast("dict[str, object]", body), "model": self.fallback_model})

    @override
    def handle(self, request: APIRequest, call_next: CallNext) -> Any:
        response = call_next(request)
        if response.status_code == 529:
            return call_next(self._fallback(request))
        return response

    @override
    async def handle_async(self, request: APIRequest, call_next: AsyncCallNext) -> Any:
        response = await call_next(request)
        if response.status_code == 529:
            return await call_next(self._fallback(request))
        return response


class ShrinkMaxTokensOnTooLarge(Middleware):