EnCompass: Enhancing Agent Programming with Search Over Program Execution Paths

AI summary

Overview

Research area: Artificial intelligence — specifically LLM-based agent programming and inference-time scaling (search over agent execution paths).

Technical level: Advanced. The paper combines a programming-language design (a probabilistic angelic nondeterminism model, a Python compiler/decorator, a search-space API) with empirical agent benchmarks. Readers will get the most from it with prior exposure to LLM agents, tree search, and probabilistic programming.

One-sentence scope: The paper introduces the PAN programming model and its Python implementation, EnCompass, which compiles an agent workflow written in ordinary Python into a search space so that inference-time strategies (best-of-N, beam search, refinement, backtracking, self-consistency) can be swapped in by changing a few inputs rather than rewriting the agent.

What This Paper Is About

Current agent programming entangles two concerns: the core workflow logic of the agent and the inference-time strategy (e.g., tree search) used to make it reliable. This coupling is inflexible, hurts readability, and makes sophisticated strategies hard to implement. The authors propose a programming model, probabilistic angelic nondeterminism (PAN), and a Python framework, EnCompass, that separate the two: the programmer marks unreliable operations with branchpoint(), the decorator @encompass.compile compiles the function into a search space object, and search is then applied over the program's nondeterministic execution paths.

Key Contributions

1. The PAN programming model. Using angelic nondeterminism, PAN separates the inference-time algorithm (search policy) from the underlying logic of the agent (the specification of the search space). The programmer writes the program as if unreliable operations always produced good outputs, and the runtime searches for an execution path where they did.

2. EnCompass, a Python library implementing PAN. It provides (1) primitives like branchpoint() usable inside an EnCompass function, (2) a Python function decorator that compiles an EnCompass function into a search space object at run time, and (3) common search algorithms plus an interface for custom search algorithms.

3. A unifying account of inference-time strategies. The paper shows that best-of-N sampling, beam search, refinement, backtracking with memory, self-consistency, and CodeT-style group evaluation are all special cases of search over nondeterministic execution paths, and that EnCompass naturally generalizes them.

4. Three case studies demonstrating that EnCompass makes it easier to implement inference-time strategies, easier to experiment across strategies, and able to surface better-scaling strategies that would otherwise be too cumbersome to implement.

Main Findings

• EnCompass reduces code churn substantially. Table 1 compares code modifications to implement search with and without EnCompass across the three case studies. On average, EnCompass saves 3–6x of coding in terms of the number of lines/words added or changed. For example, in the code repository translation case study, the non-EnCompass implementation added +423 lines (+2735 words), changed 24 lines (-62/+186 words), added +20 new function definitions, and changed the indentation level on 189 lines, versus +75 lines (+514 words), 8 changed lines (-0/+40 words), +1 new function, and 0 indentation changes with EnCompass (core logic LoC = 597 for the base agent).

• Beam search beats simpler sampling strategies on the translation agent. On the ps0 repository, "beam (coarse), beam (fine)" achieved the best scaling, outperforming the second-best strategy "beam (coarse)" with a p-value of 0.2 and all other strategies with statistical significance (p < 0.03).

• Performance scales linearly with the logarithm of cost. Consistent with prior work on inference-time scaling, on ps0 the performance of the various inference-time strategies scaled linearly with log cost (all chi-squared p-values > 0.3).

• The best-performing strategy is the hardest to implement without EnCompass. "Beam (coarse), beam (fine)" requires breaking the entire workflow into individual LLM-calling steps, each explicitly storing and retrieving variables from a frame dictionary in the plain-Python equivalent.

• Beam search generalizes across repositories. For ps1 to ps4, "beam (coarse), beam (fine)" (file-level beam width 2, method-level beam width 3) continued to outperform global best-of-N and local best-of-N (N = 16 for both), while controlling for cost.

• Global and local best-of-N are the two limiting cases of beam search. Global best-of-N is beam search with beam width N and branching factor 1 (except the root node has branching factor N); local best-of-N is beam search with beam width 1 and branching factor N.

• Hypothesis Search case study (Case Study 2). Starting from a simple agent with two LLM calls, adding a branchpoint before each LLM call and applying multithreaded BFS out of the box reproduces a parallelized version of Hypothesis Search; BFS and global best-of-N performed equally well on a subset of the ARC benchmark. Non-EnCompass modifications: +21 lines (+120 words), 3 changed lines, +2 new functions, 10 indentation changes (core logic LoC = 11); with EnCompass: +8 lines (+27 words), 1 changed line, 0 new functions, 0 indentation changes.

• Reflexion case study (Case Study 3). On LeetCodeHard, increasing N in best-of-N or the number of search steps in best-first search scaled better than increasing the number of refinement iterations in vanilla Reflexion. Non-EnCompass modifications: +27 lines (+181 words), 6 changed lines, +2 new functions, 8 indentation changes (core logic LoC = 20); with EnCompass: +9 lines (+32 words), 3 changed lines, 0 new functions, 0 indentation changes.

• Evaluation specifics for Case Study 1. The base validation repository ps0 contains 622 lines of Java code (homework ps0, Spring 2016 MIT Software Construction class on MIT OpenCourseWare). The metric is self-validation (%), the percentage match of Python and Java outputs on automatically generated test inputs, averaged across all translated non-test methods; if any validation step failed, the match percentage was counted as 0. Ps1 to ps4 each contain between 1100 and 1900 lines of code, and all four combined contain 5756 lines of code. LLM temperature was 0.0 for the base agent and 0.5 for the EnCompass agent. Error bars show standard errors of the mean over 5 runs. Average cost per run was $20–$20.5 for ps1, $27–$30 for ps2, $36–$39 for ps3, and $13.5–$14 for ps4.

Methodology in Plain English

The authors start from the observation that agent code usually mixes what the agent does with how hard the system tries. They model an agent program as a Markov chain: certain marked locations in the code (branchpoints, or the end of the program) define program states, which pair a location with a mapping from variables to values, and the code between marked locations defines a probabilistic transition function. Executing the program normally samples one trajectory; PAN instead searches over the space of possible trajectories.

This differs from standard graph search because the search cannot enumerate all children of a node — it can only stochastically sample them. Existing graph search algorithms are adapted by specifying each node's branching factor, i.e., how many children to sample. For example, DFS with branching factor 3 samples 3 next states from the current state and recurses on each.

In the EnCompass implementation, the programmer writes a normal Python function, decorates it with @encompass.compile, marks unreliable operations with branchpoint(), and records scores with record_score(score). Calling func(...).search(algo, **search_config) with a string such as "dfs" or "beam" returns the function's return value on the best execution path the algorithm found. Custom algorithms can be written against the Checkpoint class, whose step() method samples a next program state and whose score attribute holds the recorded score.

The case studies take 3 program-in-control style agents from the literature, reimplement them in EnCompass, add branchpoints before LLM calls, and compare against equivalent plain-Python implementations that express the search graph as an explicit state machine with a frame dictionary of local variables. The paper argues the state-machine approach obscures control flow, risks KeyErrors, defeats linters and static type checkers, and requires structural changes for trivial branchpoint edits.

Why This Matters

Impact on research. EnCompass reframes inference-time scaling for agents as a search-space design problem, analogous to how probabilistic programming separates model specification from inference algorithms. This gives researchers a common substrate for comparing strategies and makes it practical to implement strategies previously avoided as too cumbersome, which the authors argue could unlock better scaling laws.

Real-world applications (drawn from the paper's case studies):

• Automated translation of an entire code repository from Java to Python, with validation of each translated method against the original.
• Parallelized hypothesis-and-test agents for abstract reasoning tasks of the kind found in the ARC benchmark.
• Code-generation agents that refine their own attempts, evaluated on hard programming problems such as LeetCodeHard.
• General "program-in-control" agents built on Python frameworks (e.g., LangChain or DSPy) that call LLMs at specific subtasks and need verifiable intermediate steps.

Industry relevance. EnCompass is a Python library that can be dropped into existing agents built with Python agent frameworks, since its use consists of adding statements like branchpoint() to existing code. The reported cost figures for the translation case study, and the reported 3–6x reduction in lines/words added or changed, speak to practical engineering overhead. The framework targets developers who want to tune reliability without rewriting their agent's control flow.

Future Directions

• Extending beyond program-in-control agents. The authors note interest in "LLM+program-in-control" hybrid agents, where an LLM writes a program-in-control style agent, and suggest exploring whether EnCompass can make it easier for the LLM to implement inference-time strategies in code it writes.

• Eliminating source code modifications entirely. The authors identify as an engineering challenge choosing where to add branchpoints, supplying sufficient and good-quality intermediate reward/verification signal, and designing a good search algorithm, and suggest EnCompass could be improved to remove the need for code changes, potentially through a flexible LLM-based search strategy.

• Search strategies beyond fixed branching factors. The paper argues it is worth exploring strategies beyond fixing the branching factor in an existing graph search algorithm, citing Case Study 3's simple strategy of repeatedly choosing the highest-scoring program state and sampling one next state.

• Wider empirical validation. The case studies are limited to a small number of repositories and a subset of the ARC benchmark; scaling the comparison of strategies across more tasks and more agent architectures remains open.

Target Audience

AI agent researchers and engineers who build LLM-based, program-in-control agents and want principled control over inference-time compute. It is also relevant to programming-language and probabilistic-programming researchers interested in angelic nondeterminism and search over program execution paths, and to practitioners who want a concrete Python tool for swapping search strategies in existing agents without rewriting their workflows.

Authors’ abstract

We introduce a new approach to agent programming, the development of LLM-based agents. Current approaches to agent programming often entangle two aspects of agent design: the core workflow logic and the inference-time strategy (e.g., tree search). We introduce "probabilistic angelic nondeterminism" ("PAN"), a programming model that disentangles these two concerns, allowing the programmer to describe the agent workflow and independently experiment with different inference-time strategies by simply changing a few inputs. We provide an implementation of PAN in Python as the EnCompass framework, which uses a Python decorator to compile agent workflow programs into a search space. We present three case studies that demonstrate how the framework lets the programmer quickly improve the reliability of an agent and easily switch between different inference-time strategies, all with little additional coding.