diff --git a/ai.js b/ai.js
index c16f447..1fec7ff 100644
--- a/ai.js
+++ b/ai.js
@@ -37,10 +37,13 @@ document.addEventListener('DOMContentLoaded', () => {
             fullPrompt = prompt.replace('@code', `\n\n\`\`\`\n${editor.value}\n\`\`\`\n\n`);
         }
 
-        appendMessage('ai', 'Thinking...', true);
+        const model = document.getElementById('ai-model-selector').value;
+        const streaming = true; // For this example, we'll always stream
+
+        appendMessage('ai', '', true); // Add a placeholder for the AI response
 
         try {
-            const response = await fetch(`https://generativelanguage.googleapis.com/v1beta/models/gemini-3-flash-preview:generateContent?key=${apiKey}`, {
+            const response = await fetch(`https://generativelanguage.googleapis.com/v1beta/models/${model}:${streaming ? 'streamGenerateContent' : 'generateContent'}?key=${apiKey}`, {
                 method: 'POST',
                 headers: {
                     'Content-Type': 'application/json',
@@ -58,13 +61,45 @@ document.addEventListener('DOMContentLoaded', () => {
                 throw new Error(`API request failed with status ${response.status}`);
             }
 
-            const data = await response.json();
-            const aiResponse = data.candidates[0].content.parts[0].text;
-            updateLastMessage(aiResponse);
+            const reader = response.body.getReader();
+            const decoder = new TextDecoder();
+            let accumulatedText = '';
+
+            while (true) {
+                const {
+                    done,
+                    value
+                } = await reader.read();
+                if (done) break;
+
+                const chunk = decoder.decode(value, {
+                    stream: true
+                });
+                const lines = chunk.split('\n');
+
+                for (const line of lines) {
+                    if (line.startsWith('data: ')) {
+                        try {
+                            const json = JSON.parse(line.substring(6));
+                            if (json.candidates && json.candidates[0].content.parts[0].text) {
+                                const text = json.candidates[0].content.parts[0].text;
+                                accumulatedText += text;
+                                updateLastMessage(accumulatedText, false);
+                            }
+                        } catch (error) {
+                            console.error('Error parsing streaming JSON:', error);
+                        }
+                    }
+                }
+            }
+
+            // Final update to add the "Replace Code" button if applicable
+            updateLastMessage(accumulatedText, true);
+
 
         } catch (error) {
             console.error('Error fetching AI response:', error);
-            updateLastMessage('Sorry, I encountered an error. Please try again.');
+            updateLastMessage('Sorry, I encountered an error. Please try again.', true);
         }
     }
 
@@ -80,23 +115,30 @@ document.addEventListener('DOMContentLoaded', () => {
         aiChatContainer.scrollTop = aiChatContainer.scrollHeight;
     }
 
-    function updateLastMessage(text) {
+    function updateLastMessage(text, isFinal) {
         const lastMessage = aiChatContainer.querySelector('.ai-message:last-child');
         if (lastMessage) {
-            lastMessage.innerHTML = ''; // Clear spinner
-            lastMessage.textContent = text;
-
-            const codeBlocks = text.match(/```(\w+)?\n([\s\S]*?)```/g);
-            if (codeBlocks) {
-                const replaceBtn = document.createElement('button');
-                replaceBtn.textContent = 'Replace Code';
-                replaceBtn.classList.add('replace-code-btn');
-                replaceBtn.onclick = () => {
-                    const codeToInsert = codeBlocks.map(block => block.replace(/```(\w+)?\n|```/g, '')).join('\n');
-                    editor.value = codeToInsert;
-                    aiModalOverlay.classList.add('hidden');
-                };
-                lastMessage.appendChild(replaceBtn);
+            // Render markdown content
+            lastMessage.innerHTML = marked.parse(text);
+
+            // Highlight code blocks
+            lastMessage.querySelectorAll('pre code').forEach((block) => {
+                hljs.highlightElement(block);
+            });
+
+            if (isFinal) {
+                const codeBlocks = text.match(/```(\w+)?\n([\s\S]*?)```/g);
+                if (codeBlocks) {
+                    const replaceBtn = document.createElement('button');
+                    replaceBtn.textContent = 'Replace Code';
+                    replaceBtn.classList.add('replace-code-btn');
+                    replaceBtn.onclick = () => {
+                        const codeToInsert = codeBlocks.map(block => block.replace(/```(\w+)?\n|```/g, '')).join('\n');
+                        editor.value = codeToInsert;
+                        aiModalOverlay.classList.add('hidden');
+                    };
+                    lastMessage.appendChild(replaceBtn);
+                }
             }
         }
     }
diff --git a/editor.html b/editor.html
index 458dd69..d0cb780 100644
--- a/editor.html
+++ b/editor.html
@@ -138,7 +138,16 @@ <h3 id="modal-title"></h3>
     <div id="ai-modal-overlay" class="hidden">
         <div id="ai-modal">
             <div id="ai-modal-header">
-                <h2>AI Assistant</h2>
+                <div class="header-left-content">
+                    <h2>AI Assistant</h2>
+                    <div class="model-selector-container">
+                        <i class="fa-solid fa-robot"></i>
+                        <select id="ai-model-selector">
+                            <option value="models/gemini-3-flash-preview">Gemini 3.0 Flash</option>
+                            <option value="models/gemini-flash-latest">Gemini 2.5 Flash</option>
+                        </select>
+                    </div>
+                </div>
                 <button id="ai-modal-close-btn"><i class="fa-solid fa-xmark"></i></button>
             </div>
             <div id="ai-modal-body">
diff --git a/style.css b/style.css
index c73c51a..ebe7498 100644
--- a/style.css
+++ b/style.css
@@ -1261,4 +1261,52 @@ body {
     border-radius: 20px;
     padding: 0.5rem 1rem;
     cursor: pointer;
+}
+
+.model-selector-container {
+    display: flex;
+    align-items: center;
+    gap: 0.5rem;
+    margin-left: 1rem;
+}
+
+.model-selector-container i {
+    color: var(--text-secondary);
+}
+
+#ai-model-selector {
+    background-color: var(--bg-secondary);
+    color: var(--text-primary);
+    border: 1px solid var(--border-color);
+    border-radius: 12px;
+    padding: 0.25rem 0.5rem;
+    font-family: var(--font-sans);
+    font-size: 0.875rem;
+}
+
+#ai-model-selector:focus {
+    outline: none;
+    border-color: var(--accent-color);
+}
+
+.header-left-content {
+    display: flex;
+    align-items: center;
+}
+
+.ai-message pre {
+    white-space: pre-wrap;
+    word-wrap: break-word;
+    background-color: var(--code-bg);
+    border-radius: 8px;
+    padding: 1rem;
+    margin: 1rem 0;
+}
+
+.ai-message code {
+    font-family: var(--font-mono);
+}
+
+.ai-message p {
+    margin-bottom: 1rem;
 }
\ No newline at end of file