Update app.py

app.py

@@ -2,10 +2,14 @@ import gradio as gr
 from openai import OpenAI
 import os
 import json
-from novita_sandbox.code_interpreter import Sandbox
+import threading
+import time
 import atexit
+from novita_sandbox.code_interpreter import Sandbox
 
-# --- Initialization ---
+# -------------------------
+# Global State
+# -------------------------
 client = OpenAI(
     base_url="https://api.novita.ai/openai",
     api_key=os.environ["NOVITA_API_KEY"],
@@ -13,12 +17,55 @@ client = OpenAI(
 
 model = "meta-llama/llama-3.3-70b-instruct"
 
-# Create sandbox
-sandbox = Sandbox.create(timeout=1200)
+sandbox = None
+sandbox_timer = None
+
+# -------------------------
+# Sandbox Management
+# -------------------------
+def create_sandbox():
+    global sandbox
+    if sandbox is None:
+        sandbox = Sandbox.create(timeout=1200)
+        print("[DEBUG] Sandbox created.")
+        return "🟢 Sandbox ON"
+    return "Sandbox already running."
+
+def kill_sandbox():
+    global sandbox
+    if sandbox is not None:
+        try:
+            sandbox.kill()
+        except Exception:
+            pass
+        sandbox = None
+        print("[DEBUG] Sandbox killed.")
+        return "🔴 Sandbox OFF"
+    return "Sandbox already off."
 
-# --- Tool functions ---
+def sandbox_auto_off():
+    print("[DEBUG] Auto-off countdown started (1200 sec).")
+    time.sleep(1200)
+    print("[DEBUG] Auto-off triggered.")
+    kill_sandbox()
+
+# -------------------------
+# Sandbox Usage Guard
+# -------------------------
+def require_sandbox():
+    if sandbox is None:
+        return "❌ Sandbox is OFF. Turn it ON to use this feature."
+    return None
+
+# -------------------------
+# Tool Functions
+# -------------------------
 def read_file(path: str):
     print(f"[DEBUG] read_file called with path: {path}")
+    err = require_sandbox()
+    if err:
+        return err
+
     try:
         content = sandbox.files.read(path)
         print(f"[DEBUG] read_file result: {content}")
@@ -29,6 +76,10 @@ def read_file(path: str):
 
 def write_file(path: str, data: str):
     print(f"[DEBUG] write_file called with path: {path}")
+    err = require_sandbox()
+    if err:
+        return err
+
     try:
         sandbox.files.write(path, data)
         msg = f"File created successfully at {path}"
@@ -40,6 +91,10 @@ def write_file(path: str, data: str):
 
 def write_files(files: list):
     print(f"[DEBUG] write_files called with {len(files)} files")
+    err = require_sandbox()
+    if err:
+        return err
+
     try:
         sandbox.files.write_files(files)
         msg = f"{len(files)} file(s) created successfully"
@@ -51,15 +106,21 @@ def write_files(files: list):
 
 def run_commands(command: str):
     print(f"[DEBUG] run_commands called with command: {command}")
+    err = require_sandbox()
+    if err:
+        return err
+
     try:
         result = sandbox.commands.run(command)
-        print(f"[DEBUG] run_commands result: {result}")
+        print(f"[DEBUG] run_commands result: {result.stdout}")
         return result.stdout
     except Exception as e:
         print(f"[DEBUG] run_commands error: {e}")
         return f"Error running command: {e}"
 
-# --- Register tools ---
+# -------------------------
+# Register tools
+# -------------------------
 tools = [
     {
         "type": "function",
@@ -116,33 +177,31 @@ tools = [
         "type": "function",
         "function": {
             "name": "run_commands",
-            "description": "Run a single shell command inside the sandbox working directory",
+            "description": "Run a shell command inside the sandbox",
             "parameters": {
                 "type": "object",
-                "properties": {
-                    "command": {"type": "string"},
-                },
+                "properties": {"command": {"type": "string"}},
                 "required": ["command"],
             },
         },
     },
 ]
 
-# --- Persistent chat messages ---
+# -------------------------
+# Chat + Tool Call Debug
+# -------------------------
 messages = []
 
-# --- Global model setter ---
 def set_model(selected_model):
     global model
     model = selected_model
-    print(f"[DEBUG] Model switched to: {model}")
     return f"✅ Model switched to **{model}**"
 
 def chat_fn(user_message, history):
     global messages, model
+
     messages.append({"role": "user", "content": user_message})
 
-    # Send to model
     response = client.chat.completions.create(
         model=model,
         messages=messages,
@@ -152,72 +211,89 @@ def chat_fn(user_message, history):
     assistant_msg = response.choices[0].message
     messages.append(assistant_msg)
 
-    output_text = ""
-
+    # DEBUG tool call logging
     if assistant_msg.tool_calls:
         print(f"[DEBUG] Assistant requested {len(assistant_msg.tool_calls)} tool call(s).")
 
-        for tool_call in assistant_msg.tool_calls:
-            fn_name = tool_call.function.name
-            fn_args = json.loads(tool_call.function.arguments)
-            print(f"[DEBUG] Tool call detected: {fn_name} with args {fn_args}")
+        for tc in assistant_msg.tool_calls:
+            print(f"[DEBUG] Tool call detected: {tc.function.name} with args {tc.function.arguments}")
+
+            fn_name = tc.function.name
+            fn_args = json.loads(tc.function.arguments)
 
             if fn_name == "read_file":
-                fn_result = read_file(**fn_args)
+                result = read_file(**fn_args)
             elif fn_name == "write_file":
-                fn_result = write_file(**fn_args)
+                result = write_file(**fn_args)
             elif fn_name == "write_files":
-                fn_result = write_files(**fn_args)
+                result = write_files(**fn_args)
             elif fn_name == "run_commands":
-                fn_result = run_commands(**fn_args)
+                result = run_commands(**fn_args)
             else:
-                fn_result = f"Error: Unknown tool {fn_name}"
+                result = f"Unknown tool {fn_name}"
+
+            print(f"[DEBUG] Tool call result: {result}")
 
             messages.append({
-                "tool_call_id": tool_call.id,
+                "tool_call_id": tc.id,
                 "role": "tool",
-                "content": str(fn_result),
+                "content": str(result),
             })
 
-        follow_up = client.chat.completions.create(
+        followup = client.chat.completions.create(
             model=model,
             messages=messages,
         )
-        final_answer = follow_up.choices[0].message
-        messages.append(final_answer)
-        output_text = final_answer.content
-    else:
-        output_text = assistant_msg.content
+        final_msg = followup.choices[0].message
+        messages.append(final_msg)
+        return final_msg.content
 
-    return output_text
+    return assistant_msg.content
 
-# --- Command Interface function ---
+# -------------------------
+# Command Interface
+# -------------------------
 def execute_command(command):
     if not command.strip():
         return "⚠️ Please enter a command."
-    print(f"[DEBUG] Executing command from interface: {command}")
     output = run_commands(command)
-    return f"```bash\n{output}\n```" if output else "✅ Command executed (no output)."
+    return f"```bash\n{output}\n```"
 
-# --- Gradio UI ---
+# -------------------------
+# UI
+# -------------------------
 with gr.Blocks(title="Novita Sandbox App") as demo:
     gr.Markdown("## 🧠 Novita Sandbox Agent")
-    gr.Markdown(
-    "This app is an AI-powered **code agent** that lets you chat with intelligent assistants backed by **Novita AI LLMs**. These agents can write, read, and execute code safely inside a **Novita sandbox**, providing a secure environment for running commands, testing scripts, and managing files, all through an intuitive chat interface with model selection and command execution built right in."
-)
-
+    gr.Markdown("Interact with a Novita sandbox-enabled LLM with code execution capability.")
 
     with gr.Row(equal_height=True):
-        # Left: Chat Interface
+        # Chat
         with gr.Column(scale=2):
             gr.Markdown("### 💬 Chat Interface")
             gr.ChatInterface(chat_fn)
 
-        # Right: Command Interface
+        # Controls
         with gr.Column(scale=1):
-            gr.Markdown("### 💻 Command Interface")
-            
-            # Model selector
+            gr.Markdown("### ⚙️ Controls")
+
+            # Sandbox Switch
+            sandbox_switch = gr.Checkbox(label="Sandbox On/Off", value=False)
+            sandbox_status = gr.Markdown("🔴 Sandbox OFF")
+
+            def toggle_sandbox(is_on):
+                global sandbox_timer
+                if is_on:
+                    msg = create_sandbox()
+                    sandbox_timer = threading.Thread(target=sandbox_auto_off, daemon=True)
+                    sandbox_timer.start()
+                    return "🟢 Sandbox ON"
+                else:
+                    msg = kill_sandbox()
+                    return "🔴 Sandbox OFF"
+
+            sandbox_switch.change(toggle_sandbox, inputs=sandbox_switch, outputs=sandbox_status)
+
+            # Model Selector
             model_selector = gr.Dropdown(
                 label="Select Model",
                 choices=[
@@ -228,26 +304,18 @@ with gr.Blocks(title="Novita Sandbox App") as demo:
                     "moonshotai/kimi-k2-instruct",
                 ],
                 value=model,
-                interactive=True,
             )
-
-            model_status = gr.Markdown(f"✅ Current model: **{model}**")
+            model_status = gr.Markdown(f"Current model: **{model}**")
             model_selector.change(set_model, inputs=model_selector, outputs=model_status)
 
-            command_input = gr.Textbox(
-                label="Command",
-                placeholder="e.g., ls, python main.py",
-                lines=1,
-            )
-            with gr.Row():
-                run_btn = gr.Button("Run", variant="primary", scale=0)
-            command_output = gr.Markdown("Command output will appear here...")
-
+            # Command Runner
+            command_input = gr.Textbox(label="Command", placeholder="e.g., ls")
+            run_btn = gr.Button("Run", variant="primary")
+            command_output = gr.Markdown("Command output here...")
             run_btn.click(execute_command, inputs=command_input, outputs=command_output)
 
-# --- Cleanup on exit ---
-atexit.register(lambda: (sandbox.kill(), print("[DEBUG] Sandbox terminated. 👋")))
+# Cleanup
+atexit.register(lambda: (kill_sandbox(), print("[DEBUG] Sandbox terminated.")))
 
 if __name__ == "__main__":
     demo.launch()
-