Choice Chain

Summary: this basically replicates if-else statements for the LLM. For a real example of this chain in use, check out the SafeTerminal chain.

Sometimes you reach a junction where you need the LLM to make a decision, and then go down a different set of tracks depending on which decision the LLM made. This is where the ChoiceChain comes in. Unlike an Agent, which repeatedly makes decisions in a loop, the ChoiceChain merely involves a single decision made for a singular junction in time. However, because the subsequent Chain will differ depending on the decision made, the ChoiceChain does not produce a stable set of output keys like most other Chains do.

The ChoiceChain involves two main components:

• the picker Chain that picks what choice Chain we will go down next

• the choice Chains that are available to us as track to go down next

As an example, suppose you have a chain that evaluates the current state of LLM operations. We’ll mock this chain out in lieu of using a real LLM, but let’s pretend like an LLM is actually being used to evaluate server logs here.

from langchain_contrib.chains.testing import FakeChain

def check_status(inputs):
    if inputs["logs"].startswith("MAJOR OUTAGE"):
        return {"status": "Needs human intervention", "note": "MAJOR OUTAGE"}
    elif inputs["logs"].startswith("ERROR"):
        return {"status": "Retry command", "note": "Permission denied"}
    else:
        return {"status": "Proceed", "note": "Nothing"}

check_status_chain = FakeChain(
    expected_inputs=["logs"],
    expected_outputs=["status", "note"],
    inputs_to_outputs=check_status,
)

Suppose you also have chains that alternatively:

• pages a human

• runs a terminal command

• continues with the next step in the deployment process

page_human = FakeChain(expected_inputs=["note"], output={"email_to": "oncall@zamm.dev"})
rerun_terminal_command = FakeChain(expected_inputs=["note"], output={"command": "chmod +x test.py"})
continue_to_next_step = FakeChain(expected_inputs=["note"], output={"rest": "of", "deployment": "process"})

Let’s now glue all of this together using ChoiceChain:

from langchain_contrib.chains import ChoiceChain

decide_based_on_status = ChoiceChain(
    choice_picker=check_status_chain,
    choice_key="status",
    choices={
        "Needs human intervention": page_human,
        "Retry command": rerun_terminal_command,
        "Proceed": continue_to_next_step,
    },
)

If the server logs mention something scary, we run down the chain where the LLM asks for human intervention:

decide_based_on_status({"logs": "MAJOR OUTAGE AT us-east-2"}, return_only_outputs=True)

{'status': 'Needs human intervention',
 'note': 'MAJOR OUTAGE',
 'email_to': 'oncall@zamm.dev'}

If instead the server logs mention an error that the LLM thinks it can fix, we run down the chain where the LLM runs a different terminal command. Note the difference in output from the last invocation of the same decide_based_on_status chain!

decide_based_on_status({"logs": "ERROR: Permission denied"}, return_only_outputs=True)

{'status': 'Retry command',
 'note': 'Permission denied',
 'command': 'chmod +x test.py'}

Finally, if everything looks good, we run down the chain where everything else happens.

decide_based_on_status({"logs": "Nothing of note"}, return_only_outputs=True)

{'status': 'Proceed', 'note': 'Nothing', 'rest': 'of', 'deployment': 'process'}

Choice chains with different inputs

We’ve seen that choice chains can produce different outputs. However, if there’s existing chains we want to make use of, it’s also possible that the choice chains can take in different inputs.

Suppose that in the above example, we are gluing together chains that have radically different inputs:

page_human = FakeChain(expected_inputs=["emergency"], output={"email_to": "oncall@zamm.dev"})
rerun_terminal_command = FakeChain(expected_inputs=["error_type", "note"], output={"command": "chmod +x test.py"})
continue_to_next_step = FakeChain(expected_inputs=[], output={"rest": "of", "deployment": "process"})

decide_based_on_status.choices = {
    "Needs human intervention": page_human,
    "Retry command": rerun_terminal_command,
    "Proceed": continue_to_next_step,
}

Because the inputs have changed, the output from the choice picker chain can no longer be used directly. Some chains are missing arguments and others have extra ones.

try:
    decide_based_on_status({"logs": "MAJOR OUTAGE AT us-east-2"}, return_only_outputs=True)
except Exception as e:
    print(repr(e))

KeyError("Extra input keys for choice: {'note'}")

try:
    decide_based_on_status({"logs": "ERROR: Permission Denied"}, return_only_outputs=True)
except Exception as e:
    print(repr(e))

ValueError("Missing some input keys: {'error_type'}")

However, we could still glue these together by setting prep_picker_output to customize the arguments being passed into each step:

from typing import Dict

def prep_output(picker_output: Dict[str, str]) -> Dict[str, str]:
    status = picker_output["status"]
    if status == "Needs human intervention":
        args = {"emergency": picker_output["note"]}
    elif status == "Retry command":
        args = {"error_type": "Permissions", "note": picker_output["note"]}
    else:
        args = {}
    return {"status": status, **args}

decide_based_on_status.prep_picker_output = prep_output

Now we can invoke each chain with the inputs it expects.

decide_based_on_status({"logs": "MAJOR OUTAGE AT us-east-2"}, return_only_outputs=True)

{'status': 'Needs human intervention',
 'emergency': 'MAJOR OUTAGE',
 'email_to': 'oncall@zamm.dev'}

decide_based_on_status({"logs": "ERROR: Permission Denied"}, return_only_outputs=True)

{'status': 'Retry command',
 'error_type': 'Permissions',
 'note': 'Permission denied',
 'command': 'chmod +x test.py'}

Retaining outputs

If you have sub-chains that deal with multiple inputs/outputs, you can set emit_io_info to retain information about exactly what was passed in as input or output:

decide_based_on_status.emit_io_info = True
results = decide_based_on_status({"logs": "ERROR: Permission Denied"}, return_only_outputs=True)
decide_based_on_status.chosen_inputs(results)

{'error_type': 'Permissions', 'note': 'Permission denied'}

decide_based_on_status.chosen_outputs(results)

{'command': 'chmod +x test.py'}

Initializing from tools

To create a new ChoiceChain from tools, use ChoiceChain.from_tools. This will also set a color for each tool to use when logging.

from langchain.tools.python.tool import PythonREPLTool
from langchain_contrib.tools import TerminalTool
from langchain_contrib.chains.testing import FakePicker

chain = ChoiceChain.from_tools(
    FakePicker(input_key="choice"),
    [PythonREPLTool(), TerminalTool()],
    verbose=True,
)
chain({"choice": "Python REPL", "action_input": "print(2 ** 16)"})

> Entering new ChoiceChain chain...

65536


> Finished chain.

{'choice': 'Python REPL',
 'action_input': 'print(2 ** 16)',
 'action_result': '65536\n'}

The chain remembers colors for tools on subsequent runs:

chain({"choice": "Terminal", "action_input": "ls"})

> Entering new ChoiceChain chain...

chains.rst  choice.ipynb  dummy.ipynb  fake.ipynb  tool.ipynb


> Finished chain.

{'choice': 'Terminal',
 'action_input': 'ls',
 'action_result': 'chains.rst  choice.ipynb  dummy.ipynb  fake.ipynb  tool.ipynb\n'}

Prompting

To build prompts for the ChoiceChain, you can use the ChoicePromptTemplate for help in doing so.

Testing

To fake the picker chain that picks which choice to use for the ChoiceChain, you can use the Fake Picker.