Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import torch
 import time
 import sys
+import subprocess
 from io import StringIO
 import contextlib
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
@@ -49,6 +50,115 @@ class FileSystem:
 
 fs = FileSystem()
 
+# --- Package Manager ---
+class PackageManager:
+    """Handles pip install/uninstall operations for the HF Space runtime."""
+    
+    @staticmethod
+    def install_package(package_name: str) -> str:
+        """Install a package using pip at runtime."""
+        if not package_name or not package_name.strip():
+            return "Error: No package name provided."
+        
+        package_name = package_name.strip()
+        
+        # Basic safety check - block obviously dangerous inputs
+        dangerous = [";", "&&", "||", "|", ">", "<", "`", "$", "(", ")"]
+        if any(ch in package_name for ch in dangerous):
+            return "Error: Invalid characters in package name."
+        
+        try:
+            result = subprocess.run(
+                [sys.executable, "-m", "pip", "install", package_name],
+                capture_output=True,
+                text=True,
+                timeout=120  # 2 min timeout for large packages
+            )
+            
+            output = ""
+            if result.stdout:
+                # Filter to show only the important lines
+                lines = result.stdout.strip().split("\n")
+                important = [l for l in lines if any(kw in l.lower() for kw in 
+                    ["successfully", "already satisfied", "installing", "requirement", "downloading", "error", "warning"])]
+                output = "\n".join(important) if important else lines[-1]
+            
+            if result.returncode != 0:
+                error_msg = result.stderr.strip() if result.stderr else "Unknown error"
+                # Show last few lines of error (most relevant)
+                error_lines = error_msg.split("\n")[-5:]
+                return f"pip install {package_name} failed:\n" + "\n".join(error_lines)
+            
+            return f"✅ {output}"
+            
+        except subprocess.TimeoutExpired:
+            return f"Error: Installation of '{package_name}' timed out (>120s). Package may be too large."
+        except Exception as e:
+            return f"Error: {str(e)}"
+    
+    @staticmethod
+    def uninstall_package(package_name: str) -> str:
+        """Uninstall a package using pip."""
+        if not package_name or not package_name.strip():
+            return "Error: No package name provided."
+        
+        package_name = package_name.strip()
+        
+        # Protect critical packages from uninstall
+        protected = ["gradio", "torch", "transformers", "pip", "setuptools", "wheel"]
+        if package_name.lower() in protected:
+            return f"Error: '{package_name}' is a protected system package and cannot be uninstalled."
+        
+        try:
+            result = subprocess.run(
+                [sys.executable, "-m", "pip", "uninstall", "-y", package_name],
+                capture_output=True,
+                text=True,
+                timeout=60
+            )
+            
+            if result.returncode != 0:
+                return f"Error: {result.stderr.strip()}"
+            
+            return f"✅ Uninstalled '{package_name}'"
+            
+        except Exception as e:
+            return f"Error: {str(e)}"
+    
+    @staticmethod
+    def list_packages() -> str:
+        """List installed packages."""
+        try:
+            result = subprocess.run(
+                [sys.executable, "-m", "pip", "list", "--format=columns"],
+                capture_output=True,
+                text=True,
+                timeout=30
+            )
+            return result.stdout if result.stdout else "No packages found."
+        except Exception as e:
+            return f"Error: {str(e)}"
+    
+    @staticmethod
+    def search_package_info(package_name: str) -> str:
+        """Show info about an installed package."""
+        if not package_name or not package_name.strip():
+            return "Error: No package name provided."
+        try:
+            result = subprocess.run(
+                [sys.executable, "-m", "pip", "show", package_name.strip()],
+                capture_output=True,
+                text=True,
+                timeout=15
+            )
+            if result.returncode != 0:
+                return f"Package '{package_name.strip()}' is not installed."
+            return result.stdout
+        except Exception as e:
+            return f"Error: {str(e)}"
+
+pkg_mgr = PackageManager()
+
 # --- Model Loading ---
 @lru_cache(maxsize=1)
 def load_model():
@@ -80,24 +190,54 @@ def load_model():
 # --- Logic Functions ---
 
 def run_code(code: str) -> Tuple[str, str]:
-    """Actually executes Python code and captures stdout."""
-    start_time = time.time()
-    output_buffer = StringIO()
+    """Executes Python code and captures stdout. Supports !pip commands."""
     
-    try:
-        # Capture standard output
-        with contextlib.redirect_stdout(output_buffer):
-            # dicts for globals/locals to keep environment somewhat contained
-            exec(code, {}, {})
+    # Check if code contains !pip install lines (Jupyter-style)
+    lines = code.strip().split("\n")
+    pip_lines = [l for l in lines if l.strip().startswith("!pip ") or l.strip().startswith("!pip3 ")]
+    code_lines = [l for l in lines if not l.strip().startswith("!pip ") and not l.strip().startswith("!pip3 ")]
+    
+    output_parts = []
+    
+    # Handle any !pip commands first
+    for pip_line in pip_lines:
+        cmd = pip_line.strip().lstrip("!")  # Remove the "!" prefix
+        parts = cmd.split()
+        if len(parts) >= 3 and parts[1] == "install":
+            pkg = " ".join(parts[2:])  # Support version specifiers like "pygame==2.5.0"
+            result = pkg_mgr.install_package(pkg)
+            output_parts.append(f">>> {pip_line.strip()}\n{result}")
+        elif len(parts) >= 3 and parts[1] == "uninstall":
+            pkg = parts[2]
+            result = pkg_mgr.uninstall_package(pkg)
+            output_parts.append(f">>> {pip_line.strip()}\n{result}")
+        elif len(parts) >= 2 and parts[1] == "list":
+            result = pkg_mgr.list_packages()
+            output_parts.append(f">>> {pip_line.strip()}\n{result}")
+        else:
+            output_parts.append(f">>> {pip_line.strip()}\nUnsupported pip command. Use: !pip install <pkg>, !pip uninstall <pkg>, !pip list")
+    
+    # Run remaining Python code
+    remaining_code = "\n".join(code_lines).strip()
+    if remaining_code:
+        start_time = time.time()
+        output_buffer = StringIO()
         
-        exec_time = time.time() - start_time
-        result = output_buffer.getvalue()
-        if not result:
-            result = "(No output)"
+        try:
+            with contextlib.redirect_stdout(output_buffer):
+                exec(remaining_code, {}, {})
             
-        return f"{result}\n\n[Execution time: {exec_time:.4f}s]", ""
-    except Exception as e:
-        return output_buffer.getvalue(), f"Traceback (most recent call last):\n{str(e)}"
+            exec_time = time.time() - start_time
+            result = output_buffer.getvalue()
+            if not result:
+                result = "(No output)"
+                
+            output_parts.append(f"{result}\n[Execution time: {exec_time:.4f}s]")
+        except Exception as e:
+            output_parts.append(f"{output_buffer.getvalue()}\n=== ERROR ===\nTraceback (most recent call last):\n{str(e)}")
+    
+    combined = "\n\n".join(output_parts)
+    return combined, ""
 
 def generate_completion(code: str) -> str:
     model = load_model()
@@ -149,7 +289,6 @@ def create_diff_view(original: str, modified: str) -> str:
 
 # --- Gradio UI ---
 
-# Theme & CSS Definition
 theme = gr.themes.Default(
     font=[gr.themes.GoogleFont('JetBrains Mono'), 'monospace'],
     primary_hue="blue",
@@ -162,15 +301,14 @@ theme = gr.themes.Default(
 )
 
 css = """
-    /* Force square corners */
     * { border-radius: 0 !important; } 
     .editor-container { height: 50vh; }
     .terminal { height: 15vh; background-color: #1e1e1e; color: #d4d4d4; font-family: 'JetBrains Mono', monospace; }
     .diff-view { height: 40vh; }
     .code-wrap { white-space: pre-wrap !important; }
+    .pkg-output { font-family: 'JetBrains Mono', monospace; font-size: 0.85em; }
 """
 
-# FIX: Removed theme/css from Blocks constructor
 with gr.Blocks(title="Axon Pro - Python IDE") as demo:
     
     gr.Markdown("# 🐍 Axon Pro — Python AI IDE")
@@ -213,7 +351,7 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
                     with gr.Tabs():
                         with gr.Tab("Terminal", id="terminal-tab"):
                             terminal = gr.Textbox(
-                                value=">>> Ready to run Python code...", 
+                                value=">>> Ready to run Python code...\n>>> Tip: Use !pip install <package> in your code to install packages", 
                                 lines=6, 
                                 interactive=False,
                                 elem_classes="terminal",
@@ -227,13 +365,42 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
                                 chat_input = gr.Textbox(placeholder="Ask to generate Python code...", scale=7, container=False)
                                 send_btn = gr.Button("Generate", variant="primary", scale=1)
         
+        # --- Package Manager Tab ---
+        with gr.Tab("📦 Packages", id="pkg-tab"):
+            gr.Markdown("### Package Manager")
+            gr.Markdown("Install, uninstall, and manage Python packages at runtime.")
+            
+            with gr.Row():
+                with gr.Column(scale=3):
+                    pkg_name_input = gr.Textbox(
+                        placeholder="e.g. pygame, requests, numpy==1.24.0",
+                        label="Package Name",
+                        container=True
+                    )
+                with gr.Column(scale=1):
+                    with gr.Row():
+                        install_btn = gr.Button("📥 Install", variant="primary")
+                        uninstall_btn = gr.Button("🗑️ Uninstall", variant="secondary")
+            
+            with gr.Row():
+                list_btn = gr.Button("📋 List Installed", variant="secondary")
+                info_btn = gr.Button("ℹ️ Package Info", variant="secondary")
+            
+            pkg_output = gr.Textbox(
+                value="Package manager ready. Install any pip package here.\n\nYou can also use !pip install <package> directly in the editor.",
+                label="Output",
+                lines=12,
+                interactive=False,
+                elem_classes="pkg-output"
+            )
+        
         with gr.Tab("Diff View", id="diff-tab"):
             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** | Python 3.x | CPU Mode | TinyLlama-1.1B", elem_classes="status-bar")
+    status_bar = gr.Markdown(f"**AXON PRO v1.1** | Python 3.x | CPU Mode | TinyLlama-1.1B | pip enabled", elem_classes="status-bar")
     
     # State
     current_file_state = gr.State(fs.current_file)
@@ -241,7 +408,7 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
     diff_modified_state = gr.State("")
     diff_mode_state = gr.State(False)
     
-    # Functions
+    # --- Editor Functions ---
     def update_file(content):
         fs.save_file(content)
         return fs.get_current_file_content()
@@ -286,7 +453,26 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
         new_hist = history + [[prompt, "Generated code available in Diff View"]]
         return diff, "", generated, True, new_hist, ""
 
-    # Wiring
+    # --- Package Manager Functions ---
+    def install_wrapper(pkg_name):
+        if not pkg_name:
+            return "Please enter a package name."
+        return pkg_mgr.install_package(pkg_name)
+    
+    def uninstall_wrapper(pkg_name):
+        if not pkg_name:
+            return "Please enter a package name."
+        return pkg_mgr.uninstall_package(pkg_name)
+    
+    def list_wrapper():
+        return pkg_mgr.list_packages()
+    
+    def info_wrapper(pkg_name):
+        if not pkg_name:
+            return "Please enter a package name."
+        return pkg_mgr.search_package_info(pkg_name)
+
+    # --- Wiring: Editor ---
     editor.change(update_file, editor, None)
     file_list.change(load_file, file_list, [editor, current_file_state])
     new_file_btn.click(create_file, new_file_name, [file_list, editor])
@@ -295,7 +481,6 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
     complete_btn.click(complete_wrapper, editor, editor)
     clear_btn.click(lambda: ">>> Ready...", None, terminal)
     
-    # Tabs logic
     explain_btn.click(
         explain_wrapper, editor, 
         [diff_view, diff_original_state, diff_modified_state, diff_mode_state]
@@ -318,13 +503,18 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
     discard_btn.click(
         lambda: (gr.update(), False), None, [editor, diff_mode_state]
     ).then(lambda: gr.Tabs(selected="editor-tab"), None, main_tabs)
+    
+    # --- Wiring: Package Manager ---
+    install_btn.click(install_wrapper, pkg_name_input, pkg_output)
+    uninstall_btn.click(uninstall_wrapper, pkg_name_input, pkg_output)
+    list_btn.click(list_wrapper, None, pkg_output)
+    info_btn.click(info_wrapper, pkg_name_input, pkg_output)
 
 if __name__ == "__main__":
-    # FIX: Passed theme and css here instead of Blocks()
     demo.launch(
         server_name="0.0.0.0", 
         server_port=7860,
         theme=theme,
         css=css,
-        ssr_mode=False  # Disables SSR to prevent experimental warnings/errors
+        ssr_mode=False
     )