[Deepseek] Properly handle over-quota keys

This commit is contained in:
user
2025-01-25 11:59:41 +00:00
parent bcc83f30d9
commit 45576db441
4 changed files with 20 additions and 2 deletions
+10 -1
View File
@@ -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;
}
}
@@ -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;
@@ -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);
}
}
+2 -1
View File
@@ -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" });
}