retriever

RetrieverLike

Bases: Protocol

A protocol that defines the interface for a tree retriever.

retrieve

retrieve(
    query: str, tree: Tree, *args, **kwargs
) -> tuple[list[Node], str]

Retrieve relevant nodes from a tree given a query.

Parameters:

• query (str) –

  The query to retrieve nodes for.

• tree (Tree) –

  The tree to retrieve nodes from.

• *args –

  Additional positional arguments.

• **kwargs –

  Additional keyword arguments.

Returns:

• list[Node] –

  The retrieved nodes.

• str –

  The concatenated text of the retrieved nodes.

TreeRetriever

TreeRetriever(config: TreeRetrieverConfig)

A tree retriever that retrieves relevant nodes from a tree given a query.

It implements the RetrieverLike protocol.

Attributes:

• config (TreeRetrieverConfig) –

  The configuration for the retriever.

Parameters:

• config (TreeRetrieverConfig) –

  The configuration for the retriever.

get_relevant_node_indices

get_relevant_node_indices(
    target_embedding: list[float],
    candidate_nodes: list[Node],
) -> NDArray[int64]

Get the relevant node indices given a target embedding and candidate nodes.

Nodes are selected in the following manner

• The cosine similarity between the target embedding and the candidate node embeddings is computed.
• The cosine similarities are sorted in descending order.
• If the selection mode is SelectionMode.TOP_K, the top k nodes are selected.
• If the selection mode is SelectionMode.THRESHOLD, the nodes with cosine similarity greater than the threshold are selected.

Parameters:

• target_embedding (list[float]) –

  The target embedding to compare against.

• candidate_nodes (list[Node]) –

  The candidate nodes to compare against.

Returns:

• NDArray[int64] –

  The relevant node indices.

get_nodes_within_context

get_nodes_within_context(
    candidate_nodes: list[Node],
) -> list[Node]

Filter candidate nodes to those that fit within the maximum token length.

Parameters:

• candidate_nodes (list[Node]) –

  The candidate nodes to filter.

Returns:

• list[Node] –

  The filtered candidate nodes.

retrieve_collapse

retrieve_collapse(
    query: str, tree: Tree
) -> tuple[list[Node], str]

Retrieve relevant nodes from a tree given a query using the collapsed tree strategy.

For more details about the collapsed tree strategy, refer to the RAPTOR paper: https://arxiv.org/pdf/2401.18059.

The retrieved nodes are selected in the following manner

• The query is embedded using the embedding model.
• The candidate nodes are all the nodes in the tree.
• Relevant nodes are retrieved using get_relevant_node_indices().
• The nodes are filtered to fit within the maximum token length using get_nodes_within_context().

Parameters:

• query (str) –

  The query to retrieve nodes for.

• tree (Tree) –

  The tree to retrieve nodes from.

Returns:

• list[Node] –

  The retrieved nodes.

• str –

  The concatenated text of the retrieved nodes.

retrieve_no_collapse

retrieve_no_collapse(
    query: str, tree: Tree, start_layer: int, end_layer: int
) -> tuple[list[Node], str]

Retrieve relevant nodes from a tree given a query using the tree traversal strategy.

For more details about the collapsed tree strategy, refer to the RAPTOR paper: https://arxiv.org/pdf/2401.18059.

The retrieved nodes are selected in the following manner

• The query is embedded using the embedding model.
• The candidate nodes are all the nodes in the tree from the start layer to the end layer.
• For each layer
  • Relevant nodes are retrieved using get_relevant_node_indices().
  • The process is repeated on the children of the relevant nodes until the end layer is reached or no more children are available.
• The nodes are filtered to fit within the maximum token length using

Parameters:

• query (str) –

  The query to retrieve nodes for.

• tree (Tree) –

  The tree to retrieve nodes from.

• start_layer (int) –

  The layer to start retrieving nodes from.

• end_layer (int) –

  The layer to stop retrieving nodes at.

Returns:

• list[Node] –

  The retrieved nodes.

• str –

  The concatenated text of the retrieved nodes.

retrieve

retrieve(
    query: str,
    tree: Tree,
    start_layer: int | None = None,
    end_layer: int | None = None,
    collapse_tree: bool = True,
) -> tuple[list[Node], str]

Retrieve relevant nodes from a tree given a query.

Parameters:

• query (str) –

  The query to retrieve nodes for.

• tree (Tree) –

  The tree to retrieve nodes from.

• start_layer (int | None, default: None ) –

  The layer to start retrieving nodes from. When None, the root layer is used.

• end_layer (int | None, default: None ) –

  The layer to stop retrieving nodes at. When None, the leaf layer is used.

• collapse_tree (bool, default: True ) –

  Whether to use the collapsed tree strategy.

Returns:

• list[Node] –

  The retrieved nodes.

• str –

  The concatenated text of the retrieved nodes