TreeQuest

In TreeQuest, a node is associated with a user-definable state. This state can be a simple type like str or a complex dataclass containing LLM responses and metadata.

To integrate with an LLM, your generation function should follow this pattern:

@dataclasses.dataclass
class State:
    llm_answer: str
    score: float

def generate(parent_state: State | None) -> tuple[State, float]:
    if parent_state is None:
        state = initial_generation()
    else:
        state = refine_answer(parent_state.llm_answer, parent_state.score)
    return state, state.score

TreeQuest supports multiple action types by passing a dictionary of generation functions to the step method. This allows you to represent different LLM models or different prompting strategies as distinct actions in the search tree.

import dataclasses
import treequest as tq

@dataclasses.dataclass
class State:
    llm_answer: str
    score: float

def generate(parent_state: State | None) -> tuple[State, float]:
    # ... logic to call LLM ...
    return state, state.score

algo = tq.ABMCTSM()
search_tree = algo.init_tree()
search_tree = algo.step(search_tree, {'Action Label': generate})

TreeQuest provides visualization utilities to render the search tree. You can use the tq.render function to export the tree to different formats like html or pdf.

To use HTML visualization with custom state representations (e.g., including images or complex text), provide a state_formatter function. This function should take a state object and return an HTML string representation.

Security Warning: When using format="html" with a custom state_formatter that returns raw HTML, ensure the content is trusted to avoid XSS (cross-site scripting) attacks.

import treequest as tq
from pathlib import Path

# Example custom formatter for a state containing text and images
def state_formatter_html(state) -> str:
    # ... logic to convert state to HTML string ...
    return f'<p>{state.text}</p><br/><img src="data:image/webp;base64,{img_str}" width=100% />'

# Render the tree to HTML
tq.render(
    search_tree,
    output_basename=Path("search_tree/latest"),
    format="html",
    state_formatter=state_formatter_html,
)

# Render the tree to PDF
tq.render(
    search_tree,
    output_basename=Path("search_tree/final"),
    format="pdf"
)

For algorithms like ABMCTSM, using step can be slow. Instead, use the ask_batch and tell interface to run sampling steps in parallel.

1. algo.ask_batch(search_tree, batch_size, actions): Returns a tuple of (updated_search_tree, trials). Each trial in trials is a Trial object containing action, parent_state, and trial_id.
2. Generate Results: Iterate through the trials and call your generation function using the trial's parent_state.
3. algo.tell(search_tree, trial_id, (new_state, score)): Updates the search tree with the result of a specific trial using its trial_id.

Note: tell is order-independent and idempotent. It is recommended to keep batch_size <= 5 to avoid skewing the search tree shape.

import treequest as tq

# Setup
# ... (define generate function and actions) ...

algo = tq.ABMCTSM(max_process_workers=5)
search_tree = algo.init_tree()

# Batch loop
for _ in range(num_steps):
    search_tree, trials = algo.ask_batch(search_tree, batch_size, actions)

    for trial in trials:
        result = generate_fns[trial.action](trial.parent_state)
        search_tree = algo.tell(search_tree, trial.trial_id, result)

To perform a basic search, follow these steps:

1. Define a generation function: A function that takes a parent_state (or None for the root) and returns a (new_state, score) tuple. Scores should be normalized to the [0, 1] range.
2. Initialize: Instantiate an algorithm (e.g., tq.ABMCTSA()) and call algo.init_tree() to get the initial search tree.
3. Search: Use a loop to call algo.step(search_tree, actions_dict) where actions_dict maps action names to generation functions.
4. Extract Results: Use tq.top_k(search_tree, algo, k=1) to find the best state and score.
5. Visualize: Use tq.render(search_tree, basename, format="html") to generate a visualization file.

import random
import treequest as tq
from pathlib import Path

State = str

def generate(parent_state: State | None) -> tuple[State, float]:
    if parent_state is None:
        new_state = "Initial state"
    else:
        new_state = f"State after {parent_state}"
    score = random.random()
    return new_state, score

algo = tq.ABMCTSA()
search_tree = algo.init_tree()

for _ in range(10):
    search_tree = algo.step(search_tree, {'Action A': generate})

best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
print(f"Best state: {best_state}, Score: {best_node_score}")

tq.render(search_tree, Path("ab_mcts_a_search_tree"), format="html")

A Trial object represents a single attempt to expand a node in the search tree. It is designed to be self-contained so that experiments can be resumed even if the process is interrupted between the ask and tell phases.

Each Trial tracks:

• trial_id: A unique identifier.
• node_to_expand: The ID of the node being targeted.
• action: The specific action being taken.
• parent_state: The state of the parent node.
• score: The result of the trial (populated after completion).
• trial_status: One of RUNNING, INVALID, or COMPLETE.

If an algorithm (like standard MCTS) determines that a trial is no longer valid (e.g., because a different action was prioritized), the trial_status is set to INVALID.

from treequest.trial import Trial

# A Trial is a dataclass containing:
# trial_id: TrialId
# node_to_expand: NodeId
# action: str
# score: float | None
# parent_state: StateT | None
# created_at: str
# completed_at: str | None
# trial_status: Literal["RUNNING", "INVALID", "COMPLETE"]

The TreeOfThoughtsBFSAlgo uses a priority-based selection mechanism to decide which nodes to expand next. When the trial_store queue is empty, the algorithm performs the following:

1. Identifies Leaf Nodes: It finds all nodes in the tree that have no children and are not the root.
2. Finds Deepest Level: It identifies the maximum depth among these leaf nodes.
3. Prioritizes Nodes: It uses a heap (TreeOfThoughtsBFSHeapItem) to select the best nodes at that maximum depth. The priority is determined by:
   • Depth (Primary): Deeper nodes are prioritized to ensure the algorithm progresses through levels.
   • Score (Secondary): Among nodes at the same depth, nodes with higher scores are prioritized.
4. Distributes Samples: Once the top breadth_limit nodes are selected, the algorithm distributes the size_limit across the provided actions. Each selected node will have max(1, size_limit // len(actions)) expansions queued for each action.