Adding README
This commit is contained in:
@@ -0,0 +1,106 @@
|
||||
# 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."
|
||||
}
|
||||
|
||||
```
|
||||
Reference in New Issue
Block a user