From b3dba5f6767a3c1edc6fe5cc97f629084702c573 Mon Sep 17 00:00:00 2001
From: Sam Becker <sam@sambecker.com>
Date: Sun, 20 Apr 2025 11:48:50 -0500
Subject: [PATCH] Increase AI rate limit window for batch requests

---
 src/photo/actions.ts    |  5 +++--
 src/photo/ai/server.ts  |  6 ++++++
 src/platforms/openai.ts | 24 +++++++++++++++---------
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/photo/actions.ts b/src/photo/actions.ts
index 3ec326bd..0768c850 100644
--- a/src/photo/actions.ts
+++ b/src/photo/actions.ts
@@ -374,7 +374,7 @@ export const getExifDataAction = async (
 // - strip GPS data if necessary
 // - update blur data (or destroy if blur is disabled)
 // - generate AI text data, if enabled, and auto-generated fields are empty
-export const syncPhotoAction = async (photoId: string) =>
+export const syncPhotoAction = async (photoId: string, isBatch?: boolean) =>
   runAuthenticatedAdminServerAction(async () => {
     const photo = await getPhoto(photoId ?? '', true);
 
@@ -415,6 +415,7 @@ export const syncPhotoAction = async (photoId: string) =>
         } = await generateAiImageQueries(
           imageResizedBase64,
           photo.syncStatus.missingAiTextFields,
+          isBatch,
         );
 
         const formDataFromPhoto = convertPhotoToFormData(photo);
@@ -451,7 +452,7 @@ export const syncPhotoAction = async (photoId: string) =>
 export const syncPhotosAction = async (photoIds: string[]) =>
   runAuthenticatedAdminServerAction(async () => {
     for (const photoId of photoIds) {
-      await syncPhotoAction(photoId);
+      await syncPhotoAction(photoId, true);
     }
     revalidateAllKeysAndPaths();
   });
diff --git a/src/photo/ai/server.ts b/src/photo/ai/server.ts
index e1ac4130..d053a355 100644
--- a/src/photo/ai/server.ts
+++ b/src/photo/ai/server.ts
@@ -9,6 +9,7 @@ import { getUniqueTags } from '../db/query';
 export const generateAiImageQueries = async (
   imageBase64?: string,
   textFieldsToGenerate: AiAutoGeneratedField[] = [],
+  isBatch?: boolean,
 ): Promise<{
   title?: string
   caption?: string
@@ -31,6 +32,7 @@ export const generateAiImageQueries = async (
         const titleAndCaption = await generateOpenAiImageQuery(
           imageBase64,
           getAiImageQuery('title-and-caption'),
+          isBatch,
         );
         if (titleAndCaption) {
           const titleAndCaptionParsed = parseTitleAndCaption(titleAndCaption);
@@ -42,12 +44,14 @@ export const generateAiImageQueries = async (
           title = await generateOpenAiImageQuery(
             imageBase64,
             getAiImageQuery('title'),
+            isBatch,
           );
         }
         if (textFieldsToGenerate.includes('caption')) {
           caption =  await generateOpenAiImageQuery(
             imageBase64,
             getAiImageQuery('caption'),
+            isBatch,
           );
         }
       }
@@ -57,6 +61,7 @@ export const generateAiImageQueries = async (
         tags = await generateOpenAiImageQuery(
           imageBase64,
           getAiImageQuery('tags', existingTags),
+          isBatch,
         );
       }
   
@@ -64,6 +69,7 @@ export const generateAiImageQueries = async (
         semanticDescription = await generateOpenAiImageQuery(
           imageBase64,
           getAiImageQuery('description-small'),
+          isBatch,
         );
       }
     }
diff --git a/src/platforms/openai.ts b/src/platforms/openai.ts
index 23fb55c5..64f1d7a0 100644
--- a/src/platforms/openai.ts
+++ b/src/platforms/openai.ts
@@ -13,7 +13,6 @@ import { cleanUpAiTextResponse } from '@/photo/ai';
 const redis = HAS_REDIS_STORAGE ? Redis.fromEnv() : undefined;
 
 const RATE_LIMIT_IDENTIFIER = 'openai-image-query';
-const RATE_LIMIT_MAX_QUERIES_PER_HOUR = 100;
 const MODEL = 'gpt-4o';
 
 const openai = AI_TEXT_GENERATION_ENABLED
@@ -21,18 +20,24 @@ const openai = AI_TEXT_GENERATION_ENABLED
   : undefined;
 
 const ratelimit = redis
-  ? new Ratelimit({
-    redis,
-    limiter: Ratelimit.slidingWindow(RATE_LIMIT_MAX_QUERIES_PER_HOUR, '1h'),
-  })
+  ? {
+    basic: new Ratelimit({
+      redis,
+      limiter: Ratelimit.slidingWindow(100, '1h'),
+    }),
+    batch: new Ratelimit({
+      redis,
+      limiter: Ratelimit.slidingWindow(1200, '1d'),
+    }),
+  }
   : undefined;
 
-// Allows 100 requests per hour
-const checkRateLimitAndThrow = async () => {
+const checkRateLimitAndThrow = async (isBatch?: boolean) => {
   if (ratelimit) {
     let success = false;
     try {
-      success = (await ratelimit.limit(RATE_LIMIT_IDENTIFIER)).success;
+      const limiter = isBatch ? ratelimit.batch : ratelimit.basic;
+      success = (await limiter.limit(RATE_LIMIT_IDENTIFIER)).success;
     } catch (e: any) {
       console.error('Failed to rate limit OpenAI', e);
       throw new Error('Failed to rate limit OpenAI');
@@ -92,8 +97,9 @@ export const streamOpenAiImageQuery = async (
 export const generateOpenAiImageQuery = async (
   imageBase64: string,
   query: string,
+  isBatch?: boolean,
 ) => {
-  await checkRateLimitAndThrow();
+  await checkRateLimitAndThrow(isBatch);
 
   const args = getImageTextArgs(imageBase64, query);