From a3620db5919127fe8ac3544733972fee2a4d92ad Mon Sep 17 00:00:00 2001 From: yukianon Date: Sat, 20 May 2023 19:10:21 +0000 Subject: [PATCH] Implement Test Generation for Trial Keys (khanon/oai-reverse-proxy!12) --- package-lock.json | 18 ++++++++ package.json | 1 + src/key-management/key-checker.ts | 70 +++++++++++++++++++++---------- 3 files changed, 66 insertions(+), 23 deletions(-) diff --git a/package-lock.json b/package-lock.json index ba3875d..e0be5f6 100644 --- a/package-lock.json +++ b/package-lock.json @@ -16,6 +16,7 @@ "firebase-admin": "^11.8.0", "googleapis": "^117.0.0", "http-proxy-middleware": "^3.0.0-beta.1", + "openai": "^3.2.1", "pino": "^8.11.0", "pino-http": "^8.3.3", "showdown": "^2.1.0", @@ -3138,6 +3139,23 @@ "wrappy": "1" } }, + "node_modules/openai": { + "version": "3.2.1", + "resolved": "https://registry.npmjs.org/openai/-/openai-3.2.1.tgz", + "integrity": "sha512-762C9BNlJPbjjlWZi4WYK9iM2tAVAv0uUp1UmI34vb0CN5T2mjB/qM6RYBmNKMh/dN9fC+bxqPwWJZUTWW052A==", + "dependencies": { + "axios": "^0.26.0", + "form-data": "^4.0.0" + } + }, + "node_modules/openai/node_modules/axios": { + "version": "0.26.1", + "resolved": "https://registry.npmjs.org/axios/-/axios-0.26.1.tgz", + "integrity": "sha512-fPwcX4EvnSHuInCMItEhAGnaSEXRBjtzh9fOtsE6E1G6p7vl7edEeZe11QHf18+6+9gR5PbKV/sGKNaD8YaMeA==", + "dependencies": { + "follow-redirects": "^1.14.8" + } + }, "node_modules/optionator": { "version": "0.8.3", "resolved": "https://registry.npmjs.org/optionator/-/optionator-0.8.3.tgz", diff --git a/package.json b/package.json index 487de44..9679c58 100644 --- a/package.json +++ b/package.json @@ -25,6 +25,7 @@ "firebase-admin": "^11.8.0", "googleapis": "^117.0.0", "http-proxy-middleware": "^3.0.0-beta.1", + "openai": "^3.2.1", "pino": "^8.11.0", "pino-http": "^8.3.3", "showdown": "^2.1.0", diff --git a/src/key-management/key-checker.ts b/src/key-management/key-checker.ts index 38a7b73..6b52505 100644 --- a/src/key-management/key-checker.ts +++ b/src/key-management/key-checker.ts @@ -1,19 +1,15 @@ import axios, { AxiosError } from "axios"; +import { Configuration, OpenAIApi } from "openai"; import { logger } from "../logger"; import type { Key, KeyPool } from "./key-pool"; const MIN_CHECK_INTERVAL = 3 * 1000; // 3 seconds const KEY_CHECK_PERIOD = 5 * 60 * 1000; // 5 minutes -const GET_MODELS_URL = "https://api.openai.com/v1/models"; const GET_SUBSCRIPTION_URL = "https://api.openai.com/dashboard/billing/subscription"; const GET_USAGE_URL = "https://api.openai.com/dashboard/billing/usage"; -type GetModelsResponse = { - data: [{ id: string }]; -}; - type GetSubscriptionResponse = { plan: { title: string }; has_payment_method: boolean; @@ -26,6 +22,10 @@ type GetUsageResponse = { total_usage: number; }; +type OpenAIError = { + error: { type: string; code: string; param: unknown; message: string }; +}; + type UpdateFn = typeof KeyPool.prototype.update; export class KeyChecker { @@ -127,6 +127,13 @@ export class KeyChecker { if (isInitialCheck) { const subscription = await this.getSubscription(key); this.updateKey(key.hash, { isTrial: !subscription.has_payment_method }); + if (key.isTrial) { + this.log.debug( + { key: key.hash }, + "Attempting generation on trial key." + ); + await this.assertCanGenerate(key); + } const [provisionedModels, usage] = await Promise.all([ this.getProvisionedModels(key), this.getUsage(key), @@ -175,11 +182,10 @@ export class KeyChecker { private async getProvisionedModels( key: Key ): Promise<{ turbo: boolean; gpt4: boolean }> { - const { data } = await axios.get(GET_MODELS_URL, { - headers: { Authorization: `Bearer ${key.key}` }, - }); - const turbo = data.data.some(({ id }) => id.startsWith("gpt-3.5")); - const gpt4 = data.data.some(({ id }) => id.startsWith("gpt-4")); + const openai = new OpenAIApi(new Configuration({ apiKey: key.key })); + const models = (await openai.listModels()!).data.data; + const turbo = models.some(({ id }) => id.startsWith("gpt-3.5")); + const gpt4 = models.some(({ id }) => id.startsWith("gpt-4")); return { turbo, gpt4 }; } @@ -201,7 +207,7 @@ export class KeyChecker { } private handleAxiosError(key: Key, error: AxiosError) { - if (error.response) { + if (error.response && KeyChecker.errorIsOpenAiError(error)) { const { status, data } = error.response; if (status === 401) { this.log.warn( @@ -209,27 +215,38 @@ export class KeyChecker { "Key is invalid or revoked. Disabling key." ); this.updateKey(key.hash, { isDisabled: true }); + } else if (status === 429 && data.error.type === "insufficient_quota") { + this.log.warn( + { key: key.hash, isTrial: key.isTrial, error: data }, + "Key is out of quota. Disabling key." + ); + this.updateKey(key.hash, { isDisabled: true }); } else { this.log.error( { key: key.hash, status, error: data }, "Encountered API error while checking key." ); } - } else { - this.log.error( - { key: key.hash, error }, - "Network error while checking key." - ); + return; } + this.log.error( + { key: key.hash, error }, + "Network error while checking key; trying again later." + ); } - // TODO: Trial key usage reporting is very unreliable and keys with supposedly - // no usage are already exhausted. Instead we should try generating some text - // on the first check to quickly determine if the key is alive. - private async doTestGeneration(key: Key) { - // Generate only a single token with a very short prompt to avoid using - // too much of the key's quota. - // NYI + /** + * Trial key usage reporting is inaccurate, so we need to run an actual + * completion to test them for liveness. + */ + private async assertCanGenerate(key: Key): Promise { + const openai = new OpenAIApi(new Configuration({ apiKey: key.key })); + // This will throw an AxiosError if the key is invalid or out of quota. + await openai.createChatCompletion({ + model: "gpt-3.5-turbo", + messages: [{ role: "user", content: "Hello" }], + max_tokens: 1, + }); } static getUsageQuerystring(isTrial: boolean) { @@ -251,4 +268,11 @@ export class KeyChecker { endDate.toISOString().split("T")[0] }`; } + + static errorIsOpenAiError( + error: AxiosError + ): error is AxiosError { + const data = error.response?.data as any; + return data?.error?.type; + } }