How to Use Code Execution for Mathematical Reasoning: A Sandbox Approach

Language models can achieve high-precision mathematical reasoning by offloading calculations to a sandboxed Python interpreter through a "think-code-check" loop that isolates execution, enforces timeouts, and feeds results back into the reasoning chain.

The bojieli/ai-agent-book repository demonstrates how to enhance large language model (LLM) capabilities for complex mathematics by integrating safe code execution into the inference pipeline. By combining a restricted Python sandbox with structured tool calling, this approach eliminates arithmetic hallucinations while maintaining conversational flow. This guide explains how to implement code execution for mathematical reasoning using the repository's three-layer architecture.

The Three-Layer Architecture for Mathematical Code Execution

The system separates concerns into distinct components that handle execution safety, interface definition, and conversation management.

Sandbox Runner (run_python)

Located in chapter5/code-for-math/sandbox.py, the run_python function spawns an isolated subprocess to execute Python code safely. According to the source code, it pre-loads scientific libraries including math, sympy, numpy, and scipy via a preamble injected at lines 15-23. The implementation writes user code to a temporary file (lines 34-51), captures stdout and stderr streams, enforces a configurable timeout mechanism (lines 58-61), and guarantees cleanup by deleting the temporary file after execution (lines 64-67).

LLM Function Definition (RUN_PYTHON_TOOL)

In chapter5/code-for-math/demo.py (lines 15-33), the RUN_PYTHON_TOOL JSON schema defines the run_python interface for the language model. This structured description tells the model when and how to invoke code execution, requiring a code string parameter that must use print() statements to output results. When the model encounters a calculation requiring precision, it emits a tool call containing the Python snippet as an argument.

Orchestrator (solve routine)

The solve function in demo.py manages the conversation flow between the LLM and the sandbox. Lines 77-87 and 98-35 implement the logic that checks for tool calls, extracts the code parameter, forwards it to run_python, and appends the execution results back into the chat history. After the model completes its reasoning, the extract_answer function (lines 58-70) parses the final integer from the model's output.

The "Think-Code-Check" Execution Loop

This design implements an iterative reasoning pattern that maximizes accuracy for mathematical problems:

1. Think – The model analyzes the mathematical problem and drafts a natural-language solution strategy.
2. Code – When encountering numeric sub-tasks, the model generates Python code and invokes run_python through the defined tool interface.
3. Check – The sandbox returns exact computational results, which the model incorporates into subsequent reasoning steps to derive the final answer.

This loop continues until the model determines it has sufficient information to provide a definitive answer without further calculation.

Implementation Examples

Direct Sandbox Execution

To execute Python code directly without LLM orchestration:

from sandbox import run_python

result = run_python("print(2**10 + 7)")
print(result)          # → 1031

Full LLM-Driven Mathematical Solving

For complete integration with language model reasoning:

from sandbox import run_python
from openai import OpenAI

client, model = build_client_and_model()
question = "What is the sum of the first 100 prime numbers?"
answer, used_codes, full_text = solve(client, model,
                                      question,
                                      use_code=True,
                                      verbose=True)
print("Final answer:", answer)

Function Schema Definition

The tool definition provided to the model (from demo.py):

{
  "type": "function",
  "function": {
    "name": "run_python",
    "description": "Execute Python code with sympy/numpy/scipy pre‑loaded. Use print() to output results.",
    "parameters": {
      "type": "object",
      "properties": {
        "code": {
          "type": "string",
          "description": "Python source to execute, must print the result."
        }
      },
      "required": ["code"]
    }
  }
}

Safety and Configuration Features

The sandbox implementation prioritizes security and reliability through several mechanisms:

• Process Isolation: Code execution occurs in a separate subprocess, preventing infinite loops or malicious operations from affecting the main application.
• Timeout Enforcement: The run_python function in sandbox.py (lines 58-61) implements configurable execution timeouts to halt runaway computations.
• Library Preamble: Pre-loaded imports (math, sympy, numpy, scipy) provide deterministic mathematical primitives without requiring repetitive import statements in generated code.
• Automatic Cleanup: Temporary files created during execution are deleted immediately after capture (lines 64-67), preventing disk accumulation.

Summary

• Sandboxed Execution: The run_python utility in chapter5/code-for-math/sandbox.py isolates mathematical calculations in temporary subprocesses with pre-loaded scientific libraries.
• Tool Interface: The RUN_PYTHON_TOOL schema in demo.py enables structured communication between the LLM and the Python interpreter.
• Orchestration Logic: The solve routine manages the "think-code-check" loop, feeding execution results back into the model's reasoning chain.
• Safety Controls: Configurable timeouts, process isolation, and automatic cleanup protect against resource exhaustion and security risks.
• Integration Path: The system extends to multi-language support through chapter4/execution-tools/multilang_executor.py, which registers the Python runner for broader use cases.

Frequently Asked Questions

How does the sandbox prevent infinite loops during mathematical calculations?

The run_python implementation in sandbox.py enforces a configurable timeout mechanism (lines 58-61) that terminates the subprocess if execution exceeds the specified duration. This prevents runaway calculations from blocking the main application while allowing legitimate complex mathematics to complete within defined limits.

Which Python libraries are available in the code execution environment?

The sandbox automatically pre-loads math, sympy, numpy, and scipy through a preamble injected at lines 15-23 of sandbox.py. This provides the language model with deterministic access to symbolic mathematics, numerical arrays, and scientific computing functions without requiring explicit import statements in generated code.

Can this architecture handle multi-step mathematical reasoning problems?

Yes, the solve function in demo.py implements an iterative loop (lines 77-87 and 98-35) that allows the model to execute multiple code blocks sequentially. After each run_python invocation returns results, the orchestrator appends the output to the conversation history, enabling the model to refine its approach based on intermediate calculations until reaching the final answer.

How does the system extract the final numerical answer from model outputs?

The extract_answer function defined at lines 58-70 of demo.py parses the model's final response to isolate integer values. This extraction logic processes the conversational text to identify the definitive numerical result, separating it from explanatory reasoning or intermediate calculation steps.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}