Update app.py

app.py

@@ -1,12 +1,20 @@
-import gradio as gr
-import subprocess
 import os
+import re
+import subprocess
+import gradio as gr
+from transformers import AutoTokenizer
+from optimum.onnxruntime import ORTModelForCausalLM
 
-# Use container’s home directory to avoid permission errors
-WORKDIR = os.path.join(os.environ['HOME'], "workspace")
+# 1. Prepare a writable workspace
+WORKDIR = os.path.join(os.environ["HOME"], "workspace")
 os.makedirs(WORKDIR, exist_ok=True)
 
-# HTML iframe to embed noVNC viewer
+# 2. Load DialoGPT‐medium (chat‐tuned) via ONNXRuntime
+MODEL_ID = "microsoft/DialoGPT-medium"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = ORTModelForCausalLM.from_pretrained(MODEL_ID, provider="CPUExecutionProvider")
+
+# 3. Virtual desktop iframe
 NO_VNC_IFRAME = """
 <iframe src="http://localhost:6080/vnc.html?autoconnect=true&resize=remote"
         style="width:100%; height:600px; border:none;">
@@ -16,45 +24,92 @@ NO_VNC_IFRAME = """
 def start_desktop():
     return NO_VNC_IFRAME
 
-def run_shell(cmd):
-    result = subprocess.run(
-        cmd,
-        shell=True,
-        capture_output=True,
-        text=True,
-        env={**os.environ, "DISPLAY": ":1"}  # Supports GUI apps if needed
+def run_shell(cmd: str):
+    proc = subprocess.run(
+        cmd, shell=True, capture_output=True, text=True,
+        env={**os.environ, "DISPLAY": ":1"}
     )
-    return result.stdout + result.stderr
+    return proc.stdout + proc.stderr
 
 def upload_file(file):
-    dest_path = os.path.join(WORKDIR, os.path.basename(file.name))
-    with open(dest_path, "wb") as f:
+    dest = os.path.join(WORKDIR, os.path.basename(file.name))
+    with open(dest, "wb") as f:
         f.write(file.read())
-    return f"✅ Uploaded to: {dest_path}"
+    return f"✅ Uploaded to: {dest}"
+
+def agent_chat(user_msg, history):
+    # System prompt and conversation framing
+    system = (
+        "You are a helpful shell assistant. "
+        "When you need to run a shell command, wrap it in [CMD]...[/CMD]."
+    )
+    # Build the full conversation string
+    conv = system + tokenizer.eos_token + "\n"
+    for u, a in history:
+        conv += f"User: {u}{tokenizer.eos_token}Assistant: {a}{tokenizer.eos_token}\n"
+    conv += f"User: {user_msg}{tokenizer.eos_token}Assistant:"
+
+    # Tokenize + generate
+    inputs = tokenizer(conv, return_tensors="pt")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=150,
+        do_sample=True,
+        temperature=0.7,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    full_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    reply = full_reply.split("Assistant:")[-1].strip()
 
-with gr.Blocks(css="@import url('https://cdn.simplecss.org/simple.min.css');") as demo:
-    gr.Markdown("# 🖥️ Offline Visual Sandbox (no API, no cost)")
+    # Extract & execute shell commands
+    cmds = re.findall(r"
+
+\[CMD\]
+
+(.*?)
+
+\[/CMD\]
+
+", reply, re.DOTALL)
+    if cmds:
+        out = ""
+        for cmd in cmds:
+            out += f"$ {cmd}\n"
+            p = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+            out += p.stdout + p.stderr + "\n"
+        reply += f"\n\n**Command Output:**\n```\n{out}```"
+
+    history.append((user_msg, reply))
+    return reply, history
+
+# 4. Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖💬 ONNX Chat (DialoGPT) + Shell Agent")
 
     with gr.Tab("Desktop"):
-        launch_btn = gr.Button("Launch Desktop")
-        desktop_view = gr.HTML()
-        launch_btn.click(start_desktop, outputs=desktop_view)
+        gr.Markdown("### 🖥️ Virtual Desktop (noVNC)")
+        btn = gr.Button("Launch Desktop")
+        desktop = gr.HTML()
+        btn.click(start_desktop, outputs=desktop)
 
     with gr.Tab("Shell"):
-        cmd_input = gr.Textbox(label="Shell Command", placeholder="Try 'ls /home/user/workspace'")
-        run_btn = gr.Button("Run")
-        cmd_output = gr.Textbox(label="Command Output", lines=10)
-        run_btn.click(run_shell, inputs=cmd_input, outputs=cmd_output)
+        gr.Markdown("### 🛠️ Direct Shell")
+        cmd_in = gr.Textbox(label="Command")
+        cmd_out = gr.Textbox(label="Output", lines=6)
+        gr.Button("Run").click(run_shell, inputs=cmd_in, outputs=cmd_out)
+
+    with gr.Tab("Chat"):
+        gr.Markdown("### 💬 Chat with AI (can run shell)")
+        chat = gr.ChatInterface(
+            fn=agent_chat,
+            state=[],
+            placeholder="Ask me to list files: [CMD]ls workspace[/CMD]"
+        )
 
     with gr.Tab("Upload"):
-        file_input = gr.File(label="Upload File")
-        status_output = gr.Textbox(label="Upload Status")
-        file_input.change(upload_file, inputs=file_input, outputs=status_output)
-
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        ssr_mode=False
-    )
+        gr.Markdown("### 📂 Upload File to Sandbox")
+        uploader = gr.File(label="Choose File")
+        status = gr.Textbox(label="Status")
+        uploader.change(upload_file, inputs=uploader, outputs=status)
 
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)