106 lines
3.8 KiB
Markdown
106 lines
3.8 KiB
Markdown
# 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
|
|
|
|
1. **Selection (`UCT`):** Traverses existing paths to find the most promising code candidate using the Upper Confidence bound applied to Trees (UCT).
|
|
2. **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.
|
|
|
|
|
|
3. **Simulation (Execution):** Executes the generated Python code in an isolated dictionary-based sandboxed environment, capturing stdout and runtime errors as observations.
|
|
4. **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.
|
|
5. **Reflection:** Stores the critic's feedback directly on the node to act as debug hints for future iterations.
|
|
6. **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`:
|
|
|
|
```python
|
|
# --- 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:
|
|
|
|
```bash
|
|
pip install requests rich
|
|
|
|
```
|
|
|
|
### Running the Search
|
|
|
|
To test the script, execute it directly:
|
|
|
|
```bash
|
|
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:
|
|
|
|
```json
|
|
{
|
|
"reward": 1.0,
|
|
"reasoning": "The code runs successfully, passes the test case logic, and correctly outputs 'bcde' as the shortest substring."
|
|
}
|
|
|
|
``` |