How do I use a RouterChain with multiple chains that have different inputs?

[kells1986]

This is my sample code:

import langchain
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # read local .env file

from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.chains.router import MultiPromptChain
from langchain.chains import ConversationChain
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate


llm = ChatOpenAI(temperature=0.9, model="gpt-3.5-turbo")

refuse_template = """Politely refuse to answer the question because it's not related to soccer or football"""
refuse_prompt = PromptTemplate(template=refuse_template, input_variables=[])
refuse_chain = LLMChain(prompt=refuse_prompt, llm=llm)

chat_template = [
    ("system", "you are an expert in football (what they call soccer in the US"),
    ("system", "you are English and therefore use english terminology when talking about football"),
    ("human", "answer this {input}")
]
chat_prompt = ChatPromptTemplate.from_messages(chat_template)
chat_chain = LLMChain(prompt=chat_prompt, llm=llm)

team_template = [
    ("system", "You are an expert in English football and can answer any question about any English team"),
    ("human", "answer this {input} about this {team}")
]
team_prompt = ChatPromptTemplate.from_messages(team_template)
team_chain = LLMChain(prompt=team_prompt, llm=llm)



destination_chains = {}
destination_chains["football"] = chat_chain
destination_chains["football"] = {"prompt":chat_prompt, "llm":llm}
destination_chains["refuse"] = refuse_chain
destination_chains["refuse"] = {"prompt": refuse_prompt, "llm": llm}
destination_chains["team"] = team_chain
destination_chains["team"] = {"prompt": team_prompt, "llm": llm}

descriptions = {
    "football": "Answers association football questions",
    "refuse": "Refuse to answer the question because it's not related to football / soccer",
    "team": "Answer a question specifically about an English football team. Extract the term 'team' as an output for this chain"
}

default_chain = ConversationChain(llm=llm, output_key="text")

from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser

destinations = [f"{k}: {descriptions[k]}" for k in descriptions]
destinations_str = "\n".join(destinations)

print(destinations_str)

router_template = """
Given a raw text input to a language model select the model prompt best suited for the given input.
Extract relevant information and additional inputs for the next chain if necessary.  

<< FORMATTING >>
Return a markdown code snippet with a JSON object formatted to look like:
``json
{{
    "destination": string \ name of the prompt to use or "DEFAULT"
    "next_inputs": string \ String representation of a JSON containing the original input. Additional relevant keys extracted from the input should be included.
}}
``

REMEMBER: "destination" MUST be one of the candidate prompt names specified below OR it can be "DEFAULT" if the input is not well suited for any of the candidate prompts.
REMEMBER: "next_inputs" is a string. It is just the original input, and may contain "team" if a team is found in the text.

<< CANDIDATE PROMPTS >>
{destinations}

<< INPUT >>
{input}

"""

print(router_template)

router_prompt = PromptTemplate(
    template=router_template,
    input_variables=[
        "destinations",
        "input"
    ],
    output_parser=RouterOutputParser(),
)

router_chain = LLMRouterChain.from_llm(llm, router_prompt)


chain = MultiPromptChain(
    router_chain=router_chain,
    destination_chains=destination_chains,
    default_chain=default_chain,
    verbose=True,
)

langchain.debug = True

question = "How many league titles have Manchester United won?"
output3 = chain.run({"input":question, "destinations": destinations_str})

question = "How many players are on a football team?"
output1 = chain.run({"input":question, "destinations":destinations_str})

question = "What do you think of Boris Johnson?"
output2 = chain.run({"input":question, "destinations": destinations_str})

print("---")
print("---")
print("---")
print("---")
print(output1)
print("---")
print(output2)
print("---")
print(output3)

My problem is that sometimes I need output from the router in a different format depending on the destination chain. Is there a way to handle this?

In my case for question3 I get this error:

{
  "destination": "team",
  "next_inputs": {
    "input": "{\"text\": \"How many league titles have Manchester United won?\", \"team\": \"Manchester United\"}"
  }
}
[chain/error] [1:chain:MultiPromptChain] [2.82s] Chain run errored with error:
"ValueError(\"Missing some input keys: {'team'}\")"

Is there a more standard way to achieve my goal?

[dosubot[bot]]

🤖

Hello,

Based on the context provided, it seems like you're trying to route different inputs to different destination chains using the MultiPromptChain class in LangChain. However, you're encountering an issue where some destination chains require different input formats.

To handle this, you can create a custom chain that adapts the input variables for the destination chain. This custom chain will take the input variable from the router chain and convert it to the expected input variable for the destination chain. Here's an example of how you can create such a custom chain:

from langchain.chains.base import Chain

class InputAdapterChain(Chain):
    def __init__(self, destination_chain, input_key_map):
        self.destination_chain = destination_chain
        self.input_key_map = input_key_map

    def run(self, inputs):
        adapted_inputs = {self.input_key_map[k]: v for k, v in inputs.items()}
        return self.destination_chain.run(adapted_inputs)

Now, you can use this InputAdapterChain to wrap the destination chains and adapt the input variable:

input_key_map = {"input": "team"}
adapted_team_chain = InputAdapterChain(team_chain, input_key_map)

Finally, use the adapted chains in your MultiPromptChain setup:

destination_chains["team"] = adapted_team_chain

chain = MultiPromptChain(
    router_chain=router_chain,
    destination_chains=destination_chains,
    default_chain=default_chain,
    verbose=True,
)

This way, the InputAdapterChain will convert the input variable from the router chain to the expected input variable for the destination chain, avoiding the ValueError you encountered.

This solution was suggested in a similar issue in the LangChain repository: Issue: How to handle RouterChain when 1 or more destination chain(s) which is expecting a different input variable?

I hope this helps! Let me know if you have any other questions.

Sources

• Issue: How to handle RouterChain when 1 or more destination chain(s) which is expecting a different input variable?
• libs/langchain/langchain/chains/router/base.py
• docs/extras/expression_language/how_to/routing.ipynb
• libs/langchain/langchain/chains/router/multi_prompt.py

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[kells1986]

This didn't work for me.

I get this error:

Traceback (most recent call last):
  File "/Users/me/Code/llm-stuff/router.py", line 46, in <module>
    adapted_team_chain = InputAdapterChain(team_chain, input_key_map)
TypeError: Can't instantiate abstract class InputAdapterChain with abstract methods _call, input_keys, output_keys

I also don't think it fully addresses the error that I got.

I'd have to remap the output from:

"input": "{\"text\": \"How many league titles have Manchester United won?\", \"team\": \"Manchester United\"}"

to:

"input": "How many league titles have Manchester United won?"
"team": "Manchester United"