com.cohere.api.core.RetryInterceptor Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of cohere-java Show documentation

The official Java library for Cohere's API.

There is a newer version: 1.4.1

/**
 * This file was auto-generated by Fern from our API Definition.
 */
package com.cohere.api.core;

import java.io.IOException;
import java.time.Duration;
import java.util.Optional;
import java.util.Random;
import okhttp3.Interceptor;
import okhttp3.Response;

public class RetryInterceptor implements Interceptor {

    private static final Duration ONE_SECOND = Duration.ofSeconds(1);
    private final ExponentialBackoff backoff;
    private final Random random = new Random();

    public RetryInterceptor(int maxRetries) {
        this.backoff = new ExponentialBackoff(maxRetries);
    }

    @Override
    public Response intercept(Chain chain) throws IOException {
        Response response = chain.proceed(chain.request());

        if (shouldRetry(response.code())) {
            return retryChain(response, chain);
        }

        return response;
    }

    private Response retryChain(Response response, Chain chain) throws IOException {
        Optional nextBackoff = this.backoff.nextBackoff();
        while (nextBackoff.isPresent()) {
            try {
                Thread.sleep(nextBackoff.get().toMillis());
            } catch (InterruptedException e) {
                throw new IOException("Interrupted while trying request", e);
            }
            response.close();
            response = chain.proceed(chain.request());
            if (shouldRetry(response.code())) {
                nextBackoff = this.backoff.nextBackoff();
            } else {
                return response;
            }
        }

        return response;
    }

    private static boolean shouldRetry(int statusCode) {
        return statusCode == 408 || statusCode == 409 || statusCode == 429 || statusCode >= 500;
    }

    private final class ExponentialBackoff {

        private final int maxNumRetries;

        private int retryNumber = 0;

        ExponentialBackoff(int maxNumRetries) {
            this.maxNumRetries = maxNumRetries;
        }

        public Optional nextBackoff() {
            retryNumber += 1;
            if (retryNumber > maxNumRetries) {
                return Optional.empty();
            }

            int upperBound = (int) Math.pow(2, retryNumber);
            return Optional.of(ONE_SECOND.multipliedBy(random.nextInt(upperBound)));
        }
    }
}