From f6edaf07f193418e66994632b597e3a22164f936 Mon Sep 17 00:00:00 2001
From: Sam Becker <sam@sambecker.com>
Date: Thu, 21 Mar 2024 22:47:21 -0500
Subject: [PATCH] Add additional auth check to openai streaming requests

---
 src/services/openai.ts | 71 ++++++++++++++++++++++--------------------
 1 file changed, 37 insertions(+), 34 deletions(-)

diff --git a/src/services/openai.ts b/src/services/openai.ts
index 00f88191..c71bc108 100644
--- a/src/services/openai.ts
+++ b/src/services/openai.ts
@@ -5,6 +5,7 @@ import { createStreamableValue, render } from 'ai/rsc';
 import { kv } from '@vercel/kv';
 import { Ratelimit } from '@upstash/ratelimit';
 import { AI_TEXT_GENERATION_ENABLED, HAS_VERCEL_KV } from '@/site/config';
+import { safelyRunAdminServerAction } from '@/auth';
 
 const RATE_LIMIT_IDENTIFIER = 'openai-image-query';
 const RATE_LIMIT_MAX_QUERIES_PER_HOUR = 100;
@@ -25,43 +26,45 @@ export const streamOpenAiImageQuery = async (
   imageBase64: string,
   query: string,
 ) => {
-  if (ratelimit) {
-    const { success } = await ratelimit.limit(RATE_LIMIT_IDENTIFIER);
-    if (!success) {
-      console.error('OpenAI rate limit exceeded');
-      throw new Error('OpenAI rate limit exceeded');
+  return safelyRunAdminServerAction(async () => {
+    if (ratelimit) {
+      const { success } = await ratelimit.limit(RATE_LIMIT_IDENTIFIER);
+      if (!success) {
+        console.error('OpenAI rate limit exceeded');
+        throw new Error('OpenAI rate limit exceeded');
+      }
     }
-  }
 
-  const stream = createStreamableValue('');
+    const stream = createStreamableValue('');
 
-  if (provider) {
-    render({
-      provider,
-      model: 'gpt-4-vision-preview',
-      messages: [{
-        'role': 'user',
-        'content': [
-          {
-            'type': 'text',
-            'text': query,
-          }, {
-            'type': 'image_url',
-            'image_url': {
-              'url': imageBase64,
+    if (provider) {
+      render({
+        provider,
+        model: 'gpt-4-vision-preview',
+        messages: [{
+          'role': 'user',
+          'content': [
+            {
+              'type': 'text',
+              'text': query,
+            }, {
+              'type': 'image_url',
+              'image_url': {
+                'url': imageBase64,
+              },
             },
-          },
-        ],
-      }],
-      text: ({ content, done }): any => {
-        if (done) {
-          stream.done(content);
-        } else {
-          stream.update(content);
-        }
-      },
-    });
-  }
+          ],
+        }],
+        text: ({ content, done }): any => {
+          if (done) {
+            stream.done(content);
+          } else {
+            stream.update(content);
+          }
+        },
+      });
+    }
 
-  return stream.value;
+    return stream.value;
+  });
 };