Update app/actions/rewrite-prompt.ts

app/actions/rewrite-prompt.ts

@@ -3,33 +3,35 @@ import { GoogleGenerativeAI } from "@google/generative-ai";
 const START_REWRITE_PROMPT = ">>>>>>> START PROMPT >>>>>>";
 const END_REWRITE_PROMPT = ">>>>>>> END PROMPT >>>>>>";
 
-export const callAiRewritePrompt = async (prompt: string, { token, billTo }: { token: string, billTo?: string | null }) => {
+export const callAiRewritePrompt = async (
+  prompt: string,
+  { token, billTo }: { token: string; billTo?: string | null }
+) => {
   try {
-    // Get Gemini API key from environment variables
-    const geminiApiKey = process.env.GEMINI_API_KEY;
+    // Use the provided token as the API key, or fall back to environment variable
+    const geminiApiKey = token || process.env.GEMINI_API_KEY;
+
     if (!geminiApiKey) {
-      console.warn("Gemini API key not found, returning original prompt");
+      console.warn("Gemini API key not found. Returning original prompt.");
       return prompt;
     }
 
-    // Initialize Google Gemini AI
     const genAI = new GoogleGenerativeAI(geminiApiKey);
     const model = genAI.getGenerativeModel({ model: "gemini-2.5-flash" });
 
-    // Prepare the prompt for Gemini
-    const systemMessage = `You are a helpful assistant that rewrites prompts to make them better. All the prompts will be about creating a website or app.
+    const systemPrompt = `You are a helpful assistant that rewrites prompts to make them better. All the prompts will be about creating a website or app.
 Try to make the prompt more detailed and specific to create a good UI/UX Design and good code.
 Format the result by following this format:
 ${START_REWRITE_PROMPT}
 new prompt here
 ${END_REWRITE_PROMPT}
 If you don't rewrite the prompt, return the original prompt.
-Make sure to return the prompt in the same language as the prompt you are given. Also IMPORTANT: Make sure to keep the original intent of the prompt. Improve it it needed, but don't change the original intent.
+Make sure to return the prompt in the same language as the prompt you are given. Also IMPORTANT: Make sure to keep the original intent of the prompt. Improve it if needed, but don't change the original intent.
+`;
 
-User prompt to rewrite: ${prompt}`;
+    const fullPrompt = `${systemPrompt}\n\nUser prompt to rewrite: ${prompt}`;
 
-    // Generate content using Gemini
-    const result = await model.generateContent(systemMessage);
+    const result = await model.generateContent(fullPrompt);
     const response = result.response;
     const responseContent = response.text();
 
@@ -37,26 +39,23 @@ User prompt to rewrite: ${prompt}`;
       return prompt;
     }
 
-    // Extract the rewritten prompt between markers
     const startIndex = responseContent.indexOf(START_REWRITE_PROMPT);
     const endIndex = responseContent.indexOf(END_REWRITE_PROMPT);
-    
-    if (startIndex === -1 || endIndex === -1) {
-      // If markers not found, return original prompt
-      return prompt;
-    }
-
-    const rewrittenPrompt = responseContent.substring(
-      startIndex + START_REWRITE_PROMPT.length, 
-      endIndex
-    ).trim();
-
-    // Return rewritten prompt or fallback to original
-    return rewrittenPrompt || prompt;
 
+    if (startIndex !== -1 && endIndex !== -1) {
+      const rewrittenPrompt = responseContent
+        .substring(startIndex + START_REWRITE_PROMPT.length, endIndex)
+        .trim();
+      
+      return rewrittenPrompt || prompt;
+    }
+    
+    // If markers are not found, return the original prompt as a fallback
+    return prompt;
+    
   } catch (error) {
     console.error("Error rewriting prompt with Gemini:", error);
-    // Return original prompt if any error occurs
+    // In case of any API error, return the original prompt
     return prompt;
   }
 };