diff --git a/app/(main)/page.tsx b/app/(main)/page.tsx
index 15bb0f0e8..d77868eac 100644
--- a/app/(main)/page.tsx
+++ b/app/(main)/page.tsx
@@ -12,7 +12,9 @@ import { FormEvent, useEffect, useState } from "react";
 import LoadingDots from "../../components/loading-dots";
 
 function removeCodeFormatting(code: string): string {
-  return code.replace(/```(?:typescript|javascript|tsx)?\n([\s\S]*?)```/g, '$1').trim();
+  return code
+    .replace(/```(?:typescript|javascript|tsx)?\n([\s\S]*?)```/g, "$1")
+    .trim();
 }
 
 export default function Home() {
@@ -32,7 +34,7 @@ export default function Home() {
     {
       label: "gemini-1.5-flash",
       value: "gemini-1.5-flash",
-    }
+    },
   ];
   let [model, setModel] = useState(models[0].value);
   let [modification, setModification] = useState("");
@@ -94,6 +96,54 @@ export default function Home() {
     setStatus("created");
   }
 
+  async function modifyApp(e: FormEvent<HTMLFormElement>) {
+    e.preventDefault();
+    scrollTo({ delay: 0.5 });
+    setStatus("updating");
+
+    const newMessages = [
+      ...messages,
+      { role: "assistant", content: generatedCode },
+      { role: "user", content: modification },
+    ];
+
+    let res = await fetch("/api/generateCode", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        model,
+        messages: newMessages,
+      }),
+    });
+
+    if (!res.ok) {
+      throw new Error(res.statusText);
+    }
+
+    if (!res.body) {
+      throw new Error("No response body");
+    }
+
+    const reader = res.body.getReader();
+    let receivedData = "";
+
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) {
+        break;
+      }
+      receivedData += new TextDecoder().decode(value);
+      const cleanedData = removeCodeFormatting(receivedData);
+      setGeneratedCode(cleanedData);
+    }
+
+    setMessages(newMessages);
+    setModification("");
+    setStatus("updated");
+  }
+
   useEffect(() => {
     let el = document.querySelector(".cm-scroller");
     if (el && loading) {
@@ -105,7 +155,7 @@ export default function Home() {
   return (
     <main className="mt-12 flex w-full flex-1 flex-col items-center px-4 text-center sm:mt-1">
       <a
-        className="mb-4 inline-flex h-7 shrink-0 items-center gap-[9px] rounded-[50px] border-[0.5px] border-solid border-[#E6E6E6] bg-[rgba(234,238,255,0.65)] dark:bg-[rgba(30,41,59,0.5)] dark:border-gray-700 px-7 py-5 shadow-[0px_1px_1px_0px_rgba(0,0,0,0.25)]"
+        className="mb-4 inline-flex h-7 shrink-0 items-center gap-[9px] rounded-[50px] border-[0.5px] border-solid border-[#E6E6E6] bg-[rgba(234,238,255,0.65)] px-7 py-5 shadow-[0px_1px_1px_0px_rgba(0,0,0,0.25)] dark:border-gray-700 dark:bg-[rgba(30,41,59,0.5)]"
         href="https://ai.google.dev/gemini-api/docs"
         target="_blank"
       >
@@ -122,7 +172,7 @@ export default function Home() {
         <fieldset disabled={loading} className="disabled:opacity-75">
           <div className="relative mt-5">
             <div className="absolute -inset-2 rounded-[32px] bg-gray-300/50 dark:bg-gray-800/50" />
-            <div className="relative flex rounded-3xl bg-white dark:bg-[#1E293B] shadow-sm">
+            <div className="relative flex rounded-3xl bg-white shadow-sm dark:bg-[#1E293B]">
               <div className="relative flex flex-grow items-stretch focus-within:z-10">
                 <textarea
                   rows={3}
@@ -149,27 +199,29 @@ export default function Home() {
           </div>
           <div className="mt-6 flex flex-col justify-center gap-4 sm:flex-row sm:items-center sm:gap-8">
             <div className="flex items-center justify-between gap-3 sm:justify-center">
-              <p className="text-gray-500 dark:text-gray-400 sm:text-xs">Model:</p>
+              <p className="text-gray-500 dark:text-gray-400 sm:text-xs">
+                Model:
+              </p>
               <Select.Root
                 name="model"
                 disabled={loading}
                 value={model}
                 onValueChange={(value) => setModel(value)}
               >
-                <Select.Trigger className="group flex w-60 max-w-xs items-center rounded-2xl border-[6px] border-gray-300 dark:border-gray-700 bg-white dark:bg-[#1E293B] px-4 py-2 text-sm focus-visible:outline focus-visible:outline-2 focus-visible:outline-blue-500">
+                <Select.Trigger className="group flex w-60 max-w-xs items-center rounded-2xl border-[6px] border-gray-300 bg-white px-4 py-2 text-sm focus-visible:outline focus-visible:outline-2 focus-visible:outline-blue-500 dark:border-gray-700 dark:bg-[#1E293B]">
                   <Select.Value />
                   <Select.Icon className="ml-auto">
                     <ChevronDownIcon className="size-6 text-gray-300 group-focus-visible:text-gray-500 group-enabled:group-hover:text-gray-500 dark:text-gray-600 dark:group-focus-visible:text-gray-400 dark:group-enabled:group-hover:text-gray-400" />
                   </Select.Icon>
                 </Select.Trigger>
                 <Select.Portal>
-                  <Select.Content className="overflow-hidden rounded-md bg-white dark:bg-[#1E293B] shadow-lg">
+                  <Select.Content className="overflow-hidden rounded-md bg-white shadow-lg dark:bg-[#1E293B]">
                     <Select.Viewport className="p-2">
                       {models.map((model) => (
                         <Select.Item
                           key={model.value}
                           value={model.value}
-                          className="flex cursor-pointer items-center rounded-md px-3 py-2 text-sm data-[highlighted]:bg-gray-100 dark:data-[highlighted]:bg-gray-800 data-[highlighted]:outline-none"
+                          className="flex cursor-pointer items-center rounded-md px-3 py-2 text-sm data-[highlighted]:bg-gray-100 data-[highlighted]:outline-none dark:data-[highlighted]:bg-gray-800"
                         >
                           <Select.ItemText asChild>
                             <span className="inline-flex items-center gap-2 text-gray-500 dark:text-gray-400">
@@ -211,6 +263,41 @@ export default function Home() {
           <div className="relative mt-8 w-full overflow-hidden">
             <div className="isolate">
               <CodeViewer code={generatedCode} showEditor />
+
+              {(status === "created" || status === "updated") && (
+                <form
+                  onSubmit={modifyApp}
+                  className="mx-auto mt-8 w-full max-w-xl"
+                >
+                  <fieldset disabled={loading} className="disabled:opacity-75">
+                    <div className="relative">
+                      <div className="absolute rounded-[32px] bg-gray-300/50 dark:bg-gray-800/50" />
+                      <div className="relative flex rounded-3xl bg-white shadow-sm dark:bg-[#1E293B]">
+                        <textarea
+                          rows={3}
+                          required
+                          value={modification}
+                          onChange={(e) => setModification(e.target.value)}
+                          name="modification"
+                          className="w-full resize-none rounded-l-3xl bg-transparent px-6 py-5 text-lg focus-visible:outline focus-visible:outline-2 focus-visible:outline-blue-500 dark:text-gray-100 dark:placeholder-gray-400"
+                          placeholder="Modify the code... (e.g. make the calculator have a dark theme)"
+                        />
+                        <button
+                          type="submit"
+                          disabled={loading}
+                          className="relative -ml-px inline-flex items-center gap-x-1.5 rounded-r-3xl px-3 py-2 text-sm font-semibold text-blue-500 hover:text-blue-400 focus-visible:outline focus-visible:outline-2 focus-visible:outline-blue-500 disabled:text-gray-900 dark:disabled:text-gray-400"
+                        >
+                          {status === "updating" ? (
+                            <LoadingDots color="black" style="large" />
+                          ) : (
+                            <ArrowLongRightIcon className="size-6" />
+                          )}
+                        </button>
+                      </div>
+                    </div>
+                  </fieldset>
+                </form>
+              )}
             </div>
 
             <AnimatePresence>
@@ -225,7 +312,7 @@ export default function Home() {
                     duration: 0.85,
                     delay: 0.5,
                   }}
-                  className="absolute inset-x-0 bottom-0 top-1/2 flex items-center justify-center rounded-r border border-gray-400 dark:border-gray-700 bg-gradient-to-br from-gray-100 to-gray-300 dark:from-[#1E293B] dark:to-gray-800 md:inset-y-0 md:left-1/2 md:right-0"
+                  className="absolute inset-x-0 bottom-0 top-1/2 flex items-center justify-center rounded-r border border-gray-400 bg-gradient-to-br from-gray-100 to-gray-300 dark:border-gray-700 dark:from-[#1E293B] dark:to-gray-800 md:inset-y-0 md:left-1/2 md:right-0"
                 >
                   <p className="animate-pulse text-3xl font-bold dark:text-gray-100">
                     {status === "creating"
diff --git a/app/api/generateCode/route.ts b/app/api/generateCode/route.ts
index f57949798..345382ae6 100644
--- a/app/api/generateCode/route.ts
+++ b/app/api/generateCode/route.ts
@@ -28,11 +28,13 @@ export async function POST(req: Request) {
 
   const geminiModel = genAI.getGenerativeModel({model: model});
 
-  const geminiStream = await geminiModel.generateContentStream(
-    messages[0].content + systemPrompt + "\nPlease ONLY return code, NO backticks or language names. Don't start with \`\`\`typescript or \`\`\`javascript or \`\`\`tsx or \`\`\`."
-  );
+  // Combine all messages with system prompt for context
+  const fullPrompt = messages.map(msg => msg.content).join("\n") + systemPrompt + 
+    "\nPlease ONLY return code, NO backticks or language names. Don't start with ```typescript or ```javascript or ```tsx or ```."
 
-  console.log(messages[0].content + systemPrompt + "\nPlease ONLY return code, NO backticks or language names. Don't start with \`\`\`typescript or \`\`\`javascript or \`\`\`tsx or \`\`\`.")
+  const geminiStream = await geminiModel.generateContentStream(fullPrompt);
+
+  console.log(fullPrompt)
 
   const readableStream = new ReadableStream({
     async start(controller) {
@@ -61,12 +63,13 @@ function getSystemPrompt() {
 - Please ONLY return the full React code starting with the imports, nothing else. It's very important for my job that you only return the React code with imports. DO NOT START WITH \`\`\`typescript or \`\`\`javascript or \`\`\`tsx or \`\`\`.
 - ONLY IF the user asks for a dashboard, graph or chart, the recharts library is available to be imported, e.g. \`import { LineChart, XAxis, ... } from "recharts"\` & \`<LineChart ...><XAxis dataKey="name"> ...\`. Please only use this when needed.
 - For placeholder images, please use a <div className="bg-gray-200 border-2 border-dashed rounded-xl w-16 h-16" />
+- For icons you can only use react-icons package or make custom icons yourself.
   `;
 
   systemPrompt += `
     NO OTHER LIBRARIES (e.g. zod, hookform) ARE INSTALLED OR ABLE TO BE IMPORTED.
   `;
-
+  
   return dedent(systemPrompt);
 }
 
diff --git a/components/code-viewer.tsx b/components/code-viewer.tsx
index 47d8daa7b..789f873a4 100644
--- a/components/code-viewer.tsx
+++ b/components/code-viewer.tsx
@@ -91,6 +91,8 @@ let sharedProps = {
       "tailwind-merge": "^2.4.0",
       "tailwindcss-animate": "^1.0.7",
       vaul: "^0.9.1",
+      "react-icons": "^5.4.0",
+      "@heroicons/react": "^2.2.0",
     },
   },
 } as const;