/odra28/docetl

Wiki directory

Browse AI-generated documentation for any repository.

Optimization Engine Overview

Relevant source files

The following files were used as context for generating this wiki page:

The DocETL Optimization Engine is an advanced system designed to automatically rewrite and enhance data processing pipelines for improved performance, cost-efficiency, and accuracy. It employs an agentic, LLM-driven approach using a Monte Carlo Tree Search (MCTS) algorithm, named MOAR (Monte Carlo Optimization with Agentic Rewrites), to explore a vast space of possible pipeline configurations. The engine intelligently applies a series of transformations, called "Directives," to find optimal versions of a user-defined pipeline.

This engine is integrated into the DocETL ecosystem, accessible both via the command line and through the DocWrangler UI. Features like "Improve Prompt" and "Generate Pipeline from Natural Language" leverage this underlying optimization framework to provide AI-assisted development capabilities.

Core Architecture

The optimization engine is built around three main components: the MOAR Search algorithm, a system of Rewrite Directives, and an LLM-based agent that guides the search process.

Generating architecture diagram...

This diagram shows the high-level architecture of the optimization engine, where the MOAR Search Agent uses an LLM and a set of Rewrite Directives to transform an initial pipeline into an optimized one.

MOAR Search Agent: Implements a Monte Carlo Tree Search to explore different pipeline versions. Each version is a Node in the search tree.
Rewrite Directives: These are predefined transformation rules that modify a pipeline's operations. Examples include breaking a complex operation into a chain of simpler ones (Chaining) or adding validation loops (Gleaning).
LLM-Powered Expansion: The engine uses an LLM to decide which directive to apply to a given pipeline state (a Node). It constructs a detailed prompt with the current pipeline, action history, and performance statistics to get an expert recommendation from the LLM.

Sources: docetl/reasoning_optimizer/AGENT.md, docetl/moar/search_utils.py, docetl/moar/Node.py

MOAR Search Process

The core of the optimization is the MOAR (Monte Carlo Optimization with Agentic Rewrites) search algorithm. It iteratively builds and evaluates a tree of possible pipeline configurations. Each node in the tree represents a specific version of the pipeline YAML.

Sources: docetl/moar/Node.py, docetl/moar/search_utils.py

Search Tree Node

A Node is the fundamental data structure in the search tree, representing a single state (a pipeline configuration).

Generating architecture diagram...

This diagram illustrates the key attributes and methods of the Node class, which represents a state in the MCTS search tree.

Key attributes of a Node:

yaml_file_path: Path to the YAML file representing this pipeline configuration.
visits: Number of times this node has been visited during the search.
value: The cumulative reward (e.g., accuracy score) backpropagated to this node.
cost: The execution cost of this pipeline version.
parent / children: Pointers to build the tree structure.
memo: A list of (directive, target_operator) pairs, tracking the optimization path from the root.
message_history: The accumulated history of LLM conversations that led to this node.

Sources: docetl/moar/Node.py:18-77

The MCTS Loop

The search process follows the classic MCTS cycle: Selection, Expansion, Simulation, and Backpropagation.

Generating architecture diagram...

This flowchart shows the four main stages of the Monte Carlo Tree Search loop used by the MOAR optimizer.

Selection: Starting from the root, the algorithm recursively selects the child with the highest UCB1 (Upper Confidence Bound) score until a leaf node is reached. The UCB1 formula balances exploration and exploitation.
UCB1 = (node.value / node.visits) + exploration_constant * sqrt(log(parent_visits) / node.visits)
Sources: docetl/moar/search_utils.py:201-211, docetl/moar/Node.py:161-177
Expansion: At the leaf node, the LLM agent is invoked to suggest a new rewrite directive. This creates one or more new child nodes, expanding the search tree. The LLM is provided with a detailed context to make an informed decision.
Sources: docetl/moar/search_utils.py:270-481
Simulation: The pipeline configuration corresponding to the new node is executed. This step, node.execute_plan(), runs the DocETL pipeline and measures its execution cost. The output is then evaluated against a ground truth or a validator to determine its accuracy.
Sources: docetl/moar/Node.py:80-132
Backpropagation: The result (a score based on accuracy and/or cost) of the simulation is 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.
Sources: docetl/moar/Node.py:214-231

This loop continues for a configured number of iterations or until a time limit is reached.

Sources: docetl/moar/search_utils.py:189-199

Rewrite Directives

Directives are the core mechanism for transforming pipelines. Each directive is a Python class that defines a specific rewrite rule. The engine uses these directives to modify the list of operations in a pipeline.

Sources: docetl/reasoning_optimizer/directives/README.md

Key Directive Properties

Property	Description
`name`	A unique identifier for the directive (e.g., `cascade_filtering`).
`formal_description`	A concise representation of the transformation (e.g., `Filter => (Code Filter* -> Filter(gpt-5-nano)* ->) Filter`).
`nl_description`	A natural language explanation of what the directive does.
`when_to_use`	Guidance on the specific scenarios where the directive is most effective.
`instantiate_schema_type`	The Pydantic model the LLM must output to configure the directive.
`apply()`	The method that performs the actual transformation on the pipeline's operations list.
`llm_instantiate()`	The method that prompts the LLM to generate the configuration (`InstantiateSchema`).

Sources: docetl/reasoning_optimizer/directives/cascade_filtering.py:16-39, docetl/reasoning_optimizer/directives/README.md

Example Directives

The system includes a variety of directives for different optimization goals.

Directive	Purpose	Transformation Example	Source File
`cascade_filtering`	Optimizes expensive `filter` operations by injecting a series of cheaper code-based and LLM-based pre-filters.	`filter` -> `code_filter` -> `filter(gpt-5-nano)` -> `filter(gpt-4o)`	`cascade_filtering.py`
`deterministic_doc_compression`	Reduces token cost by adding a `code_map` operation to compress documents with regex before they reach an LLM.	`map` -> `code_map` -> `map`	`deterministic_doc_compression.py`
`chaining`	Decomposes a single complex `map` operation into a sequence of simpler `map` operations.	`map` -> `map_1` -> `map_2`	`test_directive_apply.py:16-43`
`gleaning`	Adds a validation loop to an operation to iteratively improve its output quality.	`map` -> `map(with gleaning config)`	`test_directive_apply.py:46-66`
`operator_fusion`	Fuses two consecutive operations (e.g., `map` -> `map`) into a single, more efficient operation.	`map` -> `filter` -> `map` -> `code_filter`	`test_directive_apply.py:120-143`
`isolating_subtasks`	Replaces a complex `map` operation with a `parallel_map` and a final `reduce` to combine results.	`map` -> `parallel_map` -> `reduce`	`test_directive_apply.py:207-236`

LLM-Powered Expansion and Plan Generation

A key feature of the optimization engine is its use of LLMs to guide the search. When expanding the search tree, the engine doesn't just pick a random directive; it asks an LLM for an expert recommendation.

Expansion Prompts

The create_expansion_prompt_acc and create_expansion_prompt_cost functions build a detailed prompt for the LLM. This prompt provides comprehensive context to help the LLM make a high-quality decision.

The prompt includes:

Pipeline & Operator Definitions: Descriptions of all available operators (map, filter, etc.).
Rewrite Directive Definitions: Descriptions of all available rewrite directives.
Valid Choices: A list of (operator, directive) pairs that can be applied to the current pipeline.
Action Performance History: Statistics on how each directive has performed across previous runs (e.g., average cost and accuracy change).
Model Performance Reference: Statistics on the accuracy and cost of different LLM models on the specific dataset.
Exploration History: A summary of the optimization paths already explored in the current search tree, including their costs and accuracies.
Current Pipeline State: The full YAML of the current pipeline being considered.
Data Context: The input document schema, token statistics, and a sample of the input data.

Sources: docetl/moar/search_utils.py:270-481

Plan Generation

For certain complex operations like map, the engine can generate and evaluate multiple sub-plans. The MapOptimizer uses a PlanGenerator to create different strategies.

Chunk Size Plans: Generates plans with varying chunk_size for splitting large documents. It also determines if peripheral context (chunks before and after the current one) is needed.
Parallel Plans: Decomposes a task into independent subtasks that can be run in parallel using a parallel_map operation.
Chain Plans: Decomposes a task into a sequence of simpler, dependent steps.

The PlanGenerator can also recursively call the optimizer on sub-tasks, creating a hierarchical optimization process.

Sources: docetl/optimizers/map_optimizer/plan_generators.py

API and UI Integration

The optimization engine is exposed through a WebSocket API and integrated into the DocWrangler UI.

WebSocket API

The primary endpoint for running optimizations is /ws/run_pipeline/{client_id}.

Generating architecture diagram...

This sequence diagram shows the flow of an optimization job initiated from the DocWrangler UI via the WebSocket endpoint.

When a client connects and sends a pipeline configuration with "optimize": true, the server initiates the runner.optimize() process, which in turn starts the MOAR search. Progress, rationales, and the final optimized operations are streamed back to the client.

Sources: server/app/routes/pipeline.py:165-258

DocWrangler UI

The DocWrangler frontend provides several AI-powered features that utilize the optimization engine's concepts:

Natural Language Pipeline Generator: A user can describe their goal in natural language, and the UI sends a detailed prompt to an LLM to generate a complete DocETL pipeline YAML. This mirrors the agentic decision-making process of the optimizer.
Sources: website/src/components/NaturalLanguagePipelineDialog.tsx:219-322
Prompt Improvement Dialog: This feature allows a user to get AI-driven suggestions for improving an operation's prompt. It sends the current pipeline state, the selected operation, and any user feedback to an LLM, which then returns a revised prompt. This is a micro-level application of the same principles used in the main optimization engine.
Sources: website/src/components/PromptImprovementDialog.tsx:109-216

Summary

The DocETL Optimization Engine is a sophisticated system that leverages an LLM-guided Monte Carlo Tree Search to automatically discover more efficient and accurate versions of data processing pipelines. By combining a structured search algorithm with the flexible reasoning capabilities of LLMs and a modular system of rewrite directives, it can navigate a vast space of potential pipeline configurations. Its integration into the DocWrangler UI makes these powerful optimization capabilities accessible to developers, enabling both fully automated pipeline enhancement and interactive, AI-assisted prompt engineering.

Optimization Engine Overview

Relevant source files

The following files were used as context for generating this wiki page:

Core Architecture

The optimization engine is built around three main components: the MOAR Search algorithm, a system of Rewrite Directives, and an LLM-based agent that guides the search process.

Generating architecture diagram...

MOAR Search Agent: Implements a Monte Carlo Tree Search to explore different pipeline versions. Each version is a Node in the search tree.
Rewrite Directives: These are predefined transformation rules that modify a pipeline's operations. Examples include breaking a complex operation into a chain of simpler ones (Chaining) or adding validation loops (Gleaning).
LLM-Powered Expansion: The engine uses an LLM to decide which directive to apply to a given pipeline state (a Node). It constructs a detailed prompt with the current pipeline, action history, and performance statistics to get an expert recommendation from the LLM.

Sources: docetl/reasoning_optimizer/AGENT.md, docetl/moar/search_utils.py, docetl/moar/Node.py

MOAR Search Process

Sources: docetl/moar/Node.py, docetl/moar/search_utils.py

Search Tree Node

A Node is the fundamental data structure in the search tree, representing a single state (a pipeline configuration).

Generating architecture diagram...

This diagram illustrates the key attributes and methods of the Node class, which represents a state in the MCTS search tree.

Key attributes of a Node:

yaml_file_path: Path to the YAML file representing this pipeline configuration.
visits: Number of times this node has been visited during the search.
value: The cumulative reward (e.g., accuracy score) backpropagated to this node.
cost: The execution cost of this pipeline version.
parent / children: Pointers to build the tree structure.
memo: A list of (directive, target_operator) pairs, tracking the optimization path from the root.
message_history: The accumulated history of LLM conversations that led to this node.

Sources: docetl/moar/Node.py:18-77

The MCTS Loop

The search process follows the classic MCTS cycle: Selection, Expansion, Simulation, and Backpropagation.

Generating architecture diagram...

This flowchart shows the four main stages of the Monte Carlo Tree Search loop used by the MOAR optimizer.

Selection: Starting from the root, the algorithm recursively selects the child with the highest UCB1 (Upper Confidence Bound) score until a leaf node is reached. The UCB1 formula balances exploration and exploitation.
UCB1 = (node.value / node.visits) + exploration_constant * sqrt(log(parent_visits) / node.visits)
Sources: docetl/moar/search_utils.py:201-211, docetl/moar/Node.py:161-177
Expansion: At the leaf node, the LLM agent is invoked to suggest a new rewrite directive. This creates one or more new child nodes, expanding the search tree. The LLM is provided with a detailed context to make an informed decision.
Sources: docetl/moar/search_utils.py:270-481
Simulation: The pipeline configuration corresponding to the new node is executed. This step, node.execute_plan(), runs the DocETL pipeline and measures its execution cost. The output is then evaluated against a ground truth or a validator to determine its accuracy.
Sources: docetl/moar/Node.py:80-132
Backpropagation: The result (a score based on accuracy and/or cost) of the simulation is 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.
Sources: docetl/moar/Node.py:214-231

This loop continues for a configured number of iterations or until a time limit is reached.

Sources: docetl/moar/search_utils.py:189-199

Rewrite Directives

Sources: docetl/reasoning_optimizer/directives/README.md

Key Directive Properties

Property	Description
`name`	A unique identifier for the directive (e.g., `cascade_filtering`).
`formal_description`	A concise representation of the transformation (e.g., `Filter => (Code Filter* -> Filter(gpt-5-nano)* ->) Filter`).
`nl_description`	A natural language explanation of what the directive does.
`when_to_use`	Guidance on the specific scenarios where the directive is most effective.
`instantiate_schema_type`	The Pydantic model the LLM must output to configure the directive.
`apply()`	The method that performs the actual transformation on the pipeline's operations list.
`llm_instantiate()`	The method that prompts the LLM to generate the configuration (`InstantiateSchema`).

Sources: docetl/reasoning_optimizer/directives/cascade_filtering.py:16-39, docetl/reasoning_optimizer/directives/README.md

Example Directives

The system includes a variety of directives for different optimization goals.

Directive	Purpose	Transformation Example	Source File
`cascade_filtering`	Optimizes expensive `filter` operations by injecting a series of cheaper code-based and LLM-based pre-filters.	`filter` -> `code_filter` -> `filter(gpt-5-nano)` -> `filter(gpt-4o)`	`cascade_filtering.py`
`deterministic_doc_compression`	Reduces token cost by adding a `code_map` operation to compress documents with regex before they reach an LLM.	`map` -> `code_map` -> `map`	`deterministic_doc_compression.py`
`chaining`	Decomposes a single complex `map` operation into a sequence of simpler `map` operations.	`map` -> `map_1` -> `map_2`	`test_directive_apply.py:16-43`
`gleaning`	Adds a validation loop to an operation to iteratively improve its output quality.	`map` -> `map(with gleaning config)`	`test_directive_apply.py:46-66`
`operator_fusion`	Fuses two consecutive operations (e.g., `map` -> `map`) into a single, more efficient operation.	`map` -> `filter` -> `map` -> `code_filter`	`test_directive_apply.py:120-143`
`isolating_subtasks`	Replaces a complex `map` operation with a `parallel_map` and a final `reduce` to combine results.	`map` -> `parallel_map` -> `reduce`	`test_directive_apply.py:207-236`

LLM-Powered Expansion and Plan Generation

Expansion Prompts

The prompt includes:

Pipeline & Operator Definitions: Descriptions of all available operators (map, filter, etc.).
Rewrite Directive Definitions: Descriptions of all available rewrite directives.
Valid Choices: A list of (operator, directive) pairs that can be applied to the current pipeline.
Action Performance History: Statistics on how each directive has performed across previous runs (e.g., average cost and accuracy change).
Model Performance Reference: Statistics on the accuracy and cost of different LLM models on the specific dataset.
Exploration History: A summary of the optimization paths already explored in the current search tree, including their costs and accuracies.
Current Pipeline State: The full YAML of the current pipeline being considered.
Data Context: The input document schema, token statistics, and a sample of the input data.

Sources: docetl/moar/search_utils.py:270-481

Plan Generation

For certain complex operations like map, the engine can generate and evaluate multiple sub-plans. The MapOptimizer uses a PlanGenerator to create different strategies.

Chunk Size Plans: Generates plans with varying chunk_size for splitting large documents. It also determines if peripheral context (chunks before and after the current one) is needed.
Parallel Plans: Decomposes a task into independent subtasks that can be run in parallel using a parallel_map operation.
Chain Plans: Decomposes a task into a sequence of simpler, dependent steps.

The PlanGenerator can also recursively call the optimizer on sub-tasks, creating a hierarchical optimization process.

Sources: docetl/optimizers/map_optimizer/plan_generators.py

API and UI Integration

The optimization engine is exposed through a WebSocket API and integrated into the DocWrangler UI.

WebSocket API

The primary endpoint for running optimizations is /ws/run_pipeline/{client_id}.

Generating architecture diagram...

This sequence diagram shows the flow of an optimization job initiated from the DocWrangler UI via the WebSocket endpoint.

Sources: server/app/routes/pipeline.py:165-258

DocWrangler UI

The DocWrangler frontend provides several AI-powered features that utilize the optimization engine's concepts:

Natural Language Pipeline Generator: A user can describe their goal in natural language, and the UI sends a detailed prompt to an LLM to generate a complete DocETL pipeline YAML. This mirrors the agentic decision-making process of the optimizer.
Sources: website/src/components/NaturalLanguagePipelineDialog.tsx:219-322
Prompt Improvement Dialog: This feature allows a user to get AI-driven suggestions for improving an operation's prompt. It sends the current pipeline state, the selected operation, and any user feedback to an LLM, which then returns a revised prompt. This is a micro-level application of the same principles used in the main optimization engine.
Sources: website/src/components/PromptImprovementDialog.tsx:109-216