From 871b52ca9596e87001eea81b4fe68c14d76627fa Mon Sep 17 00:00:00 2001
From: John Smith <john@johnjcsmith.com>
Date: Sun, 8 Dec 2024 14:46:18 +1030
Subject: [PATCH] feat: Retry model provider throttling

---
 .../src/modules/models/index.test.ts          | 44 +++++++++++++++++++
 control-plane/src/utilities/errors.test.ts    | 14 ++++++
 control-plane/src/utilities/errors.ts         | 24 +++++++---
 3 files changed, 77 insertions(+), 5 deletions(-)
 create mode 100644 control-plane/src/utilities/errors.test.ts

diff --git a/control-plane/src/modules/models/index.test.ts b/control-plane/src/modules/models/index.test.ts
index f49d6486..3b96a5bd 100644
--- a/control-plane/src/modules/models/index.test.ts
+++ b/control-plane/src/modules/models/index.test.ts
@@ -53,6 +53,7 @@ describe("buildModel", () => {
       });
     });
 
+
     it("should not retry other errors", async () => {
       const error = new Error("");
       mockCreate.mockImplementationOnce(() => {
@@ -77,5 +78,48 @@ describe("buildModel", () => {
         identifier: "claude-3-haiku",
       });
     });
+
+    it.skip("should throw after exhausting retries", async () => {
+      mockCreate.mockImplementation(() => {
+        throw new RetryableError("");
+      });
+
+      const model = buildModel({
+        identifier: "claude-3-haiku",
+      });
+
+      await expect(
+        () => model.call({
+          messages: [],
+        })
+      ).rejects.toThrow(RetryableError);
+
+      expect(getRouting).toHaveBeenCalledTimes(6);
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 0,
+        identifier: "claude-3-haiku",
+      });
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 1,
+        identifier: "claude-3-haiku",
+      });
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 2,
+        identifier: "claude-3-haiku",
+      });
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 3,
+        identifier: "claude-3-haiku",
+      });
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 4,
+        identifier: "claude-3-haiku",
+      });
+    }, 60_000);
   });
 });
diff --git a/control-plane/src/utilities/errors.test.ts b/control-plane/src/utilities/errors.test.ts
new file mode 100644
index 00000000..61b4bbc1
--- /dev/null
+++ b/control-plane/src/utilities/errors.test.ts
@@ -0,0 +1,14 @@
+import { isRetryableError, RetryableError } from "./errors";
+import { RateLimitError } from "@anthropic-ai/sdk";
+
+describe("isRetryableError", () => {
+  it("should return true for retryable errors", async () => {
+    expect(isRetryableError(new Error())).toBe(false);
+
+    expect(isRetryableError(new RateLimitError(429, new Error(), "", {}))).toBe(true);
+    expect(isRetryableError(new RetryableError(""))).toBe(true);
+
+
+    expect(isRetryableError(new Error("429 Too many requests, please wait before trying again."))).toBe(true);
+  });
+})
diff --git a/control-plane/src/utilities/errors.ts b/control-plane/src/utilities/errors.ts
index 9aac3127..1e10765b 100644
--- a/control-plane/src/utilities/errors.ts
+++ b/control-plane/src/utilities/errors.ts
@@ -1,3 +1,5 @@
+import { RateLimitError, InternalServerError } from "@anthropic-ai/sdk";
+
 export class RetryableError extends Error {
   constructor(message: string) {
     super(message);
@@ -5,18 +7,30 @@ export class RetryableError extends Error {
   }
 }
 
+const retryableErrors = [RetryableError, RateLimitError, InternalServerError]
 const retryableErrorMessages = [
-  "Connection terminated due to connection timeout",
+  // DB Connection Errors
+  "connection terminated due to connection timeout",
   "timeout exceeded when trying to connect",
-  "Connection terminated unexpectedly",
-  "account does not have an agreement to this model",
+  "connection terminated unexpectedly",
+  // DB Connection Pool Exhaustion
+  "remaining connection slots are reserved for roles with the SUPERUSER attribute",
+  "too many clients already",
+  // Bedrock Errors
+  "503 bedrock is unable to process your request",
+  "429 too many requests"
 ];
 
 export const isRetryableError = (error: unknown) => {
-  if (error instanceof Error && retryableErrorMessages.includes(error.message))
+  if (error instanceof Error && retryableErrorMessages.find((message) => error.message.toLowerCase().includes(message))) {
+    return true
+  }
+
+  if (error instanceof Error && retryableErrors.find((type) => error instanceof type)) {
     return true;
+  }
 
-  return error instanceof RetryableError;
+  return false;
 };
 
 export class AuthenticationError extends Error {