Allow limited imports in execute_python

README.md

@@ -1,6 +1,14 @@
 # llm-backend
 
-This project provides a simple async interface to interact with an Ollama model 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.
+This project provides a simple async interface to interact with an Ollama model
+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. Two example tools are included:
+
+* **add_two_numbers** – Adds two integers.
+* **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.
 
 ## Usage
 

run.py

@@ -7,7 +7,9 @@ from src.chat import ChatSession
 
 async def _main() -> None:
     async with ChatSession(user="demo_user", session="demo_session") as chat:
-        answer = await chat.chat("What did you just say?")
+        answer = await chat.chat(
+            "Run this Python code: import math\nresult = math.factorial(5)"
+        )
         print("\n>>>", answer)
 
 

src/__init__.py

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

src/chat.py

@@ -9,7 +9,7 @@ from .config import MAX_TOOL_CALL_DEPTH, MODEL_NAME, OLLAMA_HOST
 from .db import Conversation, Message as DBMessage, User, _db, init_db
 from .log import get_logger
 from .schema import Msg
-from .tools import add_two_numbers
+from .tools import add_two_numbers, execute_python
 
 _LOG = get_logger(__name__)
 
@@ -77,7 +77,7 @@ class ChatSession:
             self._model,
             messages=messages,
             think=think,
-            tools=[add_two_numbers],
+            tools=[add_two_numbers, execute_python],
         )
 
     async def _handle_tool_calls(
@@ -93,23 +93,28 @@ class ChatSession:
         for call in response.message.tool_calls:
             if call.function.name == "add_two_numbers":
                 result = add_two_numbers(**call.function.arguments)
-                messages.append(
-                    {
-                        "role": "tool",
-                        "name": call.function.name,
-                        "content": str(result),
-                    }
-                )
-                DBMessage.create(
-                    conversation=conversation,
-                    role="tool",
-                    content=str(result),
-                )
-                nxt = await self.ask(messages, think=True)
-                self._store_assistant_message(conversation, nxt.message)
-                return await self._handle_tool_calls(
-                    messages, nxt, conversation, depth + 1
-                )
+            elif call.function.name == "execute_python":
+                result = execute_python(**call.function.arguments)
+            else:
+                continue
+
+            messages.append(
+                {
+                    "role": "tool",
+                    "name": call.function.name,
+                    "content": str(result),
+                }
+            )
+            DBMessage.create(
+                conversation=conversation,
+                role="tool",
+                content=str(result),
+            )
+            nxt = await self.ask(messages, think=True)
+            self._store_assistant_message(conversation, nxt.message)
+            return await self._handle_tool_calls(
+                messages, nxt, conversation, depth + 1
+            )
 
         return response
 

src/tools.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-__all__ = ["add_two_numbers"]
+__all__ = ["add_two_numbers", "execute_python"]
 
 
 def add_two_numbers(a: int, b: int) -> int:  # noqa: D401
@@ -14,3 +14,61 @@ def add_two_numbers(a: int, b: int) -> int:  # noqa: D401
         int: The sum of the two numbers.
     """
     return a + b
+
+
+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,
+    }
+
+    safe_globals: dict[str, object] = {"__builtins__": allowed_builtins}
+    safe_locals: dict[str, object] = {}
+
+    stdout = StringIO()
+    original_stdout = sys.stdout
+    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()