From a2d64e281e5fc6c829df7b56243880d1437f2ada Mon Sep 17 00:00:00 2001 From: nai-degen Date: Sat, 10 Aug 2024 14:42:52 -0500 Subject: [PATCH] minor KeyProvider#getLockoutPeriod refactor --- .../key-management/anthropic/provider.ts | 23 ++-------------- src/shared/key-management/azure/provider.ts | 27 ++----------------- .../key-management/google-ai/provider.ts | 23 ++-------------- src/shared/key-management/index.ts | 25 +++++++++++++++++ .../key-management/mistral-ai/provider.ts | 23 ++-------------- src/shared/key-management/openai/provider.ts | 3 +-- 6 files changed, 34 insertions(+), 90 deletions(-) diff --git a/src/shared/key-management/anthropic/provider.ts b/src/shared/key-management/anthropic/provider.ts index 057686e..bcd8dfb 100644 --- a/src/shared/key-management/anthropic/provider.ts +++ b/src/shared/key-management/anthropic/provider.ts @@ -1,5 +1,5 @@ import crypto from "crypto"; -import { Key, KeyProvider } from ".."; +import { createGenericGetLockoutPeriod, Key, KeyProvider } from ".."; import { config } from "../../../config"; import { logger } from "../../../logger"; import { AnthropicModelFamily, getClaudeModelFamily } from "../../models"; @@ -23,10 +23,6 @@ type AnthropicKeyUsage = { export interface AnthropicKey extends Key, AnthropicKeyUsage { readonly service: "anthropic"; readonly modelFamilies: AnthropicModelFamily[]; - /** The time at which this key was last rate limited. */ - rateLimitedAt: number; - /** The time until which this key is rate limited. */ - rateLimitedUntil: number; /** * Whether this key requires a special preamble. For unclear reasons, some * Anthropic keys will throw an error if the prompt does not begin with a @@ -217,22 +213,7 @@ export class AnthropicKeyProvider implements KeyProvider { key[`${getClaudeModelFamily(model)}Tokens`] += tokens; } - public getLockoutPeriod() { - const activeKeys = this.keys.filter((k) => !k.isDisabled); - // Don't lock out if there are no keys available or the queue will stall. - // Just let it through so the add-key middleware can throw an error. - if (activeKeys.length === 0) return 0; - - const now = Date.now(); - const rateLimitedKeys = activeKeys.filter((k) => now < k.rateLimitedUntil); - const anyNotRateLimited = rateLimitedKeys.length < activeKeys.length; - - if (anyNotRateLimited) return 0; - - // If all keys are rate-limited, return the time until the first key is - // ready. - return Math.min(...activeKeys.map((k) => k.rateLimitedUntil - now)); - } + getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys); /** * This is called when we receive a 429, which means there are already five diff --git a/src/shared/key-management/azure/provider.ts b/src/shared/key-management/azure/provider.ts index 681a6ed..8a7e48a 100644 --- a/src/shared/key-management/azure/provider.ts +++ b/src/shared/key-management/azure/provider.ts @@ -1,5 +1,5 @@ import crypto from "crypto"; -import { Key, KeyProvider } from ".."; +import { createGenericGetLockoutPeriod, Key, KeyProvider } from ".."; import { config } from "../../../config"; import { PaymentRequiredError } from "../../errors"; import { logger } from "../../../logger"; @@ -14,10 +14,6 @@ type AzureOpenAIKeyUsage = { export interface AzureOpenAIKey extends Key, AzureOpenAIKeyUsage { readonly service: "azure"; readonly modelFamilies: AzureOpenAIModelFamily[]; - /** The time at which this key was last rate limited. */ - rateLimitedAt: number; - /** The time until which this key is rate limited. */ - rateLimitedUntil: number; contentFiltering: boolean; } @@ -156,26 +152,7 @@ export class AzureOpenAIKeyProvider implements KeyProvider { key[`${getAzureOpenAIModelFamily(model)}Tokens`] += tokens; } - // TODO: all of this shit is duplicate code - - public getLockoutPeriod(family: AzureOpenAIModelFamily) { - const activeKeys = this.keys.filter( - (key) => !key.isDisabled && key.modelFamilies.includes(family) - ); - - // Don't lock out if there are no keys available or the queue will stall. - // Just let it through so the add-key middleware can throw an error. - if (activeKeys.length === 0) return 0; - - const now = Date.now(); - const rateLimitedKeys = activeKeys.filter((k) => now < k.rateLimitedUntil); - const anyNotRateLimited = rateLimitedKeys.length < activeKeys.length; - - if (anyNotRateLimited) return 0; - - // If all keys are rate-limited, return time until the first key is ready. - return Math.min(...activeKeys.map((k) => k.rateLimitedUntil - now)); - } + getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys); /** * This is called when we receive a 429, which means there are already five diff --git a/src/shared/key-management/google-ai/provider.ts b/src/shared/key-management/google-ai/provider.ts index 47695f5..e94460b 100644 --- a/src/shared/key-management/google-ai/provider.ts +++ b/src/shared/key-management/google-ai/provider.ts @@ -1,5 +1,5 @@ import crypto from "crypto"; -import { Key, KeyProvider } from ".."; +import { createGenericGetLockoutPeriod, Key, KeyProvider } from ".."; import { config } from "../../../config"; import { logger } from "../../../logger"; import { getGoogleAIModelFamily, type GoogleAIModelFamily } from "../../models"; @@ -28,10 +28,6 @@ type GoogleAIKeyUsage = { export interface GoogleAIKey extends Key, GoogleAIKeyUsage { readonly service: "google-ai"; readonly modelFamilies: GoogleAIModelFamily[]; - /** The time at which this key was last rate limited. */ - rateLimitedAt: number; - /** The time until which this key is rate limited. */ - rateLimitedUntil: number; /** All detected model IDs on this key. */ modelIds: string[]; } @@ -162,22 +158,7 @@ export class GoogleAIKeyProvider implements KeyProvider { key[`${getGoogleAIModelFamily(model)}Tokens`] += tokens; } - public getLockoutPeriod() { - const activeKeys = this.keys.filter((k) => !k.isDisabled); - // Don't lock out if there are no keys available or the queue will stall. - // Just let it through so the add-key middleware can throw an error. - if (activeKeys.length === 0) return 0; - - const now = Date.now(); - const rateLimitedKeys = activeKeys.filter((k) => now < k.rateLimitedUntil); - const anyNotRateLimited = rateLimitedKeys.length < activeKeys.length; - - if (anyNotRateLimited) return 0; - - // If all keys are rate-limited, return the time until the first key is - // ready. - return Math.min(...activeKeys.map((k) => k.rateLimitedUntil - now)); - } + getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys); /** * This is called when we receive a 429, which means there are already five diff --git a/src/shared/key-management/index.ts b/src/shared/key-management/index.ts index 67dfad4..c1b3f6a 100644 --- a/src/shared/key-management/index.ts +++ b/src/shared/key-management/index.ts @@ -30,6 +30,10 @@ export interface Key { lastChecked: number; /** Hash of the key, for logging and to find the key in the pool. */ hash: string; + /** The time at which this key was last rate limited. */ + rateLimitedAt: number; + /** The time until which this key is rate limited. */ + rateLimitedUntil: number; } /* @@ -58,6 +62,27 @@ export interface KeyProvider { recheck(): void; } +export function createGenericGetLockoutPeriod( + getKeys: () => T[] +) { + return function (this: unknown, family?: ModelFamily): number { + const keys = getKeys(); + const activeKeys = keys.filter( + (k) => !k.isDisabled && (!family || k.modelFamilies.includes(family)) + ); + + if (activeKeys.length === 0) return 0; + + const now = Date.now(); + const rateLimitedKeys = activeKeys.filter((k) => now < k.rateLimitedUntil); + const anyNotRateLimited = rateLimitedKeys.length < activeKeys.length; + + if (anyNotRateLimited) return 0; + + return Math.min(...activeKeys.map((k) => k.rateLimitedUntil - now)); + }; +} + export const keyPool = new KeyPool(); export { AnthropicKey } from "./anthropic/provider"; export { OpenAIKey } from "./openai/provider"; diff --git a/src/shared/key-management/mistral-ai/provider.ts b/src/shared/key-management/mistral-ai/provider.ts index 83785f8..20f9dae 100644 --- a/src/shared/key-management/mistral-ai/provider.ts +++ b/src/shared/key-management/mistral-ai/provider.ts @@ -1,5 +1,5 @@ import crypto from "crypto"; -import { Key, KeyProvider } from ".."; +import { createGenericGetLockoutPeriod, Key, KeyProvider } from ".."; import { config } from "../../../config"; import { logger } from "../../../logger"; import { MistralAIModelFamily, getMistralAIModelFamily } from "../../models"; @@ -13,10 +13,6 @@ type MistralAIKeyUsage = { export interface MistralAIKey extends Key, MistralAIKeyUsage { readonly service: "mistral-ai"; readonly modelFamilies: MistralAIModelFamily[]; - /** The time at which this key was last rate limited. */ - rateLimitedAt: number; - /** The time until which this key is rate limited. */ - rateLimitedUntil: number; } /** @@ -150,22 +146,7 @@ export class MistralAIKeyProvider implements KeyProvider { key[`${family}Tokens`] += tokens; } - public getLockoutPeriod() { - const activeKeys = this.keys.filter((k) => !k.isDisabled); - // Don't lock out if there are no keys available or the queue will stall. - // Just let it through so the add-key middleware can throw an error. - if (activeKeys.length === 0) return 0; - - const now = Date.now(); - const rateLimitedKeys = activeKeys.filter((k) => now < k.rateLimitedUntil); - const anyNotRateLimited = rateLimitedKeys.length < activeKeys.length; - - if (anyNotRateLimited) return 0; - - // If all keys are rate-limited, return the time until the first key is - // ready. - return Math.min(...activeKeys.map((k) => k.rateLimitedUntil - now)); - } + getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys); /** * This is called when we receive a 429, which means there are already five diff --git a/src/shared/key-management/openai/provider.ts b/src/shared/key-management/openai/provider.ts index 809262c..528f029 100644 --- a/src/shared/key-management/openai/provider.ts +++ b/src/shared/key-management/openai/provider.ts @@ -26,8 +26,6 @@ export interface OpenAIKey extends Key, OpenAIKeyUsage { isTrial: boolean; /** Set when key check returns a non-transient 429. */ isOverQuota: boolean; - /** The time at which this key was last rate limited. */ - rateLimitedAt: number; /** * Last known X-RateLimit-Requests-Reset header from OpenAI, converted to a * number. @@ -111,6 +109,7 @@ export class OpenAIKeyProvider implements KeyProvider { .digest("hex") .slice(0, 8)}`, rateLimitedAt: 0, + rateLimitedUntil: 0, rateLimitRequestsReset: 0, rateLimitTokensReset: 0, turboTokens: 0,