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

Copy Markdown View Source

Pure state machine for the TRM (Tiny-Recursive-Model) reasoning pattern.

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

Overview

TRM uses recursive reasoning to iteratively improve answers through a reason-supervise-improve cycle. Each iteration:

  1. Reasoning: Generate insights about the current answer
  2. Supervision: Evaluate the answer and provide feedback
  3. Improvement: Apply feedback to generate a better answer

The cycle continues until either:

  • Maximum supervision steps reached
  • ACT (Adaptive Computational Time) threshold exceeded (high confidence)
  • Convergence detected (improvements have plateaued)

States

  • :idle - Initial state, waiting for a question
  • :reasoning - Generating reasoning insights
  • :supervising - Evaluating current answer quality
  • :improving - Applying feedback to improve answer
  • :completed - Final state, best answer found
  • :error - Error state

Usage

The machine is used by the TRM strategy:

machine = Machine.new(max_supervision_steps: 5, act_threshold: 0.9)
{machine, directives} = Machine.update(machine, {:start, question, call_id}, env)

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

Telemetry Events

The machine emits the following telemetry events:

[:jido, :ai, :trm, :start]

Emitted when reasoning starts with a new question.

  • Measurements: (none)
  • Metadata: %{call_id: String.t(), question_length: non_neg_integer()}

[:jido, :ai, :trm, :step]

Emitted at each supervision step.

  • Measurements: %{step: pos_integer(), quality_score: float()}
  • Metadata: %{call_id: String.t(), issues_count: non_neg_integer(), suggestions_count: non_neg_integer()}

[:jido, :ai, :trm, :act_triggered]

Emitted when ACT early stopping triggers.

  • Measurements: %{confidence: float(), threshold: float()}
  • Metadata: %{call_id: String.t(), step: pos_integer()}

[:jido, :ai, :trm, :error]

Emitted when an error occurs.

  • Measurements: (none)
  • Metadata: %{call_id: String.t(), error: term(), usage: map()}

[:jido, :ai, :trm, :complete]

Emitted when reasoning completes successfully.

  • Measurements: %{steps: non_neg_integer(), best_score: float(), duration_ms: non_neg_integer()}
  • Metadata: %{call_id: String.t(), termination_reason: atom(), usage: map()}

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

act_state()
directive()
external_status()

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

internal_status()

Internal machine status (string) - required by Fsmx library

latent_state()
msg()
t()
termination_reason()
usage()

Functions

before_transition(struct, from, to, state_field)
check_act_condition(machine)

Checks the ACT condition for early stopping.

extract_confidence(latent_state)

Extracts the current confidence score from latent state.

from_map(map)

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

generate_call_id()

Generates a unique call ID for LLM requests.

initialize_latent_state(question, answer)

Initializes the latent state from a question and optional initial answer.

merge_reasoning_trace(trace, new_entry)

Merges a new entry into the reasoning trace, keeping recent history.

new()

Creates a new machine in the idle state with default configuration.

new(opts)

Creates a new machine in the idle state with custom configuration.

should_terminate_max_steps?(machine)

Checks if the machine should terminate due to max steps.

to_map(machine)

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

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

Updates the machine state based on a message.

update_latent_state(latent_state, phase, content)

Updates the latent state with new reasoning insights.

Types

act_state()

@type act_state() :: %{
  threshold: float(),
  current_confidence: float(),
  history: [float()]
}

directive()

@type directive() ::
  {:reason, id :: String.t(), context :: map()}
  | {:supervise, id :: String.t(), context :: map()}
  | {:improve, id :: String.t(), context :: map()}

external_status()

@type external_status() ::
  :idle | :reasoning | :supervising | :improving | :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

latent_state()

@type latent_state() :: %{
  question_context: String.t() | nil,
  answer_context: String.t() | nil,
  reasoning_trace: [String.t()],
  confidence_score: float(),
  step_count: non_neg_integer()
}

msg()

@type msg() ::
  {:start, question :: String.t(), call_id :: String.t()}
  | {:reasoning_result, call_id :: String.t(), result :: term()}
  | {:supervision_result, call_id :: String.t(), result :: term()}
  | {:improvement_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.TRM.Machine{
  act_state: act_state(),
  act_threshold: float(),
  act_triggered: boolean(),
  answer_history: [String.t()],
  best_answer: String.t() | nil,
  best_score: float(),
  current_answer: String.t() | nil,
  current_call_id: String.t() | nil,
  latent_state: latent_state(),
  max_supervision_steps: pos_integer(),
  parsed_feedback: map() | nil,
  question: String.t() | nil,
  result: term(),
  started_at: integer() | nil,
  status: internal_status(),
  streaming_text: String.t(),
  supervision_feedback: String.t() | nil,
  supervision_step: non_neg_integer(),
  termination_reason: termination_reason(),
  usage: usage()
}

termination_reason()

@type termination_reason() ::
  :max_steps | :act_threshold | :convergence_detected | :error | nil

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, state_field)

check_act_condition(machine)

@spec check_act_condition(t()) :: boolean()

Checks the ACT condition for early stopping.

extract_confidence(latent_state)

@spec extract_confidence(latent_state()) :: float()

Extracts the current confidence score from latent state.

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.

initialize_latent_state(question, answer)

@spec initialize_latent_state(String.t(), String.t() | nil) :: latent_state()

Initializes the latent state from a question and optional initial answer.

merge_reasoning_trace(trace, new_entry)

@spec merge_reasoning_trace([String.t()], String.t()) :: [String.t()]

Merges a new entry into the reasoning trace, keeping recent history.

new()

@spec new() :: t()

Creates a new machine in the idle state with default configuration.

new(opts)

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

Creates a new machine in the idle state with custom configuration.

Options

  • :max_supervision_steps - Maximum iterations before termination (default: 5)
  • :act_threshold - Confidence threshold for early stopping (default: 0.9)

should_terminate_max_steps?(machine)

@spec should_terminate_max_steps?(t()) :: boolean()

Checks if the machine should terminate due to max steps.

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, question, call_id} - Start TRM reasoning
  • {:reasoning_result, call_id, result} - Handle reasoning response
  • {:supervision_result, call_id, result} - Handle supervision feedback
  • {:improvement_result, call_id, result} - Handle improved answer
  • {:llm_partial, call_id, delta, chunk_type} - Handle streaming chunk

Directives

  • {:reason, id, context} - Request reasoning LLM call
  • {:supervise, id, context} - Request supervision LLM call
  • {:improve, id, context} - Request improvement LLM call

update_latent_state(latent_state, phase, content)

@spec update_latent_state(latent_state(), atom(), String.t()) :: latent_state()

Updates the latent state with new reasoning insights.