From ffeef657c3e56f5db68e09fc2e631d62e1e8ab9a Mon Sep 17 00:00:00 2001
From: Sam Becker <sam@sambecker.com>
Date: Wed, 12 Jun 2024 11:41:02 -0500
Subject: [PATCH] Streamline AI text formatting

---
 src/photo/ai/index.ts           |  7 +++--
 src/photo/ai/server.ts          | 46 ++++++++++++++-------------------
 src/photo/ai/useAiImageQuery.ts |  6 ++---
 src/services/openai.ts          |  5 ++--
 4 files changed, 28 insertions(+), 36 deletions(-)

diff --git a/src/photo/ai/index.ts b/src/photo/ai/index.ts
index 66b458f6..79287d5b 100644
--- a/src/photo/ai/index.ts
+++ b/src/photo/ai/index.ts
@@ -64,9 +64,8 @@ export const parseTitleAndCaption = (text: string) => {
   };
 };
 
-export const cleanUpAiTextResponse = (text?: string) => text
-  ? text
+export const cleanUpAiTextResponse = (text: string) =>
+  text
     .replaceAll('\n', ' ')
     .replaceAll('"', '')
-    .replace(/\.$/, '')
-  : undefined;
+    .replace(/\.$/, '');
diff --git a/src/photo/ai/server.ts b/src/photo/ai/server.ts
index f57d1ff8..163845c1 100644
--- a/src/photo/ai/server.ts
+++ b/src/photo/ai/server.ts
@@ -2,7 +2,6 @@ import { generateOpenAiImageQuery } from '@/services/openai';
 import {
   AI_IMAGE_QUERIES,
   AiAutoGeneratedField,
-  cleanUpAiTextResponse,
   parseTitleAndCaption,
 } from '.';
 
@@ -28,11 +27,10 @@ export const generateAiImageQueries = async (
         textFieldsToGenerate.includes('title') &&
         textFieldsToGenerate.includes('caption')
       ) {
-        const titleAndCaption = cleanUpAiTextResponse(
-          await generateOpenAiImageQuery(
-            imageBase64,
-            AI_IMAGE_QUERIES['title-and-caption'],
-          ));
+        const titleAndCaption = await generateOpenAiImageQuery(
+          imageBase64,
+          AI_IMAGE_QUERIES['title-and-caption'],
+        );
         if (titleAndCaption) {
           const titleAndCaptionParsed = parseTitleAndCaption(titleAndCaption);
           title = titleAndCaptionParsed.title;
@@ -40,35 +38,31 @@ export const generateAiImageQueries = async (
         }
       } else {
         if (textFieldsToGenerate.includes('title')) {
-          title = cleanUpAiTextResponse(
-            await generateOpenAiImageQuery(
-              imageBase64,
-              AI_IMAGE_QUERIES['title'],
-            ));
+          title = await generateOpenAiImageQuery(
+            imageBase64,
+            AI_IMAGE_QUERIES['title'],
+          );
         }
         if (textFieldsToGenerate.includes('caption')) {
-          caption = cleanUpAiTextResponse(
-            await generateOpenAiImageQuery(
-              imageBase64,
-              AI_IMAGE_QUERIES['caption'],
-            ));
+          caption =  await generateOpenAiImageQuery(
+            imageBase64,
+            AI_IMAGE_QUERIES['caption'],
+          );
         }
       }
   
       if (textFieldsToGenerate.includes('tags')) {
-        tags = cleanUpAiTextResponse(
-          await generateOpenAiImageQuery(
-            imageBase64,
-            AI_IMAGE_QUERIES['tags'],
-          ));
+        tags = await generateOpenAiImageQuery(
+          imageBase64,
+          AI_IMAGE_QUERIES['tags'],
+        );
       }
   
       if (textFieldsToGenerate.includes('semantic')) {
-        semanticDescription = cleanUpAiTextResponse(
-          await generateOpenAiImageQuery(
-            imageBase64,
-            AI_IMAGE_QUERIES['description-small'],
-          ));
+        semanticDescription = await generateOpenAiImageQuery(
+          imageBase64,
+          AI_IMAGE_QUERIES['description-small'],
+        );
       }
     }
   } catch (e: any) {
diff --git a/src/photo/ai/useAiImageQuery.ts b/src/photo/ai/useAiImageQuery.ts
index f75c7395..fcb8e10e 100644
--- a/src/photo/ai/useAiImageQuery.ts
+++ b/src/photo/ai/useAiImageQuery.ts
@@ -1,7 +1,7 @@
 import { useCallback, useState } from 'react';
 import { streamAiImageQueryAction } from '../actions';
 import { readStreamableValue } from 'ai/rsc';
-import { AiImageQuery, cleanUpAiTextResponse } from '.';
+import { AiImageQuery } from '.';
 
 export default function useAiImageQuery(
   imageBase64: string | undefined,
@@ -21,9 +21,7 @@ export default function useAiImageQuery(
           query,
         );
         for await (const text of readStreamableValue(textStream)) {
-          setText(current =>
-            cleanUpAiTextResponse(`${current}${text ?? ''}`) ?? ''
-          );
+          setText(current => `${current}${text ?? ''}`);
         }
         setIsLoading(false);
       } catch (e) {
diff --git a/src/services/openai.ts b/src/services/openai.ts
index 443ea7d4..bd9681f6 100644
--- a/src/services/openai.ts
+++ b/src/services/openai.ts
@@ -5,6 +5,7 @@ import { kv } from '@vercel/kv';
 import { Ratelimit } from '@upstash/ratelimit';
 import { AI_TEXT_GENERATION_ENABLED, HAS_VERCEL_KV } from '@/site/config';
 import { removeBase64Prefix } from '@/utility/image';
+import { cleanUpAiTextResponse } from '@/photo/ai';
 
 const RATE_LIMIT_IDENTIFIER = 'openai-image-query';
 const RATE_LIMIT_MAX_QUERIES_PER_HOUR = 100;
@@ -73,7 +74,7 @@ export const streamOpenAiImageQuery = async (
     (async () => {
       const { textStream } = await streamText(args);
       for await (const delta of textStream) {
-        stream.update(delta);
+        stream.update(cleanUpAiTextResponse(delta));
       }
       stream.done();
     })();
@@ -92,7 +93,7 @@ export const generateOpenAiImageQuery = async (
 
   if (args) {
     return generateText(args)
-      .then(({ text }) => text);
+      .then(({ text }) => cleanUpAiTextResponse(text));
   }
 };