Merge pull request #10 from penguineer/rate-limit-exception

Introduce and throw a Rate Limit Exception on API calls

src/main/java/com/penguineering/hareairis/ai/AIChatService.java

@@ -60,6 +60,9 @@ public ChatResponse handleChatRequest(ChatRequest chatRequest) {
             throw new ChatException(ChatException.Code.CODE_BAD_REQUEST, e.getMessage());
         } catch (HttpResponseException e) {
             var response = e.getResponse();
+            if (response.getStatusCode() == 429)
+                throw RateLimitException.fromHttpResponse(response, logger::warn);
+
             throw new ChatException(response.getStatusCode(), e.getMessage());
         }
     }

src/main/java/com/penguineering/hareairis/ai/RateLimitException.java

@@ -0,0 +1,94 @@
+package com.penguineering.hareairis.ai;
+
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpResponse;
+import com.penguineering.hareairis.model.ChatException;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.function.Consumer;
+
+/**
+ * Exception thrown when the OpenAI service returns a rate-limiting error.
+ */
+public class RateLimitException extends ChatException {
+    /**
+     * Creates a RateLimitException from the specified HttpResponse.
+     *
+     * @param response    The HttpResponse to create the exception from.
+     * @param logConsumer The consumer to log any parsing errors.
+     * @return The RateLimitException created from the HttpResponse.
+     * @throws IllegalArgumentException If the response status code is not 429.
+     */
+    public static RateLimitException fromHttpResponse(HttpResponse response, Consumer<String> logConsumer) {
+        if (response.getStatusCode() != 429)
+            throw new IllegalArgumentException("Expected status code 429 for RateLimitException, but received " + response.getStatusCode() + ".");
+
+        String errorMessage = response.getBodyAsString().block();
+        String retryAfterHeader = response.getHeaders().getValue(HttpHeaderName.RETRY_AFTER);
+
+        if (Objects.isNull(retryAfterHeader))
+            return new RateLimitException(errorMessage);
+
+        // Try parsing as Instant
+        try {
+            Instant retryAfter = Instant.parse(retryAfterHeader);
+            return new RateLimitException(errorMessage, retryAfter);
+        } catch (Exception ex) {
+            // Ignore and try parsing as Duration
+        }
+
+        // Try parsing as seconds
+        try {
+            long seconds = Long.parseLong(retryAfterHeader);
+            return new RateLimitException(errorMessage, Duration.ofSeconds(seconds));
+        } catch (NumberFormatException e) {
+            logConsumer.accept("Failed to parse Retry-After header and returning an exception with message only: " + retryAfterHeader);
+        }
+
+        return new RateLimitException(errorMessage);
+    }
+
+    private final Instant retryAfter;
+
+    /**
+     * Constructs a new RateLimitException with the specified detail message.
+     *
+     * @param message The detail message.
+     */
+    public RateLimitException(String message) {
+        this(message, (Instant) null);
+    }
+
+    /**
+     * Constructs a new RateLimitException with the specified detail message and retry after time.
+     *
+     * @param message    The detail message.
+     * @param retryAfter The optional point in time to hint when the service can be called again.
+     */
+    public RateLimitException(String message, Instant retryAfter) {
+        super(Code.CODE_TOO_MANY_REQUESTS, message);
+        this.retryAfter = retryAfter;
+    }
+
+    /**
+     * Constructs a new RateLimitException with the specified detail message and retry after duration.
+     *
+     * @param message  The detail message.
+     * @param duration The duration after which the service can be called again.
+     */
+    public RateLimitException(String message, Duration duration) {
+        this(message, Instant.now().plus(duration));
+    }
+
+    /**
+     * Returns the optional point in time to hint when the service can be called again.
+     *
+     * @return The optional retry after time.
+     */
+    public Optional<Instant> getRetryAfter() {
+        return Optional.ofNullable(retryAfter);
+    }
+}

src/test/java/com/penguineering/hareairis/ai/RateLimitExceptionTest.java

@@ -0,0 +1,132 @@
+package com.penguineering.hareairis.ai;
+
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpResponse;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import reactor.core.publisher.Mono;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.function.Consumer;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+@ExtendWith(MockitoExtension.class)
+class RateLimitExceptionTest {
+
+    @Test
+    void testConstructorWithMessage() {
+        String message = "Rate limit exceeded";
+        RateLimitException exception = new RateLimitException(message);
+
+        assertEquals(message, exception.getMessage(), "The exception message should match the provided message.");
+        assertTrue(exception.getRetryAfter().isEmpty(), "The retryAfter should be empty when not provided.");
+    }
+
+    @Test
+    void testConstructorWithMessageAndInstant() {
+        String message = "Rate limit exceeded";
+        Instant retryAfter = Instant.now().plus(Duration.ofMinutes(1));
+        RateLimitException exception = new RateLimitException(message, retryAfter);
+
+        assertEquals(message, exception.getMessage(), "The exception message should match the provided message.");
+        assertTrue(exception.getRetryAfter().isPresent(), "The retryAfter should be present when provided.");
+        assertEquals(retryAfter, exception.getRetryAfter().get(), "The retryAfter should match the provided Instant.");
+    }
+
+    @Test
+    void testConstructorWithMessageAndDuration() {
+        String message = "Rate limit exceeded";
+        Duration duration = Duration.ofMinutes(1);
+        Instant beforeCreation = Instant.now();
+        RateLimitException exception = new RateLimitException(message, duration);
+        Instant afterCreation = Instant.now();
+
+        assertEquals(message, exception.getMessage(), "The exception message should match the provided message.");
+        assertTrue(exception.getRetryAfter().isPresent(), "The retryAfter should be present when duration is provided.");
+        Instant retryAfter = exception.getRetryAfter().get();
+        assertTrue(retryAfter.isAfter(beforeCreation), "The retryAfter should be after the time before creation.");
+        assertTrue(retryAfter.isBefore(afterCreation.plus(duration)), "The retryAfter should be within the expected duration range.");
+    }
+
+
+    @Test
+    void testFromHttpResponseWithoutRetryAfterHeader() {
+        HttpResponse response = mock(HttpResponse.class);
+        when(response.getStatusCode()).thenReturn(429);
+        when(response.getBodyAsString()).thenReturn(Mono.just("Rate limit exceeded"));
+        when(response.getHeaders()).thenReturn(new HttpHeaders());
+
+        RateLimitException exception = RateLimitException.fromHttpResponse(response, System.out::println);
+
+        assertEquals("Rate limit exceeded", exception.getMessage());
+        assertTrue(exception.getRetryAfter().isEmpty());
+    }
+
+    // Apply similar changes to other test methods
+    @Test
+    void testFromHttpResponseWithRetryAfterDuration() {
+        HttpResponse response = mock(HttpResponse.class);
+        when(response.getStatusCode()).thenReturn(429);
+        when(response.getBodyAsString()).thenReturn(Mono.just("Rate limit exceeded"));
+        HttpHeaders headers = new HttpHeaders().set(HttpHeaderName.RETRY_AFTER, "60");
+        when(response.getHeaders()).thenReturn(headers);
+
+        RateLimitException exception = RateLimitException.fromHttpResponse(response, System.out::println);
+
+        assertEquals("Rate limit exceeded", exception.getMessage());
+        assertTrue(exception.getRetryAfter().isPresent());
+        assertTrue(exception.getRetryAfter().get().isAfter(Instant.now().minusSeconds(60)));
+    }
+
+    @Test
+    void testFromHttpResponseWithRetryAfterInstant() {
+        HttpResponse response = mock(HttpResponse.class);
+        when(response.getStatusCode()).thenReturn(429);
+        when(response.getBodyAsString()).thenReturn(Mono.just("Rate limit exceeded"));
+        Instant retryAfter = Instant.now().plus(Duration.ofMinutes(1));
+        HttpHeaders headers = new HttpHeaders().set(HttpHeaderName.RETRY_AFTER, retryAfter.toString());
+        when(response.getHeaders()).thenReturn(headers);
+
+        RateLimitException exception = RateLimitException.fromHttpResponse(response, System.out::println);
+
+        assertEquals("Rate limit exceeded", exception.getMessage());
+        assertTrue(exception.getRetryAfter().isPresent());
+        assertEquals(retryAfter, exception.getRetryAfter().get());
+    }
+
+    @Mock
+    Consumer<String> logConsumer;
+
+    @Test
+    void testFromHttpResponseWithInvalidRetryAfterHeader() {
+        HttpResponse response = mock(HttpResponse.class);
+        when(response.getStatusCode()).thenReturn(429);
+        when(response.getBodyAsString()).thenReturn(Mono.just("Rate limit exceeded"));
+        HttpHeaders headers = new HttpHeaders().set(HttpHeaderName.RETRY_AFTER, "invalid");
+        when(response.getHeaders()).thenReturn(headers);
+
+        RateLimitException exception = RateLimitException.fromHttpResponse(response, logConsumer);
+
+        assertEquals("Rate limit exceeded", exception.getMessage());
+        assertTrue(exception.getRetryAfter().isEmpty());
+        verify(logConsumer).accept("Failed to parse Retry-After header and returning an exception with message only: invalid");
+    }
+
+    @Test
+    void testFromHttpResponseWithNon429StatusCode() {
+        HttpResponse response = mock(HttpResponse.class);
+        when(response.getStatusCode()).thenReturn(400); // Any status code other than 429
+
+        IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> {
+            throw RateLimitException.fromHttpResponse(response, logConsumer);
+        });
+
+        assertEquals("Expected status code 429 for RateLimitException, but received " + response.getStatusCode() + ".", exception.getMessage());
+    }
+}