diff --git a/src/admin/web/manage.ts b/src/admin/web/manage.ts index de52288..3b37a6d 100644 --- a/src/admin/web/manage.ts +++ b/src/admin/web/manage.ts @@ -274,6 +274,7 @@ router.post("/maintenance", (req, res) => { "aws", "gcp", "azure", + "google-ai" ]; checkable.forEach((s) => keyPool.recheck(s)); const keyCount = keyPool diff --git a/src/proxy/middleware/response/index.ts b/src/proxy/middleware/response/index.ts index d0fd678..932ec5c 100644 --- a/src/proxy/middleware/response/index.ts +++ b/src/proxy/middleware/response/index.ts @@ -602,16 +602,37 @@ async function handleGoogleAIBadRequestError( } //{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED"} +// async function handleGoogleAIRateLimitError( req: Request, errorPayload: ProxiedErrorPayload ) { const status = errorPayload.error?.status; + const text = JSON.stringify(errorPayload.error); + + // sometimes they block keys by rate limiting them to 0 requests per minute + // for some indefinite period of time + const keyDeadMsgs = [ + /GenerateContentRequestsPerMinutePerProjectPerRegion/i, + /"quota_limit_value":"0"/i, + ]; + switch (status) { - case "RESOURCE_EXHAUSTED": + case "RESOURCE_EXHAUSTED": { + if (keyDeadMsgs.every((msg) => text.match(msg))) { + req.log.warn( + { key: req.key?.hash, error: text }, + "Google API key appears to be temporarily inoperative and will be disabled." + ); + keyPool.disable(req.key!, "revoked"); + errorPayload.proxy_note = `Assigned API key cannot be used.`; + return; + } + keyPool.markRateLimited(req.key!); await reenqueueRequest(req); throw new RetryableError("Rate-limited request re-enqueued."); + } default: errorPayload.proxy_note = `Unrecognized rate limit error from Google AI (${status}). Please report this.`; break; diff --git a/src/shared/key-management/key-pool.ts b/src/shared/key-management/key-pool.ts index eccde21..cab1eae 100644 --- a/src/shared/key-management/key-pool.ts +++ b/src/shared/key-management/key-pool.ts @@ -178,8 +178,9 @@ export class KeyPool { const job = schedule.scheduleJob(crontab, () => { const next = job.nextInvocation(); - logger.info({ next }, "Performing periodic recheck of OpenAI keys"); + logger.info({ next }, "Performing periodic recheck."); this.recheck("openai"); + this.recheck("google-ai"); }); logger.info( { rule: crontab, next: job.nextInvocation() },