Reference Agent Agents

LATS Reasoning Tree Search

A reusable agent pattern for replacing single-path generation with scored reasoning trees that can be inspected, compared, and improved.

View implementation ↗

Problem Greedy single-path agents often commit too early and provide little evidence about why one reasoning path was better than another.

Outcome A Language Agent Tree Search loop with structured state, UCT selection, dimensional scoring, and trace-based refinement.

Implementation evidence

The solution is backed by inspectable code

https://github.com/ernanhughes/co-ai

This solves multi-path reasoning for agents: represent reasoning as a search tree, score candidate paths, and preserve traces for later refinement.

Code

from collections import defaultdict


class LATSAgent:
    def __init__(self, cfg, scorer, generator):
        self.max_depth = cfg.get("max_depth", 5)
        self.ucb_weight = cfg.get("ucb_weight", 1.41)
        self.N = defaultdict(int)
        self.W = defaultdict(float)
        self.children = {}
        self.scorer = scorer
        self.generator = generator

    def update_state(self, state, action):
        return {
            **state,
            "current": state.get("current", "") + "\n" + action,
            "trace": state.get("trace", []) + [action],
        }

    def score_hypothesis(self, hypothesis, context, metrics="lats_node"):
        dimension_scores = self.scorer.evaluate(
            hypothesis=hypothesis,
            context=context,
            metrics=metrics,
        )
        weighted_total = sum(
            score["score"] * score.get("weight", 1.0)
            for score in dimension_scores.values()
        )
        weight_sum = sum(score.get("weight", 1.0) for score in dimension_scores.values())
        return {
            "score": round(weighted_total / weight_sum, 2) if weight_sum else 0.0,
            "scores": dimension_scores,
        }

    def resolve_node(self, node):
        if isinstance(node, str):
            return {"current": node, "trace": node.splitlines()}
        return node

    def expand(self, node, context):
        node = self.resolve_node(node)
        actions = self.generator.generate_next_steps(
            state=node.get("current", ""),
            trace=node.get("trace", []),
            context=context,
        )
        self.children[id(node)] = [self.update_state(node, action) for action in actions]
        return self.children[id(node)]

    def best_child(self, node):
        candidates = self.children.get(id(node), [])
        if not candidates:
            return None

        def uct(child):
            child_id = id(child)
            if self.N[child_id] == 0:
                return float("inf")
            exploitation = self.W[child_id] / self.N[child_id]
            exploration = self.ucb_weight * (self.N[id(node)] ** 0.5 / (1 + self.N[child_id]))
            return exploitation + exploration

        return max(candidates, key=uct)

    def backpropagate(self, path, reward):
        for node in path:
            node_id = id(node)
            self.N[node_id] += 1
            self.W[node_id] += reward

Usage

Use this pattern when you need the agent to explore alternatives:

root = {"current": "Goal: compare two research strategies", "trace": []}
children = lats.expand(root, context)
for child in children:
    scored = lats.score_hypothesis({"text": child["current"]}, context)
    lats.backpropagate([root, child], scored["score"])

Configuration

Keep early runs small:

lats:
  max_depth: 5
  ucb_weight: 1.41
  scoring_dimensions:
    correctness: 1.2
    feasibility: 1.0
    insightfulness: 0.8

Source

The implementation is part of the Stephanie system in ernanhughes/co-ai.

Full explanation

For the full architecture, DSPy integration, pitfalls, and symbolic refinement loop, read: Learning to Learn: A LATS-Based Framework for Self-Aware AI Pipelines.