From 2b7c901951bfaab0548811c151962df5d0598d1b Mon Sep 17 00:00:00 2001 From: reanon <85157-reanon@users.noreply.gitgud.io> Date: Fri, 4 Apr 2025 02:48:25 +0000 Subject: [PATCH] Update 3 files - /src/shared/key-management/xai/checker.ts - /src/shared/key-management/key-pool.ts - /src/service-info.ts --- src/service-info.ts | 5 ++- src/shared/key-management/key-pool.ts | 7 +-- src/shared/key-management/xai/checker.ts | 55 ++++++++++++++---------- 3 files changed, 40 insertions(+), 27 deletions(-) diff --git a/src/service-info.ts b/src/service-info.ts index 281172c..91d6aaa 100644 --- a/src/service-info.ts +++ b/src/service-info.ts @@ -481,6 +481,9 @@ function getInfoForFamily(family: ModelFamily): BaseFamilyInfo { break; case "deepseek": info.overQuotaKeys = familyStats.get(`${family}__overQuota`) || 0; + case "xai": + info.overQuotaKeys = familyStats.get(`${family}__overQuota`) || 0; + break; } } @@ -505,4 +508,4 @@ function getQueueInformation(partition: ModelFamily) { proomptersInQueue: getQueueLength(partition), estimatedQueueTime: waitMs > 2000 ? waitTime : "no wait", }; -} +} \ No newline at end of file diff --git a/src/shared/key-management/key-pool.ts b/src/shared/key-management/key-pool.ts index ec8566b..0ca6955 100644 --- a/src/shared/key-management/key-pool.ts +++ b/src/shared/key-management/key-pool.ts @@ -16,7 +16,7 @@ import { MistralAIKeyProvider } from "./mistral-ai/provider"; import { DeepseekKeyProvider } from "./deepseek/provider"; import { XaiKeyProvider } from "./xai/provider"; -type AllowedPartial = OpenAIKeyUpdate | AnthropicKeyUpdate | Partial; +type AllowedPartial = OpenAIKeyUpdate | AnthropicKeyUpdate | Partial | Partial; export class KeyPool { private keyProviders: KeyProvider[] = []; @@ -74,7 +74,8 @@ export class KeyPool { if ( service instanceof OpenAIKeyProvider || service instanceof AnthropicKeyProvider || - service instanceof DeepseekKeyProvider + service instanceof DeepseekKeyProvider || + service instanceof XaiKeyProvider ) { service.update(key.hash, { isOverQuota: reason === "quota" }); } @@ -197,4 +198,4 @@ export class KeyPool { ); this.recheckJobs.openai = job; } -} +} \ No newline at end of file diff --git a/src/shared/key-management/xai/checker.ts b/src/shared/key-management/xai/checker.ts index c0f3e4f..4ef269a 100644 --- a/src/shared/key-management/xai/checker.ts +++ b/src/shared/key-management/xai/checker.ts @@ -7,23 +7,36 @@ const CHECK_TIMEOUT = 10000; export class XaiKeyChecker { private log = logger.child({ module: "key-checker", service: "xai" }); - constructor(private readonly update: (hash: string, key: Partial) => void) {} + constructor(private readonly update: (hash: string, key: Partial) => void) { + this.log.info("XaiKeyChecker initialized"); + } public async checkKey(key: XaiKey): Promise { + this.log.info({ hash: key.hash }, "Starting key validation check"); try { const result = await this.validateKey(key); this.handleCheckResult(key, result); } catch (error) { - this.log.warn( - { error, hash: key.hash }, - "Failed to check key status" - ); + if (error instanceof Error) { + this.log.warn( + { error: error.message, stack: error.stack, hash: key.hash }, + "Failed to check key status" + ); + } else { + this.log.warn( + { error, hash: key.hash }, + "Failed to check key status with unknown error" + ); + } } } private async validateKey(key: XaiKey): Promise<"valid" | "invalid" | "quota"> { const controller = new AbortController(); - const timeout = setTimeout(() => controller.abort(), CHECK_TIMEOUT); + const timeout = setTimeout(() => { + controller.abort(); + this.log.warn({ hash: key.hash }, "Key validation timed out after " + CHECK_TIMEOUT + "ms"); + }, CHECK_TIMEOUT); try { const response = await fetch("https://api.x.ai/v1/chat/completions", { @@ -40,32 +53,27 @@ export class XaiKeyChecker { signal: controller.signal, }); - const rateLimit = { - limit: parseInt(response.headers.get("x-ratelimit-limit") || "200"), - remaining: parseInt(response.headers.get("x-ratelimit-remaining") || "199"), - }; - switch (response.status) { case 400: - this.log.debug( - { key: key.hash, rateLimit }, - "Key check successful, updating rate limit info" - ); return "valid"; case 401: return "invalid"; - case 402: - return "quota"; + case 403: + return "invalid"; case 429: - this.log.warn({ key: key.hash }, "Key is rate limited"); - return "valid"; + return "quota"; default: this.log.warn( { status: response.status, hash: key.hash }, "Unexpected status code while checking key" ); - return "valid"; + return "invalid"; } + } catch (error) { + if (error instanceof Error && error.name === 'AbortError') { + this.log.warn({ hash: key.hash }, "Key validation aborted"); + } + throw error; } finally { clearTimeout(timeout); } @@ -77,13 +85,14 @@ export class XaiKeyChecker { ): void { switch (result) { case "valid": + this.log.info({ hash: key.hash }, "Key is valid and enabled"); this.update(key.hash, { isDisabled: false, lastChecked: Date.now(), }); break; case "invalid": - this.log.warn({ hash: key.hash }, "Key is invalid"); + this.log.warn({ hash: key.hash }, "Key is invalid, marking as revoked"); this.update(key.hash, { isDisabled: true, isRevoked: true, @@ -91,7 +100,7 @@ export class XaiKeyChecker { }); break; case "quota": - this.log.warn({ hash: key.hash }, "Key has exceeded its quota"); + this.log.warn({ hash: key.hash }, "Key has exceeded its quota, disabling"); this.update(key.hash, { isDisabled: true, isOverQuota: true, @@ -102,4 +111,4 @@ export class XaiKeyChecker { assertNever(result); } } -} \ No newline at end of file +}