Omni.Session.Tree (Omni Agent v0.3.1)

Copy Markdown View Source

A branching message tree used by Omni.Session to store conversation history.

Each node holds a message and an optional parent pointer, forming a tree where linear conversations are the common case and branches represent alternate replies or edits.

Active path and cursors

An active path acts as a cursor through the tree — push/3 always appends to the head of this path, and navigate/2 moves it to a different branch.

Both push/3 and navigate/2 record cursors that track which child was most recently selected at a given parent. push/3 sets a single cursor at the new node's immediate parent. navigate/2 sets cursors for every parent-child pair along the path from root to the target, so the branch that was just navigated to is fully pinned. extend/1 follows these cursors to walk from the current head to a leaf, reconstructing a full path after navigating to a mid-tree branch point.

Example

A fresh tree is empty. Push messages to build up an active path; node ids are auto-assigned starting at 1:

tree =
  %Tree{}
  |> Tree.push(%Message{role: :user, content: "..."})
  |> Tree.push(%Message{role: :assistant, content: "..."})
  |> Tree.push(%Message{role: :user, content: "..."})
  |> Tree.push(%Message{role: :assistant, content: "..."})

Tree.messages(tree)   #=> [user, assistant, user, assistant]
Tree.size(tree)       #=> 4

Navigate back to an earlier node, then push — this creates a branch:

{:ok, tree} = Tree.navigate(tree, 3)
tree = Tree.push(tree, %Message{role: :assistant, content: "..."})

Tree.size(tree)         #=> 5
Tree.children(tree, 3)  #=> [4, 5]   # two branches from node 3

Navigate between siblings to make either branch the live conversation:

{:ok, tree} = Tree.navigate(tree, 4)   # original branch is live
{:ok, tree} = Tree.navigate(tree, 5)   # switch to the new branch

When navigate/2 lands on an interior node, extend/1 walks down to a leaf, following cursors left by previous pushes and navigations:

{:ok, tree} = Tree.navigate(tree, 2)
tree = Tree.extend(tree)               # path extended to a leaf

Enumerable

Implements Enumerable, yielding tree nodes (maps with :id, :parent_id, :message, and :usage keys) along the active path in root-to-leaf order. Enum.count/1 returns the active path length.

Summary

Types

Integer node identifier, assigned sequentially.

t()

A tree of conversation messages with an active path cursor.

A tree node wrapping a message with tree metadata.

Functions

Returns the IDs of all nodes whose parent is the given node.

Extends the active path from head to a leaf node.

Returns the message for a given ID, or nil if not found.

Returns the full tree node for a given ID, or nil if not found.

Returns the ID of the last node in the active path, or nil if empty.

Returns a flat list of all messages along the active path, in order.

Sets the active path by walking parent pointers from node_id back to root.

Creates a new tree from a keyword list or map.

Walks parent pointers from node_id to root, returns the path in root-first order.

Appends a message to the head of the active path. Pipe-safe.

Like push/3, but returns {node_id, tree} for when you need the new node's ID.

Returns IDs of all nodes with parent_id: nil.

Returns other children of the same parent, excluding the given node.

Returns the total number of nodes in the tree.

Returns the cumulative usage across all nodes in the tree.

Types

node_id()

@type node_id() :: non_neg_integer()

Integer node identifier, assigned sequentially.

t()

@type t() :: %Omni.Session.Tree{
  cursors: %{required(node_id()) => node_id()},
  nodes: %{required(node_id()) => tree_node()},
  path: [node_id()]
}

A tree of conversation messages with an active path cursor.

tree_node()

@type tree_node() :: %{
  id: node_id(),
  parent_id: node_id() | nil,
  message: Omni.Message.t(),
  usage: Omni.Usage.t() | nil
}

A tree node wrapping a message with tree metadata.

Functions

children(tree, node_id)

@spec children(t(), node_id()) :: [node_id()]

Returns the IDs of all nodes whose parent is the given node.

extend(tree)

@spec extend(t()) :: t()

Extends the active path from head to a leaf node.

At each level, follows the cursor if one exists for the current head, otherwise falls back to the last (most recent) child. Stops when reaching a node with no children.

get_message(tree, id)

@spec get_message(t(), node_id()) :: Omni.Message.t() | nil

Returns the message for a given ID, or nil if not found.

get_node(tree, id)

@spec get_node(t(), node_id()) :: tree_node() | nil

Returns the full tree node for a given ID, or nil if not found.

head(tree)

@spec head(t()) :: node_id() | nil

Returns the ID of the last node in the active path, or nil if empty.

messages(tree)

@spec messages(t()) :: [Omni.Message.t()]

Returns a flat list of all messages along the active path, in order.

new(attrs)

@spec new(Enumerable.t()) :: t()

Creates a new tree from a keyword list or map.

Accepts :nodes (list of tree nodes or a pre-built node_id => node map), :path, and :cursors. Primarily used for hydration from a persisted store.

path_to(tree, node_id)

@spec path_to(t(), node_id()) :: {:ok, [node_id()]} | {:error, :not_found}

Walks parent pointers from node_id to root, returns the path in root-first order.

Useful for UIs that need to show the full path to a specific branch point.

push(tree, message, usage \\ nil)

@spec push(t(), Omni.Message.t(), Omni.Usage.t() | nil) :: t()

Appends a message to the head of the active path. Pipe-safe.

Also sets the cursor for the parent node to point at the new node, so that extend/1 will follow this branch by default.

push_node(tree, message, usage \\ nil)

@spec push_node(t(), Omni.Message.t(), Omni.Usage.t() | nil) :: {node_id(), t()}

Like push/3, but returns {node_id, tree} for when you need the new node's ID.

Sets the cursor for the parent node, same as push/3.

roots(tree)

@spec roots(t()) :: [node_id()]

Returns IDs of all nodes with parent_id: nil.

siblings(tree, node_id)

@spec siblings(t(), node_id()) :: [node_id()]

Returns other children of the same parent, excluding the given node.

size(tree)

@spec size(t()) :: non_neg_integer()

Returns the total number of nodes in the tree.

usage(tree)

@spec usage(t()) :: Omni.Usage.t()

Returns the cumulative usage across all nodes in the tree.