Jido.AI.Reasoning.TreeOfThoughts.Machine (Jido AI v2.2.0)

Copy Markdown View Source

Pure state machine for the Tree-of-Thoughts (ToT) reasoning pattern.

This module implements state transitions for a ToT agent without any side effects. It uses Fsmx for state machine management and returns directives that describe what external effects should be performed.

Overview

Tree-of-Thoughts extends Chain-of-Thought by generating multiple candidate thoughts at each step, evaluating them, and exploring the most promising branches. This approach is effective for problems requiring search, like puzzles, planning, and creative writing.

States

  • :idle - Initial state, waiting for a prompt
  • :generating - Generating candidate thoughts for current node
  • :evaluating - Evaluating candidate thoughts
  • :expanding - Selecting next node to expand
  • :completed - Final state, solution found
  • :error - Error state

Tree Structure

The tree is stored as a map of nodes:

%{
  "node_1" => %{
    id: "node_1",
    parent_id: nil,
    content: "Initial problem...",
    score: nil,
    children: ["node_2", "node_3"],
    depth: 0
  },
  ...
}

Usage

The machine is used by the ToT strategy:

machine = Machine.new()
{machine, directives} = Machine.update(machine, {:start, prompt, call_id}, env)

All state transitions are pure - side effects are described in directives.

Status Type Boundary

Internal (Machine struct): Status is stored as strings ("idle", "completed") due to Fsmx library requirements.

External (Strategy state, Snapshots): Status is converted to atoms (:idle, :completed) via to_map/1 before storage in agent state.

Never compare machine.status directly with atoms - use Machine.to_map/1 first.

Summary

Types

External status (atom) - used in strategy state after to_map/1 conversion

Internal machine status (string) - required by Fsmx library

t()

Functions

Returns the default system prompt for thought evaluation.

Returns the default system prompt for thought generation.

Finds the best leaf node by score.

Finds all leaf nodes.

Creates a machine from a map (e.g., from strategy state storage).

Generates a unique call ID for LLM requests.

Generates a unique node ID.

Gets all children of a node.

Gets a node by ID from the machine's node map.

Gets the path from root to a given node.

Creates a new machine in the idle state.

Parses evaluation scores from LLM response text.

Parses numbered thoughts from LLM response text.

Converts the machine state to a map suitable for strategy state storage.

Updates the machine state based on a message.

Types

directive()

@type directive() ::
  {:generate_thoughts, id :: String.t(), context :: list(),
   count :: pos_integer()}
  | {:evaluate_thoughts, id :: String.t(), thoughts :: [thought_entry()]}
  | {:call_llm_stream, id :: String.t(), context :: list()}
  | {:request_error, id :: String.t(), atom(), String.t()}

external_status()

@type external_status() ::
  :idle | :generating | :evaluating | :expanding | :completed | :error

External status (atom) - used in strategy state after to_map/1 conversion

internal_status()

@type internal_status() :: String.t()

Internal machine status (string) - required by Fsmx library

msg()

@type msg() ::
  {:start, prompt :: String.t(), call_id :: String.t()}
  | {:thoughts_generated, call_id :: String.t(), thoughts :: [String.t()]}
  | {:thoughts_evaluated, call_id :: String.t(),
     scores :: %{required(String.t()) => float()}}
  | {:llm_result, call_id :: String.t(), result :: term()}
  | {:llm_partial, call_id :: String.t(), delta :: String.t(),
     chunk_type :: atom()}

t()

@type t() :: %Jido.AI.Reasoning.TreeOfThoughts.Machine{
  beam_width: pos_integer() | nil,
  branching_factor: pos_integer(),
  convergence_window: pos_integer(),
  current_call_id: String.t() | nil,
  current_node_id: String.t() | nil,
  early_success_threshold: float(),
  frontier: [String.t()],
  max_depth: pos_integer(),
  max_duration_ms: pos_integer() | nil,
  max_nodes: pos_integer(),
  max_parse_retries: non_neg_integer(),
  min_depth: non_neg_integer(),
  min_score_improvement: float(),
  nodes: %{required(String.t()) => thought_node()},
  parse_retries: %{generation: non_neg_integer(), evaluation: non_neg_integer()},
  parser_errors: [atom()],
  parser_mode: atom() | nil,
  pending_scores: %{required(String.t()) => float()},
  pending_thoughts: [thought_entry()],
  prompt: String.t() | nil,
  recent_best_scores: [float()],
  result: term(),
  root_id: String.t() | nil,
  solution_path: [String.t()],
  started_at: integer() | nil,
  status: internal_status(),
  streaming_text: String.t(),
  termination_reason: termination_reason(),
  top_k: pos_integer(),
  traversal_strategy: traversal_strategy(),
  usage: usage()
}

termination_reason()

@type termination_reason() ::
  :success
  | :threshold
  | :max_depth
  | :max_nodes
  | :max_duration
  | :converged
  | :error
  | nil

thought_entry()

@type thought_entry() :: %{id: String.t(), content: String.t()}

thought_node()

@type thought_node() :: %{
  id: String.t(),
  parent_id: String.t() | nil,
  content: String.t(),
  score: float() | nil,
  children: [String.t()],
  depth: non_neg_integer()
}

traversal_strategy()

@type traversal_strategy() :: :bfs | :dfs | :best_first

usage()

@type usage() :: %{
  optional(:input_tokens) => non_neg_integer(),
  optional(:output_tokens) => non_neg_integer(),
  optional(:total_tokens) => non_neg_integer()
}

Functions

before_transition(struct, from, to)

@spec before_transition(t(), internal_status(), internal_status()) :: {:ok, t()}

before_transition(struct, from, to, state_field)

default_evaluation_prompt()

@spec default_evaluation_prompt() :: String.t()

Returns the default system prompt for thought evaluation.

default_generation_prompt()

@spec default_generation_prompt() :: String.t()

Returns the default system prompt for thought generation.

find_best_leaf(machine)

@spec find_best_leaf(t()) :: thought_node() | nil

Finds the best leaf node by score.

find_leaves(machine)

@spec find_leaves(t()) :: [thought_node()]

Finds all leaf nodes.

from_map(map)

@spec from_map(map()) :: t()

Creates a machine from a map (e.g., from strategy state storage).

generate_call_id()

@spec generate_call_id() :: String.t()

Generates a unique call ID for LLM requests.

generate_node_id()

@spec generate_node_id() :: String.t()

Generates a unique node ID.

get_children(machine, node_id)

@spec get_children(t(), String.t()) :: [thought_node()]

Gets all children of a node.

get_node(machine, node_id)

@spec get_node(t(), String.t()) :: thought_node() | nil

Gets a node by ID from the machine's node map.

get_path_to_node(machine, node_id)

@spec get_path_to_node(t(), String.t()) :: [thought_node()]

Gets the path from root to a given node.

new(opts \\ [])

@spec new(keyword()) :: t()

Creates a new machine in the idle state.

Options

  • :branching_factor - Number of thoughts to generate at each node (default: 3)
  • :max_depth - Maximum depth of the tree (default: 3)
  • :traversal_strategy - :bfs, :dfs, or :best_first (default: :best_first)
  • :top_k - Number of ranked candidates in final result (default: 3)
  • :min_depth - Minimum depth before early success completion (default: 2)
  • :max_nodes - Hard cap on explored nodes (default: 100)
  • :max_duration_ms - Optional wall-time cap in milliseconds
  • :beam_width - Optional frontier cap for best-first expansion
  • :early_success_threshold - Score threshold for early completion (default: 1.0)
  • :convergence_window - Number of recent best scores for convergence check (default: 2)
  • :min_score_improvement - Minimum best-score improvement required across convergence window (default: 0.02)
  • :max_parse_retries - Number of parser repair retries per phase (default: 1)

parse_scores(text, thoughts)

@spec parse_scores(String.t(), [thought_entry() | String.t()]) ::
  {%{required(String.t()) => float()}, atom()}

Parses evaluation scores from LLM response text.

parse_thoughts(text)

@spec parse_thoughts(String.t()) :: {[String.t()], atom()}

Parses numbered thoughts from LLM response text.

to_map(machine)

@spec to_map(t()) :: map()

Converts the machine state to a map suitable for strategy state storage.

update(machine, msg, env \\ %{})

@spec update(t(), msg(), map()) :: {t(), [directive()]}

Updates the machine state based on a message.

Returns the updated machine and a list of directives describing external effects to be performed.

Messages

  • {:start, prompt, call_id} - Start ToT exploration
  • {:thoughts_generated, call_id, thoughts} - Handle generated thoughts
  • {:thoughts_evaluated, call_id, scores} - Handle evaluation scores
  • {:llm_result, call_id, result} - Handle LLM response
  • {:llm_partial, call_id, delta, chunk_type} - Handle streaming chunk

Directives

  • {:generate_thoughts, id, context, count} - Request thought generation
  • {:evaluate_thoughts, id, thoughts} - Request thought evaluation
  • {:call_llm_stream, id, context} - Request LLM call