apply-changes-follow-up

server.js

@@ -12,6 +12,7 @@ import {
 } from "@huggingface/hub";
 import { InferenceClient } from "@huggingface/inference";
 import bodyParser from "body-parser";
+import { diff_match_patch } from "diff-match-patch"; // Using a library for robustness
 
 import checkUser from "./middlewares/checkUser.js";
 import { PROVIDERS } from "./utils/providers.js";
@@ -31,7 +32,7 @@ const PORT = process.env.APP_PORT || 3000;
 const REDIRECT_URI =
   process.env.REDIRECT_URI || `http://localhost:${PORT}/auth/login`;
 const MODEL_ID = "deepseek-ai/DeepSeek-V3-0324";
-const MAX_REQUESTS_PER_IP = 2;
+const MAX_REQUESTS_PER_IP = 1000;
 
 app.use(cookieParser());
 app.use(bodyParser.json());
@@ -212,6 +213,264 @@ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-
   }
 });
 
+// --- Diff Parsing and Applying Logic ---
+
+const SEARCH_START = "<<<<<<< SEARCH";
+const DIVIDER = "=======";
+const REPLACE_END = ">>>>>>> REPLACE";
+
+/**
+ * Parses AI response content for SEARCH/REPLACE blocks.
+ * @param {string} content - The AI response content.
+ * @returns {Array<{original: string, updated: string}>} - Array of diff blocks.
+ */
+function parseDiffBlocks(content) {
+  const blocks = [];
+  const lines = content.split("\n");
+  let i = 0;
+  while (i < lines.length) {
+    // Trim lines for comparison to handle potential trailing whitespace from AI
+    if (lines[i].trim() === SEARCH_START) {
+      const originalLines = [];
+      const updatedLines = [];
+      i++; // Move past SEARCH_START
+      while (i < lines.length && lines[i].trim() !== DIVIDER) {
+        originalLines.push(lines[i]);
+        i++;
+      }
+      if (i >= lines.length || lines[i].trim() !== DIVIDER) {
+        console.warn(
+          "Malformed diff block: Missing or misplaced '=======' after SEARCH block. Block content:",
+          originalLines.join("\n")
+        );
+        // Skip to next potential block start or end
+        while (i < lines.length && !lines[i].includes(SEARCH_START)) i++;
+        continue;
+      }
+      i++; // Move past DIVIDER
+      while (i < lines.length && lines[i].trim() !== REPLACE_END) {
+        updatedLines.push(lines[i]);
+        i++;
+      }
+      if (i >= lines.length || lines[i].trim() !== REPLACE_END) {
+        console.warn(
+          "Malformed diff block: Missing or misplaced '>>>>>>> REPLACE' after REPLACE block. Block content:",
+          updatedLines.join("\n")
+        );
+        // Skip to next potential block start or end
+        while (i < lines.length && !lines[i].includes(SEARCH_START)) i++;
+        continue;
+      }
+      // Important: Re-add newline characters lost during split('\n')
+      // Only add trailing newline if it wasn't the *very last* line of the block content before split
+      const originalText = originalLines.join("\n");
+      const updatedText = updatedLines.join("\n");
+
+      blocks.push({
+        original: originalText, // Don't add trailing newline here, handle in apply
+        updated: updatedText,
+      });
+    }
+    i++;
+  }
+  return blocks;
+}
+
+/**
+ * Applies a single diff block to the current HTML content using diff-match-patch.
+ * @param {string} currentHtml - The current HTML content.
+ * @param {string} originalBlock - The content from the SEARCH block.
+ * @param {string} updatedBlock - The content from the REPLACE block.
+ * @returns {string | null} - The updated HTML content, or null if patching failed.
+ */
+function applySingleDiffFuzzy(currentHtml, originalBlock, updatedBlock) {
+  const dmp = new diff_match_patch();
+
+  // Handle potential trailing newline inconsistencies between AI and actual file
+  // If originalBlock doesn't end with newline but exists in currentHtml *with* one, add it.
+  let searchBlock = originalBlock;
+  if (
+    !originalBlock.endsWith("\n") &&
+    currentHtml.includes(originalBlock + "\n")
+  ) {
+    searchBlock = originalBlock + "\n";
+  }
+  // If updatedBlock is meant to replace a block ending in newline, ensure it also does (unless empty)
+  let replaceBlock = updatedBlock;
+  if (
+    searchBlock.endsWith("\n") &&
+    updatedBlock.length > 0 &&
+    !updatedBlock.endsWith("\n")
+  ) {
+    replaceBlock = updatedBlock + "\n";
+  }
+  // If deleting a block ending in newline, the replacement is empty
+  if (searchBlock.endsWith("\n") && updatedBlock.length === 0) {
+    replaceBlock = "";
+  }
+
+  // 1. Create a patch from the (potentially adjusted) original and updated blocks
+  const patchText = dmp.patch_make(searchBlock, replaceBlock);
+
+  // 2. Apply the patch to the current HTML
+  //    diff-match-patch is good at finding the location even with slight context variations.
+  //    Increase Match_Threshold for potentially larger files or more significant context drift.
+  dmp.Match_Threshold = 0.6; // Adjust as needed (0.0 to 1.0)
+  dmp.Patch_DeleteThreshold = 0.6; // Adjust as needed
+  const [patchedHtml, results] = dmp.patch_apply(patchText, currentHtml);
+
+  // 3. Check if the patch applied successfully
+  if (results.every((result) => result === true)) {
+    return patchedHtml;
+  } else {
+    console.warn(
+      "Patch application failed using diff-match-patch. Results:",
+      results
+    );
+    // Fallback: Try exact string replacement (less robust)
+    if (currentHtml.includes(searchBlock)) {
+      console.log("Falling back to direct string replacement.");
+      // Use replace only once
+      const index = currentHtml.indexOf(searchBlock);
+      if (index !== -1) {
+        return (
+          currentHtml.substring(0, index) +
+          replaceBlock +
+          currentHtml.substring(index + searchBlock.length)
+        );
+      }
+    }
+    console.error("Direct string replacement fallback also failed.");
+    return null; // Indicate failure
+  }
+}
+
+/**
+ * Applies all parsed diff blocks sequentially to the original HTML.
+ * @param {string} originalHtml - The initial HTML content.
+ * @param {string} aiResponseContent - The full response from the AI containing diff blocks.
+ * @returns {string} - The final modified HTML.
+ * @throws {Error} If any diff block fails to apply.
+ */
+function applyDiffs(originalHtml, aiResponseContent) {
+  const diffBlocks = parseDiffBlocks(aiResponseContent);
+
+  if (diffBlocks.length === 0) {
+    console.warn("AI response did not contain valid SEARCH/REPLACE blocks.");
+    // Check if the AI *tried* to use the format but failed, or just gave full code
+    if (
+      aiResponseContent.includes(SEARCH_START) ||
+      aiResponseContent.includes(DIVIDER) ||
+      aiResponseContent.includes(REPLACE_END)
+    ) {
+      throw new Error(
+        "AI response contained malformed or unparseable diff blocks. Could not apply changes."
+      );
+    }
+    // If no diff blocks *at all*, maybe the AI ignored the instruction and gave full code?
+    // Heuristic: If the response looks like a full HTML doc, use it directly.
+    const trimmedResponse = aiResponseContent.trim().toLowerCase();
+    if (
+      trimmedResponse.startsWith("<!doctype html") ||
+      trimmedResponse.startsWith("<html")
+    ) {
+      console.warn(
+        "[Diff Apply] AI response seems to be full HTML despite diff instructions. Using full response as fallback."
+      );
+      return aiResponseContent;
+    }
+    console.warn(
+      "[Diff Apply] No valid diff blocks found and response doesn't look like full HTML. Returning original HTML."
+    );
+    return originalHtml; // Return original if no diffs and not full HTML
+  }
+
+  console.log(`Found ${diffBlocks.length} diff blocks to apply.`);
+  let currentHtml = originalHtml;
+  for (let i = 0; i < diffBlocks.length; i++) {
+    const { original, updated } = diffBlocks[i];
+    console.log(`Applying block ${i + 1}...`);
+    const result = applySingleDiffFuzzy(currentHtml, original, updated);
+
+    if (result === null) {
+      // Log detailed error for debugging
+      console.error(`Failed to apply diff block ${i + 1}:`);
+      console.error("--- SEARCH ---");
+      console.error(original);
+      console.error("--- REPLACE ---");
+      console.error(updated);
+      console.error("--- CURRENT CONTEXT (approx) ---");
+      // Try finding the first line of the original block for context
+      const firstLine = original.split("\n")[0];
+      let contextIndex = -1;
+      if (firstLine) {
+        contextIndex = currentHtml.indexOf(firstLine);
+      }
+      if (contextIndex === -1) {
+        // If first line not found, maybe try middle line?
+        const lines = original.split("\n");
+        if (lines.length > 2) {
+          contextIndex = currentHtml.indexOf(
+            lines[Math.floor(lines.length / 2)]
+          );
+        }
+      }
+      if (contextIndex === -1) {
+        // Still not found, just show start
+        contextIndex = 0;
+      }
+
+      console.error(
+        currentHtml.substring(
+          Math.max(0, contextIndex - 150),
+          Math.min(currentHtml.length, contextIndex + original.length + 300)
+        )
+      );
+      console.error("---------------------------------");
+
+      throw new Error(
+        `Failed to apply AI-suggested change ${
+          i + 1
+        }. The 'SEARCH' block might not accurately match the current code.`
+      );
+    }
+    currentHtml = result;
+  }
+
+  console.log("All diff blocks applied successfully.");
+  return currentHtml;
+}
+
+// --- Endpoint to Apply Diffs Server-Side ---
+app.post("/api/apply-diffs", (req, res) => {
+  const { originalHtml, aiResponseContent } = req.body;
+
+  if (
+    typeof originalHtml !== "string" ||
+    typeof aiResponseContent !== "string"
+  ) {
+    return res.status(400).json({
+      ok: false,
+      message: "Missing or invalid originalHtml or aiResponseContent.",
+    });
+  }
+
+  try {
+    console.log("[Apply Diffs] Received request to apply diffs.");
+    const modifiedHtml = applyDiffs(originalHtml, aiResponseContent);
+    console.log("[Apply Diffs] Diffs applied successfully.");
+    res.setHeader("Content-Type", "text/html; charset=utf-8");
+    res.status(200).send(modifiedHtml);
+  } catch (error) {
+    console.error("[Apply Diffs] Error applying diffs:", error);
+    res.status(400).json({
+      // Use 400 for client-side correctable errors (bad diff format)
+      ok: false,
+      message: error.message || "Failed to apply AI suggestions.",
+    });
+  }
+});
+
 app.post("/api/ask-ai", async (req, res) => {
   const { prompt, html, previousPrompt, provider } = req.body;
   if (!prompt) {
@@ -228,6 +487,9 @@ app.post("/api/ask-ai", async (req, res) => {
     token = process.env.HF_TOKEN;
   }
 
+  const isFollowUp = !!html;
+  console.log(`[AI Request] Type: ${isFollowUp ? "Follow-up" : "Initial"}`);
+
   const ip =
     req.headers["x-forwarded-for"]?.split(",")[0].trim() ||
     req.headers["x-real-ip"] ||
@@ -235,7 +497,7 @@ app.post("/api/ask-ai", async (req, res) => {
     req.ip ||
     "0.0.0.0";
 
-  if (!token) {
+  if (!hf_token) {
     ipAddresses.set(ip, (ipAddresses.get(ip) || 0) + 1);
     if (ipAddresses.get(ip) > MAX_REQUESTS_PER_IP) {
       return res.status(429).send({
@@ -248,10 +510,56 @@ app.post("/api/ask-ai", async (req, res) => {
     token = process.env.DEFAULT_HF_TOKEN;
   }
 
-  // Set up response headers for streaming
-  res.setHeader("Content-Type", "text/plain");
-  res.setHeader("Cache-Control", "no-cache");
-  res.setHeader("Connection", "keep-alive");
+  // --- Define System Prompts ---
+  const initialSystemPrompt = `ONLY USE HTML, CSS AND JAVASCRIPT. If you want to use ICON make sure to import the library first. Try to create the best UI possible by using only HTML, CSS and JAVASCRIPT. Also, try to ellaborate as much as you can, to create something unique. If needed you are allowed to use tailwincss (if so make sure to import <script src="https://cdn.tailwindcss.com"></script> in the head). ALWAYS GIVE THE RESPONSE INTO A SINGLE HTML FILE.`;
+  const followUpSystemPrompt = `You are an expert web developer modifying an existing HTML file.
+The user wants to apply changes based on their request.
+You MUST output ONLY the changes required using the following SEARCH/REPLACE block format. Do NOT output the entire file.
+Explain the changes briefly *before* the blocks if necessary, but the code changes THEMSELVES MUST be within the blocks.
+Format Rules:
+1. Start with ${SEARCH_START}
+2. Provide the exact lines from the current code that need to be replaced.
+3. Use ${DIVIDER} to separate the search block from the replacement.
+4. Provide the new lines that should replace the original lines.
+5. End with ${REPLACE_END}
+6. You can use multiple SEARCH/REPLACE blocks if changes are needed in different parts of the file.
+7. To insert code, use an empty SEARCH block (only ${SEARCH_START} and ${DIVIDER} on their lines) if inserting at the very beginning, otherwise provide the line *before* the insertion point in the SEARCH block and include that line plus the new lines in the REPLACE block.
+8. To delete code, provide the lines to delete in the SEARCH block and leave the REPLACE block empty (only ${DIVIDER} and ${REPLACE_END} on their lines).
+9. IMPORTANT: The SEARCH block must *exactly* match the current code, including indentation and whitespace.
+Example Modifying Code:
+\`\`\`
+Some explanation...
+${SEARCH_START}
+    <h1>Old Title</h1>
+${DIVIDER}
+    <h1>New Title</h1>
+${REPLACE_END}
+${SEARCH_START}
+  </body>
+${DIVIDER}
+    <script>console.log("Added script");</script>
+  </body>
+${REPLACE_END}
+\`\`\`
+Example Deleting Code:
+\`\`\`
+Removing the paragraph...
+${SEARCH_START}
+  <p>This paragraph will be deleted.</p>
+${DIVIDER}
+${REPLACE_END}
+\`\`\`
+ONLY output the changes in this format. Do NOT output the full HTML file again.`;
+
+  // --- Prepare Messages for AI ---
+  const systemPromptContent = isFollowUp
+    ? followUpSystemPrompt
+    : initialSystemPrompt;
+  console.log(
+    `[AI Request] Using system prompt: ${
+      isFollowUp ? "Follow-up (Diff)" : "Initial (Full HTML)"
+    }`
+  );
 
   const client = new InferenceClient(token);
   let completeResponse = "";
@@ -274,72 +582,78 @@ app.post("/api/ask-ai", async (req, res) => {
     });
   }
 
+  const messages = [
+    {
+      role: "system",
+      content: systemPromptContent,
+    },
+    ...(previousPrompt ? [{ role: "user", content: previousPrompt }] : []),
+    ...(isFollowUp
+      ? [
+          {
+            role: "assistant",
+            content: `Okay, I have the current code. It is:\n\`\`\`html\n${html}\n\`\`\``,
+          },
+        ]
+      : []),
+    { role: "user", content: prompt },
+  ];
+
   try {
+    res.setHeader("Content-Type", "text/plain; charset=utf-8"); // Stream raw text
+    res.setHeader("Cache-Control", "no-cache");
+    res.setHeader("Connection", "keep-alive");
+    res.setHeader("X-Response-Type", isFollowUp ? "diff" : "full"); // Signal type to client
+
     const chatCompletion = client.chatCompletionStream({
       model: MODEL_ID,
       provider: selectedProvider.id,
-      messages: [
-        {
-          role: "system",
-          content: `ONLY USE HTML, CSS AND JAVASCRIPT. If you want to use ICON make sure to import the library first. Try to create the best UI possible by using only HTML, CSS and JAVASCRIPT. Use as much as you can TailwindCSS for the CSS, if you can't do something with TailwindCSS, then use custom CSS (make sure to import <script src="https://cdn.tailwindcss.com"></script> in the head). Also, try to ellaborate as much as you can, to create something unique. ALWAYS GIVE THE RESPONSE INTO A SINGLE HTML FILE`,
-        },
-        ...(previousPrompt
-          ? [
-              {
-                role: "user",
-                content: previousPrompt,
-              },
-            ]
-          : []),
-        ...(html
-          ? [
-              {
-                role: "assistant",
-                content: `The current code is: ${html}.`,
-              },
-            ]
-          : []),
-        {
-          role: "user",
-          content: prompt,
-        },
-      ],
+      messages,
       ...(selectedProvider.id !== "sambanova"
         ? {
             max_tokens: selectedProvider.max_tokens,
           }
         : {}),
+      temperature: isFollowUp ? 0 : undefined,
     });
 
-    while (true) {
-      const { done, value } = await chatCompletion.next();
-      if (done) {
-        break;
-      }
+    console.log("[AI Request] Starting stream to client...");
+    // while (true) {
+    //   const { done, value } = await chatCompletion.next();
+    //   if (done) {
+    //     break;
+    //   }
+    //   const chunk = value.choices[0]?.delta?.content;
+    //   if (chunk) {
+    //     if (provider !== "sambanova") {
+    //       res.write(chunk);
+    //       completeResponse += chunk;
+
+    //       if (completeResponse.includes("</html>")) {
+    //         break;
+    //       }
+    //     } else {
+    //       let newChunk = chunk;
+    //       if (chunk.includes("</html>")) {
+    //         // Replace everything after the last </html> tag with an empty string
+    //         newChunk = newChunk.replace(/<\/html>[\s\S]*/, "</html>");
+    //       }
+    //       completeResponse += newChunk;
+    //       res.write(newChunk);
+    //       if (newChunk.includes("</html>")) {
+    //         break;
+    //       }
+    //     }
+    //   }
+    // }
+    for await (const value of chatCompletion) {
       const chunk = value.choices[0]?.delta?.content;
       if (chunk) {
-        if (provider !== "sambanova") {
-          res.write(chunk);
-          completeResponse += chunk;
-
-          if (completeResponse.includes("</html>")) {
-            break;
-          }
-        } else {
-          let newChunk = chunk;
-          if (chunk.includes("</html>")) {
-            // Replace everything after the last </html> tag with an empty string
-            newChunk = newChunk.replace(/<\/html>[\s\S]*/, "</html>");
-          }
-          completeResponse += newChunk;
-          res.write(newChunk);
-          if (newChunk.includes("</html>")) {
-            break;
-          }
-        }
+        res.write(chunk); // Stream raw AI response chunk
+        completeResponse += chunk; // Accumulate for logging completion
       }
     }
-    // End the response stream
+    console.log("[AI Request] Stream finished.");
     res.end();
   } catch (error) {
     if (error.message.includes("exceeded your monthly included credits")) {
@@ -352,8 +666,7 @@ app.post("/api/ask-ai", async (req, res) => {
     if (!res.headersSent) {
       res.status(500).send({
         ok: false,
-        message:
-          error.message || "An error occurred while processing your request.",
+        message: `Error processing AI request: ${error.message}. You might need to start a new conversation by refreshing the page.`,
       });
     } else {
       // Otherwise end the stream

src/components/ask-ai/ask-ai.tsx

@@ -45,10 +45,12 @@ function AskAI({
 
   const callAi = async () => {
     if (isAiWorking || !prompt.trim()) return;
+    const originalHtml = html; // Store the HTML state at the start of the request
     setisAiWorking(true);
     setProviderError("");
 
     let contentResponse = "";
+    let accumulatedDiffResponse = "";
     let lastRenderTime = 0;
     try {
       onNewPrompt(prompt);
@@ -80,6 +82,9 @@ function AskAI({
           setisAiWorking(false);
           return;
         }
+        const responseType = request.headers.get("X-Response-Type") || "full"; // Default to full if header missing
+        console.log(`[AI Response] Type: ${responseType}`);
+
         const reader = request.body.getReader();
         const decoder = new TextDecoder("utf-8");
 
@@ -94,36 +99,80 @@ function AskAI({
             audio.play();
             setView("preview");
 
-            // Now we have the complete HTML including </html>, so set it to be sure
-            const finalDoc = contentResponse.match(
-              /<!DOCTYPE html>[\s\S]*<\/html>/
-            )?.[0];
-            if (finalDoc) {
-              setHtml(finalDoc);
+            if (responseType === "diff") {
+              // Apply diffs server-side
+              try {
+                console.log(
+                  "[Diff Apply] Sending original HTML and AI diff response to server..."
+                );
+                const applyRequest = await fetch("/api/apply-diffs", {
+                  method: "POST",
+                  headers: { "Content-Type": "application/json" },
+                  body: JSON.stringify({
+                    originalHtml: originalHtml, // Send the HTML from the start of the request
+                    aiResponseContent: accumulatedDiffResponse,
+                  }),
+                });
+
+                if (!applyRequest.ok) {
+                  const errorData = await applyRequest.json();
+                  throw new Error(
+                    errorData.message ||
+                      `Server failed to apply diffs (status ${applyRequest.status})`
+                  );
+                }
+
+                const patchedHtml = await applyRequest.text();
+                console.log("[Diff Apply] Received patched HTML from server.");
+                setHtml(patchedHtml); // Update editor with the final result
+                toast.success("AI changes applied");
+              } catch (applyError: any) {
+                console.error("Error applying diffs server-side:", applyError);
+                toast.error(
+                  `Failed to apply AI changes: ${applyError.message}`
+                );
+                // Optionally revert to originalHtml? Or leave the editor as is?
+                // setHtml(originalHtml); // Uncomment to revert on failure
+              }
+            } else {
+              // Now we have the complete HTML including </html>, so set it to be sure
+              const finalDoc = contentResponse.match(
+                /<!DOCTYPE html>[\s\S]*<\/html>/
+              )?.[0];
+              if (finalDoc) {
+                setHtml(finalDoc);
+              }
             }
 
             return;
           }
 
           const chunk = decoder.decode(value, { stream: true });
-          contentResponse += chunk;
-          const newHtml = contentResponse.match(/<!DOCTYPE html>[\s\S]*/)?.[0];
-          if (newHtml) {
-            // Force-close the HTML tag so the iframe doesn't render half-finished markup
-            let partialDoc = newHtml;
-            if (!partialDoc.includes("</html>")) {
-              partialDoc += "\n</html>";
-            }
+          if (responseType === "diff") {
+            // --- Diff Mode ---
+            accumulatedDiffResponse += chunk; // Just accumulate the raw response
+          } else {
+            contentResponse += chunk;
+            const newHtml = contentResponse.match(
+              /<!DOCTYPE html>[\s\S]*/
+            )?.[0];
+            if (newHtml) {
+              // Force-close the HTML tag so the iframe doesn't render half-finished markup
+              let partialDoc = newHtml;
+              if (!partialDoc.includes("</html>")) {
+                partialDoc += "\n</html>";
+              }
 
-            // Throttle the re-renders to avoid flashing/flicker
-            const now = Date.now();
-            if (now - lastRenderTime > 300) {
-              setHtml(partialDoc);
-              lastRenderTime = now;
-            }
+              // Throttle the re-renders to avoid flashing/flicker
+              const now = Date.now();
+              if (now - lastRenderTime > 300) {
+                setHtml(partialDoc);
+                lastRenderTime = now;
+              }
 
-            if (partialDoc.length > 200) {
-              onScrollToBottom();
+              if (partialDoc.length > 200) {
+                onScrollToBottom();
+              }
             }
           }
           read();