Using RunnableBranch to Route Execution to Different Prompts in LangChain.js

DanielJuly 8, 2024July 9, 2024

By routing in LangChain.js, we refer to the ability to create chains where the output of a previous step determines the action from the chain.

Let’s say we have two different chains with similar steps:

// use this chain for Star Wars related questions
const starWarsChain = PromptTemplate.fromTemplate(
  `You are an expert in the Star Wars universe.
  Speak like Yoda. Always answer questions starting with "As Yoda told me:".
  Respond to the following question:
  {question}`
).pipe(model).pipe(stringParser)

// use this chain for any other question
const generalChain = PromptTemplate.fromTemplate(
  `You are a helpful assistant.
  Respond to the following question: 
  {question}`
).pipe(model).pipe(stringParser)

We want to use the starWarsChain when the user asks something related to the Star Wars universe and the generalChain otherwise.

I’ve written an article about how we can solve a similar routing situation using a custom function made with RunnableLambda, but this time we will use the RunnableBranch.

Introduction to RunnableBranch

The RunnableBranch expects the following parameters:

• A list of (condition, runnable) pairs
• A default runnable

RunnableBranch first checks if one of the conditions is met and calls the runnable from the same pair. If no provided conditions are met, it runs the default runnable.

You can think of it as a switch-case JavaScript statement.

Let’s take a look at what our RunnableBranch will look like:

const branch = RunnableBranch.from([
  // the list of (condition, runnable) pairs 
  [
    data => data.topic.toUpperCase().includes("STAR-WARS"),
    starWarsChain
  ],
  // the default runnable
  generalChain
]);

Making a Classification Chain for RunnableBranch

The data.topic.toUpperCase().includes(some_value_here) is the main condition in our RunnableBranch.

To decide if we need to call the starWarsChain or the generalChain, we will first figure out if the user asked a Star Wars-related question or not.

For this, we will need to make a classification chain that will look at the question passed in the invoke() method of the main chain. The classification chain will return STAR-WARS or OTHER as the result:

const classificationChain = PromptTemplate.fromTemplate(
`You are good at classifying a question.
Given the user question below, classify it as either being about:
- STAR-WARS
- or "OTHER".
Do not respond with more than one word.

<question>
{question}
</question>

Classification:`
).pipe(model).pipe(stringParser);

await classificationChain.invoke("What is 2 + 2?"); 
// returns OTHER

await classificationChain.invoke("Who is Darth Vader?"); 
// returns STAR-WARS

Putting It All Together

Now that we have all the pieces of the puzzle, it’s time to put it all together.

This is how the full code of our example will look:

import { ChatOpenAI } from "@langchain/openai"
import { StringOutputParser } from "@langchain/core/output_parsers"
import { PromptTemplate } from "@langchain/core/prompts"
import { RunnableBranch, RunnableSequence } from "@langchain/core/runnables"
import * as dotenv from "dotenv"

dotenv.config()
const model = new ChatOpenAI({})
const stringParser = new StringOutputParser()

const classificationChain = PromptTemplate.fromTemplate(
`You are good at classifying a question.
Given the user question below, classify it as either being about:
- STAR-WARS
- or "OTHER".
Do not respond with more than one word.

<question>
{question}
</question>

Classification:`
).pipe(model).pipe(stringParser)

const starWarsChain = PromptTemplate.fromTemplate(
`You are an expert in the Star Wars universe.
Speak like Yoda. Always answer questions starting with "As Yoda told me:".
Respond to the following question:
{question}`
).pipe(model).pipe(stringParser)

const generalChain = PromptTemplate.fromTemplate(
`You are a helpful assistant.
Respond to the following question: 
{question}`
).pipe(model).pipe(stringParser)

const branch = RunnableBranch.from([
  [
    data => data.topic.toUpperCase().includes("STAR-WARS"),
    starWarsChain
  ],
  generalChain
])

const fullChain = RunnableSequence.from([
  {
    topic: classificationChain,
    question: new RunnablePassthrough(),
  },
  branch,
])

Let’s give it a test drive. We will first want to use the starWarsChain chain:

const result1 = await fullChain.invoke({
    question: "What is the age of Chewbacca?",
})
// will use the generalChain 
console.log(result1)
// outputs - As Yoda told me: Age of Chewbacca, a mystery it is. 
// Much older than seems, he is. Wisdom of many years, in him it resides.

And now let’s test the default generalChain runnable:

const result2 = await fullChain.invoke({
    question: "What is 2 + 2?",
})
// will use the generalChain 
console.log(result2)
// outputs - 2 + 2 equals 4

RunnableBranch vs RunnableLambda

As you can see in the other article I’ve written about routing in LangChain.js, the RunnableBranch overlaps with using a RunnableLambda as a router.

Furthermore, the RunnableBranch is marked as legacy in the documentation:

I’ve asked on Reddit about this, and someone from the LangChain team was kind enough to reply. It seems that RunnableBranch will not be removed (yet), but in the future, using a custom lambda function is the way to go.

Nevertheless, I think it’s good to know about RunnableBranch in case you see it in some codebase.

And there you have it. This is how you can make a simple routing classification example with RunnableBranch in LangChain.js.

As usual, you can see the full code on my GitHub repo. Happy coding! 👋