Planning and Search with LLMs: Tree-of-Thought at Inference Scale

Most LLM reasoning follows a straight line: you give the model a prompt, it produces a chain of thought, and it lands on an answer. This works well enough for many tasks, but planning problems are different. Planning requires the model to consider multiple possible futures, evaluate them, and commit to the best path. A single linear generation does not naturally support backtracking or comparison across branches.

Tree-of-Thought (ToT) is an approach that structures LLM reasoning as a search process over a tree of partial solutions. Each node in the tree is a reasoning step or candidate plan. The model generates multiple continuations from each node, evaluates them, and expands the most promising ones. The result is more like a minimax search than a forward pass.

The catch is that this approach multiplies your inference calls. A tree with branching factor b and depth d requires up to b^d generation steps before you apply any pruning, plus additional calls to score each node. For any agent where planning actually matters, inference speed is not a footnote -- it is the constraint that determines whether your search budget is useful or negligible.

How Tree-of-Thought Works

The original Tree of Thoughts paper from Princeton and Google DeepMind describes three components:

Thought decomposition: break the problem into a sequence of intermediate steps, where each step is small enough to evaluate independently.
Thought generation: at each node, generate multiple candidate next steps (either by sampling the same prompt k times, or by asking the model to propose k options in a single call).
State evaluation: score each candidate node -- either with a value heuristic ("rate this partial solution from 1-10") or by asking the model to vote across candidates.

Search strategy sits on top of these components. Breadth-first search (BFS) expands all nodes at the current depth before going deeper. Depth-first search (DFS) commits to one branch and backtracks when it hits a dead end. Beam search keeps only the top-k nodes at each level.

Here is the basic structure in Python:

from typing import Callable
from dataclasses import dataclass, field

@dataclass
class ThoughtNode:
    state: str
    score: float = 0.0
    children: list = field(default_factory=list)
    depth: int = 0

def tree_of_thought_bfs(
    problem: str,
    generate_fn: Callable[[str], list[str]],
    evaluate_fn: Callable[[str, str], float],
    max_depth: int = 3,
    beam_width: int = 3,
) -> ThoughtNode:
    root = ThoughtNode(state=problem)
    beam = [root]

    for depth in range(max_depth):
        candidates = []
        for node in beam:
            thoughts = generate_fn(node.state)
            for thought in thoughts:
                child = ThoughtNode(
                    state=thought,
                    depth=depth + 1,
                )
                child.score = evaluate_fn(problem, thought)
                node.children.append(child)
                candidates.append(child)

        # Keep the top beam_width nodes for the next level
        candidates.sort(key=lambda n: n.score, reverse=True)
        beam = candidates[:beam_width]

    return max(beam, key=lambda n: n.score)

The generate_fn and evaluate_fn are both LLM calls. With beam_width=3 and max_depth=3, that is at minimum 9 generation calls and 9 scoring calls -- 18 calls total for a modest search tree. Deeper trees or higher branching factors scale this up quickly.

The Inference Math

To make ToT practical, you need to think carefully about how inference latency compounds.

If each LLM call takes 2 seconds, a tree with depth 4 and beam width 3 requires roughly:

Generation: 3 × 4 = 12 calls × 2s = 24 seconds
Evaluation: 12 calls × 2s = 24 seconds
Total (sequential): ~48 seconds

Most of these calls are independent and can run in parallel. The generation calls at each depth level can fire simultaneously once you have the parent nodes. The scoring calls for a given set of candidates can also run in parallel. With full parallelism, the wall-clock time collapses to roughly max_depth × (generation_latency + evaluation_latency) -- in this case, 4 × (2s + 2s) = 16 seconds.

But at 500ms per call (achievable with fast inference), the same tree takes 4 × (0.5s + 0.5s) = 4 seconds. That difference matters a lot in interactive applications, and it determines how much search you can afford before a user gives up waiting.

Here is a parallel version of the inner loop using asyncio:

import asyncio

async def tree_of_thought_bfs_async(
    problem: str,
    generate_fn,  # async callable
    evaluate_fn,  # async callable
    max_depth: int = 3,
    beam_width: int = 3,
) -> ThoughtNode:
    root = ThoughtNode(state=problem)
    beam = [root]

    for depth in range(max_depth):
        # Generate candidates for all beam nodes in parallel
        generation_tasks = [generate_fn(node.state) for node in beam]
        all_thoughts_per_node = await asyncio.gather(*generation_tasks)

        # Flatten and score all candidates in parallel
        candidate_pairs = [
            (node, thought)
            for node, thoughts in zip(beam, all_thoughts_per_node)
            for thought in thoughts
        ]

        score_tasks = [
            evaluate_fn(problem, thought)
            for _, thought in candidate_pairs
        ]
        scores = await asyncio.gather(*score_tasks)

        # Build candidate nodes
        candidates = []
        for (node, thought), score in zip(candidate_pairs, scores):
            child = ThoughtNode(state=thought, depth=depth + 1, score=score)
            node.children.append(child)
            candidates.append(child)

        candidates.sort(key=lambda n: n.score, reverse=True)
        beam = candidates[:beam_width]

    return max(beam, key=lambda n: n.score)

Even with full parallelism, each depth level still has a sequential dependency: you cannot generate depth-2 nodes until you have scored and selected the depth-1 beam. This is why per-call latency (time-to-first-token plus generation time) still matters, even when you parallelize aggressively within a level.

Where ToT Beats Chain-of-Thought

Tree-of-Thought is not universally better than chain-of-thought. It adds overhead and complexity, so you want to use it where the planning structure actually helps.

Combinatorial tasks with clear intermediate states. The original paper uses the Game of 24 (combine four numbers with arithmetic to reach 24) and crossword puzzle solving. These problems have well-defined intermediate states and a clear evaluation function.

Code planning before writing. Before generating code for a complex function, you can run a shallow ToT search over possible approaches: generate 3 high-level plans, score them on feasibility and simplicity, then generate the implementation from the best one. This is a 2-level tree with a branching factor of 3, totaling about 6 calls. Fast inference makes this feasible to do inline.

Multi-step reasoning with backtracking. Tasks like "write a proof sketch for X" or "design a schema that satisfies constraints A, B, and C" benefit from being able to try partial solutions and abandon dead ends. Linear chain-of-thought cannot backtrack; it just keeps generating forward even when earlier reasoning was wrong.

Agent task decomposition. When an agent receives a high-level task, it can use a ToT step to generate multiple decompositions, score each on estimated complexity and risk, and pick the one that looks most tractable before starting execution.

Practical ToT with an OpenAI-Compatible API

Here is a self-contained example using the OpenAI SDK against a fast inference endpoint. The task is planning a sequence of steps to refactor a codebase module.

import asyncio
from openai import AsyncOpenAI

client = AsyncOpenAI(
    base_url="https://api.generalcompute.com/v1",
    api_key="your-api-key",
)
MODEL = "llama-4-maverick"

async def generate_thoughts(state: str, n: int = 3) -> list[str]:
    """Generate n candidate next steps from the current plan state."""
    response = await client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a software architect. Given a partial refactoring plan, "
                    "propose the next concrete step. Be specific and actionable."
                ),
            },
            {"role": "user", "content": f"Current plan:\n{state}\n\nPropose the next step."},
        ],
        n=n,
        temperature=0.8,
        max_tokens=200,
    )
    return [choice.message.content for choice in response.choices]


async def score_thought(problem: str, state: str) -> float:
    """Score a candidate plan state from 0.0 to 1.0."""
    response = await client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Rate the quality of this partial refactoring plan from 0 to 10. "
                    "Consider correctness, safety, and completeness. "
                    "Reply with a single integer and nothing else."
                ),
            },
            {
                "role": "user",
                "content": f"Original task: {problem}\n\nPlan so far:\n{state}",
            },
        ],
        max_tokens=5,
        temperature=0,
    )
    try:
        score = int(response.choices[0].message.content.strip())
        return max(0.0, min(10.0, score)) / 10.0
    except ValueError:
        return 0.5


async def plan_with_tot(task: str, depth: int = 3, beam: int = 3) -> str:
    current_state = f"Task: {task}\n\nPlan steps:"
    beam_nodes = [current_state]

    for _ in range(depth):
        # Generate candidates in parallel
        gen_tasks = [generate_thoughts(s, n=beam) for s in beam_nodes]
        all_thoughts = await asyncio.gather(*gen_tasks)

        candidates = [
            f"{parent}\n- {thought}"
            for parent, thoughts in zip(beam_nodes, all_thoughts)
            for thought in thoughts
        ]

        # Score all candidates in parallel
        score_tasks = [score_thought(task, c) for c in candidates]
        scores = await asyncio.gather(*score_tasks)

        ranked = sorted(zip(scores, candidates), reverse=True)
        beam_nodes = [plan for _, plan in ranked[:beam]]

    return beam_nodes[0]


if __name__ == "__main__":
    task = (
        "Refactor the authentication module to use JWT tokens instead of "
        "session cookies, without breaking existing user sessions."
    )
    result = asyncio.run(plan_with_tot(task))
    print(result)

A few notes on this implementation:

Passing n=3 to the completions API generates three candidates in a single request, which is more efficient than three separate calls when the model and provider support it.
The scoring prompt asks for a single integer, which keeps output tokens low and evaluation fast.
The beam parameter controls the breadth-versus-depth trade-off. A beam of 1 is effectively greedy best-first search. A beam of 5 with depth 4 produces a wider, slower search.

Controlling the Search Budget

In production, you usually cannot give ToT an unlimited compute budget. There are a few knobs worth building into your implementation:

Time budget. Set a wall-clock deadline. If the search is not complete by a threshold, return the best node found so far. This requires running the evaluation scores as they complete rather than waiting for the full gather.

Token budget. Track total tokens across all calls in the tree. Once you hit a cap, stop expanding and return the current best leaf.

Early termination. If a node's score exceeds a threshold (say, 0.9 out of 1.0), treat it as good enough and stop expanding siblings. This is most useful when your scoring function is reliable.

Depth-adaptive beam width. Start with a wider beam at shallow depths (exploring broadly) and narrow it as you go deeper (committing to the best path). This is closer to how humans approach planning: consider many options early, then commit.

The Inference Bottleneck

The reason ToT is underused in production is not that it lacks theoretical merit -- it is that most inference APIs are too slow to make the approach feel responsive. If each call takes 2 to 3 seconds, even a modest 3-level tree with a beam of 3 takes 18 to 27 seconds total in the sequential case. Users will not wait for that.

With sub-500ms inference, the same tree completes in 3 to 5 seconds end-to-end when parallelized properly. That is a usable latency budget for a coding agent or research assistant where the user expects to wait a moment for a plan before execution begins.

The faster your inference, the deeper you can search within any given time budget, and the better your planning agent's decisions will be. A 10x speedup in per-call latency translates to roughly 10x more search depth (or equivalently, a much larger beam at the same depth) within the same wall-clock constraint.

What This Means for Agent Design

If you are building agents that need to make decisions -- not just respond to prompts -- the inference speed of your provider directly shapes how sophisticated your planning can be.

A coding agent that plans with a 2-level ToT before writing code will produce better-structured implementations than one that jumps straight to generation. A research agent that explores multiple lines of inquiry in parallel will surface better answers than one following a single chain of thought. But these agents only become practical when each LLM call is fast enough to keep the total latency within what users will tolerate.

Start simple: a 2-level tree with branching factor 3 and full parallelism requires only 6 LLM calls. With fast inference, that is a few seconds of wall clock time and already meaningfully better than greedy generation for complex planning tasks. Once you have that baseline working, you can tune depth, beam width, and evaluation prompts to fit your specific problem.

If you want to try ToT on a fast inference endpoint, General Compute's API is OpenAI-compatible and designed for the kind of parallel, high-throughput workloads that planning agents require. Check out the docs to get started with your API key.