structures

Node dataclass

Node(
    text: str,
    index: int,
    children: set[int],
    embeddings: list[float],
    metadata: dict[str, Any] | None = None,
    layer: int = 0,
)

A node in the RAPTOR tree.

Attributes:

• text (str) –

  The text of the node.

• index (int) –

  The global index of the node in the tree.

• children (set[int]) –

  The global indices of the children nodes.

• embeddings (list[float]) –

  Embeddings of the node’s text.

• metadata (dict[str, Any] | None) –

  Metadata about the node’s text.

• layer (int) –

  Tree layer the node belongs to.

num_children property

num_children: int

Returns:

• int –

  Number of children nodes.

from_text classmethod

from_text(
    index: int,
    text: str,
    embedding_model: EmbeddingModelLike,
    children_indices: set[int] | None = None,
    metadata: dict[str, str] | None = None,
    layer: int = 0,
) -> Node

Create a node from text by embedding it using the given embedding model.

Parameters:

• index (int) –

  The global index of the node in the tree.

• text (str) –

  The text of the node.

• embedding_model (EmbeddingModelLike) –

  The embedding model to use for embedding the text.

• children_indices (set[int] | None, default: None ) –

  The global indices of the children nodes.

• metadata (dict[str, str] | None, default: None ) –

  Metadata about the node’s text.

• layer (int, default: 0 ) –

  Tree layer the node belongs to.

Returns:

• Node –

  A node created from the text.

from_children classmethod

from_children(
    children: list[Node],
    embedding_model: EmbeddingModelLike,
    summarization_model: SummarizationModelLike,
    index: int,
    layer: int = 0,
) -> Node

Create a node from a list of children nodes by summarizing their texts.

The text of the children nodes is concatenated using concatenate_node_texts() and passed to the summarization model to generate a summary.

Parameters:

• children (list[Node]) –

  A list of children nodes.

• embedding_model (EmbeddingModelLike) –

  The embedding model to use for embedding the summarized text.

• summarization_model (SummarizationModelLike) –

  The summarization model to use for summarizing the text of the children nodes.

• index (int) –

  The global index of the node in the tree.

• layer (int, default: 0 ) –

  Tree layer the node belongs to.

Returns:

• Node –

  A node created from the children nodes.

Tree dataclass

Tree(
    all_nodes: dict[int, Node],
    root_nodes: dict[int, Node],
    leaf_nodes: dict[int, Node],
    num_layers: int,
    layer_to_nodes: dict[int, list[Node]],
)

A RAPTOR tree.

Attributes:

• all_nodes (dict[int, Node]) –

  All nodes in the tree, mapped to their global indices.

• root_nodes (dict[int, Node]) –

  Root nodes in the tree, mapped to their global indices.

• leaf_nodes (dict[int, Node]) –

  Leaf nodes in the tree, mapped to their global indices.

• num_layers (int) –

  Number of layers in the tree.

• layer_to_nodes (dict[int, list[Node]]) –

  Nodes in a layer, mapped to their layer index.

num_nodes property

num_nodes: int

Returns:

• int –

  The number of nodes in the tree.

top_layer property

top_layer: int

Returns:

• int –

  Index of the root layer of the tree.

tolist

tolist() -> list[Node]

Returns:

• list[Node] –

  List of all nodes in the tree.

get_node

get_node(index: int) -> Node

Fetch a node in the tree by its global index.

Parameters:

• index (int) –

  The global index of the node to fetch.

Returns:

• Node –

  The node with the given global index.

fetch_layer

fetch_layer(layer: int) -> list[Node]

Fetch all nodes in a layer of the tree.

Parameters:

• layer (int) –

  The layer index to fetch nodes from.

Returns:

• list[Node] –

  List of nodes in the given layer.

fetch_node_layer

fetch_node_layer(node_idx: int) -> int

Fetch the index of the layer a node belongs to in the tree.

Parameters:

• node_idx (int) –

  The global index of the node.

Returns:

• int –

  The index of the layer the node belongs to.

concatenate_node_texts

concatenate_node_texts(nodes: list[Node]) -> str

Concatenate the texts of a list of nodes.