Replace python execution tool with terminal

README.md

@@ -5,9 +5,8 @@ and demonstrates basic tool usage. Chat histories are stored in a local SQLite
 database using Peewee. Histories are persisted per user and session so
 conversations can be resumed with context. One example tool is included:
 
-* **execute_python** – Executes Python code in a sandbox with selected built-ins
-  and allows importing safe modules like ``math``. The result is returned from a
-  ``result`` variable or captured output.
+* **execute_terminal** – Executes a shell command in a Linux VM with network
+  access. Output from ``stdout`` and ``stderr`` is captured and returned.
 
 The application now injects a system prompt that instructs the model to chain
 multiple tools when required. This prompt ensures the assistant can orchestrate
@@ -19,4 +18,4 @@ tool calls in sequence to satisfy the user's request.
 python run.py
 ```
 
-The script will ask the model to compute an arithmetic expression and print the answer. Conversations are automatically persisted to `chat.db` and are now associated with a user and session.
+The script will instruct the model to run a simple shell command and print the result. Conversations are automatically persisted to `chat.db` and are now associated with a user and session.

run.py

@@ -8,7 +8,7 @@ from src.chat import ChatSession
 async def _main() -> None:
     async with ChatSession(user="demo_user", session="demo_session") as chat:
         answer = await chat.chat(
-            "Run this Python code: import math\nresult = math.factorial(5)"
+            "Execute this command: echo Hello from the VM"
         )
         print("\n>>>", answer)
 

src/__init__.py

@@ -1,4 +1,4 @@
 from .chat import ChatSession
-from .tools import execute_python
+from .tools import execute_terminal
 
-__all__ = ["ChatSession", "execute_python"]
+__all__ = ["ChatSession", "execute_terminal"]

src/chat.py

@@ -15,7 +15,7 @@ from .config import (
 from .db import Conversation, Message as DBMessage, User, _db, init_db
 from .log import get_logger
 from .schema import Msg
-from .tools import execute_python
+from .tools import execute_terminal
 
 _LOG = get_logger(__name__)
 
@@ -93,7 +93,7 @@ class ChatSession:
             self._model,
             messages=messages,
             think=think,
-            tools=[execute_python],
+            tools=[execute_terminal],
             options={"num_ctx": NUM_CTX},
         )
 
@@ -108,8 +108,8 @@ class ChatSession:
             return response
 
         for call in response.message.tool_calls:
-            if call.function.name == "execute_python":
-                result = execute_python(**call.function.arguments)
+            if call.function.name == "execute_terminal":
+                result = execute_terminal(**call.function.arguments)
             else:
                 continue
 

src/tools.py

@@ -1,61 +1,32 @@
 from __future__ import annotations
 
-__all__ = ["execute_python"]
+__all__ = ["execute_terminal"]
 
+import subprocess
+from typing import Final
 
-def execute_python(code: str) -> str:
-    """Execute Python code in a sandbox with a broader set of built-ins.
 
-    The code is executed with restricted but useful built-ins and can import a
-    small whitelist of standard library modules. Results should be stored in a
-    variable named ``result`` or printed. The value of ``result`` is returned if
-    present; otherwise any standard output captured during execution is
-    returned.
-    """
-    import sys
-    from io import StringIO
-
-    allowed_modules = {"math", "random", "statistics"}
-
-    def _safe_import(name: str, globals=None, locals=None, fromlist=(), level=0):
-        if name in allowed_modules:
-            return __import__(name, globals, locals, fromlist, level)
-        raise ImportError(f"Import of '{name}' is not allowed")
-
-    allowed_builtins = {
-        "abs": abs,
-        "min": min,
-        "max": max,
-        "sum": sum,
-        "len": len,
-        "range": range,
-        "sorted": sorted,
-        "enumerate": enumerate,
-        "map": map,
-        "filter": filter,
-        "list": list,
-        "dict": dict,
-        "set": set,
-        "tuple": tuple,
-        "float": float,
-        "int": int,
-        "str": str,
-        "bool": bool,
-        "print": print,
-        "__import__": _safe_import,
-    }
+def execute_terminal(command: str, *, timeout: int = 30) -> str:
+    """Execute a shell command inside an isolated Linux VM.
 
-    safe_globals: dict[str, object] = {"__builtins__": allowed_builtins}
-    safe_locals: dict[str, object] = {}
-
-    stdout = StringIO()
-    original_stdout = sys.stdout
+    The command is executed with network access enabled. Output from both
+    ``stdout`` and ``stderr`` is captured and returned. Commands are killed if
+    they exceed ``timeout`` seconds.
+    """
     try:
-        sys.stdout = stdout
-        exec(code, safe_globals, safe_locals)
-    finally:
-        sys.stdout = original_stdout
-
-    if "result" in safe_locals:
-        return str(safe_locals["result"])
-    return stdout.getvalue().strip()
+        completed = subprocess.run(
+            command,
+            shell=True,
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+        )
+    except subprocess.TimeoutExpired as exc:
+        return f"Command timed out after {timeout}s: {exc.cmd}"
+    except Exception as exc:  # pragma: no cover - unforeseen errors
+        return f"Failed to execute command: {exc}"
+
+    output = completed.stdout
+    if completed.stderr:
+        output = f"{output}\n{completed.stderr}" if output else completed.stderr
+    return output.strip()