diff --git a/README.md b/README.md new file mode 100644 index 0000000..1d3dc51 --- /dev/null +++ b/README.md @@ -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." +} + +``` \ No newline at end of file