elm.tree.DecisionTree

class DecisionTree(graph)[source]

Bases: object

Class to traverse a directed graph of LLM prompts. Nodes are prompts and edges are transitions between prompts based on conditions being met in the LLM response.

Class to traverse a directed graph of LLM prompts. Nodes are prompts and edges are transitions between prompts based on conditions being met in the LLM response.

Examples

Here’s a simple example to setup a decision tree graph and run with the DecisionTree class:

>>> import logging
>>> import networkx as nx
>>> from rex import init_logger
>>> from elm.base import ApiBase
>>> from elm.tree import DecisionTree
>>>
>>> init_logger('elm.tree')
>>>
>>> G = nx.DiGraph(text='hello', name='Grant',
                   api=ApiBase(model='gpt-35-turbo'))
>>>
>>> G.add_node('init', prompt='Say {text} to {name}')
>>> G.add_edge('init', 'next', condition=lambda x: 'Grant' in x)
>>> G.add_node('next', prompt='How are you?')
>>>
>>> tree = DecisionTree(G)
>>> out = tree.run()
>>>
>>> print(tree.all_messages_txt)
Parameters:

graph (nx.DiGraph) – Directed acyclic graph where nodes are LLM prompts and edges are logical transitions based on the response. Must have high-level graph attribute “api” which is an ApiBase instance. Nodes should have attribute “prompt” which can have {format} named arguments that will be filled from the high-level graph attributes. Edges can have attribute “condition” that is a callable to be executed on the LLM response text. An edge from a node without a condition acts as an “else” statement if no other edge conditions are satisfied. A single edge from node to node does not need a condition.

Methods

call_node(node0)

Call the LLM with the prompt from the input node and search the successor edges for a valid transition condition

run([node0])

Traverse the decision tree starting at the input node.

Attributes

all_messages_txt

Get a printout of the full conversation with the LLM

api

Get the ApiBase object.

graph

Get the networkx graph object

history

Get a record of the nodes traversed in the tree

messages

Get a list of the conversation messages with the LLM.

property api

Get the ApiBase object.

Returns:

ApiBase

property messages

Get a list of the conversation messages with the LLM.

Returns:

list

property all_messages_txt

Get a printout of the full conversation with the LLM

Returns:

str

property history

Get a record of the nodes traversed in the tree

Returns:

list

property graph

Get the networkx graph object

Returns:

nx.DiGraph

call_node(node0)[source]

Call the LLM with the prompt from the input node and search the successor edges for a valid transition condition

Parameters:

node0 (str) – Name of node being executed.

Returns:

out (str) – Next node or LLM response if at a leaf node.

run(node0='init')[source]

Traverse the decision tree starting at the input node.

Parameters:

node0 (str) – Name of starting node in the graph. This is typically called “init”

Returns:

out (str) – Final response from LLM at the leaf node.