Axon-Pro-IDE / app.py
AIencoder's picture
Update app.py
d36b3a0 verified
raw
history blame
27.2 kB
import warnings
warnings.filterwarnings("ignore", category=FutureWarning, message=".*torch.distributed.reduce_op.*")
import gradio as gr
import torch
import time
import sys
import subprocess
import os
import pty
import select
import signal
import fcntl
import struct
import termios
import threading
import re
import tempfile
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, GenerationConfig
from typing import List
from functools import lru_cache
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# --- Configuration ---
MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
DEVICE = "cpu"
DTYPE = torch.float32
DISPLAY_NUM = ":99"
SCREEN_W, SCREEN_H = 800, 600
# --- Virtual Display ---
class VirtualDisplay:
def __init__(self):
self.xvfb_proc = None
self.display = DISPLAY_NUM
self._start_xvfb()
def _start_xvfb(self):
try:
subprocess.run(["pkill", "-f", f"Xvfb {self.display}"],
capture_output=True, timeout=5)
time.sleep(0.2)
self.xvfb_proc = subprocess.Popen(
["Xvfb", self.display, "-screen", "0",
f"{SCREEN_W}x{SCREEN_H}x24", "-ac", "-nolisten", "tcp"],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
time.sleep(0.5)
if self.xvfb_proc.poll() is not None:
print("[Xvfb] Failed, installing...")
subprocess.run(["apt-get", "update", "-qq"], capture_output=True, timeout=60)
subprocess.run(["apt-get", "install", "-y", "-qq", "xvfb", "x11-utils",
"imagemagick"], capture_output=True, timeout=120)
self.xvfb_proc = subprocess.Popen(
["Xvfb", self.display, "-screen", "0",
f"{SCREEN_W}x{SCREEN_H}x24", "-ac", "-nolisten", "tcp"],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
time.sleep(0.5)
if self.xvfb_proc.poll() is None:
print(f"[Xvfb] Running on {self.display}")
os.environ["DISPLAY"] = self.display
else:
print("[Xvfb] Could not start")
self.xvfb_proc = None
except FileNotFoundError:
print("[Xvfb] Not found, installing...")
try:
subprocess.run(["apt-get", "update", "-qq"], capture_output=True, timeout=60)
subprocess.run(["apt-get", "install", "-y", "-qq", "xvfb", "x11-utils",
"imagemagick"], capture_output=True, timeout=120)
self._start_xvfb()
except Exception as e:
print(f"[Xvfb] Install failed: {e}")
self.xvfb_proc = None
except Exception as e:
print(f"[Xvfb] Error: {e}")
self.xvfb_proc = None
def capture_screenshot(self) -> str | None:
if not self.is_running:
return None
try:
tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
tmp.close()
result = subprocess.run(
["import", "-window", "root", "-display", self.display, tmp.name],
capture_output=True, timeout=5)
if result.returncode == 0 and os.path.getsize(tmp.name) > 100:
return tmp.name
os.unlink(tmp.name)
except Exception:
pass
return None
@property
def is_running(self) -> bool:
return self.xvfb_proc is not None and self.xvfb_proc.poll() is None
def cleanup(self):
if self.xvfb_proc:
self.xvfb_proc.terminate()
try:
self.xvfb_proc.wait(timeout=3)
except subprocess.TimeoutExpired:
self.xvfb_proc.kill()
vdisplay = VirtualDisplay()
# --- File System ---
class FileSystem:
def __init__(self):
self.files = {
"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.py"
def save_file(self, content: str) -> None:
if self.current_file:
self.files[self.current_file] = content
def get_current_file_content(self) -> str:
return self.files.get(self.current_file, "")
def set_current_file(self, filename: str) -> None:
if filename in self.files:
self.current_file = filename
def create_file(self, filename: str, content: str = "") -> None:
if filename not in self.files:
self.files[filename] = content
self.current_file = filename
def get_all_files(self) -> List[str]:
return list(self.files.keys())
fs = FileSystem()
# --- PTY Terminal ---
class PTYTerminal:
STRIP_ANSI = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07|\x1b\[.*?[@-~]|\r')
def __init__(self):
self.master_fd = None
self.pid = None
self.lock = threading.Lock()
self.log_lines: List[str] = []
self.max_lines = 800
self._spawn_shell()
self._append("AXON TERMINAL v3.0")
self._append("══════════════════════════════════════")
self._append("Full PTY shell — everything works:")
self._append(" pip/npm/apt-get/git/curl/wget/make")
self._append(" python/node — GUI apps render to Display")
self._append("══════════════════════════════════════")
def _spawn_shell(self):
try:
pid, fd = pty.openpty()
self.pid = os.fork()
if self.pid == 0:
os.close(pid)
os.setsid()
fcntl.ioctl(fd, termios.TIOCSCTTY, 0)
os.dup2(fd, 0); os.dup2(fd, 1); os.dup2(fd, 2)
if fd > 2: os.close(fd)
env = os.environ.copy()
env.update({"TERM": "dumb", "PS1": "$ ",
"DEBIAN_FRONTEND": "noninteractive", "DISPLAY": DISPLAY_NUM})
os.execvpe("/bin/bash", ["/bin/bash", "--norc", "--noprofile", "-i"], env)
else:
os.close(fd)
self.master_fd = pid
flags = fcntl.fcntl(self.master_fd, fcntl.F_GETFL)
fcntl.fcntl(self.master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ,
struct.pack("HHHH", 40, 120, 0, 0))
time.sleep(0.3)
self._read(0.2)
except Exception as e:
self._append(f"[Shell Error] {e}")
self.master_fd = None
def _read(self, timeout=0.1) -> str:
if self.master_fd is None: return ""
out = []
deadline = time.time() + timeout
while True:
rem = deadline - time.time()
if rem <= 0: break
try:
r, _, _ = select.select([self.master_fd], [], [], min(rem, 0.05))
if r:
c = os.read(self.master_fd, 4096)
if c:
out.append(c.decode("utf-8", errors="replace"))
deadline = time.time() + 0.15
else: break
elif out: break
except OSError: break
return "".join(out)
def _clean(self, text):
c = self.STRIP_ANSI.sub("", text)
return "".join(ch for ch in c if ch in "\n\t" or ord(ch) >= 32)
def _append(self, text):
self.log_lines.append(text)
while len(self.log_lines) > self.max_lines:
self.log_lines.pop(0)
def get_log(self) -> str:
return "\n".join(self.log_lines)
def run_command(self, cmd: str) -> str:
cmd = cmd.strip()
if not cmd: return self.get_log()
with self.lock:
if cmd.lower() == "clear":
self.log_lines = []
return ""
if self.master_fd is None:
self._append(f"$ {cmd}")
self._append("[Error] No shell. Restart Space.")
return self.get_log()
self._read(0.05)
self._append(f"$ {cmd}")
try:
os.write(self.master_fd, (cmd + "\n").encode())
except OSError as e:
self._append(f"[Write Error] {e}")
return self.get_log()
parts = cmd.split()
base = parts[0].lower() if parts else ""
long = ["pip","pip3","npm","npx","apt-get","apt","git",
"wget","curl","make","cmake","cargo","yarn","conda"]
wait = 180 if base in long else (60 if base in ("python","python3","node") else 15)
chunks = []
start = time.time()
idle = 0
while time.time() - start < wait:
c = self._read(0.3)
if c:
idle = 0; chunks.append(c)
if "".join(chunks).rstrip().endswith("$"): break
else:
idle += 1
if base not in long and idle >= 3: break
if base in long and idle >= 10: break
raw = self._clean("".join(chunks))
lines = raw.split("\n")
filtered = []
skip = True
for l in lines:
s = l.strip()
if skip and s == cmd.strip():
skip = False; continue
if s in ("$", "$ "): continue
filtered.append(l)
result = "\n".join(filtered).strip()
if result: self._append(result)
return self.get_log()
def run_editor_code(self, code: str) -> str:
with self.lock:
self._append("$ python [editor]")
tmp = "/tmp/_axon_run.py"
try:
with open(tmp, "w") as f: f.write(code)
except Exception as e:
with self.lock: self._append(f"[Error] {e}")
return self.get_log()
try:
env = os.environ.copy()
env["DISPLAY"] = DISPLAY_NUM
r = subprocess.run([sys.executable, tmp], capture_output=True,
text=True, timeout=30, env=env)
with self.lock:
if r.stdout.strip(): self._append(r.stdout.rstrip())
if r.stderr.strip(): self._append(r.stderr.rstrip())
if not r.stdout.strip() and not r.stderr.strip():
self._append("(No output)")
except subprocess.TimeoutExpired:
with self.lock: self._append("[Timed out 30s]")
except Exception as e:
with self.lock: self._append(f"[Error] {e}")
return self.get_log()
def cleanup(self):
if self.pid and self.pid > 0:
try: os.kill(self.pid, signal.SIGTERM)
except ProcessLookupError: pass
if self.master_fd is not None:
try: os.close(self.master_fd)
except OSError: pass
terminal = PTYTerminal()
# --- Model ---
@lru_cache(maxsize=1)
def load_model():
print("Loading TinyLlama...")
t0 = time.time()
try:
tok = AutoTokenizer.from_pretrained(MODEL_NAME)
mdl = AutoModelForCausalLM.from_pretrained(
MODEL_NAME, dtype=DTYPE, device_map=DEVICE, low_cpu_mem_usage=True)
gc = GenerationConfig(max_new_tokens=256, temperature=0.2, top_p=0.95, do_sample=True)
pipe = pipeline("text-generation", model=mdl, tokenizer=tok, generation_config=gc)
print(f"Loaded in {time.time()-t0:.1f}s")
return pipe
except Exception as e:
print(f"Model error: {e}")
return None
def _llm(system, user, max_tok=512):
m = load_model()
if not m: return "Error: model not loaded"
prompt = f"<|system|>{system}</s><|user|>\n{user}\n</s><|assistant|>"
try:
return m(prompt, max_new_tokens=max_tok)[0]["generated_text"].strip()
except Exception as e:
return str(e)
def generate_completion(code): return _llm("Complete this Python code. No explanations.", code, 256)
def explain_code(code): return _llm("Explain this Python code concisely.", code)
def refactor_code(code): return _llm("Refactor for PEP 8 and best practices.", code)
def generate_code(p): return _llm(f"Write Python code for: {p}", "")
def create_diff_view(orig, mod):
import difflib
return "\n".join(difflib.unified_diff(orig.splitlines(), mod.splitlines(), lineterm=""))
# ═══════════════════════════════════════════════════════════
# THEME — only valid params
# ═══════════════════════════════════════════════════════════
theme = gr.themes.Default(
font=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
primary_hue="cyan",
neutral_hue="gray",
).set(
body_background_fill="#080808",
body_text_color="#e0e0e0",
button_primary_background_fill="#0ff",
button_primary_background_fill_hover="#0dd",
button_primary_text_color="#000",
button_secondary_background_fill="#1a1a1a",
button_secondary_background_fill_hover="#2a2a2a",
button_secondary_text_color="#0ff",
button_secondary_border_color="#0ff4",
block_background_fill="#0d0d0d",
block_border_color="#1a1a1a",
block_label_text_color="#0f8",
input_background_fill="#0a0a0a",
input_border_color="#1a3a2a",
input_placeholder_color="#0f83",
)
# ═══════════════════════════════════════════════════════════
# CSS
# ═══════════════════════════════════════════════════════════
css = """
* { border-radius: 0 !important; }
.gradio-container { max-width: 100% !important; }
/* ─── Header ─── */
.axon-header {
background: linear-gradient(90deg, #080808 0%, #0a1a1a 50%, #080808 100%);
border-bottom: 1px solid #0ff3;
padding: 12px 20px !important;
margin-bottom: 8px;
}
.axon-header h1 {
font-family: 'IBM Plex Mono', monospace !important;
background: linear-gradient(90deg, #0ff, #0f8);
-webkit-background-clip: text;
-webkit-text-fill-color: transparent;
font-size: 1.4em !important;
letter-spacing: 3px;
margin: 0 !important;
}
.axon-header p { color: #0f84 !important; font-size: 0.75em; letter-spacing: 1px; }
/* ─── Terminal ─── */
.terminal-box textarea {
background: #050505 !important;
color: #00ff41 !important;
font-family: 'IBM Plex Mono', monospace !important;
font-size: 12px !important;
line-height: 1.5 !important;
border: 1px solid #0f31 !important;
text-shadow: 0 0 4px #0f32;
}
.terminal-input input {
background: #050505 !important;
color: #00ff41 !important;
font-family: 'IBM Plex Mono', monospace !important;
border: 1px solid #0f32 !important;
font-size: 12px !important;
caret-color: #0f8 !important;
}
.terminal-input input::placeholder { color: #0f83 !important; }
.terminal-input input:focus { border-color: #0ff !important; box-shadow: 0 0 8px #0ff2 !important; }
/* ─── Editor ─── */
.code-editor textarea, .code-editor .cm-editor {
font-family: 'IBM Plex Mono', monospace !important;
font-size: 13px !important;
}
/* ─── File explorer ─── */
.file-explorer { border-right: 1px solid #1a2a2a !important; }
/* ─── Display ─── */
.display-panel { border: 1px solid #1a2a2a !important; background: #050505 !important; }
.display-panel img { image-rendering: pixelated; }
/* ─── Toolbar buttons ─── */
.toolbar-btn {
font-size: 0.8em !important;
padding: 6px 12px !important;
letter-spacing: 0.5px;
text-transform: uppercase;
font-weight: 600 !important;
border: 1px solid transparent !important;
transition: all 0.2s ease !important;
}
.toolbar-btn:hover { border-color: #0ff !important; box-shadow: 0 0 12px #0ff2 !important; }
/* ─── Accordion ─── */
.gr-accordion { border: 1px solid #1a2a2a !important; background: #0a0a0a !important; }
.gr-accordion > .label-wrap {
background: #0d0d0d !important;
border-bottom: 1px solid #1a2a2a !important;
transition: background 0.2s ease;
}
.gr-accordion > .label-wrap:hover { background: #121a1a !important; }
.gr-accordion > .label-wrap span {
color: #0ff !important;
font-family: 'IBM Plex Mono', monospace !important;
letter-spacing: 1px;
font-size: 0.85em;
}
/* ─── Tabs ─── */
.tabs > .tab-nav > button {
font-family: 'IBM Plex Mono', monospace !important;
letter-spacing: 1px; font-size: 0.8em; text-transform: uppercase;
color: #888 !important; border-bottom: 2px solid transparent !important;
transition: all 0.2s ease;
}
.tabs > .tab-nav > button.selected {
color: #0ff !important; border-bottom-color: #0ff !important;
text-shadow: 0 0 8px #0ff4;
}
/* ─── Status ─── */
.status-bar {
background: #0a0a0a !important; border-top: 1px solid #1a2a2a;
padding: 4px 12px !important; font-size: 11px !important; color: #0f84 !important;
}
.status-bar strong { color: #0ff !important; }
/* ─── Scrollbar ─── */
::-webkit-scrollbar { width: 6px; height: 6px; }
::-webkit-scrollbar-track { background: #0a0a0a; }
::-webkit-scrollbar-thumb { background: #0f83; }
::-webkit-scrollbar-thumb:hover { background: #0ff; }
/* ─── Generic input text color (covers what input_text_color would do) ─── */
input, textarea, select { color: #0f8 !important; }
"""
# ═══════════════════════════════════════════════════════════
# UI LAYOUT
# ═══════════════════════════════════════════════════════════
with gr.Blocks(title="Axon Pro — Python AI IDE") as demo:
with gr.Row(elem_classes="axon-header"):
gr.Markdown("# ⬡ AXON PRO\n\nPYTHON AI IDE — v3.0")
with gr.Row(equal_height=False):
# ── LEFT: FILE EXPLORER ──
with gr.Column(scale=1, min_width=180, elem_classes="file-explorer"):
with gr.Accordion("📁 EXPLORER", open=True):
file_list = gr.Dropdown(
choices=fs.get_all_files(), value=fs.current_file,
label="", interactive=True, container=False)
new_file_name = gr.Textbox(placeholder="filename.py", label="", container=False)
with gr.Row():
new_file_btn = gr.Button("+ NEW", variant="secondary", size="sm",
elem_classes="toolbar-btn")
save_btn = gr.Button("SAVE", variant="secondary", size="sm",
elem_classes="toolbar-btn")
# ── CENTER: EDITOR + TERMINAL ──
with gr.Column(scale=4):
editor = gr.Code(
value=fs.get_current_file_content(),
label=f" {fs.current_file}",
language="python", lines=22, interactive=True,
elem_classes="code-editor")
with gr.Row():
run_btn = gr.Button("▶ RUN", variant="primary", elem_classes="toolbar-btn")
complete_btn = gr.Button("✦ COMPLETE", variant="secondary", elem_classes="toolbar-btn")
explain_btn = gr.Button("◈ EXPLAIN", variant="secondary", elem_classes="toolbar-btn")
refactor_btn = gr.Button("⟲ REFACTOR", variant="secondary", elem_classes="toolbar-btn")
with gr.Tabs() as bottom_tabs:
with gr.Tab("⌘ TERMINAL", id="terminal-tab"):
terminal_output = gr.Textbox(
value=terminal.get_log(), lines=10, max_lines=25,
interactive=False, elem_classes="terminal-box",
label="", show_label=False)
with gr.Row():
terminal_input = gr.Textbox(
placeholder="$ pip install pygame | npm i | git clone | ls ...",
scale=9, container=False, elem_classes="terminal-input")
terminal_run_btn = gr.Button("⏎", variant="primary", scale=1, min_width=45)
clear_btn = gr.Button("CLEAR", variant="secondary", size="sm",
elem_classes="toolbar-btn")
with gr.Tab("💬 AI CHAT", id="chat-tab"):
chat_history = gr.Chatbot(label="", height=250)
with gr.Row():
chat_input = gr.Textbox(placeholder="Describe the code you want...",
scale=8, container=False)
send_btn = gr.Button("GEN", variant="primary", scale=1,
elem_classes="toolbar-btn")
with gr.Tab("⇄ DIFF", id="diff-tab"):
diff_view = gr.Code(label="AI Changes", language="python",
interactive=False, lines=15)
with gr.Row():
apply_btn = gr.Button("✓ APPLY", variant="primary",
elem_classes="toolbar-btn")
discard_btn = gr.Button("✗ DISCARD", variant="secondary",
elem_classes="toolbar-btn")
# ── RIGHT: DISPLAY ──
with gr.Column(scale=2, min_width=250):
with gr.Accordion("🖥 DISPLAY OUTPUT", open=True):
gr.Markdown(
"<small style='color:#0f84'>GUI apps (pygame, tkinter, etc.) "
"render here via virtual display</small>")
display_image = gr.Image(
label="", type="filepath", interactive=False,
elem_classes="display-panel", height=400)
with gr.Row():
capture_btn = gr.Button("📸 CAPTURE", variant="secondary", size="sm",
elem_classes="toolbar-btn")
auto_refresh = gr.Checkbox(label="Auto-refresh", value=False)
display_status = gr.Markdown(
f"<small style='color:#0f84'>Xvfb: "
f"{'● running' if vdisplay.is_running else '○ stopped'} "
f"| {SCREEN_W}x{SCREEN_H}</small>")
status_bar = gr.Markdown(
f"**AXON PRO v3.0** │ Python {sys.version.split()[0]} │ CPU │ "
f"TinyLlama-1.1B │ PTY Shell │ Xvfb {'ON' if vdisplay.is_running else 'OFF'}",
elem_classes="status-bar")
# ── State ──
diff_original_state = gr.State("")
diff_modified_state = gr.State("")
# ── Handlers ──
def update_file(content):
fs.save_file(content)
def load_file(filename):
fs.set_current_file(filename)
return fs.get_current_file_content(), gr.update(label=f" {filename}")
def create_file_fn(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(), gr.update(label=f" {name}"))
return gr.update(choices=fs.get_all_files()), fs.get_current_file_content(), gr.update()
def terminal_cmd(cmd):
return terminal.run_command(cmd), ""
def clear_term():
terminal.log_lines = []
return ""
def run_editor(content):
return terminal.run_editor_code(content)
def complete_fn(content):
return content + "\n" + generate_completion(content)
def explain_fn(content):
exp = explain_code(content)
fs.save_file(content)
mod = f'"""\nEXPLANATION:\n{exp}\n"""\n\n{content}'
return create_diff_view(content, mod), content, mod
def refactor_fn(content):
ref = refactor_code(content)
fs.save_file(content)
return create_diff_view(content, ref), content, ref
def generate_fn(prompt_text, history):
gen = generate_code(prompt_text)
diff = create_diff_view("", gen)
new_h = history + [
{"role": "user", "content": prompt_text},
{"role": "assistant", "content": "Code generated → check DIFF tab"},
]
return diff, "", gen, new_h, ""
def capture_display():
return vdisplay.capture_screenshot()
# ── Wiring ──
editor.change(update_file, editor, None)
file_list.change(load_file, file_list, [editor, editor])
new_file_btn.click(create_file_fn, new_file_name, [file_list, editor, editor])
terminal_input.submit(terminal_cmd, terminal_input, [terminal_output, terminal_input])
terminal_run_btn.click(terminal_cmd, terminal_input, [terminal_output, terminal_input])
clear_btn.click(clear_term, None, terminal_output)
run_btn.click(run_editor, editor, terminal_output)
complete_btn.click(complete_fn, editor, editor)
explain_btn.click(explain_fn, editor,
[diff_view, diff_original_state, diff_modified_state]
).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
refactor_btn.click(refactor_fn, editor,
[diff_view, diff_original_state, diff_modified_state]
).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
chat_input.submit(generate_fn, [chat_input, chat_history],
[diff_view, diff_original_state, diff_modified_state, chat_history, chat_input]
).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
send_btn.click(generate_fn, [chat_input, chat_history],
[diff_view, diff_original_state, diff_modified_state, chat_history, chat_input]
).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
apply_btn.click(lambda mod: mod, diff_modified_state, editor
).then(lambda: gr.Tabs(selected="terminal-tab"), None, bottom_tabs)
discard_btn.click(lambda: gr.update(), None, editor
).then(lambda: gr.Tabs(selected="terminal-tab"), None, bottom_tabs)
capture_btn.click(capture_display, None, display_image)
if __name__ == "__main__":
import atexit
atexit.register(terminal.cleanup)
atexit.register(vdisplay.cleanup)
demo.launch(server_name="0.0.0.0", server_port=7860,
theme=theme, css=css, ssr_mode=False)