Update app.py

app.py

@@ -1,39 +1,35 @@
 import gradio as gr
 import torch
 import time
-import re
+import sys
+from io import StringIO
+import contextlib
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from typing import Dict, List, Tuple, Optional
+from typing import List, Tuple
 import os
-import json
 from functools import lru_cache
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# CPU-optimized configuration
+# --- Configuration ---
 MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 DEVICE = "cpu"
 DTYPE = torch.float32 
 
-# --- File System Class ---
+# --- File System (Python Version) ---
 class FileSystem:
     def __init__(self):
         self.files = {
-            "main.js": "// Start coding here\nconsole.log('Hello Axon Pro');",
-            "utils.js": "// Utility functions\nfunction add(a, b) {\n  return a + b;\n}",
-            "style.css": "/* Add your styles here */\nbody {\n  font-family: monospace;\n}"
+            "main.py": "# Start coding here\nprint('Hello Axon Pro (Python Edition)')",
+            "utils.py": "# Utility functions\ndef add(a, b):\n    return a + b",
+            "notes.txt": "Project requirements:\n- Build a cool AI app\n- Use Gradio"
         }
-        self.current_file = "main.js"
-        self.history = []
-        self.max_history = 50
+        self.current_file = "main.py"
     
     def save_file(self, content: str) -> None:
         if self.current_file:
             self.files[self.current_file] = content
     
-    def get_file(self, filename: str) -> str:
-        return self.files.get(filename, "")
-    
     def get_current_file_content(self) -> str:
         return self.files.get(self.current_file, "")
     
@@ -48,19 +44,13 @@ class FileSystem:
     
     def get_all_files(self) -> List[str]:
         return list(self.files.keys())
-    
-    def get_context(self) -> str:
-        context = ""
-        for filename, content in self.files.items():
-            context += f"// File: {filename}\n{content}\n\n"
-        return context
 
 fs = FileSystem()
 
 # --- Model Loading ---
 @lru_cache(maxsize=1)
 def load_model():
-    print("Loading TinyLlama model for CPU... This may take a moment.")
+    print("Loading TinyLlama model for CPU...")
     start_time = time.time()
     try:
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
@@ -79,37 +69,49 @@ def load_model():
             top_p=0.95,
             do_sample=True
         )
-        print(f"Model loaded successfully in {time.time() - start_time:.2f} seconds!")
+        print(f"Model loaded in {time.time() - start_time:.2f}s")
         return pipe
     except Exception as e:
-        print(f"Error loading model: {str(e)}")
+        print(f"Error loading model: {e}")
         return None
 
 # --- Logic Functions ---
+
 def run_code(code: str) -> Tuple[str, str]:
+    """Actually executes Python code and captures stdout."""
     start_time = time.time()
+    output_buffer = StringIO()
+    
     try:
-        # Placeholder for actual execution logic
-        output = f"✓ Code executed successfully\nExecution time: {time.time() - start_time:.4f}s"
-        return output, ""
+        # Capture standard output
+        with contextlib.redirect_stdout(output_buffer):
+            # dicts for globals/locals to keep environment somewhat contained
+            exec(code, {}, {})
+        
+        exec_time = time.time() - start_time
+        result = output_buffer.getvalue()
+        if not result:
+            result = "(No output)"
+            
+        return f"{result}\n\n[Execution time: {exec_time:.4f}s]", ""
     except Exception as e:
-        return "", f"✗ Error: {str(e)}"
+        return output_buffer.getvalue(), f"Traceback (most recent call last):\n{str(e)}"
 
 def generate_completion(code: str) -> str:
     model = load_model()
-    if not model: return "Error: Model failed to load."
+    if not model: return "# Error: Model failed to load."
     
-    prompt = f"<|system|>You are a JS expert. Complete the code. No explanations.</s><|user|>\n{code}\n</s><|assistant|>"
+    prompt = f"<|system|>You are a Python expert. Complete the code. No explanations.</s><|user|>\n{code}\n</s><|assistant|>"
     try:
         result = model(prompt, num_return_sequences=1, return_full_text=False)
         return result[0]['generated_text'].strip()
     except Exception as e:
-        return f"/* Error: {str(e)} */"
+        return f"# Error: {str(e)}"
 
 def explain_code(code: str) -> str:
     model = load_model()
     if not model: return "Error loading model."
-    prompt = f"<|system|>Explain this code concisely.</s><|user|>\n{code}\n</s><|assistant|>"
+    prompt = f"<|system|>Explain this Python code concisely.</s><|user|>\n{code}\n</s><|assistant|>"
     try:
         result = model(prompt, max_new_tokens=512)
         return result[0]['generated_text'].strip()
@@ -119,7 +121,7 @@ def explain_code(code: str) -> str:
 def refactor_code(code: str) -> str:
     model = load_model()
     if not model: return "Error loading model."
-    prompt = f"<|system|>Refactor this JS code for best practices.</s><|user|>\n{code}\n</s><|assistant|>"
+    prompt = f"<|system|>Refactor this Python code for PEP 8 and best practices.</s><|user|>\n{code}\n</s><|assistant|>"
     try:
         result = model(prompt, max_new_tokens=512)
         return result[0]['generated_text'].strip()
@@ -129,7 +131,7 @@ def refactor_code(code: str) -> str:
 def generate_code(prompt_text: str) -> str:
     model = load_model()
     if not model: return "Error loading model."
-    prompt = f"<|system|>Write JavaScript code for: {prompt_text}</s><|assistant|>"
+    prompt = f"<|system|>Write Python code for: {prompt_text}</s><|assistant|>"
     try:
         result = model(prompt, max_new_tokens=512)
         return result[0]['generated_text'].strip()
@@ -137,20 +139,17 @@ def generate_code(prompt_text: str) -> str:
         return str(e)
 
 def create_diff_view(original: str, modified: str) -> str:
-    lines_original = original.split('\n')
-    lines_modified = modified.split('\n')
-    diff = []
+    lines_original = original.splitlines()
+    lines_modified = modified.splitlines()
     import difflib
-    for line in difflib.unified_diff(lines_original, lines_modified, lineterm=""):
-        diff.append(line)
+    diff = difflib.unified_diff(lines_original, lines_modified, lineterm="")
     return "\n".join(diff)
 
-# --- Gradio UI Configuration ---
+# --- Gradio UI ---
 
-# 1. THEME: Cleaned up to avoid errors
+# Theme & CSS
 theme = gr.themes.Default(
     font=[gr.themes.GoogleFont('JetBrains Mono'), 'monospace'],
-    font_mono=[gr.themes.GoogleFont('JetBrains Mono'), 'monospace'],
     primary_hue="blue",
     neutral_hue="gray",
 ).set(
@@ -160,36 +159,25 @@ theme = gr.themes.Default(
     body_text_color="*neutral_50",
 )
 
-# 2. CSS: Styling
 css = """
-    /* Force square corners globally */
+    /* Force square corners */
     * { border-radius: 0 !important; } 
-    
     .editor-container { height: 50vh; }
-    .terminal { height: 15vh; background-color: #1e1e1e; color: #d4d4d4; overflow: auto; }
-    .diff-view { height: 40vh; overflow: auto; }
-    .file-explorer { height: 50vh; background-color: #252526; overflow: auto; }
-    .ai-chat { height: 40vh; overflow: auto; }
-    .status-bar { background-color: #007acc; color: white; padding: 5px; text-align: center; }
-    .btn { min-width: 100px; }
-    .monaco-editor { height: 100% !important; }
-    .theme-dark { --body-bg: #1e1e1e; --editor-bg: #1e1e1e; }
-    .button-large { min-height: 40px !important; }
-    .tab-nav { background-color: #252526 !important; }
-    .tabitem { background-color: #1e1e1e !important; }
-    .code-wrap { white-space: pre-wrap !important; word-break: break-word !important; }
+    .terminal { height: 15vh; background-color: #1e1e1e; color: #d4d4d4; font-family: 'JetBrains Mono', monospace; }
+    .diff-view { height: 40vh; }
+    .code-wrap { white-space: pre-wrap !important; }
 """
 
-with gr.Blocks(title="Axon Pro - Free AI IDE", theme=theme, css=css) as demo:
+with gr.Blocks(title="Axon Pro - Python IDE", theme=theme, css=css) as demo:
     
-    gr.Markdown("# ⚡ Axon Pro — Free AI-Powered Code IDE")
+    gr.Markdown("# 🐍 Axon Pro — Python AI IDE")
     
     with gr.Tabs() as main_tabs:
         with gr.Tab("Editor", id="editor-tab"):
             with gr.Row(equal_height=True):
-                # File Explorer
+                # Sidebar
                 with gr.Column(scale=1, min_width=200):
-                    gr.Markdown("### 📁 File Explorer")
+                    gr.Markdown("### 📁 Explorer")
                     file_list = gr.Dropdown(
                         choices=fs.get_all_files(), 
                         value=fs.current_file,
@@ -200,14 +188,15 @@ with gr.Blocks(title="Axon Pro - Free AI IDE", theme=theme, css=css) as demo:
                     with gr.Row():
                         new_file_btn = gr.Button("➕ New", variant="secondary")
                         save_btn = gr.Button("💾 Save", variant="secondary")
-                    new_file_name = gr.Textbox(placeholder="filename.js", label="New File Name", container=False)
+                    new_file_name = gr.Textbox(placeholder="script.py", label="New File", container=False)
                 
-                # Main Editor Area
+                # Editor
                 with gr.Column(scale=4):
+                    # Changed language to 'python'
                     editor = gr.Code(
                         value=fs.get_current_file_content(),
                         label="Code Editor",
-                        language="javascript",
+                        language="python",
                         lines=20,
                         interactive=True,
                         elem_classes="code-wrap"
@@ -222,36 +211,37 @@ with gr.Blocks(title="Axon Pro - Free AI IDE", theme=theme, css=css) as demo:
                     with gr.Tabs():
                         with gr.Tab("Terminal", id="terminal-tab"):
                             terminal = gr.Textbox(
-                                value="$ _", 
-                                lines=5, 
+                                value=">>> Ready to run Python code...", 
+                                lines=6, 
                                 interactive=False,
-                                elem_classes="terminal"
+                                elem_classes="terminal",
+                                label="Output"
                             )
-                            clear_btn = gr.Button("CLEAR TERMINAL", variant="secondary", size="sm")
+                            clear_btn = gr.Button("Clear Output", variant="secondary", size="sm")
                         
                         with gr.Tab("AI Chat", id="chat-tab"):
                             chat_history = gr.Chatbot(label="Axon AI", height=300)
                             with gr.Row():
-                                chat_input = gr.Textbox(placeholder="Ask to generate code...", scale=7, container=False)
+                                chat_input = gr.Textbox(placeholder="Ask to generate Python code...", scale=7, container=False)
                                 send_btn = gr.Button("Generate", variant="primary", scale=1)
         
         with gr.Tab("Diff View", id="diff-tab"):
-            # FIX: Changed language="diff" to language="javascript" to prevent crash
-            diff_view = gr.Code(label="AI Changes", language="javascript", elem_classes="diff-view code-wrap", interactive=False)
+            # Changed language to 'python' for diff view as well
+            diff_view = gr.Code(label="AI Changes", language="python", elem_classes="diff-view code-wrap", interactive=False)
             with gr.Row():
                 apply_btn = gr.Button("Apply Changes", variant="primary")
                 discard_btn = gr.Button("Discard Changes", variant="secondary")
     
-    status_bar = gr.Markdown(f"**AXON PRO v1.0** | CPU Mode | TinyLlama-1.1B", elem_classes="status-bar")
+    status_bar = gr.Markdown(f"**AXON PRO v1.0** | Python 3.x | CPU Mode | TinyLlama-1.1B", elem_classes="status-bar")
     
-    # State management
+    # State
     current_file_state = gr.State(fs.current_file)
     diff_original_state = gr.State("")
     diff_modified_state = gr.State("")
     diff_mode_state = gr.State(False)
     
-    # Event Handlers
-    def update_file_content(content):
+    # Functions
+    def update_file(content):
         fs.save_file(content)
         return fs.get_current_file_content()
     
@@ -259,84 +249,74 @@ with gr.Blocks(title="Axon Pro - Free AI IDE", theme=theme, css=css) as demo:
         fs.set_current_file(filename)
         return fs.get_current_file_content(), filename
     
-    def create_new_file(name):
+    def create_file(name):
         if name and "." in name:
             fs.create_file(name)
             return gr.update(choices=fs.get_all_files(), value=name), fs.get_current_file_content()
         return gr.update(choices=fs.get_all_files()), fs.get_current_file_content()
     
-    def run_code_wrapper(content):
+    def run_wrapper(content):
         out, err = run_code(content)
-        return out if not err else err
+        if err:
+            return f"{out}\n\n=== ERROR ===\n{err}"
+        return out
     
-    def complete_code_wrapper(content):
+    def complete_wrapper(content):
         comp = generate_completion(content)
-        return content + comp
+        return content + "\n" + comp
     
-    def explain_code_wrapper(content):
+    def explain_wrapper(content):
         explanation = explain_code(content)
         fs.save_file(content)
         diff_orig = content
-        diff_mod = f"/* EXPLANATION:\n{explanation}\n*/\n\n{content}"
+        diff_mod = f'"""\nEXPLANATION:\n{explanation}\n"""\n\n{content}'
         diff = create_diff_view(diff_orig, diff_mod)
         return diff, diff_orig, diff_mod, True
     
-    def refactor_code_wrapper(content):
+    def refactor_wrapper(content):
         refactored = refactor_code(content)
         fs.save_file(content)
         diff = create_diff_view(content, refactored)
         return diff, content, refactored, True
 
-    def generate_code_wrapper(prompt, history):
+    def generate_wrapper(prompt, history):
         generated = generate_code(prompt)
         diff = create_diff_view("", generated)
         new_hist = history + [[prompt, "Generated code available in Diff View"]]
         return diff, "", generated, True, new_hist, ""
 
-    # Binding Events
-    editor.change(update_file_content, editor, None)
+    # Wiring
+    editor.change(update_file, editor, None)
     file_list.change(load_file, file_list, [editor, current_file_state])
-    new_file_btn.click(create_new_file, new_file_name, [file_list, editor])
+    new_file_btn.click(create_file, new_file_name, [file_list, editor])
     
-    run_btn.click(run_code_wrapper, editor, terminal)
-    complete_btn.click(complete_code_wrapper, editor, editor)
-    clear_btn.click(lambda: "$ _", None, terminal)
+    run_btn.click(run_wrapper, editor, terminal)
+    complete_btn.click(complete_wrapper, editor, editor)
+    clear_btn.click(lambda: ">>> Ready...", None, terminal)
     
-    # Tab switching actions
+    # Tabs logic
     explain_btn.click(
-        explain_code_wrapper, editor, 
+        explain_wrapper, editor, 
         [diff_view, diff_original_state, diff_modified_state, diff_mode_state]
-    ).then(
-        lambda: gr.Tabs(selected="diff-tab"), None, main_tabs
-    )
+    ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
     
     refactor_btn.click(
-        refactor_code_wrapper, editor,
+        refactor_wrapper, editor,
         [diff_view, diff_original_state, diff_modified_state, diff_mode_state]
-    ).then(
-        lambda: gr.Tabs(selected="diff-tab"), None, main_tabs
-    )
+    ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
     
     chat_input.submit(
-        generate_code_wrapper, [chat_input, chat_history], 
+        generate_wrapper, [chat_input, chat_history], 
         [diff_view, diff_original_state, diff_modified_state, diff_mode_state, chat_history, chat_input]
-    ).then(
-        lambda: gr.Tabs(selected="diff-tab"), None, main_tabs
-    )
+    ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
     
     apply_btn.click(
         lambda mod: (mod, False), diff_modified_state, [editor, diff_mode_state]
-    ).then(
-        lambda: gr.Tabs(selected="editor-tab"), None, main_tabs
-    )
+    ).then(lambda: gr.Tabs(selected="editor-tab"), None, main_tabs)
     
     discard_btn.click(
         lambda: (gr.update(), False), None, [editor, diff_mode_state]
-    ).then(
-        lambda: gr.Tabs(selected="editor-tab"), None, main_tabs
-    )
+    ).then(lambda: gr.Tabs(selected="editor-tab"), None, main_tabs)
 
 if __name__ == "__main__":
-    # If the warning about theme/css being in launch() persists, you can move them here.
-    # But for now, we keep them in Blocks to support older/standard versions as well.
     demo.launch(server_name="0.0.0.0", server_port=7860)