From 45576db441cdce96a233634a678bbf1e5f74a4ab Mon Sep 17 00:00:00 2001
From: user <user@example.com>
Date: Sat, 25 Jan 2025 11:59:41 +0000
Subject: [PATCH] [Deepseek] Properly handle over-quota keys

---
 src/service-info.ts                            | 11 ++++++++++-
 src/shared/key-management/deepseek/checker.ts  |  1 +
 src/shared/key-management/deepseek/provider.ts |  7 +++++++
 src/shared/key-management/key-pool.ts          |  3 ++-
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/service-info.ts b/src/service-info.ts
index e8df550..71126fe 100644
--- a/src/service-info.ts
+++ b/src/service-info.ts
@@ -2,6 +2,7 @@ import { config, listConfig } from "./config";
 import {
   AnthropicKey,
   AwsBedrockKey,
+  DeepseekKey,
   GcpKey,
   keyPool,
   OpenAIKey,
@@ -35,6 +36,8 @@ const keyIsAnthropicKey = (k: KeyPoolKey): k is AnthropicKey =>
   k.service === "anthropic";
 const keyIsAwsKey = (k: KeyPoolKey): k is AwsBedrockKey => k.service === "aws";
 const keyIsGcpKey = (k: KeyPoolKey): k is GcpKey => k.service === "gcp";
+const keyIsDeepseekKey = (k: KeyPoolKey): k is DeepseekKey =>
+  k.service === "deepseek";
 
 /** Stats aggregated across all keys for a given service. */
 type ServiceAggregate = "keys" | "uncheckedKeys" | "orgs";
@@ -384,7 +387,11 @@ function addKeyToAggregates(k: KeyPoolKey) {
     case "google-ai":
     case "mistral-ai":
     case "deepseek":
-      k.modelFamilies.forEach(incrementGenericFamilyStats);
+      if (!keyIsDeepseekKey(k)) throw new Error("Invalid key type");
+      k.modelFamilies.forEach((f) => {
+        incrementGenericFamilyStats(f);
+        addToFamily(`${f}__overQuota`, k.isOverQuota ? 1 : 0);
+      });
       break;
     default:
       assertNever(k.service);
@@ -452,6 +459,8 @@ function getInfoForFamily(family: ModelFamily): BaseFamilyInfo {
           info.enabledVariants = "not implemented";
         }
         break;
+      case "deepseek":
+        info.overQuotaKeys = familyStats.get(`${family}__overQuota`) || 0;
     }
   }
 
diff --git a/src/shared/key-management/deepseek/checker.ts b/src/shared/key-management/deepseek/checker.ts
index d759e09..2195800 100644
--- a/src/shared/key-management/deepseek/checker.ts
+++ b/src/shared/key-management/deepseek/checker.ts
@@ -94,6 +94,7 @@ export class DeepseekKeyChecker {
         this.log.warn({ hash: key.hash }, "Key has exceeded its quota");
         this.update(key.hash, {
           isDisabled: true,
+          isOverQuota: true,
           lastChecked: Date.now(),
         });
         break;
diff --git a/src/shared/key-management/deepseek/provider.ts b/src/shared/key-management/deepseek/provider.ts
index 2adba08..e79802c 100644
--- a/src/shared/key-management/deepseek/provider.ts
+++ b/src/shared/key-management/deepseek/provider.ts
@@ -11,6 +11,7 @@ type DeepseekKeyUsage = {
 export interface DeepseekKey extends Key, DeepseekKeyUsage {
   readonly service: "deepseek";
   readonly modelFamilies: DeepseekModelFamily[];
+  isOverQuota: boolean;
 }
 
 export class DeepseekKeyProvider implements KeyProvider<DeepseekKey> {
@@ -42,6 +43,7 @@ export class DeepseekKeyProvider implements KeyProvider<DeepseekKey> {
         rateLimitedAt: 0,
         rateLimitedUntil: 0,
         "deepseekTokens": 0,
+        isOverQuota: false,
       });
     }
   }
@@ -130,6 +132,11 @@ export class DeepseekKeyProvider implements KeyProvider<DeepseekKey> {
   public recheck(): void {
     if (!this.checker || !config.checkKeys) return;
     for (const key of this.keys) {
+      this.update(key.hash, { 
+        isOverQuota: false,
+        isDisabled: false,
+        lastChecked: 0 
+      });
       void this.checker.checkKey(key);
     }
   }
diff --git a/src/shared/key-management/key-pool.ts b/src/shared/key-management/key-pool.ts
index 5c6db18..80f02cc 100644
--- a/src/shared/key-management/key-pool.ts
+++ b/src/shared/key-management/key-pool.ts
@@ -71,7 +71,8 @@ export class KeyPool {
     service.update(key.hash, { isRevoked: reason === "revoked" });
     if (
       service instanceof OpenAIKeyProvider ||
-      service instanceof AnthropicKeyProvider
+      service instanceof AnthropicKeyProvider ||
+      service instanceof DeepseekKeyProvider
     ) {
       service.update(key.hash, { isOverQuota: reason === "quota" });
     }