feat: introduce apply button functionality (v0.1)

### Summary - Added a new "Apply" button to code snippets provided by the AI assistant. - The button is designed to seamlessly merge the AI-generated snippet into the relevant file in the editor. ### Current Issues 1. **Sticky Accept/Decline Buttons:** These activate for every snippet instead of being limited to the relevant snippet. 2. **Discard Button:** Currently non-functional. 3. **Highlight Inconsistencies:** The green-red code highlights for old and new code are inconsistent. ### To Do - Implement a toast notification when the "Apply" button is pressed on an irrelevant tab to prevent code application errors. ### Workflow Implemented 1. The "Apply" button is added alongside "Copy" and "Reply" for AI-generated code snippets. 2. Upon clicking "Apply," the code snippet and relevant file content (active file) are sent to a secondary model (GPT-4O). 3. The system prompt for GPT-4O instructs it to merge the snippet with the file content: - Ensure the original file functionality remains intact. - Integrate the code snippet seamlessly. 4. The output from GPT-4O is injected directly into the code editor. 5. Changes are visually highlighted: - Green for new code. - Red for removed code. 6. Highlights remain until the user explicitly accepts or discards the changes.
2024-11-30 20:56:56 -05:00 · 2024-11-30 20:56:56 -05:00 · aa33ad3031
commit aa33ad3031
parent f79115974c
10 changed files with 372 additions and 5 deletions
--- a/frontend/app/api/merge/route.ts
+++ b/frontend/app/api/merge/route.ts
@ -0,0 +1,67 @@
+import OpenAI from "openai"
+
+const openai = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY,
+})
+
+export async function POST(request: Request) {
+  try {
+    const { originalCode, newCode, fileName } = await request.json()
+
+    const systemPrompt = `You are a code merging assistant. Your task is to merge the new code snippet with the original file content while:
+1. Preserving the original file's functionality
+2. Ensuring proper integration of the new code
+3. Maintaining consistent style and formatting
+4. Resolving any potential conflicts
+5. Output ONLY the raw code without any:
+   - Code fence markers (\`\`\`)
+   - Language identifiers (typescript, javascript, etc.)
+   - Explanations or comments
+   - Markdown formatting
+
+The output should be the exact code that will replace the existing code, nothing more and nothing less.`
+
+    const mergedCode = `Original file (${fileName}):\n${originalCode}\n\nNew code to merge:\n${newCode}`
+
+    const response = await openai.chat.completions.create({
+      model: "gpt-4o",
+      messages: [
+        { role: "system", content: systemPrompt },
+        { role: "user", content: mergedCode },
+      ],
+      prediction: {
+        type: "content",
+        content: mergedCode,
+      },
+      stream: true,
+    })
+
+    // Clean and stream response
+    const encoder = new TextEncoder()
+    return new Response(
+      new ReadableStream({
+        async start(controller) {
+          let buffer = ""
+          for await (const chunk of response) {
+            if (chunk.choices[0]?.delta?.content) {
+              buffer += chunk.choices[0].delta.content
+              // Clean any code fence markers that might appear in the stream
+              const cleanedContent = buffer
+                .replace(/^```[\w-]*\n|```\s*$/gm, "") // Remove code fences
+                .replace(/^(javascript|typescript|python|html|css)\n/gm, "") // Remove language identifiers
+              controller.enqueue(encoder.encode(cleanedContent))
+              buffer = ""
+            }
+          }
+          controller.close()
+        },
+      })
+    )
+  } catch (error) {
+    console.error("Merge error:", error)
+    return new Response(
+      error instanceof Error ? error.message : "Failed to merge code",
+      { status: 500 }
+    )
+  }
+}
--- a/frontend/app/globals.css
+++ b/frontend/app/globals.css
@ -175,3 +175,23 @@
 .tab-scroll::-webkit-scrollbar {
  display: none;
 }
+
+.added-line-decoration {
+  background-color: rgba(0, 255, 0, 0.1);
+}
+
+.removed-line-decoration {
+  background-color: rgba(255, 0, 0, 0.1);
+}
+
+.added-line-glyph {
+  background-color: #28a745;
+  width: 4px !important;
+  margin-left: 3px;
+}
+
+.removed-line-glyph {
+  background-color: #dc3545;
+  width: 4px !important;
+  margin-left: 3px;
+}
--- a/frontend/components/editor/AIChat/ApplyButton.tsx
+++ b/frontend/components/editor/AIChat/ApplyButton.tsx
@ -0,0 +1,77 @@
+import { Check, Loader2 } from "lucide-react"
+import { useState } from "react"
+import { toast } from "sonner"
+import { Button } from "../../ui/button"
+
+interface ApplyButtonProps {
+  code: string
+  activeFileName: string
+  activeFileContent: string
+  editorRef: { current: any }
+  onApply: (mergedCode: string) => void
+}
+
+export default function ApplyButton({
+  code,
+  activeFileName,
+  activeFileContent,
+  editorRef,
+  onApply,
+}: ApplyButtonProps) {
+  const [isApplying, setIsApplying] = useState(false)
+
+  const handleApply = async () => {
+    setIsApplying(true)
+    try {
+      const response = await fetch("/api/merge", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          originalCode: activeFileContent,
+          newCode: String(code),
+          fileName: activeFileName,
+        }),
+      })
+
+      if (!response.ok) {
+        throw new Error(await response.text())
+      }
+
+      const reader = response.body?.getReader()
+      const decoder = new TextDecoder()
+      let mergedCode = ""
+
+      if (reader) {
+        while (true) {
+          const { done, value } = await reader.read()
+          if (done) break
+          mergedCode += decoder.decode(value, { stream: true })
+        }
+      }
+      onApply(mergedCode.trim())
+    } catch (error) {
+      console.error("Error applying code:", error)
+      toast.error(
+        error instanceof Error ? error.message : "Failed to apply code changes"
+      )
+    } finally {
+      setIsApplying(false)
+    }
+  }
+
+  return (
+    <Button
+      onClick={handleApply}
+      size="sm"
+      variant="ghost"
+      className="p-1 h-6"
+      disabled={isApplying}
+    >
+      {isApplying ? (
+        <Loader2 className="w-4 h-4 animate-spin" />
+      ) : (
+        <Check className="w-4 h-4" />
+      )}
+    </Button>
+  )
+}
--- a/frontend/components/editor/AIChat/ChatMessage.tsx
+++ b/frontend/components/editor/AIChat/ChatMessage.tsx
@ -13,6 +13,12 @@ export default function ChatMessage({
  setContext,
  setIsContextExpanded,
  socket,
+  handleApplyCode,
+  activeFileName,
+  activeFileContent,
+  editorRef,
+  mergeDecorationsCollection,
+  setMergeDecorationsCollection,
 }: MessageProps) {
  // State for expanded message index
  const [expandedMessageIndex, setExpandedMessageIndex] = useState<
@ -104,7 +110,13 @@ export default function ChatMessage({
  const components = createMarkdownComponents(
    renderCopyButton,
    renderMarkdownElement,
-    askAboutCode
+    askAboutCode,
+    activeFileName,
+    activeFileContent,
+    editorRef,
+    handleApplyCode,
+    mergeDecorationsCollection,
+    setMergeDecorationsCollection
  )

  return (
--- a/frontend/components/editor/AIChat/index.tsx
+++ b/frontend/components/editor/AIChat/index.tsx
@ -1,6 +1,6 @@
 import { useSocket } from "@/context/SocketContext"
 import { TFile } from "@/lib/types"
-import { X, ChevronDown } from "lucide-react"
+import { ChevronDown, X } from "lucide-react"
 import { nanoid } from "nanoid"
 import { useEffect, useRef, useState } from "react"
 import LoadingDots from "../../ui/LoadingDots"
@ -18,6 +18,9 @@ export default function AIChat({
  lastCopiedRangeRef,
  files,
  templateType,
+  handleApplyCode,
+  mergeDecorationsCollection,
+  setMergeDecorationsCollection,
 }: AIChatProps) {
  // Initialize socket and messages
  const { socket } = useSocket()
@ -176,7 +179,7 @@ export default function AIChat({
          const fileExt = tab.name.split(".").pop() || "txt"
          return {
            ...tab,
-            content: `\`\`\`${fileExt}\n${activeFileContent}\n\`\`\``
+            content: `\`\`\`${fileExt}\n${activeFileContent}\n\`\`\``,
          }
        }
        return tab
@ -214,6 +217,12 @@ export default function AIChat({
            setContext={setContext}
            setIsContextExpanded={setIsContextExpanded}
            socket={socket}
+            handleApplyCode={handleApplyCode}
+            activeFileName={activeFileName}
+            activeFileContent={activeFileContent}
+            editorRef={editorRef}
+            mergeDecorationsCollection={mergeDecorationsCollection}
+            setMergeDecorationsCollection={setMergeDecorationsCollection}
          />
        ))}
        {isLoading && <LoadingDots />}
--- a/frontend/components/editor/AIChat/lib/markdownComponents.tsx
+++ b/frontend/components/editor/AIChat/lib/markdownComponents.tsx
@ -1,15 +1,23 @@
-import { CornerUpLeft } from "lucide-react"
+import { Check, CornerUpLeft, X } from "lucide-react"
+import monaco from "monaco-editor"
 import { Components } from "react-markdown"
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"
 import { vscDarkPlus } from "react-syntax-highlighter/dist/esm/styles/prism"
 import { Button } from "../../../ui/button"
+import ApplyButton from "../ApplyButton"
 import { stringifyContent } from "./chatUtils"

 // Create markdown components for chat message component
 export const createMarkdownComponents = (
  renderCopyButton: (text: any) => JSX.Element,
  renderMarkdownElement: (props: any) => JSX.Element,
-  askAboutCode: (code: any) => void
+  askAboutCode: (code: any) => void,
+  activeFileName: string,
+  activeFileContent: string,
+  editorRef: any,
+  handleApplyCode: (mergedCode: string) => void,
+  mergeDecorationsCollection?: monaco.editor.IEditorDecorationsCollection,
+  setMergeDecorationsCollection?: (collection: undefined) => void
 ): Components => ({
  code: ({
    node,
@ -33,6 +41,57 @@ export const createMarkdownComponents = (
          <div className="flex border border-input shadow-lg bg-background rounded-md">
            {renderCopyButton(children)}
            <div className="w-px bg-input"></div>
+            {!mergeDecorationsCollection ? (
+              <ApplyButton
+                code={String(children)}
+                activeFileName={activeFileName}
+                activeFileContent={activeFileContent}
+                editorRef={editorRef}
+                onApply={handleApplyCode}
+              />
+            ) : (
+              <>
+                <Button
+                  onClick={() => {
+                    if (
+                      setMergeDecorationsCollection &&
+                      mergeDecorationsCollection &&
+                      editorRef?.current
+                    ) {
+                      mergeDecorationsCollection.clear()
+                      setMergeDecorationsCollection(undefined)
+                    }
+                  }}
+                  size="sm"
+                  variant="ghost"
+                  className="p-1 h-6"
+                  title="Accept Changes"
+                >
+                  <Check className="w-4 h-4 text-green-500" />
+                </Button>
+                <div className="w-px bg-input"></div>
+                <Button
+                  onClick={() => {
+                    if (
+                      setMergeDecorationsCollection &&
+                      mergeDecorationsCollection &&
+                      editorRef?.current
+                    ) {
+                      editorRef.current.getModel()?.setValue(activeFileContent)
+                      mergeDecorationsCollection.clear()
+                      setMergeDecorationsCollection(undefined)
+                    }
+                  }}
+                  size="sm"
+                  variant="ghost"
+                  className="p-1 h-6"
+                  title="Discard Changes"
+                >
+                  <X className="w-4 h-4 text-red-500" />
+                </Button>
+              </>
+            )}
+            <div className="w-px bg-input"></div>
            <Button
              onClick={(e) => {
                e.preventDefault()
--- a/frontend/components/editor/AIChat/types/index.ts
+++ b/frontend/components/editor/AIChat/types/index.ts
@ -56,6 +56,9 @@ export interface AIChatProps {
  files: (TFile | TFolder)[]
  templateType: string
  templateConfig?: TemplateConfig
+  handleApplyCode: (mergedCode: string) => void
+  mergeDecorationsCollection?: monaco.editor.IEditorDecorationsCollection
+  setMergeDecorationsCollection?: (collection: undefined) => void
 }

 // Chat input props interface
@ -105,6 +108,12 @@ export interface MessageProps {
  ) => void
  setIsContextExpanded: (isExpanded: boolean) => void
  socket: Socket | null
+  handleApplyCode: (mergedCode: string) => void
+  activeFileName: string
+  activeFileContent: string
+  editorRef: any
+  mergeDecorationsCollection?: monaco.editor.IEditorDecorationsCollection
+  setMergeDecorationsCollection?: (collection: undefined) => void
 }

 // Context tabs props interface
--- a/frontend/components/editor/index.tsx
+++ b/frontend/components/editor/index.tsx
@ -104,6 +104,13 @@ export default function CodeEditor({
  // Added this state to track the most recent content for each file
  const [fileContents, setFileContents] = useState<Record<string, string>>({})

+  // Apply Button merger decoration state
+  const [mergeDecorations, setMergeDecorations] = useState<
+    monaco.editor.IModelDeltaDecoration[]
+  >([])
+  const [mergeDecorationsCollection, setMergeDecorationsCollection] =
+    useState<monaco.editor.IEditorDecorationsCollection>()
+
  // Editor state
  const [editorLanguage, setEditorLanguage] = useState("plaintext")
  const [cursorLine, setCursorLine] = useState(0)
@ -375,6 +382,49 @@ export default function CodeEditor({
    })
  }, [editorRef])

+  // handle apply code
+  const handleApplyCode = useCallback(
+    (mergedCode: string) => {
+      if (!editorRef) return
+
+      const originalLines = activeFileContent.split("\n")
+      const mergedLines = mergedCode.split("\n")
+
+      const decorations: monaco.editor.IModelDeltaDecoration[] = []
+
+      for (
+        let i = 0;
+        i < Math.max(originalLines.length, mergedLines.length);
+        i++
+      ) {
+        if (originalLines[i] !== mergedLines[i]) {
+          decorations.push({
+            range: new monaco.Range(i + 1, 1, i + 1, 1),
+            options: {
+              isWholeLine: true,
+              className:
+                i < originalLines.length
+                  ? "removed-line-decoration"
+                  : "added-line-decoration",
+              glyphMarginClassName:
+                i < originalLines.length
+                  ? "removed-line-glyph"
+                  : "added-line-glyph",
+            },
+          })
+        }
+      }
+
+      // Update editor content
+      editorRef.setValue(mergedCode)
+
+      // Apply decorations
+      const newDecorations = editorRef.createDecorationsCollection(decorations)
+      setMergeDecorationsCollection(newDecorations)
+    },
+    [editorRef, activeFileContent]
+  )
+
  // Generate widget effect
  useEffect(() => {
    if (generate.show) {
@ -1234,6 +1284,9 @@ export default function CodeEditor({
                  lastCopiedRangeRef={lastCopiedRangeRef}
                  files={files}
                  templateType={sandboxData.type}
+                  handleApplyCode={handleApplyCode}
+                  mergeDecorationsCollection={mergeDecorationsCollection}
+                  setMergeDecorationsCollection={setMergeDecorationsCollection}
                />
              </ResizablePanel>
            </>
--- a/frontend/package-lock.json
+++ b/frontend/package-lock.json
@ -54,6 +54,7 @@
        "monaco-themes": "^0.4.4",
        "next": "14.1.3",
        "next-themes": "^0.3.0",
+        "openai": "^4.73.1",
        "posthog-js": "^1.147.0",
        "react": "^18.3.1",
        "react-dom": "^18",
@ -7492,6 +7493,65 @@
        "url": "https://github.com/sponsors/sindresorhus"
      }
    },
+    "node_modules/openai": {
+      "version": "4.73.1",
+      "resolved": "https://registry.npmjs.org/openai/-/openai-4.73.1.tgz",
+      "integrity": "sha512-nWImDJBcUsqrhy7yJScXB4+iqjzbUEgzfA3un/6UnHFdwWhjX24oztj69Ped/njABfOdLcO/F7CeWTI5dt8Xmg==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@types/node": "^18.11.18",
+        "@types/node-fetch": "^2.6.4",
+        "abort-controller": "^3.0.0",
+        "agentkeepalive": "^4.2.1",
+        "form-data-encoder": "1.7.2",
+        "formdata-node": "^4.3.2",
+        "node-fetch": "^2.6.7"
+      },
+      "bin": {
+        "openai": "bin/cli"
+      },
+      "peerDependencies": {
+        "zod": "^3.23.8"
+      },
+      "peerDependenciesMeta": {
+        "zod": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/openai/node_modules/@types/node": {
+      "version": "18.19.67",
+      "resolved": "https://registry.npmjs.org/@types/node/-/node-18.19.67.tgz",
+      "integrity": "sha512-wI8uHusga+0ZugNp0Ol/3BqQfEcCCNfojtO6Oou9iVNGPTL6QNSdnUdqq85fRgIorLhLMuPIKpsN98QE9Nh+KQ==",
+      "license": "MIT",
+      "dependencies": {
+        "undici-types": "~5.26.4"
+      }
+    },
+    "node_modules/openai/node_modules/@types/node-fetch": {
+      "version": "2.6.12",
+      "resolved": "https://registry.npmjs.org/@types/node-fetch/-/node-fetch-2.6.12.tgz",
+      "integrity": "sha512-8nneRWKCg3rMtF69nLQJnOYUcbafYeFSjqkw3jCRLsqkWFlHaoQrr5mXmofFGOx3DKn7UfmBMyov8ySvLRVldA==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/node": "*",
+        "form-data": "^4.0.0"
+      }
+    },
+    "node_modules/openai/node_modules/form-data": {
+      "version": "4.0.1",
+      "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.1.tgz",
+      "integrity": "sha512-tzN8e4TX8+kkxGPK8D5u0FNmjPUjw3lwC9lSLxxoB/+GtsJG91CO8bSWy73APlgAZzZbXEYZJuxjkHH2w+Ezhw==",
+      "license": "MIT",
+      "dependencies": {
+        "asynckit": "^0.4.0",
+        "combined-stream": "^1.0.8",
+        "mime-types": "^2.1.12"
+      },
+      "engines": {
+        "node": ">= 6"
+      }
+    },
    "node_modules/ora": {
      "version": "6.3.1",
      "resolved": "https://registry.npmjs.org/ora/-/ora-6.3.1.tgz",
--- a/frontend/package.json
+++ b/frontend/package.json
@ -55,6 +55,7 @@
    "monaco-themes": "^0.4.4",
    "next": "14.1.3",
    "next-themes": "^0.3.0",
+    "openai": "^4.73.1",
    "posthog-js": "^1.147.0",
    "react": "^18.3.1",
    "react-dom": "^18",