# 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."
}

```