Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
 import requests
 import json
+import time
 from faster_whisper import WhisperModel
 
 OLLAMA_URL = "http://localhost:11434"
+MAX_RETRIES = 3
+TIMEOUT = 300
 
 MODELS = {
     "⭐ Qwen3 30B-A3B (Best)": "hf.co/bartowski/Qwen_Qwen3-30B-A3B-GGUF:Q4_K_M",
@@ -19,13 +22,13 @@ MODELS = {
 }
 
 MODEL_INFO = {
-    "⭐ Qwen3 30B-A3B (Best)": "🏆 Best quality • MoE 30B/3B active • Complex tasks",
+    "⭐ Qwen3 30B-A3B (Best)": "🏆 Best quality • MoE 30B/3B active",
     "Qwen2.5 Coder 7B": "⚖️ Balanced • Great for most tasks",
     "Qwen2.5 Coder 3B": "🚀 Fast & capable • Recommended",
     "Qwen2.5 Coder 1.5B (Fast)": "⚡ Fastest • Simple tasks",
     "DeepSeek Coder 6.7B": "🧠 Complex logic • Algorithms",
     "DeepSeek Coder 1.3B (Fast)": "⚡ Quick completions",
-    "StarCoder2 7B": "🐙 GitHub trained • Real-world patterns",
+    "StarCoder2 7B": "🐙 GitHub trained",
     "StarCoder2 3B": "🐙 Fast GitHub style",
     "CodeGemma 7B": "🔷 Google • Strong docs",
     "CodeGemma 2B (Fast)": "🔷 Quick & efficient",
@@ -38,67 +41,119 @@ LANGUAGES = [
     "HTML/CSS", "SQL", "Bash", "PowerShell", "Lua"
 ]
 
-print("Loading Whisper...")
-whisper_model = WhisperModel("tiny", device="cpu", compute_type="int8")
-print("Whisper ready!")
+# Whisper init
+whisper_model = None
+try:
+    print("Loading Whisper...")
+    whisper_model = WhisperModel("tiny", device="cpu", compute_type="int8")
+    print("✅ Whisper ready!")
+except Exception as e:
+    print(f"❌ Whisper failed: {e}")
 
-# ===== HELPER FUNCTIONS =====
+# ===== HELPERS =====
+
+def check_ollama():
+    try:
+        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
+        return r.status_code == 200
+    except:
+        return False
 
 def get_status():
     try:
-        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=2)
+        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
         if r.status_code == 200:
             models = r.json().get("models", [])
             return f"🟢 Online • {len(models)} models"
     except:
         pass
-    return "🟡 Starting..."
+    return "🔴 Offline"
 
 def get_model_info(model_name):
     return MODEL_INFO.get(model_name, "")
 
+def validate_input(text, name="Input"):
+    if not text or not text.strip():
+        return False, f"⚠️ {name} cannot be empty."
+    if len(text) > 100000:
+        return False, f"⚠️ {name} too long (max 100KB)."
+    return True, None
+
 def transcribe_audio(audio):
     if audio is None:
         return ""
+    if whisper_model is None:
+        return "❌ Whisper not available."
     try:
         segments, _ = whisper_model.transcribe(audio)
-        return " ".join([seg.text for seg in segments]).strip()
+        text = " ".join([seg.text for seg in segments]).strip()
+        return text if text else "⚠️ No speech detected."
     except Exception as e:
-        return f"[STT Error: {e}]"
+        return f"❌ Transcription failed: {str(e)[:50]}"
 
 def call_ollama(model_name, prompt, temperature=0.7, max_tokens=2048):
+    if not check_ollama():
+        return "❌ **Ollama not running.** Please wait for it to start."
+    
     model = MODELS.get(model_name, "qwen2.5-coder:3b")
-    try:
-        r = requests.post(
-            f"{OLLAMA_URL}/api/generate",
-            json={"model": model, "prompt": prompt, "stream": False, "options": {"temperature": temperature, "num_predict": max_tokens}},
-            timeout=300
-        )
-        if r.status_code == 200:
-            return r.json().get("response", "")
-        return f"❌ Error: {r.text}"
-    except Exception as e:
-        return f"❌ Error: {e}"
+    
+    for attempt in range(MAX_RETRIES):
+        try:
+            r = requests.post(
+                f"{OLLAMA_URL}/api/generate",
+                json={"model": model, "prompt": prompt, "stream": False, 
+                      "options": {"temperature": temperature, "num_predict": max_tokens}},
+                timeout=TIMEOUT
+            )
+            
+            if r.status_code == 200:
+                response = r.json().get("response", "")
+                return response if response.strip() else "⚠️ Empty response. Try rephrasing."
+            elif r.status_code == 404:
+                return f"❌ **Model not found:** `{model}`"
+            else:
+                return f"❌ **Error {r.status_code}**"
+                
+        except requests.exceptions.Timeout:
+            if attempt < MAX_RETRIES - 1:
+                time.sleep(2)
+                continue
+            return "❌ **Timeout.** Try smaller model."
+        except requests.exceptions.ConnectionError:
+            return "❌ **Connection failed.**"
+        except Exception as e:
+            return f"❌ **Error:** {str(e)[:50]}"
+    
+    return "❌ **Max retries reached.**"
 
 def extract_code(text):
-    if "```" in text:
+    if not text or "```" not in text:
+        return text
+    try:
         parts = text.split("```")
         if len(parts) >= 2:
             code = parts[1]
             if "\n" in code:
                 code = code.split("\n", 1)[-1]
             return code.strip()
+    except:
+        pass
     return text
 
 # ===== CORE FUNCTIONS =====
 
 def chat_stream(message, history, model_name, temperature, max_tokens):
-    if not message.strip():
-        yield history
+    valid, error = validate_input(message, "Message")
+    if not valid:
+        yield history + [[message, error]]
+        return
+    
+    if not check_ollama():
+        yield history + [[message, "❌ **Ollama not running.**"]]
         return
     
     model = MODELS.get(model_name, "qwen2.5-coder:3b")
-    messages = [{"role": "system", "content": "You are an expert coding assistant. Provide clear, well-commented code. Always use markdown code blocks with language tags."}]
+    messages = [{"role": "system", "content": "You are an expert coding assistant. Use markdown code blocks."}]
 
     for user_msg, assistant_msg in history:
         messages.append({"role": "user", "content": user_msg})
@@ -110,9 +165,14 @@ def chat_stream(message, history, model_name, temperature, max_tokens):
     try:
         response = requests.post(
             f"{OLLAMA_URL}/api/chat",
-            json={"model": model, "messages": messages, "stream": True, "options": {"temperature": temperature, "num_predict": max_tokens}},
-            stream=True, timeout=300
+            json={"model": model, "messages": messages, "stream": True,
+                  "options": {"temperature": temperature, "num_predict": max_tokens}},
+            stream=True, timeout=TIMEOUT
         )
+        
+        if response.status_code != 200:
+            yield history + [[message, f"❌ **Error {response.status_code}**"]]
+            return
 
         full = ""
         for line in response.iter_lines():
@@ -124,936 +184,262 @@ def chat_stream(message, history, model_name, temperature, max_tokens):
                         yield history + [[message, full]]
                 except:
                     continue
+                    
+    except requests.exceptions.Timeout:
+        yield history + [[message, "❌ **Timeout.**"]]
     except Exception as e:
-        yield history + [[message, f"❌ Error: {e}"]]
+        yield history + [[message, f"❌ **Error:** {str(e)[:50]}"]]
 
 def generate_code(prompt, language, model_name, temperature, max_tokens):
-    if not prompt.strip():
-        return "⚠️ Please describe what you want to build."
+    valid, error = validate_input(prompt, "Description")
+    if not valid:
+        return error
     
     full_prompt = (
-        f"Write {language} code for the following task:\n\n"
-        f"{prompt}\n\n"
-        "Requirements:\n"
-        "- Clean, production-ready code\n"
-        "- Add helpful comments\n"
-        "- Handle edge cases\n"
-        "- Output ONLY the code in a markdown code block"
+        f"Write {language} code for:\n\n{prompt}\n\n"
+        "Requirements: Clean code, comments, handle edge cases. Output ONLY code in markdown block."
     )
-    
     result = call_ollama(model_name, full_prompt, temperature, max_tokens)
-    return extract_code(result)
+    return result if result.startswith("❌") or result.startswith("⚠️") else extract_code(result)
 
 def explain_code(code, model_name, detail_level, max_tokens):
-    if not code.strip():
-        return "⚠️ Paste code to explain."
+    valid, error = validate_input(code, "Code")
+    if not valid:
+        return error
     
-    detail_prompts = {
-        "Brief": "Give a brief 2-3 sentence explanation of what this code does.",
-        "Normal": "Explain what this code does, including the main logic and any important details.",
-        "Detailed": "Give a detailed explanation including: purpose, how it works step-by-step, time/space complexity, and potential improvements."
+    details = {
+        "Brief": "Give 2-3 sentence explanation.",
+        "Normal": "Explain the code with main logic.",
+        "Detailed": "Detailed explanation with complexity and improvements."
     }
-    
-    prompt = detail_prompts[detail_level] + "\n\nCode:\n" + code
+    prompt = details.get(detail_level, details["Normal"]) + "\n\nCode:\n" + code
     return call_ollama(model_name, prompt, 0.5, max_tokens)
 
-def fix_code(code, error, model_name, max_tokens):
-    if not code.strip():
-        return "⚠️ Paste code to fix."
+def fix_code(code, error_msg, model_name, max_tokens):
+    valid, error = validate_input(code, "Code")
+    if not valid:
+        return error
     
-    error_msg = error if error else "Code is not working as expected"
-    prompt = (
-        "Fix the following code and explain what was wrong.\n\n"
-        "Code:\n" + code + "\n\n"
-        "Error/Problem: " + error_msg + "\n\n"
-        "Provide:\n"
-        "1. The fixed code in a markdown code block\n"
-        "2. Brief explanation of what was wrong\n"
-        "3. Any suggestions to prevent similar issues"
-    )
+    err = error_msg if error_msg and error_msg.strip() else "Not working as expected"
+    prompt = f"Fix this code and explain what was wrong.\n\nCode:\n{code}\n\nError: {err}"
     return call_ollama(model_name, prompt, 0.3, max_tokens)
 
 def review_code(code, model_name, max_tokens):
-    if not code.strip():
-        return "⚠️ Paste code to review."
+    valid, error = validate_input(code, "Code")
+    if not valid:
+        return error
     
     prompt = (
-        "Review this code and provide feedback on:\n\n"
-        "1. **Code Quality** - Is it clean, readable, well-structured?\n"
-        "2. **Bugs/Issues** - Any potential bugs or problems?\n"
-        "3. **Performance** - Any performance concerns?\n"
-        "4. **Security** - Any security issues?\n"
-        "5. **Suggestions** - How could it be improved?\n\n"
+        "Review this code for:\n"
+        "1. Code Quality\n2. Bugs\n3. Performance\n4. Security\n5. Suggestions\n\n"
         "Code:\n" + code
     )
     return call_ollama(model_name, prompt, 0.4, max_tokens)
 
-def convert_code(code, source_lang, target_lang, model_name, max_tokens):
-    if not code.strip():
-        return "⚠️ Paste code to convert."
-    
-    if source_lang == target_lang:
-        return "⚠️ Source and target languages are the same."
-    
-    prompt = (
-        f"Convert this {source_lang} code to {target_lang}.\n\n"
-        "Requirements:\n"
-        f"- Write idiomatic {target_lang} code\n"
-        "- Preserve the functionality exactly\n"
-        "- Add comments explaining any language-specific differences\n"
-        "- Output ONLY the converted code in a markdown code block\n\n"
-        f"{source_lang} Code:\n" + code
-    )
+def convert_code(code, from_lang, to_lang, model_name, max_tokens):
+    valid, error = validate_input(code, "Code")
+    if not valid:
+        return error
+    if from_lang == to_lang:
+        return "⚠️ Same language selected."
     
+    prompt = f"Convert this {from_lang} to {to_lang}. Output ONLY code.\n\n{from_lang}:\n{code}"
     result = call_ollama(model_name, prompt, 0.3, max_tokens)
-    return extract_code(result)
+    return result if result.startswith("❌") or result.startswith("⚠️") else extract_code(result)
 
 def generate_tests(code, language, framework, model_name, max_tokens):
-    if not code.strip():
-        return "⚠️ Paste code to generate tests for."
-    
-    frameworks = {
-        "Python": "pytest",
-        "JavaScript": "Jest",
-        "TypeScript": "Jest",
-        "Java": "JUnit",
-        "C#": "NUnit",
-        "Go": "testing package",
-        "Rust": "built-in test framework",
-        "Ruby": "RSpec",
-        "PHP": "PHPUnit",
-    }
-    
-    fw = framework if framework else frameworks.get(language, "appropriate testing framework")
-    
-    prompt = (
-        f"Generate comprehensive unit tests for this {language} code using {fw}.\n\n"
-        "Requirements:\n"
-        "- Test all functions/methods\n"
-        "- Include edge cases\n"
-        "- Include both positive and negative tests\n"
-        "- Add descriptive test names\n"
-        "- Output ONLY the test code in a markdown code block\n\n"
-        "Code to test:\n" + code
-    )
+    valid, error = validate_input(code, "Code")
+    if not valid:
+        return error
     
+    fw = framework if framework else "pytest" if language == "Python" else "Jest"
+    prompt = f"Generate unit tests for this {language} code using {fw}. Output ONLY test code.\n\n{code}"
     result = call_ollama(model_name, prompt, 0.3, max_tokens)
-    return extract_code(result)
+    return result if result.startswith("❌") or result.startswith("⚠️") else extract_code(result)
 
 def document_code(code, language, style, model_name, max_tokens):
-    if not code.strip():
-        return "⚠️ Paste code to document."
-    
-    styles = {
-        "Docstrings": "Add comprehensive docstrings to all functions, classes, and methods",
-        "Comments": "Add inline comments explaining the logic",
-        "Both": "Add both docstrings and inline comments",
-        "README": "Generate a README.md documenting this code"
-    }
-    
-    prompt = (
-        f"Document this {language} code.\n\n"
-        f"Task: {styles[style]}\n\n"
-        "Requirements:\n"
-        "- Be clear and concise\n"
-        "- Explain parameters, return values, and exceptions\n"
-        "- Include usage examples where helpful\n"
-        "- Output the fully documented code in a markdown code block\n\n"
-        "Code:\n" + code
-    )
+    valid, error = validate_input(code, "Code")
+    if not valid:
+        return error
     
+    prompt = f"Add {style.lower()} to this {language} code.\n\n{code}"
     result = call_ollama(model_name, prompt, 0.4, max_tokens)
-    if style == "README":
-        return result
-    return extract_code(result)
+    return result if style == "README" or result.startswith("❌") else extract_code(result)
 
 def optimize_code(code, language, focus, model_name, max_tokens):
-    if not code.strip():
-        return "⚠️ Paste code to optimize."
-    
-    focus_prompts = {
-        "Performance": "Optimize for speed and efficiency. Reduce time complexity where possible.",
-        "Readability": "Refactor for better readability and maintainability. Follow best practices.",
-        "Memory": "Optimize memory usage. Reduce allocations and improve data structures.",
-        "All": "Optimize for performance, readability, and memory usage."
-    }
-    
-    prompt = (
-        f"Optimize this {language} code.\n\n"
-        f"Focus: {focus_prompts[focus]}\n\n"
-        "Requirements:\n"
-        "- Explain what you changed and why\n"
-        "- Preserve the original functionality\n"
-        "- Show before/after complexity if relevant\n"
-        "- Output the optimized code in a markdown code block\n\n"
-        "Code:\n" + code
-    )
+    valid, error = validate_input(code, "Code")
+    if not valid:
+        return error
     
+    prompt = f"Optimize this {language} code for {focus.lower()}. Explain changes.\n\n{code}"
     return call_ollama(model_name, prompt, 0.3, max_tokens)
 
 def build_regex(description, model_name, max_tokens):
-    if not description.strip():
-        return "⚠️ Describe the pattern you want to match."
-    
-    prompt = (
-        "Create a regex pattern for the following requirement:\n\n"
-        f"{description}\n\n"
-        "Provide:\n"
-        "1. The regex pattern\n"
-        "2. Explanation of each part\n"
-        "3. Example matches and non-matches\n"
-        "4. Code example in Python showing usage"
-    )
+    valid, error = validate_input(description, "Description")
+    if not valid:
+        return error
     
+    prompt = f"Create regex for: {description}\n\nProvide pattern, explanation, examples, and Python code."
     return call_ollama(model_name, prompt, 0.3, max_tokens)
 
 def build_api(description, framework, model_name, max_tokens):
-    if not description.strip():
-        return "⚠️ Describe the API endpoint you want to build."
-    
-    prompt = (
-        f"Create a REST API endpoint using {framework}.\n\n"
-        f"Requirements:\n{description}\n\n"
-        "Include:\n"
-        "- Route definition with proper HTTP methods\n"
-        "- Request validation\n"
-        "- Error handling\n"
-        "- Response formatting\n"
-        "- Brief documentation comments\n"
-        "- Output the code in a markdown code block"
-    )
+    valid, error = validate_input(description, "Description")
+    if not valid:
+        return error
     
+    prompt = f"Create REST API using {framework}:\n\n{description}\n\nInclude validation and error handling."
     result = call_ollama(model_name, prompt, 0.3, max_tokens)
-    return extract_code(result)
-
-# ===== PREMIUM CSS =====
-
-css = """
-/* ===== GLOBAL ===== */
-:root {
-    --primary: #6366f1;
-    --primary-dark: #4f46e5;
-    --secondary: #8b5cf6;
-    --accent: #06b6d4;
-    --bg-dark: #0f172a;
-    --bg-card: #1e293b;
-    --bg-hover: #334155;
-    --text-primary: #f1f5f9;
-    --text-secondary: #94a3b8;
-    --border: #334155;
-    --success: #10b981;
-    --warning: #f59e0b;
-    --error: #ef4444;
-    --gradient: linear-gradient(135deg, #6366f1 0%, #8b5cf6 50%, #06b6d4 100%);
-}
-
-.gradio-container {
-    max-width: 1500px !important;
-    margin: auto !important;
-    background: var(--bg-dark) !important;
-    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
-}
-
-.dark {
-    --bg-dark: #0f172a;
-}
-
-/* ===== HEADER ===== */
-.header-section {
-    background: var(--gradient);
-    border-radius: 20px;
-    padding: 32px 40px;
-    margin-bottom: 24px;
-    position: relative;
-    overflow: hidden;
-    box-shadow: 0 20px 40px rgba(99, 102, 241, 0.3);
-}
-
-.header-section::before {
-    content: '';
-    position: absolute;
-    top: -50%;
-    right: -50%;
-    width: 100%;
-    height: 200%;
-    background: radial-gradient(circle, rgba(255,255,255,0.1) 0%, transparent 60%);
-    animation: pulse 4s ease-in-out infinite;
-}
-
-@keyframes pulse {
-    0%, 100% { transform: scale(1); opacity: 0.5; }
-    50% { transform: scale(1.1); opacity: 0.8; }
-}
-
-.header-content {
-    position: relative;
-    z-index: 1;
-    display: flex;
-    justify-content: space-between;
-    align-items: center;
-    flex-wrap: wrap;
-    gap: 20px;
-}
-
-.header-title {
-    color: white;
-    margin: 0;
-    font-size: 2.8rem;
-    font-weight: 800;
-    letter-spacing: -0.02em;
-    text-shadow: 0 2px 10px rgba(0,0,0,0.2);
-}
-
-.header-subtitle {
-    color: rgba(255,255,255,0.9);
-    margin: 8px 0 0 0;
-    font-size: 1.1rem;
-    font-weight: 400;
-}
-
-.header-badges {
-    display: flex;
-    gap: 10px;
-    flex-wrap: wrap;
-}
-
-.badge {
-    background: rgba(255,255,255,0.2);
-    backdrop-filter: blur(10px);
-    padding: 8px 16px;
-    border-radius: 50px;
-    font-size: 0.85rem;
-    font-weight: 500;
-    color: white;
-    border: 1px solid rgba(255,255,255,0.2);
-}
-
-/* ===== STATUS BAR ===== */
-.status-bar {
-    background: var(--bg-card);
-    border: 1px solid var(--border);
-    border-radius: 16px;
-    padding: 16px 24px;
-    margin-bottom: 20px;
-    display: flex;
-    justify-content: space-between;
-    align-items: center;
-    flex-wrap: wrap;
-    gap: 16px;
-}
-
-.status-indicator {
-    display: flex;
-    align-items: center;
-    gap: 8px;
-    font-weight: 500;
-    color: var(--text-primary);
-}
-
-.status-dot {
-    width: 10px;
-    height: 10px;
-    border-radius: 50%;
-    background: var(--success);
-    animation: blink 2s ease-in-out infinite;
-}
-
-@keyframes blink {
-    0%, 100% { opacity: 1; }
-    50% { opacity: 0.5; }
-}
-
-/* ===== SETTINGS PANEL ===== */
-.settings-panel {
-    background: var(--bg-card);
-    border: 1px solid var(--border);
-    border-radius: 16px;
-    padding: 20px 24px;
-    margin-bottom: 20px;
-}
-
-.settings-panel label {
-    color: var(--text-secondary) !important;
-    font-weight: 500 !important;
-    font-size: 0.9rem !important;
-}
-
-.settings-panel input, .settings-panel select {
-    background: var(--bg-dark) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 10px !important;
-    color: var(--text-primary) !important;
-    transition: all 0.2s ease !important;
-}
-
-.settings-panel input:focus, .settings-panel select:focus {
-    border-color: var(--primary) !important;
-    box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.2) !important;
-}
-
-/* ===== MODEL INFO ===== */
-.model-info-box {
-    background: linear-gradient(135deg, rgba(99, 102, 241, 0.1) 0%, rgba(139, 92, 246, 0.1) 100%);
-    border: 1px solid rgba(99, 102, 241, 0.3);
-    border-radius: 12px;
-    padding: 12px 18px;
-    font-size: 0.9rem;
-    color: var(--text-secondary);
-    margin-top: 12px;
-}
-
-/* ===== TABS ===== */
-.tabs {
-    background: transparent !important;
-}
-
-.tab-nav {
-    background: var(--bg-card) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 16px !important;
-    padding: 8px !important;
-    gap: 6px !important;
-    margin-bottom: 20px !important;
-    flex-wrap: wrap !important;
-}
-
-.tab-nav button {
-    background: transparent !important;
-    border: none !important;
-    border-radius: 12px !important;
-    padding: 12px 20px !important;
-    font-weight: 600 !important;
-    font-size: 0.9rem !important;
-    color: var(--text-secondary) !important;
-    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
-}
-
-.tab-nav button:hover {
-    background: var(--bg-hover) !important;
-    color: var(--text-primary) !important;
-}
-
-.tab-nav button.selected {
-    background: var(--gradient) !important;
-    color: white !important;
-    box-shadow: 0 4px 15px rgba(99, 102, 241, 0.4) !important;
-}
-
-.tabitem {
-    background: transparent !important;
-    border: none !important;
-    padding: 0 !important;
-}
-
-/* ===== CARDS ===== */
-.card {
-    background: var(--bg-card);
-    border: 1px solid var(--border);
-    border-radius: 16px;
-    padding: 24px;
-    transition: all 0.3s ease;
-}
-
-.card:hover {
-    border-color: var(--primary);
-    box-shadow: 0 8px 30px rgba(99, 102, 241, 0.15);
-}
-
-/* ===== CHATBOT ===== */
-.chatbot {
-    background: var(--bg-card) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 16px !important;
-    height: 500px !important;
-}
-
-.chatbot .message {
-    border-radius: 16px !important;
-    padding: 16px 20px !important;
-    margin: 8px !important;
-    font-size: 0.95rem !important;
-    line-height: 1.6 !important;
-}
-
-.chatbot .user {
-    background: var(--gradient) !important;
-    color: white !important;
-    margin-left: 20% !important;
-}
-
-.chatbot .bot {
-    background: var(--bg-hover) !important;
-    color: var(--text-primary) !important;
-    margin-right: 20% !important;
-    border: 1px solid var(--border) !important;
-}
-
-/* ===== INPUTS ===== */
-.input-container textarea,
-.input-container input[type="text"] {
-    background: var(--bg-card) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 12px !important;
-    color: var(--text-primary) !important;
-    padding: 14px 18px !important;
-    font-size: 0.95rem !important;
-    transition: all 0.2s ease !important;
-}
-
-.input-container textarea:focus,
-.input-container input[type="text"]:focus {
-    border-color: var(--primary) !important;
-    box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.2) !important;
-    outline: none !important;
-}
-
-.input-container textarea::placeholder {
-    color: var(--text-secondary) !important;
-}
-
-/* ===== CODE BLOCKS ===== */
-.code-container {
-    border-radius: 16px !important;
-    overflow: hidden !important;
-    border: 1px solid var(--border) !important;
-}
-
-.code-container pre {
-    background: #0d1117 !important;
-    padding: 20px !important;
-    margin: 0 !important;
-    font-family: 'JetBrains Mono', 'Fira Code', monospace !important;
-    font-size: 0.9rem !important;
-    line-height: 1.6 !important;
-}
-
-/* ===== BUTTONS ===== */
-.primary-btn {
-    background: var(--gradient) !important;
-    border: none !important;
-    border-radius: 12px !important;
-    padding: 14px 28px !important;
-    font-weight: 600 !important;
-    font-size: 0.95rem !important;
-    color: white !important;
-    cursor: pointer !important;
-    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
-    box-shadow: 0 4px 15px rgba(99, 102, 241, 0.3) !important;
-}
-
-.primary-btn:hover {
-    transform: translateY(-2px) !important;
-    box-shadow: 0 8px 25px rgba(99, 102, 241, 0.4) !important;
-}
-
-.primary-btn:active {
-    transform: translateY(0) !important;
-}
-
-.secondary-btn {
-    background: var(--bg-hover) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 12px !important;
-    padding: 12px 20px !important;
-    font-weight: 500 !important;
-    color: var(--text-secondary) !important;
-    transition: all 0.2s ease !important;
-}
-
-.secondary-btn:hover {
-    background: var(--bg-card) !important;
-    border-color: var(--primary) !important;
-    color: var(--text-primary) !important;
-}
-
-/* ===== DROPDOWNS ===== */
-.dropdown-container select {
-    background: var(--bg-card) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 10px !important;
-    color: var(--text-primary) !important;
-    padding: 10px 14px !important;
-}
-
-/* ===== SLIDERS ===== */
-input[type="range"] {
-    accent-color: var(--primary) !important;
-}
-
-/* ===== RADIO & CHECKBOX ===== */
-.radio-group label {
-    background: var(--bg-hover) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 10px !important;
-    padding: 10px 16px !important;
-    color: var(--text-secondary) !important;
-    transition: all 0.2s ease !important;
-}
-
-.radio-group label.selected {
-    background: var(--primary) !important;
-    border-color: var(--primary) !important;
-    color: white !important;
-}
-
-/* ===== ACCORDION ===== */
-.accordion {
-    background: var(--bg-card) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 12px !important;
-    margin-top: 16px !important;
-}
-
-.accordion-header {
-    padding: 14px 18px !important;
-    color: var(--text-secondary) !important;
-    font-weight: 500 !important;
-}
-
-/* ===== MARKDOWN OUTPUT ===== */
-.markdown-output {
-    background: var(--bg-card);
-    border: 1px solid var(--border);
-    border-radius: 16px;
-    padding: 24px;
-    color: var(--text-primary);
-    line-height: 1.7;
-}
-
-.markdown-output h1, .markdown-output h2, .markdown-output h3 {
-    color: var(--text-primary);
-    margin-top: 1.5em;
-    margin-bottom: 0.5em;
-}
-
-.markdown-output code {
-    background: var(--bg-hover);
-    padding: 2px 6px;
-    border-radius: 4px;
-    font-family: 'JetBrains Mono', monospace;
-    font-size: 0.9em;
-}
-
-.markdown-output pre code {
-    display: block;
-    padding: 16px;
-    overflow-x: auto;
-}
-
-/* ===== AUDIO INPUT ===== */
-.audio-input {
-    background: var(--bg-card) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 12px !important;
-}
-
-/* ===== DIVIDER ===== */
-.divider {
-    height: 1px;
-    background: var(--border);
-    margin: 24px 0;
-}
-
-/* ===== TOOL SECTION ===== */
-.tool-section {
-    background: var(--bg-card);
-    border: 1px solid var(--border);
-    border-radius: 16px;
-    padding: 24px;
-    margin-bottom: 20px;
-}
+    return result if result.startswith("❌") else extract_code(result)
 
-.tool-title {
-    color: var(--text-primary);
-    font-size: 1.1rem;
-    font-weight: 600;
-    margin-bottom: 16px;
-    display: flex;
-    align-items: center;
-    gap: 8px;
-}
-
-/* ===== FOOTER ===== */
-.footer {
-    text-align: center;
-    padding: 24px;
-    color: var(--text-secondary);
-    font-size: 0.85rem;
-    border-top: 1px solid var(--border);
-    margin-top: 32px;
-}
-
-.footer a {
-    color: var(--primary);
-    text-decoration: none;
-}
-
-/* ===== SCROLLBAR ===== */
-::-webkit-scrollbar {
-    width: 8px;
-    height: 8px;
-}
-
-::-webkit-scrollbar-track {
-    background: var(--bg-dark);
-}
+# ===== UI (Gradio 6.0 compatible) =====
 
-::-webkit-scrollbar-thumb {
-    background: var(--border);
-    border-radius: 4px;
-}
-
-::-webkit-scrollbar-thumb:hover {
-    background: var(--text-secondary);
-}
-
-/* ===== HIDE FOOTER ===== */
-footer { display: none !important; }
-
-/* ===== RESPONSIVE ===== */
-@media (max-width: 768px) {
-    .header-title { font-size: 2rem; }
-    .header-content { flex-direction: column; text-align: center; }
-    .header-badges { justify-content: center; }
-    .tab-nav button { padding: 10px 14px !important; font-size: 0.8rem !important; }
-}
-"""
-
-# ===== UI =====
-
-with gr.Blocks(title="Axon v6", css=css, theme=gr.themes.Base()) as demo:
+with gr.Blocks(title="Axon v6") as demo:
     
-    # Header
-    gr.HTML("""
-    <div class="header-section">
-        <div class="header-content">
-            <div>
-                <h1 class="header-title">🔥 Axon v6</h1>
-                <p class="header-subtitle">AI-Powered Coding Assistant</p>
-            </div>
-            <div class="header-badges">
-                <span class="badge">🤖 10 Models</span>
-                <span class="badge">🛠️ 9 Tools</span>
-                <span class="badge">🔒 100% Local</span>
-                <span class="badge">⚡ No Rate Limits</span>
-            </div>
-        </div>
-    </div>
-    """)
+    gr.Markdown("# 🔥 Axon v6\n### AI Coding Assistant • 10 Models • 9 Tools • 100% Local")
     
-    # Status & Model Info
-    with gr.Row():
-        status = gr.Markdown(value=get_status, every=5)
+    status = gr.Markdown(value=get_status, every=5)
     
-    # Settings Panel
-    with gr.Row(elem_classes="settings-panel"):
-        model_dropdown = gr.Dropdown(
-            choices=list(MODELS.keys()), 
-            value="Qwen2.5 Coder 3B", 
-            label="🤖 Model",
-            scale=3
-        )
+    with gr.Row():
+        model_dropdown = gr.Dropdown(choices=list(MODELS.keys()), value="Qwen2.5 Coder 3B", label="🤖 Model", scale=3)
         temperature = gr.Slider(0, 1, value=0.7, step=0.1, label="🌡️ Creativity", scale=2)
         max_tokens = gr.Slider(256, 8192, value=2048, step=256, label="📏 Max Tokens", scale=2)
     
-    model_info_display = gr.Markdown(value="🚀 Fast & capable • Recommended", elem_classes="model-info-box")
-    model_dropdown.change(get_model_info, model_dropdown, model_info_display)
+    model_info = gr.Markdown(value="🚀 Fast & capable • Recommended")
+    model_dropdown.change(get_model_info, model_dropdown, model_info)
 
-    with gr.Tabs(elem_classes="tabs"):
+    with gr.Tabs():
         
-        # ===== CHAT =====
-        with gr.TabItem("💬 Chat", elem_classes="tabitem"):
-            chatbot = gr.Chatbot(height=500, show_copy_button=True, bubble_full_width=False, elem_classes="chatbot")
+        with gr.TabItem("💬 Chat"):
+            chatbot = gr.Chatbot(height=450)
             with gr.Row():
-                msg = gr.Textbox(
-                    placeholder="Ask anything about coding... Press Enter to send",
-                    show_label=False, scale=8, lines=1,
-                    elem_classes="input-container"
-                )
-                send = gr.Button("Send ➤", variant="primary", scale=1, elem_classes="primary-btn")
+                msg = gr.Textbox(placeholder="Ask anything...", show_label=False, scale=8)
+                send = gr.Button("Send ➤", variant="primary", scale=1)
             with gr.Row():
-                audio_input = gr.Audio(sources=["microphone"], type="filepath", label="🎤 Voice", scale=2)
-                transcribe_btn = gr.Button("🎤 Transcribe", scale=1, elem_classes="secondary-btn")
-                clear = gr.Button("🗑️ Clear", scale=1, elem_classes="secondary-btn")
-            with gr.Accordion("💡 Quick Prompts", open=False):
-                gr.Examples([
-                    "Write a Python function to find all prime numbers up to n",
-                    "Explain async/await vs promises in JavaScript",
-                    "How do I implement a binary search tree?",
-                    "Write a REST API with authentication in FastAPI"
-                ], inputs=msg)
+                audio = gr.Audio(sources=["microphone"], type="filepath", label="🎤", scale=2)
+                transcribe = gr.Button("🎤 Transcribe", scale=1)
+                clear = gr.Button("🗑️ Clear", scale=1)
         
-        # ===== GENERATE =====
-        with gr.TabItem("⚡ Generate", elem_classes="tabitem"):
+        with gr.TabItem("⚡ Generate"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    gen_prompt = gr.Textbox(
-                        label="📝 Describe what you want to build",
-                        placeholder="e.g., A function that validates email addresses with regex",
-                        lines=5, elem_classes="input-container"
-                    )
+                    gen_prompt = gr.Textbox(label="📝 Describe what to build", lines=4)
                     with gr.Row():
-                        gen_lang = gr.Dropdown(LANGUAGES, value="Python", label="🔤 Language", scale=2)
+                        gen_lang = gr.Dropdown(LANGUAGES, value="Python", label="Language", scale=2)
                         gen_temp = gr.Slider(0, 1, value=0.3, step=0.1, label="🌡️", scale=1)
-                    gen_btn = gr.Button("⚡ Generate Code", variant="primary", size="lg", elem_classes="primary-btn")
+                    gen_btn = gr.Button("⚡ Generate", variant="primary")
                 with gr.Column(scale=2):
-                    gen_output = gr.Code(label="Generated Code", language="python", lines=22, elem_classes="code-container")
+                    gen_output = gr.Code(label="Code", language="python", lines=18)
         
-        # ===== EXPLAIN =====
-        with gr.TabItem("🔍 Explain", elem_classes="tabitem"):
+        with gr.TabItem("🔍 Explain"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    explain_input = gr.Code(label="📋 Paste your code", lines=14, elem_classes="code-container")
-                    explain_detail = gr.Radio(
-                        ["Brief", "Normal", "Detailed"],
-                        value="Normal", label="📊 Detail Level"
-                    )
-                    explain_btn = gr.Button("🔍 Explain Code", variant="primary", size="lg", elem_classes="primary-btn")
+                    explain_input = gr.Code(label="📋 Code", lines=12)
+                    explain_detail = gr.Radio(["Brief", "Normal", "Detailed"], value="Normal", label="Detail")
+                    explain_btn = gr.Button("🔍 Explain", variant="primary")
                 with gr.Column(scale=1):
-                    explain_output = gr.Markdown(label="Explanation", elem_classes="markdown-output")
+                    explain_output = gr.Markdown(label="Explanation")
         
-        # ===== DEBUG =====
-        with gr.TabItem("🔧 Debug", elem_classes="tabitem"):
+        with gr.TabItem("🔧 Debug"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    fix_input = gr.Code(label="🐛 Paste buggy code", lines=12, elem_classes="code-container")
-                    fix_error = gr.Textbox(
-                        label="❌ Error message (optional)",
-                        placeholder="Paste error or describe the issue",
-                        lines=3, elem_classes="input-container"
-                    )
-                    fix_btn = gr.Button("🔧 Fix Code", variant="primary", size="lg", elem_classes="primary-btn")
+                    fix_input = gr.Code(label="🐛 Buggy Code", lines=10)
+                    fix_error = gr.Textbox(label="❌ Error", lines=2)
+                    fix_btn = gr.Button("🔧 Fix", variant="primary")
                 with gr.Column(scale=1):
-                    fix_output = gr.Markdown(label="Solution", elem_classes="markdown-output")
+                    fix_output = gr.Markdown(label="Solution")
         
-        # ===== REVIEW =====
-        with gr.TabItem("📋 Review", elem_classes="tabitem"):
+        with gr.TabItem("📋 Review"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    review_input = gr.Code(label="📋 Code to review", lines=16, elem_classes="code-container")
-                    review_btn = gr.Button("📋 Review Code", variant="primary", size="lg", elem_classes="primary-btn")
+                    review_input = gr.Code(label="📋 Code", lines=14)
+                    review_btn = gr.Button("📋 Review", variant="primary")
                 with gr.Column(scale=1):
-                    review_output = gr.Markdown(label="Code Review", elem_classes="markdown-output")
+                    review_output = gr.Markdown(label="Review")
         
-        # ===== CONVERT =====
-        with gr.TabItem("🔄 Convert", elem_classes="tabitem"):
+        with gr.TabItem("🔄 Convert"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    convert_input = gr.Code(label="📥 Source Code", lines=14, elem_classes="code-container")
+                    convert_input = gr.Code(label="📥 Source", lines=12)
                     with gr.Row():
-                        convert_from = gr.Dropdown(LANGUAGES, value="Python", label="From", scale=1)
-                        gr.HTML("<div style='display:flex;align-items:center;justify-content:center;font-size:1.5rem;'>➜</div>")
-                        convert_to = gr.Dropdown(LANGUAGES, value="JavaScript", label="To", scale=1)
-                    convert_btn = gr.Button("🔄 Convert Code", variant="primary", size="lg", elem_classes="primary-btn")
+                        convert_from = gr.Dropdown(LANGUAGES, value="Python", label="From")
+                        convert_to = gr.Dropdown(LANGUAGES, value="JavaScript", label="To")
+                    convert_btn = gr.Button("🔄 Convert", variant="primary")
                 with gr.Column(scale=1):
-                    convert_output = gr.Code(label="📤 Converted Code", lines=14, elem_classes="code-container")
+                    convert_output = gr.Code(label="📤 Result", lines=12)
         
-        # ===== TEST =====
-        with gr.TabItem("🧪 Test", elem_classes="tabitem"):
+        with gr.TabItem("🧪 Test"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    test_input = gr.Code(label="📋 Code to test", lines=14, elem_classes="code-container")
+                    test_input = gr.Code(label="📋 Code", lines=12)
                     with gr.Row():
-                        test_lang = gr.Dropdown(LANGUAGES[:12], value="Python", label="Language", scale=2)
-                        test_framework = gr.Textbox(label="Framework", placeholder="e.g., pytest", scale=2)
-                    test_btn = gr.Button("🧪 Generate Tests", variant="primary", size="lg", elem_classes="primary-btn")
+                        test_lang = gr.Dropdown(LANGUAGES[:12], value="Python", label="Language")
+                        test_fw = gr.Textbox(label="Framework", placeholder="pytest")
+                    test_btn = gr.Button("🧪 Generate Tests", variant="primary")
                 with gr.Column(scale=1):
-                    test_output = gr.Code(label="Generated Tests", lines=14, elem_classes="code-container")
+                    test_output = gr.Code(label="Tests", lines=12)
         
-        # ===== DOCUMENT =====
-        with gr.TabItem("📝 Document", elem_classes="tabitem"):
+        with gr.TabItem("📝 Document"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    doc_input = gr.Code(label="📋 Code to document", lines=14, elem_classes="code-container")
+                    doc_input = gr.Code(label="📋 Code", lines=12)
                     with gr.Row():
-                        doc_lang = gr.Dropdown(LANGUAGES, value="Python", label="Language", scale=2)
-                        doc_style = gr.Dropdown(
-                            ["Docstrings", "Comments", "Both", "README"],
-                            value="Both", label="Style", scale=2
-                        )
-                    doc_btn = gr.Button("📝 Document", variant="primary", size="lg", elem_classes="primary-btn")
+                        doc_lang = gr.Dropdown(LANGUAGES, value="Python", label="Language")
+                        doc_style = gr.Dropdown(["Docstrings", "Comments", "Both", "README"], value="Both", label="Style")
+                    doc_btn = gr.Button("📝 Document", variant="primary")
                 with gr.Column(scale=1):
-                    doc_output = gr.Code(label="Documented Code", lines=14, elem_classes="code-container")
+                    doc_output = gr.Code(label="Documented", lines=12)
         
-        # ===== OPTIMIZE =====
-        with gr.TabItem("🚀 Optimize", elem_classes="tabitem"):
+        with gr.TabItem("🚀 Optimize"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    opt_input = gr.Code(label="📋 Code to optimize", lines=14, elem_classes="code-container")
+                    opt_input = gr.Code(label="📋 Code", lines=12)
                     with gr.Row():
-                        opt_lang = gr.Dropdown(LANGUAGES, value="Python", label="Language", scale=2)
-                        opt_focus = gr.Dropdown(
-                            ["All", "Performance", "Readability", "Memory"],
-                            value="All", label="Focus", scale=2
-                        )
-                    opt_btn = gr.Button("🚀 Optimize", variant="primary", size="lg", elem_classes="primary-btn")
+                        opt_lang = gr.Dropdown(LANGUAGES, value="Python", label="Language")
+                        opt_focus = gr.Dropdown(["All", "Performance", "Readability", "Memory"], value="All", label="Focus")
+                    opt_btn = gr.Button("🚀 Optimize", variant="primary")
                 with gr.Column(scale=1):
-                    opt_output = gr.Markdown(label="Optimized Code", elem_classes="markdown-output")
+                    opt_output = gr.Markdown(label="Optimized")
         
-        # ===== TOOLS =====
-        with gr.TabItem("🛠️ Tools", elem_classes="tabitem"):
-            
-            # Regex Builder
-            gr.HTML("<div class='tool-section'><div class='tool-title'>🎯 Regex Builder</div></div>")
+        with gr.TabItem("🛠️ Tools"):
+            gr.Markdown("### 🎯 Regex Builder")
             with gr.Row():
                 with gr.Column(scale=1):
-                    regex_desc = gr.Textbox(
-                        label="Describe the pattern",
-                        placeholder="e.g., Match email addresses, validate phone numbers with country code...",
-                        lines=3, elem_classes="input-container"
-                    )
-                    regex_btn = gr.Button("🎯 Build Regex", variant="primary", elem_classes="primary-btn")
+                    regex_desc = gr.Textbox(label="Describe pattern", lines=2)
+                    regex_btn = gr.Button("🎯 Build Regex", variant="primary")
                 with gr.Column(scale=1):
-                    regex_output = gr.Markdown(label="Regex Pattern", elem_classes="markdown-output")
+                    regex_output = gr.Markdown(label="Pattern")
             
-            gr.HTML("<div class='divider'></div>")
-            
-            # API Builder
-            gr.HTML("<div class='tool-section'><div class='tool-title'>🔗 API Builder</div></div>")
+            gr.Markdown("---\n### 🔗 API Builder")
             with gr.Row():
                 with gr.Column(scale=1):
-                    api_desc = gr.Textbox(
-                        label="Describe the endpoint",
-                        placeholder="e.g., POST endpoint for user registration with email/password validation...",
-                        lines=3, elem_classes="input-container"
-                    )
-                    api_framework = gr.Dropdown(
-                        ["FastAPI (Python)", "Express (Node.js)", "Gin (Go)", "Spring Boot (Java)", "Flask (Python)", "Django REST (Python)"],
-                        value="FastAPI (Python)", label="Framework"
-                    )
-                    api_btn = gr.Button("🔗 Build API", variant="primary", elem_classes="primary-btn")
+                    api_desc = gr.Textbox(label="Describe endpoint", lines=2)
+                    api_fw = gr.Dropdown(["FastAPI", "Express", "Flask", "Gin"], value="FastAPI", label="Framework")
+                    api_btn = gr.Button("🔗 Build API", variant="primary")
                 with gr.Column(scale=1):
-                    api_output = gr.Code(label="API Code", lines=14, elem_classes="code-container")
+                    api_output = gr.Code(label="API Code", lines=12)
 
-    # Footer
-    gr.HTML("""
-    <div class="footer">
-        <p>🔒 Running 100% locally via Ollama • Your code never leaves your machine</p>
-        <p style="margin-top: 8px; opacity: 0.7;">Built with ❤️ using Gradio</p>
-    </div>
-    """)
+    gr.Markdown("<center>🔒 100% Local • Robust Error Handling</center>")
 
-    # ===== EVENT HANDLERS =====
-    
+    # Events
     def respond(message, history, model, temp, tokens):
         history = history or []
-        for updated_history in chat_stream(message, history, model, temp, tokens):
-            yield updated_history, ""
+        for updated in chat_stream(message, history, model, temp, tokens):
+            yield updated, ""
     
     msg.submit(respond, [msg, chatbot, model_dropdown, temperature, max_tokens], [chatbot, msg])
     send.click(respond, [msg, chatbot, model_dropdown, temperature, max_tokens], [chatbot, msg])
     clear.click(lambda: [], None, chatbot)
-    transcribe_btn.click(transcribe_audio, audio_input, msg)
+    transcribe.click(transcribe_audio, audio, msg)
     
     gen_btn.click(generate_code, [gen_prompt, gen_lang, model_dropdown, gen_temp, max_tokens], gen_output)
     explain_btn.click(explain_code, [explain_input, model_dropdown, explain_detail, max_tokens], explain_output)
     fix_btn.click(fix_code, [fix_input, fix_error, model_dropdown, max_tokens], fix_output)
     review_btn.click(review_code, [review_input, model_dropdown, max_tokens], review_output)
     convert_btn.click(convert_code, [convert_input, convert_from, convert_to, model_dropdown, max_tokens], convert_output)
-    test_btn.click(generate_tests, [test_input, test_lang, test_framework, model_dropdown, max_tokens], test_output)
+    test_btn.click(generate_tests, [test_input, test_lang, test_fw, model_dropdown, max_tokens], test_output)
     doc_btn.click(document_code, [doc_input, doc_lang, doc_style, model_dropdown, max_tokens], doc_output)
     opt_btn.click(optimize_code, [opt_input, opt_lang, opt_focus, model_dropdown, max_tokens], opt_output)
     regex_btn.click(build_regex, [regex_desc, model_dropdown, max_tokens], regex_output)
-    api_btn.click(build_api, [api_desc, api_framework, model_dropdown, max_tokens], api_output)
+    api_btn.click(build_api, [api_desc, api_fw, model_dropdown, max_tokens], api_output)
 
 demo.launch(server_name="0.0.0.0", server_port=7860)