Functionality for prompts based on decision tree / flow charts.
Labels: name: area: agent, name: auto:enhancement, name: area: models

### Feature request

I want to request functionality for decision tree / flow chart like prompt architecture. The idea is that there would be a Prompt Tree that starts on a specific branch then allows the LLM to select new branches as part of its toolkit. Each branch would have its own prompts meaning that the AI does not need to be given all the information up front and instead can break down its commands into bite sized chunks that it sees at each branch in the tree.

### Motivation

This would help chatbot workflows by limiting the amount of information the LLM sees at each point in time, it can collect variables through different branches of the tree that it will use later, and it would improve reliability with LLM outputs because it would be easier to implement checks. It also could eliminate the need for a scratchpad, which can become costly if abused by the LLM.

Also, this is a feature that is available in other systems such as [LLMFlows](https://github.com/stoyan-stoyanov/llmflows) and [Amazon Lex](https://aws.amazon.com/lex/). And from what I have seen it is frequently on message boards here.

### Your contribution

I have made a simple example script to show how this could work in principle. However, I do not have experience contributing to open source projects so I am not sure what formatting mistakes I may be making, nor where exactly in the object heirarchy this should belong (is this a type of Prompt? Or Agent?). I would love to learn about what is needed to incorporate this into the LangChain functionality.

In my example I make a PromptTree class which stores the state and can access the current prompt. Inside the tree are a variety of branches which point to eachother according to a dictionary. Each branch produces a tool which allows the AI to switch branches by updating the prompttree.
```python

# Import libraries
import ast
from pydantic.v1 import BaseModel, Field
from langchain.tools import Tool
from langchain.schema import HumanMessage, AIMessage, SystemMessage, FunctionMessage
from langchain.tools import format_tool_to_openai_function
from langchain.chat_models import ChatOpenAI


### Define PromptBranch ###

# Declare function name variable
SELECT_BRANCH = 'select_branch'
UPDATE_INSIGHT = 'update_insight'


# Create PromptTreeBranch class
class PromptBranch:
    """A branch in the PromptTree."""

    # Declare PromptBranch variables
    description = None  # Default description of the branch
    header = None       # Header prompt
    footer = None       # Footer prompt
    children = {}       # Dictionary of children branches with descriptions. Format={name: description (None for default)}
    initial_state = {}  # Initial state of the branch
    pass_info = {}      # Additional info to be passed to children
    insights = {}       # Dictionary of insights that the AI can update. Format={name: description}

    # Get branch ID
    @property
    def branch_id(self):
        """Get the branch ID."""
        return type(self).__name__
    
    def __init__(self, parent, **kwargs):
        """Initialize the PromptBranch."""
        self.parent = parent
        self.initialize_state(**kwargs)
        return
    
    def initialize_state(self, **kwargs):
        """Initialize the branch state."""
        # We allow kwargs to be passed in case the branch needs to be initialized with additional info
        self.state = {
            **self.initial_state,
            'insights': {x: None for x in self.insights.keys()}  # Initialize insights to None
        }
        return
    
    def __call__(self, messages):
        """Call the PromptBranch."""
        return (
            self.get_prompt(messages),
            self.get_tools(),
        )
    
    def get_pass_info(self):
        """Pass info to children."""
        return self.pass_info
    
    def get_prompt(self, messages):
        """Get the prompt."""

        # Initialze prompt
        prompt = []

        # Add preamble
        preamble = self.parent.preamble
        if preamble is not None:
            prompt.append(SystemMessage(content=preamble))

        # Add header
        header = self.get_header()
        if header is not None:
            prompt.append(SystemMessage(content=header))

        # Add messages
        prompt += messages

        # Add footer
        footer = self.get_footer()
        if footer is not None:
            prompt.append(SystemMessage(content=footer))

        # Add insights
        insights = self.get_insights()
        if insights is not None:
            prompt.append(SystemMessage(content=insights))

        # Return
        return prompt
    
    def get_header(self):
        """Get header."""
        return self.header
    
    def get_footer(self):
        """Get footer."""
        return self.footer
    
    def get_insights(self):
        """Get insights."""
        if len(self.insights) == 0:
            return None
        else:
            insights = f"Your insights so far are:"
            for name, state in self.state['insights'].items():
                insights += f"\n{name}: {state}"
            return insights

    def get_tools(self):
        """Get tools."""

        # Initialize tools
        tools = []

        # Add switch branch tool
        if len(self.children) > 0:
            tools.append(self._tool_switch_branch())

        # Add update insights tool
        if len(self.insights) > 0:
            tools.append(self._tool_update_insight())

        # Return
        return tools
    
    def _tool_switch_branch(self):
        """Create tool to select next branch."""

        # Get variables
        tool_name = SELECT_BRANCH
        children = self.children
            
        # Create tool function
        tool_func = self.switch_branch
        
        # Create tool description
        tool_description = "Select the next branch to continue the conversation. Your options are:"
        for branch_id, branch_description in children.items():
            if branch_description is None:
                branch_description = self.parent.all_branches[branch_id].description
            tool_description += f"\n{branch_id}: {branch_description}"

        # Create tool schema
        class ToolSchema(BaseModel):
            branch: str = Field(
                description="Select next branch.",
                enum=list(children.keys()),
            )

        # Create tool
        tool_obj = Tool(
            name=tool_name,
            func=tool_func,
            description=tool_description,
            args_schema=ToolSchema,
        )

        # Return
        return tool_obj
    
    def _tool_update_insight(self):
        """Create tool to update an insight."""

        # Get variables
        tool_name = UPDATE_INSIGHT
        insights = self.insights
            
        # Create tool function
        tool_func = self.update_insight
        
        # Create tool description
        tool_description = "Update an insight. You can choose to update any of the following insights:"
        for name, state in insights.items():
            tool_description += f"\n{name}: {state}"

        # Create tool schema
        class ToolSchema(BaseModel):
            insight: str = Field(
                description="Select insight to update.",
                enum=list(insights.keys()),
            )
            value: str = Field(
                description="New value of the insight.",
            )

        # Create tool
        tool_obj = Tool(
            name=tool_name,
            func=tool_func,
            description=tool_description,
            args_schema=ToolSchema,
        )

        # Return
        return tool_obj

    def switch_branch(self, branch):
        """Switch to a new branch."""
        
        # Switch parent tree branch
        self.parent.branch = self.parent.all_branches[branch](parent=self.parent, **self.get_pass_info())

        # Return function message
        message = FunctionMessage(
            name=SELECT_BRANCH,
            content=f"You have switched to the {branch} branch.",
            additional_kwargs={'internal_function': True},
        )
        return message
    
    def update_insight(self, insight, value):
        """Update an insight."""

        # Update insight
        self.state['insights'][insight] = value

        # Return function message
        message = FunctionMessage(
            name=UPDATE_INSIGHT,
            content=f"You have updated the {insight} insight to {value}.",
            additional_kwargs={'internal_function': True},
        )
        return message


### Define PromptTree ###


# Create PromptTree class
class PromptTree:
    """A decision tree for prompting the AI."""
    
    # Declare PromptTree variables
    preamble = None      # System prompt to put before each branch prompt
    first_branch = None  # Name of first branch to start the prompt tree
    all_branches = {}    # Dictionary of all branches in the tree. Format={branch_id: branch_class}

    def __init__(self):
        """Initialize the PromptTree branch state."""
        self.branch = self.all_branches[self.first_branch](parent=self)
        return

    def __call__(self, messages, **kwargs):
        """Call the PromptTree."""
        return self.branch(messages, **kwargs)
    
    def get_state(self):
        """Get the current branch state."""
        return {
            'branch_id': self.branch.branch_id,
            'branch_state': self.branch.state,
        }
    
    def load_state(self, state):
        """Load a branch from the state."""
        branch_id = state['branch_id']
        branch_state = state['branch_state']
        if branch_id not in self.all_branches:
            raise ValueError(f"Unknown branch_id: {branch_id}")
        self.branch = self.all_branches[branch_id](parent=self)
        self.branch.state = branch_state
        return


### Define TreeAgent ###


# Create TreeAgent class
class TreeAgent:
    """An AI agent based on the PromptTree class."""

    def __init__(self, tree, model):
        """Initialize the TreeAgent."""
        self.tree = tree
        self.model = model
        return
    
    def __call__(self, messages, **kwargs):
        """Call the TreeAgent."""
        return self.respond(messages, **kwargs)
    
    def get_state(self):
        """Get the current state of the TreeAgent."""
        return self.tree.get_state()
    
    def load_state(self, state):
        """Load the state of the TreeAgent."""
        self.tree.load_state(state)
        return
    
    def respond(self, messages):
        """Respond to the messages."""

        # Initialize new messages
        new_messages = []

        # Loop until no function calls
        while True:

            # Get the prompt
            prompt, tools = self.tree(messages+new_messages)

            # Get the response
            funcs = [format_tool_to_openai_function(t) for t in tools]
            response = self.model.predict_messages(prompt, functions=funcs)
            new_messages.append(response)

            # Check for function calls
            if 'function_call' in new_messages[-1].additional_kwargs:
                # Get function call
                func_call = new_messages[-1].additional_kwargs['function_call']
                func_name = func_call['name']
                func_args = ast.literal_eval(func_call['arguments'])
                func = [x.func for x in tools if x.name == func_name][0]

                # Call the function
                func_response = func(**func_args)
                new_messages.append(func_response)
                continue
            else:
                # If no function call, break
                break

        # Return
        return new_messages




####################################################################################################
####################################################################################################

### EXAMPLE ###

# Create PromptBranches
class BranchA(PromptBranch):
    header = "You love icecream, but you only like vanilla icecream."
    footer = "If you choose to respond make sure you mention icecream."
    description = "A Branch to talk about icecream."
    children = {
        'BranchB': 'If someone mentions anything fancy, be sure to switch to this branch.',
        'BranchC': None,
    }
class BranchB(PromptBranch):
    header = "You love fine wines, but only if they are over 10 years old."
    footer = "If you choose to respond make sure you mention wine."
    description = "A Branch to talk about wine."
    children = {
        'BranchA': None,
        'BranchC': None,
    }
class BranchC(PromptBranch):
    header = "You love going to the beach all the time no matter what."
    footer = "If you choose to respond make sure you mention that you love the beach."
    description = "A Branch to talk about the beach."
    children = {
        'BranchA': None,
        'BranchB': None,
    }

# Create PromptTree
class MyPromptTree(PromptTree):
    preamble = "You are an AI who is obsessed with a few things."
    first_branch = 'BranchA'
    all_branches = {
        'BranchA': BranchA,
        'BranchB': BranchB,
        'BranchC': BranchC,
    }


### CONVERSATION ###

# Initialize the AI
llm = ChatOpenAI(model="gpt-3.5-turbo-0613")
tree = MyPromptTree()
agent = TreeAgent(tree, llm)

# Create sample conversation
messages = []
while True:

    # Human input
    user_message = input("You: ")
    messages += [HumanMessage(content=user_message)]
    
    # AI response
    new_messages = agent(messages)
    for m in new_messages:
        print("AI:", m)
    messages += new_messages

```

While this may not be a perfect way to go about things, it does demonstrate that with a relatively small amount of code we can work with existing LangChain arcitecture to implement a toy model. I think that with a little bit of work this could be made into something very useful.

I would love to learn more about if/how I can help contribute to incorporate this.