Add editing functionality BEN-1079 (#29)

### TL;DR

Added code editing capabilities to the UI builder, allowing users to modify existing components through chat instructions.

### What changed?

- Added `editApp` function in `lib/openai.ts` to handle code modifications
- Created new prompt templates for editing in `lib/prompts.ts`
- Extended the component schema to support editing instructions and existing files
- Implemented file merging logic to combine updated files with existing ones
- Updated the chat interface to send edit requests with current files
- Added session storage to persist builder results
- Enhanced logging for better debugging of edit operations

### How to test?

1. Generate a component using the UI builder
2. Use the chat interface to request changes (e.g., "Add a dark mode toggle")
3. Verify that the component updates according to your instructions
4. Check that the changes persist when refreshing the page (via sessionStorage)
5. Examine the console logs to verify the edit flow is working correctly

### Why make this change?

This enhancement significantly improves the user experience by allowing iterative refinement of generated components without starting from scratch. Users can now have a conversation with the AI to progressively improve their components, making the tool more flexible and powerful for real-world development workflows.

Author: juancastano

app/api/generate/route.ts

@@ -1,10 +1,12 @@
 // app/api/generate/route.ts
 import { NextRequest, NextResponse } from 'next/server';
-import { generateApp } from '@/lib/openai';
+import { generateApp, editApp } from '@/lib/openai';
 import { createSandbox } from '@/lib/e2b';
 import { componentSchema } from '@/lib/schemas';
 import { Benchify } from 'benchify';
 import { applyPatch } from 'diff';
+import { z } from 'zod';
+import { benchifyFileSchema } from '@/lib/schemas';
 
 const benchify = new Benchify({
     apiKey: process.env.BENCHIFY_API_KEY,
@@ -32,12 +34,30 @@ export default App;`
     }
 ];
 
+// Extended schema to support editing
+const extendedComponentSchema = componentSchema.extend({
+    existingFiles: benchifyFileSchema.optional(),
+    editInstruction: z.string().optional(),
+});
+
+// Helper function to merge updated files with existing files
+function mergeFiles(existingFiles: z.infer<typeof benchifyFileSchema>, updatedFiles: z.infer<typeof benchifyFileSchema>): z.infer<typeof benchifyFileSchema> {
+    const existingMap = new Map(existingFiles.map(file => [file.path, file]));
+
+    // Apply updates
+    updatedFiles.forEach(updatedFile => {
+        existingMap.set(updatedFile.path, updatedFile);
+    });
+
+    return Array.from(existingMap.values());
+}
+
 export async function POST(request: NextRequest) {
     try {
         const body = await request.json();
 
-        // Validate the request using Zod schema
-        const validationResult = componentSchema.safeParse(body);
+        // Validate the request using extended schema
+        const validationResult = extendedComponentSchema.safeParse(body);
 
         if (!validationResult.success) {
             return NextResponse.json(
@@ -46,30 +66,53 @@ export async function POST(request: NextRequest) {
             );
         }
 
-        const { description } = validationResult.data;
+        const { description, existingFiles, editInstruction } = validationResult.data;
+
+        console.log('API Request:', {
+            isEdit: !!(existingFiles && editInstruction),
+            filesCount: existingFiles?.length || 0,
+            editInstruction: editInstruction || 'none',
+            description: description || 'none'
+        });
+
+        let filesToSandbox;
+
+        // Determine if this is an edit request or new generation
+        if (existingFiles && editInstruction) {
+            // Edit existing code
+            console.log('Processing edit request...');
+            console.log('Existing files:', existingFiles.map(f => ({ path: f.path, contentLength: f.content.length })));
+
+            const updatedFiles = await editApp(existingFiles, editInstruction);
+            console.log('Updated files from AI:', updatedFiles.map(f => ({ path: f.path, contentLength: f.content.length })));
 
-        // Generate the Vue app using OpenAI
-        let generatedFiles;
-        if (debug) {
-            generatedFiles = buggyCode;
+            // Merge the updated files with the existing files
+            filesToSandbox = mergeFiles(existingFiles, updatedFiles);
+            console.log('Final merged files:', filesToSandbox.map(f => ({ path: f.path, contentLength: f.content.length })));
         } else {
-            generatedFiles = await generateApp(description);
+            // Generate new app
+            console.log('Processing new generation request...');
+            if (debug) {
+                filesToSandbox = buggyCode;
+            } else {
+                filesToSandbox = await generateApp(description);
+            }
         }
 
         // Repair the generated code using Benchify's API
         // const { data } = await benchify.fixer.run({
-        //     files: generatedFiles.map(file => ({
+        //     files: filesToSandbox.map(file => ({
         //         path: file.path,
         //         contents: file.content
         //     }))
         // });
 
-        let repairedFiles = generatedFiles;
+        let repairedFiles = filesToSandbox;
         // if (data) {
         //     const { success, diff } = data;
 
         //     if (success && diff) {
-        //         repairedFiles = generatedFiles.map(file => {
+        //         repairedFiles = filesToSandbox.map(file => {
         //             const patchResult = applyPatch(file.content, diff);
         //             return {
         //                 ...file,
@@ -83,12 +126,13 @@ export async function POST(request: NextRequest) {
 
         // Return the results to the client
         return NextResponse.json({
-            originalFiles: generatedFiles,
+            originalFiles: filesToSandbox,
             repairedFiles: sandboxResult.allFiles, // Use the allFiles from the sandbox
             buildOutput: `Sandbox created with template: ${sandboxResult.template}, ID: ${sandboxResult.sbxId}`,
             previewUrl: sandboxResult.url,
             buildErrors: sandboxResult.buildErrors,
             hasErrors: sandboxResult.hasErrors,
+            ...(editInstruction && { editInstruction }),
         });
     } catch (error) {
         console.error('Error generating app:', error);

app/chat/page.tsx

@@ -100,7 +100,12 @@ export default function ChatPage() {
             <div className="w-1/4 min-w-80 border-r border-border bg-card flex-shrink-0">
                 <ChatInterface
                     initialPrompt={initialPrompt}
-                    onUpdateResult={setResult}
+                    currentFiles={result?.repairedFiles || result?.originalFiles}
+                    onUpdateResult={(updatedResult) => {
+                        setResult(updatedResult);
+                        // Save updated result to sessionStorage
+                        sessionStorage.setItem('builderResult', JSON.stringify(updatedResult));
+                    }}
                 />
             </div>
 

components/ui-builder/chat-interface.tsx

@@ -18,15 +18,24 @@ interface Message {
 
 interface ChatInterfaceProps {
     initialPrompt: string;
+    currentFiles?: z.infer<typeof benchifyFileSchema>;
     onUpdateResult: (result: {
         repairedFiles?: z.infer<typeof benchifyFileSchema>;
         originalFiles?: z.infer<typeof benchifyFileSchema>;
         buildOutput: string;
         previewUrl: string;
+        buildErrors?: Array<{
+            type: 'typescript' | 'build' | 'runtime';
+            message: string;
+            file?: string;
+            line?: number;
+            column?: number;
+        }>;
+        hasErrors?: boolean;
     }) => void;
 }
 
-export function ChatInterface({ initialPrompt, onUpdateResult }: ChatInterfaceProps) {
+export function ChatInterface({ initialPrompt, currentFiles, onUpdateResult }: ChatInterfaceProps) {
     const [messages, setMessages] = useState<Message[]>([
         {
             id: '1',
@@ -64,30 +73,63 @@ export function ChatInterface({ initialPrompt, onUpdateResult }: ChatInterfacePr
         };
 
         setMessages(prev => [...prev, userMessage]);
+        const editInstruction = newMessage;
         setNewMessage('');
         setIsLoading(true);
 
         try {
-            // Here you would call your API to process the new request
-            // For now, we'll add a placeholder response
-            const assistantMessage: Message = {
+            // Add thinking message
+            const thinkingMessage: Message = {
                 id: (Date.now() + 1).toString(),
                 type: 'assistant',
                 content: "I understand your request. Let me update the component for you...",
                 timestamp: new Date(),
             };
-
-            setMessages(prev => [...prev, assistantMessage]);
-
-            // TODO: Implement actual regeneration with the new prompt
-            // This would call your generate API with the conversation context
+            setMessages(prev => [...prev, thinkingMessage]);
+
+            // Call the edit API
+            const response = await fetch('/api/generate', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify({
+                    type: 'component',
+                    description: '', // Not used for edits
+                    existingFiles: currentFiles,
+                    editInstruction: editInstruction,
+                }),
+            });
+
+            console.log('Edit request:', {
+                existingFiles: currentFiles,
+                editInstruction: editInstruction,
+                filesCount: currentFiles?.length || 0
+            });
+
+            if (!response.ok) {
+                throw new Error('Failed to edit component');
+            }
+
+            const editResult = await response.json();
+            console.log('Edit response:', editResult);
+
+            // Update the result in the parent component
+            onUpdateResult(editResult);
+
+            // Update the thinking message to success
+            setMessages(prev => prev.map(msg =>
+                msg.id === thinkingMessage.id
+                    ? { ...msg, content: `Great! I've updated the component according to your request: "${editInstruction}"` }
+                    : msg
+            ));
 
         } catch (error) {
-            console.error('Error processing message:', error);
+            console.error('Error processing edit:', error);
             const errorMessage: Message = {
                 id: (Date.now() + 1).toString(),
                 type: 'assistant',
-                content: "I'm sorry, there was an error processing your request. Please try again.",
+                content: "I'm sorry, there was an error processing your edit request. Please try again.",
                 timestamp: new Date(),
             };
             setMessages(prev => [...prev, errorMessage]);
@@ -137,8 +179,8 @@ export function ChatInterface({ initialPrompt, onUpdateResult }: ChatInterfacePr
                             >
                                 <div
                                     className={`max-w-[80%] p-3 rounded-lg ${message.type === 'user'
-                                            ? 'bg-primary text-primary-foreground'
-                                            : 'bg-muted text-muted-foreground'
+                                        ? 'bg-primary text-primary-foreground'
+                                        : 'bg-muted text-muted-foreground'
                                         }`}
                                 >
                                     <p className="text-sm whitespace-pre-wrap">{message.content}</p>

lib/openai.ts

@@ -2,7 +2,8 @@
 import { streamObject } from 'ai';
 import { openai } from '@ai-sdk/openai';
 import { z } from 'zod';
-import { REACT_APP_SYSTEM_PROMPT, REACT_APP_USER_PROMPT, TEMPERATURE, MODEL } from './prompts';
+import { REACT_APP_SYSTEM_PROMPT, REACT_APP_USER_PROMPT, TEMPERATURE, MODEL, EDIT_SYSTEM_PROMPT, createEditUserPrompt } from './prompts';
+import { benchifyFileSchema } from './schemas';
 
 const OPENAI_API_KEY = process.env.OPENAI_API_KEY;
 
@@ -16,7 +17,9 @@ const fileSchema = z.object({
     content: z.string()
 });
 
-// Generate a Vue application using AI SDK
+
+
+// Generate a new application using AI SDK
 export async function generateApp(
     description: string,
 ): Promise<Array<{ path: string; content: string }>> {
@@ -50,4 +53,41 @@ export async function generateApp(
         console.error('Error generating app:', error);
         throw error;
     }
+}
+
+// Edit existing application using AI SDK
+export async function editApp(
+    existingFiles: z.infer<typeof benchifyFileSchema>,
+    editInstruction: string,
+): Promise<Array<{ path: string; content: string }>> {
+    console.log("Editing app with instruction: ", editInstruction);
+
+    try {
+        const { elementStream } = streamObject({
+            model: openai('gpt-4o-mini'),
+            output: 'array',
+            schema: fileSchema,
+            temperature: 0.3, // Lower temperature for more consistent edits
+            messages: [
+                { role: 'system', content: EDIT_SYSTEM_PROMPT },
+                { role: 'user', content: createEditUserPrompt(existingFiles, editInstruction) }
+            ]
+        });
+
+        const updatedFiles = [];
+        for await (const file of elementStream) {
+            updatedFiles.push(file);
+        }
+
+        if (!updatedFiles.length) {
+            throw new Error("Failed to generate updated files - received empty response");
+        }
+
+        console.log("Generated updated files: ", updatedFiles);
+
+        return updatedFiles;
+    } catch (error) {
+        console.error('Error editing app:', error);
+        throw error;
+    }
 }

lib/prompts.ts

@@ -1,4 +1,6 @@
 // lib/prompts.ts
+import { z } from 'zod';
+import { benchifyFileSchema } from "./schemas";
 
 export const REACT_APP_SYSTEM_PROMPT = `You are an expert React, TypeScript, and Tailwind CSS developer.
 You will be generating React application code based on the provided description.
@@ -47,5 +49,39 @@ export const REACT_APP_USER_PROMPT = (description: string) => `
 Create a React application with the following requirements:
 ${description}`;
 
+export const EDIT_SYSTEM_PROMPT = `You are an expert React/TypeScript developer. You will be given existing code files and an edit instruction. Your job is to modify the existing code according to the instruction while maintaining:
+
+1. Code quality and best practices
+2. Existing functionality that shouldn't be changed
+3. Proper TypeScript types
+4. Modern React patterns
+5. Tailwind CSS for styling
+6. shadcn/ui components where appropriate
+
+Return ONLY the files that need to be changed. Do not return unchanged files.
+
+Rules:
+- Only return files that have actual changes
+- Make minimal changes necessary to fulfill the instruction
+- Keep all imports and dependencies that are still needed
+- Add new dependencies only if absolutely necessary
+- Use Tailwind classes for styling changes
+- Follow React best practices
+- Ensure all returned files are complete and valid`;
+
+export function createEditUserPrompt(files: z.infer<typeof benchifyFileSchema>, editInstruction: string): string {
+  const filesContent = files.map(file =>
+    `### ${file.path}\n\`\`\`\n${file.content}\n\`\`\``
+  ).join('\n\n');
+
+  return `Here are the current files:
+
+${filesContent}
+
+Edit instruction: ${editInstruction}
+
+Please update the code according to this instruction and return all files with their updated content.`;
+}
+
 export const TEMPERATURE = 0.7;
 export const MODEL = 'gpt-4o';