Language Agent Tree Search (LATS) for Python Code Generation

A lightweight, local-first Python implementation of the Language Agent Tree Search (LATS) framework. This script integrates a Large Language Model (LLM) with Monte Carlo Tree Search (MCTS) to generate, execute, evaluate, and refine Python code iteratively to solve complex programming tasks.

🛠️ How It Works

LATS enhances standard LLM generation by framing code synthesis as a tree-search problem. Rather than relying on a single-shot generation, it explores different paths using an iterative MCTS process:

                  [ Root Node ] (Empty State)
                 /             \
       [Option 1: Code]     [Option 2: Code]
         /         \            /         \
   (Execute)    (Score)    (Execute)    (Score)

The 6-Step Search Loop

Selection (UCT): Traverses existing paths to find the most promising code candidate using the Upper Confidence bound applied to Trees (UCT).
Expansion: Generates multiple alternative code modifications or fixes (EXPANSION_N) using the LLM.

If at the root, it generates a first draft.
If deep in the tree, it acts as a debugging agent using previous execution logs and reflections.

Simulation (Execution): Executes the generated Python code in an isolated dictionary-based sandboxed environment, capturing stdout and runtime errors as observations.
Evaluation (Scoring): Employs a critic LLM with a structured JSON schema to score the execution result (0.0 to 1.0) and provide step-by-step reasoning.
Reflection: Stores the critic's feedback directly on the node to act as debug hints for future iterations.
Backpropagation: Propagates the highest reward score back up the selection path to update the value metrics of parent nodes.

⚙️ Configuration

You can easily configure the script parameters at the top of lats.py:

# --- CONFIGURATION ---
LLM_API_URL = "http://localhost:8090/v1/chat/completions"  # Endpoint of your LLM provider
MAX_ITERATIONS = 3                                        # Number of complete MCTS cycles
EXPANSION_N = 2                                           # Children nodes to generate per expansion
UCT_CONSTANT = 1.41                                       # Exploration constant for MCTS selection

Note: The LLM client expects an OpenAI-compatible API endpoint (e.g., LocalAI, vLLM, Ollama, or LM Studio).

🚀 Getting Started

Prerequisites

Python 3.8+
An OpenAI-compatible LLM server running locally (or pointing to a cloud provider)
Required Python libraries:

pip install requests rich

Running the Search

To test the script, execute it directly:

python lats.py

By default, the script runs a sample algorithm-heavy task:

Given strings S and T, find the shortest substring of S which has T as a subsequence. Return the substring or empty string if none.

📦 Key Functions

`LATS_Search(task_description)`

The main entry point. Executes the entire MCTS loop to search for the most optimal, bug-free Python code snippet matching the task_description.

`execute_python(code, supplied_input="")`

Safely runs the generated code string dynamically using Python's built-in exec(). It redirects stdout and traps standard runtime errors, returning them to the LLM agent as feedback.

`clean_code(text)`

Extracts raw Python scripts out of the LLM's markdown formatting block (python ... ) and normalizes code indentation before execution.

📊 Evaluation Output Format

The evaluation step relies on structured JSON outputs from the LLM, ensuring deterministic grading:

{
  "reward": 1.0,
  "reasoning": "The code runs successfully, passes the test case logic, and correctly outputs 'bcde' as the shortest substring."
}