C ode T ree: Agent-guided Tree Search for Code Generation with Large Language Models (2025)

Abstract: Pretrained on massive amounts of code and text data, large language models (LLMs) have demonstrated remarkable achievements in performing code generation tasks. With additional execution-based feedback, these models can act as agents with capabilities to self-refine and improve generated code autonomously. However, on challenging coding tasks with extremely large search space, current agentic approaches still struggle with multi-stage planning, generating, and debugging. To address this problem, we propose CodeTree, a framework for LLM agents to efficiently explore the search space in different stages of the code generation process. Specifically, we adopted a unified tree structure to explicitly explore different coding strategies, generate corresponding coding solutions, and subsequently refine the solutions. In each stage, critical decision-making (ranking, termination, expanding) of the exploration process is guided by both the environmental execution-based feedback and LLM-agent-generated feedback. We comprehensively evaluated CodeTree on 7 code generation benchmarks and demonstrated the significant performance gains of CodeTree against strong baselines. Using GPT-4o as the base model, we consistently achieved top results of 95.1% on HumanEval, 98.7% on MBPP, and 43.0% on CodeContests. On the challenging SWEBench benchmark, our approach led to significant performance gains, achieving a 31.9% solving rate.

## introduction recently, we have witnessed significant impacts of large language models (llms) beyond the nlp domain such as in coding tasks @xcite @xcite . however, different from traditional nlp tasks, coding tasks require generated code to be fully executable and functionally correct i.e. containing no programmatic syntax errors and passing all possible test cases @xcite @xcite . given the extremely large search space in code, early methods propose to sample a very large number of generation outputs (for example, @xcite generated up to 1 million samples per problem) to increase the chance of generating a correct code solution. more recently, several approaches adopted a "vertical" strategy in which llms first generate one (or very few) generation output, and then iteratively refine this output multiple times, often conditioned by some forms of external feedback @xcite @xcite . while these approaches are more cost-effective by focusing only on a small subset of the search space (i.e. starting from an initial output candidate), the performances of these approaches are bounded by the local optima of the chosen search space. related to our work, several methods for nlp reasoning tasks were introduced to control and enhance the generation procedure of llms. for example, @xcite proposed to enhance llms with chains of thought and statistically select the right solutions based on majority voting. @xcite decomposed a task into smaller sub-tasks and addressed them by increasing the order of difficulty. @xcite proposed to improve llms by adopting a tree-based structure to explicitly simulate the exploration of thoughts in a tree. we are motivated by this line of research and proposed codetree, a new generation framework to effectively explore the search space of code generation tasks through a tree-based structure. an overview of codetree is given in figure 1 . we define 3 standard agents, thinker, solver, and debugger, to equip the strategy-planning, solution implementation, and solution improvement correspondingly, posing comprehensive roles needed for code generation. a codetree starts approach explore exploit execution feedback ai feedback multiagent action coderanker (inala et al., 2022) ✓ ✓ alphacode @xcite , mbr-exec @xcite , codet @xcite ✓ ✓ lever @xcite , coder-reviewer @xcite ) ✓ ✓ ✓ self-correct @xcite , ilf @xcite , self-refine @xcite ✓ ✓ codechain @xcite ) ✓ ✓ self-debug @xcite , @xcite , reflexion @xcite ✓ ✓ ✓ camel @xcite ) ✓ ✓ chatdev @xcite , metagpt @xcite , agentverse @xcite ✓ ✓ ✓ self-collaboration @xcite from the input problem as the tree root and subsequent nodes represent code solutions. at any node of the tree, one can either explore sibling nodes (other strategies from the same parent node) or its children (refinements of this node). within code-tree, agents can interact with each other through a tree expansion guided by a critic agent, searching for the optimal code solution. rather than following heuristic rules or classic tree traversal methods, we use critic agent to selfevaluate the status of tree nodes at each tree expansion step by executing the following tasks: based on the outputs of the above evaluation tasks, critic agent can take action to either refine, abort, or accept the current solution, which automatically expands or terminates the tree search. codetree is flexible and efficient, avoiding duplicated or redundant exploration of functionally similar or unfeasible solutions. we comprehensively evaluated codetree on diverse code generation benchmarks from beginner-to competition-level coding tasks. our results demonstrated the significant and consistent performance gains of codetree against strong baselines. using gpt-4o as the base models, we achieved the top results on humaneval+, mbpp+ @xcite , and codecontests @xcite respectively. on the challenging swebench benchmark, our approach led to significant performance gains. we also conducted comprehensive ablation and qualitative analysis to derive the best practices and any limitations of the current method. ## related work our work is broadly related to the research of large language models (llms) @xcite @xcite . rozière et al. (2023 @xcite @xcite . early approaches @xcite @xcite adopted a brute-force principle by independently generating a very large number of output code candidates, so one of which might be the optimal code solution. subsequently, @xcite ; @xcite ; @xcite proposed to utilize unit test outcomes to filter for more prospective generated code samples. also related to our work is the studies of selfrefinement capabilities of llms. these studies leverage the inherent ability of llm to perceive arbitrary natural language contexts, including information such as environmental feedback, to iteratively fix problematic code and improve its correctness. for example, @xcite al. (2023a) @xcite introduced a tree structure to explore sub-functions in generated code. different from prior approaches, we introduce a tree-based structure as a unified search space for llm agents to efficiently perform exploration throughout different stages of code generation. see ## method to apply llms to code generation, we can define this task as a sequence-to-sequence task. the input sequence consists of a problem description d, usually in the form of a function docstring (including expected input and output) or the textual explanation of the problem. the output is a corresponding code solution, flattened into a sequence of tokens ŵ = ( ŵ1 , ..., ŵt ) with ŵt ∈ v. generated codes are evaluated against hidden test cases to check the functional correctness @xcite @xcite . the test cases are a set of input-output pairs {(i j , o j )} = {(i j , o j ) v } ∪ {(i j , o j ) h }. visible test cases are de-noted as {(i j , o j ) v } while hidden test cases are denoted as {(i j , o j ) h }. an output code ŵ is correct when ŵ (i j ) = o j ∀j. please refer to figure 1 for an overview of our method and figure 2 for a simplified version of instruction prompts to our llm agents. ## conclusion we introduce codetree, a new framework of agentguided tree search for code generation tasks. introducing a tree-based structure as a unified search space, codetree includes a critic agent to guide the tree search and make critical decisions such as termination, expanding and scoring of tree nodes. codetree facilitates multi-agent collaboration (among thinker, solver, and debugger agents) to find the correct solution within a limited solution generation budget. in our comprehensive experiments, codetree exhibits consistent and strong performance against many baselines across diverse code generation benchmarks.