From a6c6e4c694ffe195d290155a26dac85e2b1752b8 Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Thu, 1 Jun 2023 08:31:41 -0500
Subject: [PATCH] allows KeyPool.available to be scoped per-service

---
 src/key-management/key-pool.ts         | 15 ++++++---------
 src/proxy/middleware/response/index.ts |  2 +-
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/src/key-management/key-pool.ts b/src/key-management/key-pool.ts
index 5730e53..acdd154 100644
--- a/src/key-management/key-pool.ts
+++ b/src/key-management/key-pool.ts
@@ -13,7 +13,7 @@ export class KeyPool {
 
   public init() {
     this.keyProviders.forEach((provider) => provider.init());
-    const availableKeys = this.available();
+    const availableKeys = this.available("all");
     if (availableKeys === 0) {
       throw new Error(
         "No keys loaded. Ensure either OPENAI_KEY or ANTHROPIC_KEY is set."
@@ -35,14 +35,11 @@ export class KeyPool {
     service.disable(key);
   }
 
-  // TODO: this probably needs to be scoped to a specific provider. I think the
-  // only code calling this is the error handler which needs to know how many
-  // more keys are available for the provider the user tried to use.
-  public available(): number {
-    return this.keyProviders.reduce(
-      (sum, provider) => sum + provider.available(),
-      0
-    );
+  public available(service: AIService | "all" = "all"): number {
+    return this.keyProviders.reduce((sum, provider) => {
+      const includeProvider = service === "all" || service === provider.service;
+      return sum + (includeProvider ? provider.available() : 0);
+    }, 0);
   }
 
   public anyUnchecked(): boolean {
diff --git a/src/proxy/middleware/response/index.ts b/src/proxy/middleware/response/index.ts
index 0980424..95210a2 100644
--- a/src/proxy/middleware/response/index.ts
+++ b/src/proxy/middleware/response/index.ts
@@ -223,7 +223,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
   let errorPayload: Record<string, any>;
   // Subtract 1 from available keys because if this message is being shown,
   // it's because the key is about to be disabled.
-  const availableKeys = keyPool.available() - 1;
+  const availableKeys = keyPool.available(req.outboundApi) - 1;
   const tryAgainMessage = Boolean(availableKeys)
     ? `There are ${availableKeys} more keys available; try your request again.`
     : "There are no more keys available.";