The Reasoning Optimizer Agent
Relevant source files
The following files were used as context for generating this wiki page:- docetl/reasoning_optimizer/AGENT.md
- docetl/moar/Node.py
- docetl/moar/search_utils.py
- docetl/reasoning_optimizer/directives/README.md
- experiments/reasoning/RUNMOAR.md
- docetl/reasoning_optimizer/directives/agent_utils.py
- docetl/reasoning_optimizer/directives/cascade_filtering.py
- docetl/reasoning_optimizer/directives/swap_with_code.py
- tests/reasoning_optimizer/test_directive_apply.py
The Reasoning Optimizer is an autonomous agent that uses a graph-search algorithm, specifically Monte Carlo Tree Search (MCTS), to rewrite and optimize DocETL pipelines. It intelligently applies a system of LLM-based "directives" to transform pipeline operations, aiming to find Pareto-optimal configurations that balance cost and accuracy.
This agent automates the complex task of prompt engineering and pipeline structuring by exploring a vast space of possible pipeline configurations. It iteratively modifies a pipeline, executes the new version to measure its performance, and uses the results to guide its next decision.
Sources: docetl/reasoning_optimizer/AGENT.md:5-6, experiments/reasoning/RUNMOAR.md:30-31
Core Architecture
The Reasoning Optimizer is built around three main components: the MOARSearch algorithm, the Node data structure, and the Directive system. These components work together to explore and evaluate different pipeline configurations.
The diagram below illustrates the high-level architecture of the Reasoning Optimizer Agent.
This diagram shows the Agent using the MCTS loop to traverse the search tree of Nodes, applying Directives to create new pipeline configurations.
- MOARSearch Algorithm: The core engine that implements a multi-objective MCTS to search for optimal pipelines. It manages the search process, including selection, expansion, simulation, and backpropagation.
- Node: Represents a single state in the search tree, corresponding to a specific pipeline configuration (defined by a YAML file). Each node stores its visit count, value, cost, and relationships to parent and child nodes.
- Directives: These are transformation rules that the agent can apply to a pipeline's operations. Each directive modifies the pipeline in a specific way, such as changing an LLM model, splitting an operation into multiple steps, or adding a pre-processing step.
Sources: docetl/moar/Node.py:10-19, docetl/reasoning_optimizer/directives/README.md:9-12, experiments/reasoning/RUNMOAR.md:30-31
The MOARSearch Algorithm
MOARSearch (Multi-Objective Agentic Rewriting) is the MCTS implementation that drives the optimization process. It iteratively builds a search tree where each node is a pipeline configuration. The goal is to find configurations on the Pareto frontier for cost and accuracy. The search follows four main steps in a loop: Selection, Expansion, Simulation, and Backpropagation.
Sources: experiments/reasoning/RUNMOAR.md:30-31
The following flowchart shows the main loop of the MOARSearch algorithm.
This diagram illustrates the four stages of the MCTS loop: Selection, Expansion, Simulation, and Backpropagation.
1. Selection
The search starts at the root node and traverses down the tree by repeatedly selecting the "best" child. The selection is guided by the Upper Confidence Bound (UCB1) formula, which balances exploitation (choosing nodes with high observed values) and exploration (choosing less-visited nodes). The best_child method in the Node class, using the calculate_ucb1 utility, implements this logic.
Sources: docetl/moar/Node.py:126-148, docetl/moar/search_utils.py:124-134
2. Expansion
When a leaf node is selected, the agent expands the tree by creating a new child node. This involves prompting an LLM to choose a (directive, operator) pair to apply to the current pipeline. The prompt is constructed by functions like create_expansion_prompt_acc or create_expansion_prompt_cost, which provide the LLM with extensive context, including:
- The current pipeline structure.
- Descriptions of all available operators and directives.
- A history of previously attempted actions and their performance.
- A summary of the exploration tree.
The LLM's choice is then used to generate a new pipeline configuration, which becomes the new child node.
Sources: docetl/moar/search_utils.py:173-356, docetl/moar/Node.py:277-290
The sequence diagram below shows the expansion process.
This diagram shows how the agent prompts the LLM to get a new action, which is then used to create a new child node in the search tree.
3. Simulation
Once a new node is created, the agent simulates its performance. This involves executing the new pipeline configuration against a sample of the dataset using the DSLRunner. The execute_plan method in the Node class handles this step. The simulation yields the total execution cost and the output results, which are then used for evaluation.
Sources: docetl/moar/Node.py:73-122
4. Backpropagation
The results of the simulation (cost and a calculated accuracy/value score) are propagated back up the tree from the new node to the root. Each node in the path has its visits count incremented and its value updated. This update refines the UCB scores for future selection steps. This is handled by the update_value and update_visit methods in the Node class.
Sources: docetl/moar/Node.py:192-209
Directives System
Directives are the fundamental actions the agent can take to modify a pipeline. They are self-contained transformation rules that can be applied to one or more operations.
Sources: docetl/reasoning_optimizer/directives/README.md:9-10
Instantiate Schema
A key concept is the Instantiate Schema, a Pydantic model that defines the structured output an LLM must produce to apply a directive. This schema acts as a "configuration blueprint," ensuring the agent provides all necessary parameters in the correct format for the directive to work. For example, to apply the deterministic_doc_compression directive, the LLM must output a JSON object matching the DeterministicDocCompressionInstantiateSchema, which specifies the Python code for the new code_map operation.
Sources: docetl/reasoning_optimizer/directives/README.md:14-29, docetl/reasoning_optimizer/directives/deterministic_doc_compression.py:17-21
Directive Patterns
Directives follow several common patterns for transforming pipelines:
| Pattern | Description | Example Directives |
|---|---|---|
| Single Operation Modification | Adds or modifies configuration on an existing operation. | gleaning, change_model |
| Operation Replacement | Replaces one complex operation with multiple simpler ones. | chaining, isolating_subtasks |
| Pipeline Preprocessing | Adds a new operation at the beginning of the pipeline. | doc_summarization, deterministic_doc_compression |
| Agentic Analysis | Uses an iterative, agentic loop to analyze data before generating a configuration. | cascade_filtering, swap_with_code |
Sources: docetl/reasoning_optimizer/AGENT.md:58-70, docetl/reasoning_optimizer/directives/cascade_filtering.py, docetl/reasoning_optimizer/directives/swap_with_code.py
Agentic Directives
Some directives, like CascadeFilteringDirective and SwapWithCodeDirective, use an "agentic" runner for more complex instantiations. The AgenticDirectiveRunner class manages an iterative loop where the LLM can choose to:
read_next_docs: Analyze more sample documents to better understand the data.read_operator_doc: Read the documentation for a specific operator to understand its functionality.output_schema: Finalize its analysis and generate the required instantiate schema.
This allows the agent to gather more information before committing to a specific transformation, leading to more robust and data-aware pipeline modifications.
Sources: docetl/reasoning_optimizer/directives/agent_utils.py:13-23, 195-200, docetl/reasoning_optimizer/directives/cascade_filtering.py:138-142, docetl/reasoning_optimizer/directives/swap_with_code.py:126-130
The diagram below illustrates the agentic directive loop.
This diagram shows the iterative process where the agent can read documents multiple times before deciding to generate the final configuration schema.
Running the Optimizer
The Reasoning Optimizer is executed via command-line scripts, primarily run_moar.py. This script orchestrates the entire MCTS process.
Sources: experiments/reasoning/RUNMOAR.md:2-3
Example Command
python experiments/reasoning/run_moar.py \
--yaml_path experiments/reasoning/pipelines/cuad.yaml \
--dataset_path experiments/reasoning/data/train/cuad.json \
--experiment_name cuad_moar \
--dataset cuad
Sources: experiments/reasoning/RUNMOAR.md:37-42
Key Parameters
| Parameter | Description |
|---|---|
--yaml_path | Path to the initial user-authored YAML pipeline file. |
--dataset_path | Path to the dataset file for sample input data. |
--experiment_name | A unique name for the experiment, used for output directories. |
--dataset | The name of the dataset, used for evaluation. Required if a custom accuracy function is not provided. |
--max_iterations | The maximum number of MCTS iterations to run. Default is 100. |
--exploration_weight | The UCB exploration parameter c. Default is 1.414. |
--model | The LLM to use for directive instantiation. Default is gpt-5. |
--available_models | A space-separated list of models to consider for change_model directives. |
--accuracy_function | Path to a Python file with a custom evaluate_results function for user-defined datasets. |
Sources: experiments/reasoning/RUNMOAR.md:53-68
Conclusion
The Reasoning Optimizer Agent is a sophisticated system that automates the optimization of DocETL pipelines. By combining the strategic exploration of MCTS with the flexible, LLM-driven transformations of the Directives system, it can autonomously navigate the complex trade-offs between processing cost and output quality. This enables developers to find highly effective and efficient pipeline configurations with minimal manual intervention.