feat: Update model options and enhance AI response handling in webview

Author: TheUnknown550

server/server.py

@@ -9,7 +9,7 @@
 
 MODELS = {
     "gemma3-4b": "gemma3:4b",
-    "gemma3-8b": "gemma3:8b",
+    "gemma3-12b": "gemma3:12b",
     "deepseek-r1-8b": "deepseek-r1:8b"
 }
 
@@ -51,10 +51,21 @@ def completion():
         prompt = data.get("prompt", "")
         context = data.get("context", "")
         task_type = data.get("task_type", "chat")  # 'chat', 'code', 'edit'
-        
-        if model not in MODELS:
-            return jsonify({"error": f"Model {model} not available"}), 400
-        
+
+        # Accept both "gemma3:4b" and "gemma3-4b" formats
+        model_key = model.replace(":", "-")
+        if model_key not in MODELS and model not in MODELS.values():
+            print(f"Model not found: {model}")
+            return jsonify({"error": f"Model {model} not available. Valid models: {list(MODELS.keys()) + list(MODELS.values())}"}), 400
+
+        # Use Ollama model name
+        ollama_model = MODELS.get(model_key, model)
+
+        # Validate prompt
+        if not prompt:
+            print("Missing prompt in request.")
+            return jsonify({"error": "Missing prompt in request."}), 400
+
         # Build the prompt based on task type
         if task_type == "code":
             system_prompt = "You are a coding assistant. Provide only code without explanations unless asked. Format code properly."
@@ -69,7 +80,7 @@ def completion():
         response = requests.post(
             OLLAMA_API_URL,
             json={
-                "model": MODELS[model],
+                "model": ollama_model,
                 "prompt": full_prompt,
                 "stream": False
             },
@@ -80,11 +91,19 @@ def completion():
             return jsonify({"error": f"Ollama API error: {response.text}"}), 500
 
         result = response.json()
+        print("Ollama API response:", result)
         completion_text = result.get("response", "").strip()
-        
+
+        if not completion_text:
+            print("Ollama returned an empty response.")
+            return jsonify({
+                "error": "Ollama returned an empty response. Check if the model is running and prompt is valid.",
+                "ollama_response": result
+            }), 500
+
         # Determine if this is code or chat
         is_code = detect_code_blocks(completion_text) or task_type in ["code", "edit"]
-        
+
         # Extract code if it's in markdown format
         if is_code and '```' in completion_text:
             code_content = extract_code_from_markdown(completion_text)
@@ -93,7 +112,7 @@ def completion():
                 "raw_completion": completion_text,
                 "completion_type": "code"
             })
-        
+
         return jsonify({
             "completion": completion_text,
             "completion_type": "code" if is_code else "chat"

vscode-extension/src/extension.ts

@@ -156,7 +156,7 @@ function openChatPanel(context: vscode.ExtensionContext) {
 
 async function handleChatMessage(message: any) {
     const editor = vscode.window.activeTextEditor;
-    
+
     // Show loading state
     currentPanel?.webview.postMessage({
         type: 'aiResponse',
@@ -169,28 +169,31 @@ async function handleChatMessage(message: any) {
         let context = '';
         let taskType = 'chat';
 
-        // Get context from active editor if available
+        // Smart context gathering
         if (editor) {
             const selection = editor.selection;
             const selectedText = editor.document.getText(selection);
-            
+            const document = editor.document;
+            const cursorPosition = editor.selection.active;
+            // Get up to 40 lines before and after cursor for context
+            const startLine = Math.max(0, cursorPosition.line - 40);
+            const endLine = Math.min(document.lineCount - 1, cursorPosition.line + 40);
+            const range = new vscode.Range(startLine, 0, endLine, document.lineAt(endLine).text.length);
+            const surroundingContext = document.getText(range);
+
             if (selectedText) {
-                context = selectedText;
+                context = selectedText + '\n\n--- Surrounding Context ---\n' + surroundingContext;
                 taskType = 'code';
             } else {
-                // Get surrounding context
-                const document = editor.document;
-                const cursorPosition = editor.selection.active;
-                const startLine = Math.max(0, cursorPosition.line - 10);
-                const endLine = Math.min(document.lineCount - 1, cursorPosition.line + 10);
-                const range = new vscode.Range(startLine, 0, endLine, document.lineAt(endLine).text.length);
-                context = document.getText(range);
+                context = surroundingContext;
             }
         }
 
         const response = await callAI(message.prompt, taskType, context, message.model);
+        console.log('AI response from server:', response);
 
         if (response.error) {
+            console.error('AI error response:', response.error);
             currentPanel?.webview.postMessage({
                 type: 'error',
                 content: response.error
@@ -200,16 +203,26 @@ async function handleChatMessage(message: any) {
         }
 
         const isCode = response.completion_type === 'code';
+        console.log('Sending message to webview:', isCode ? 'codeResponse' : 'chatResponse', response.completion);
 
-        currentPanel?.webview.postMessage({
-            type: 'aiResponse',
-            content: response.completion,
-            rawContent: response.raw_completion,
-            isCode: isCode,
-            loading: false
-        });
+        if (isCode) {
+            currentPanel?.webview.postMessage({
+                type: 'codeResponse',
+                code: response.completion,
+                rawContent: response.raw_completion,
+                explanation: 'Preview and insert, replace, or diff with selection.',
+                loading: false
+            });
+        } else {
+            currentPanel?.webview.postMessage({
+                type: 'chatResponse',
+                content: response.completion,
+                loading: false
+            });
+        }
 
     } catch (err: any) {
+        console.error('Exception in handleChatMessage:', err);
         currentPanel?.webview.postMessage({
             type: 'error',
             content: `Error: ${err.message}`
@@ -249,8 +262,14 @@ function insertCodeInEditor(code: string) {
         return;
     }
 
+    // Smart insertion: insert code at cursor, auto-indent
     editor.edit((editBuilder) => {
-        editBuilder.insert(editor.selection.active, code);
+        const position = editor.selection.active;
+        // Indent code to match current line
+        const currentLine = editor.document.lineAt(position.line);
+        const indent = currentLine.text.match(/^\s*/)?.[0] || '';
+        const indentedCode = code.split('\n').map((line, i) => i === 0 ? line : indent + line).join('\n');
+        editBuilder.insert(position, indentedCode);
     });
 }
 
@@ -263,7 +282,7 @@ async function applyCodeEdit(code: string) {
 
     const selection = editor.selection;
 
-    // Show diff first
+    // Show diff preview before applying edit
     const originalDocument = editor.document;
     const originalText = originalDocument.getText(selection);
 
@@ -278,7 +297,8 @@ async function applyCodeEdit(code: string) {
         'vscode.diff',
         originalDocument.uri,
         newDoc.uri,
-        'Original ← → AI Suggestion'
+        'Original ← → AI Suggestion',
+        { preview: true }
     );
 
     // Ask user to apply

vscode-extension/src/webview.html

@@ -174,7 +174,7 @@
             <label for="modelSelect">Model:</label>
             <select id="modelSelect">
                 <option value="gemma3:4b">Gemma3 4B (Fast)</option>
-                <option value="gemma3:8b">Gemma3 8B (Balanced)</option>
+                <option value="gemma3:12b">Gemma3 12B (Large)</option>
                 <option value="deepseek-r1:8b">DeepSeek-R1 8B (Reasoning)</option>
             </select>
         </div>

vscode-extension/src/webview.js

@@ -1,3 +1,4 @@
+alert('Webview loaded!');
 (function() {
     const vscode = acquireVsCodeApi();
 
@@ -133,7 +134,8 @@
     // Listen for messages from the extension
     window.addEventListener('message', event => {
         const message = event.data;
-        
+        console.log('Received message from extension:', message);
+
         removeLoadingMessage();
         isProcessing = false;
         sendButton.disabled = false;
@@ -150,6 +152,10 @@
                 }
                 break;
 
+            case 'aiResponse':
+                addMessage(message.content, 'ai', message.isCode);
+                break;
+            
             case 'error':
                 addMessage(`Error: ${message.content}`, 'system');
                 break;