MultiRouteChain not working as expected
Labels: name: auto:bug, name: area: models

### System Info

python==3.9
langchain==0.0.246

### Who can help?

@hwchase17 

### Information

- [ ] The official example notebooks/scripts
- [X] My own modified scripts

### Related Components

- [X] LLMs/Chat Models
- [ ] Embedding Models
- [X] Prompts / Prompt Templates / Prompt Selectors
- [ ] Output Parsers
- [ ] Document Loaders
- [ ] Vector Stores / Retrievers
- [ ] Memory
- [ ] Agents / Agent Executors
- [ ] Tools / Toolkits
- [X] Chains
- [ ] Callbacks/Tracing
- [ ] Async

### Reproduction

I am trying to use GPT4ALL model with MultiRouterChain and it is throwing some weird error.

## Utility Function:
```
def create_transform_func(remove_index):
    def my_transform_func(inputs: dict):
        return transform_func(inputs, remove_index)
    return my_transform_func


def transform_func(inputs: dict, remove_index = 5) -> dict:
    text = inputs['input'].strip()
    accumulate = ""
    for i, s in enumerate(text.split(' ')):
        i = i + 1
        if i % remove_index != 0:
            accumulate += f"{s} "
    return {"input": accumulate.strip()}
```

```
"""Callback Handler that prints to std out."""
from typing import Any, Dict, List, Optional, Union

from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import AgentAction, AgentFinish

from langchain.schema import LLMResult

from pathlib import Path
from datetime import datetime
import re


class FileCallbackHandler(BaseCallbackHandler):
    """Callback Handler that prints to std out."""

    def __init__(self, 
                 path: Path, 
                 print_prompts: bool=False,
                 print_class: bool=False,
                 title: Optional[str] = "Conversation Log",
                 color: Optional[str] = None
        ) -> None:
        """Initialize callback handler."""
        self.color = color
        self.print_prompts = print_prompts
        self.print_class = print_class
        self.path = path
        self.file_handle = open(path, 'w')
        self.title = title
        self.output_keys = []
        self.output_values = []

    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Print out the prompts."""
        if self.print_prompts:
            self.file_handle.write(f"=============== PROMPTS ==================\n")
            for prompt in prompts:
                self.file_handle.write(f"{prompt}\n")
            self.file_handle.write("\n")
            self.file_handle.flush()
            self.file_handle.write(f"============ END PROMPTS =================\n\n")

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Do nothing."""
        self.file_handle.write(f"=============== LLM END ==================\n")
        pass

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Do nothing."""
        pass

    def on_llm_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Print out that we are entering a chain."""
        if self.print_class:
            self.file_handle.write(f"================ CLASS ===================\n")
            class_name = serialized["name"]
            self.file_handle.write(f">>> class: {class_name}\n")
            self.file_handle.write(f"============== END CLASS =================\n\n")
            self.file_handle.flush()
        self.output_keys.append(list(inputs.keys()))
        self.output_values.append(list(inputs.values()))

    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """Print out that we finished a chain."""
        # print("\n\033[1m> Finished chain.\033[0m")
        # self.file_handle.close()
        self.file_handle.write(f"================ OUTPUT ==================\n")
        keys = []
        values = []
        for k, v in outputs.items():
            keys.append(k)
            values.append(v)
            self.file_handle.write(f"{k}:\n")
            self.file_handle.write(f"{v}\n\n")
        self.output_keys.append(keys)
        self.output_values.append(values)
        self.file_handle.write(f"================ OUTPUT ==================\n")
        self.file_handle.flush()

    def on_chain_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        """Do nothing."""
        self.file_handle.write(datetime.today().strftime('%Y-%m-%d'))
        self.file_handle.write("\n========")
        self.file_handle.flush()
        

    def on_agent_action(
        self, action: AgentAction, color: Optional[str] = None, **kwargs: Any
    ) -> Any:
        """Run on agent action."""
        self.file_handle.write(f">>> action: {action.log}")

    def on_tool_end(
        self,
        output: str,
        color: Optional[str] = None,
        observation_prefix: Optional[str] = None,
        llm_prefix: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """If not the final action, print out observation."""
        if observation_prefix is not None:
            self.file_handle.write(f"\n{observation_prefix}")
        self.file_handle.write(output)
        if llm_prefix is not None:
            self.file_handle.write(f"\n{llm_prefix}")
        self.file_handle.flush()

    def on_tool_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_text(self, text: str, color: Optional[str] = None, end: str = "", **kwargs: Any) -> None:
        """Run when agent ends."""
        self.file_handle.write(f"================ TEXT ===================\n")
        self.file_handle.write(f"{text}\n")
        self.file_handle.flush()
        self.file_handle.write(f"============== END TEXT =================\n\n")
        agent = extract_agent(text)
        if agent != "":
            self.output_keys.append([agent])
            self.output_values.append([text])

    def on_agent_finish(
        self, finish: AgentFinish, color: Optional[str] = None, **kwargs: Any
    ) -> None:
        """Run on agent end."""
        self.file_handle.write(f"{finish.log}\n")
        self.file_handle.flush()
        self.file_handle.close()

    def create_html(self):
        table: str = """
<table class="table table-striped">
    <tr>
        <th>
            Agent
        </th>
        <th>
            Type
        </th>
        <th>
            Output
        </th>
    </tr>
        """
        dedup_hash = set()
        for keys, values in zip(self.output_keys, self.output_values):
            for key, val in zip(keys, values):
                if val not in dedup_hash:
                    dedup_hash.add(val)
                else:
                    continue
                agent = extract_agent(val)
                table += (f"""
        <tr>
            <td>
            </td>
            <td>
                <pre>{key}</pre>
            </td>
            <td>
                <pre>{val}</pre>
            </td>
        </tr>
                """ if agent == "" else f"""
        <tr>
            <td>{agent}</td>
            <td></td>
            <td></td>
        </tr>                
                """)
        table += "</table>"

        target_file = f"{self.path.stem}.html" 
        with open(target_file, "w", encoding='utf-8') as f:
            f.write(f"""
<html>
    <head>
        <meta charset="UTF-8" />
        <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-9ndCyUaIbzAi2FUVXJi0CjmCapSmO7SnpJef0486qhLnuZ2cdeRhO02iuK6FUUVM" crossorigin="anonymous">
        <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/js/bootstrap.bundle.min.js" integrity="sha384-geWF76RCwLtnZ8qwWowPQNguL3RmwHVBC9FhGdlKrxdiJJigb/j/68SIy3Te4Bkz" crossorigin="anonymous"></script>
        <style>
            pre {{
                white-space: pre-wrap;
            }}
        </style>
    </head>
    <body>
        <div class="container-fluid">
            <h1>{self.title}</h1>
            <h2>{generate_timestamp()}</h2>
            {table}
        </div>
    </body>
</html>
            """)
        print(f"Saved chat content to {target_file}")


def generate_timestamp():
    # Get the current date and time
    now = datetime.now()

    # Get the weekday, day, month, year, and time in English
    weekday = now.strftime("%A")
    day = now.strftime("%d")
    month = now.strftime("%B")
    year = now.strftime("%Y")
    time = now.strftime("%H:%M:%S")

    # Create the timestamp string
    timestamp = f"{weekday}, {day} {month} {year} {time}"

    return timestamp

def extract_input(text):
    return re.sub(r".+?'input':\s*'(.+)'}", r"\1", text)

def extract_agent(text):
    regex = r"^([a-z\s]+)\:.+"
    match = re.search(regex, text)
    if match is None:
        return ""
    return re.sub(regex, r"\1", text)
```

## Driver Code
```
from langchain.chains.router import MultiRouteChain, RouterChain
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationChain
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate

from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE
from langchain.chains import SimpleSequentialChain, TransformChain

from prompt_toolkit import HTML, prompt
import langchain.callbacks

from replace_function import create_transform_func

from langchain.callbacks import StdOutCallbackHandler

from FileCallbackHandler import FileCallbackHandler

from pathlib import Path
from typing import Mapping, List, Union
import openai, os
from langchain.llms import GPT4All, LlamaCpp, OpenAI, AzureOpenAI, SelfHostedHuggingFaceLLM, HuggingFacePipeline
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler


file_ballback_handler = FileCallbackHandler(Path('router_chain.txt'), print_prompts=True)


class GPT4_Config():
    current_path = os.path.dirname(os.path.realpath(__file__))
    llm = GPT4All(
        model=os.path.abspath(os.path.join(current_path, r"../../llms/gpt4all/ggml-gpt4all-j-v1.3-groovy.bin")), 
        # n_ctx=1000,
        verbose=True,
        backend='gptj',
        callbacks=[StreamingStdOutCallbackHandler()]
    )


cfg = GPT4_Config()


class PromptFactory():
    developer_template = """You are a very smart Python programmer. \
    You provide answers for algorithmic and computer problems in Python. \
    You explain the code in a detailed manner. \

    Here is a question:
    {input}
    
    Answer:"""

    python_test_developer_template = """You are a very smart Python programmer who writes unit tests using pytest. \
    You provide test functions written in pytest with asserts. \
    You explain the code in a detailed manner. \

    Here is a input on which you create a test:
    {input}
    
    Answer:"""

    kotlin_developer_template = """You are a very smart Kotlin programmer. \
    You provide answers for algorithmic and computer science problems in Kotlin. \
    You explain the code in a detailed manner. \

    Here is a question:
    {input}
    
    Answer:"""

    kotlin_test_developer_template = """You are a very smart Kotlin programmer who writes unit tests using JUnit 5. \
    You provide test functions written in JUnit 5 with JUnit asserts. \
    You explain the code in a detailed manner. \

    Here is a input on which you create a test:
    {input}
    
    Answer:"""

    poet_template = """You are a poet who replies to creative requests with poems in English. \
    You provide answers which are poems in the style of Lord Byron or Shakespeare. \

    Here is a question:
    {input}
    
    Answer:"""

    wiki_template = """You are a Wikipedia expert. \
    You answer common knowledge questions based on Wikipedia knowledge. \
    Your explanations are detailed and in plain English.

    Here is a question:
    {input}
    
    Answer:"""

    image_creator_template = """You create a creator of images. \
    You provide graphic representations of answers using SVG images.

    Here is a question:
    {input}
    
    Answer:"""

    legal_expert_template = """You are a UK or US legal expert. \
    You explain questions related to the UK or US legal systems in an accessible language \
    with a good number of examples.

    Here is a question:
    {input}
    
    Answer:"""

    word_filler = """Your job is to fill the words in a sentence in which words seems to be missing.
    
    Here is the input:
    {input}
    
    Answer:"""

    python_programmer = 'python programmer'
    kotlin_programmer = 'kotlin programmer'

    programmer_test_dict = {
        python_programmer: python_test_developer_template,
        kotlin_programmer: kotlin_test_developer_template
    }

    word_filler_name = 'word filler'

    prompt_infos = [
        {
            'name': python_programmer,
            'description': 'Good for questions about coding and algorithms in Python',
            'prompt_template': developer_template
        },
        {
            'name': 'python tester',
            'description': 'Good for for generating Python tests from existing Python code',
            'prompt_template': python_test_developer_template
        },
        {
            'name': kotlin_programmer,
            'description': 'Good for questions about coding and algorithms in Kotlin',
            'prompt_template': kotlin_developer_template
        },
        {
            'name': 'kotlin tester',
            'description': 'Good for for generating Kotlin tests from existing Kotlin code',
            'prompt_template': kotlin_test_developer_template
        },
        {
            'name': 'poet',
            'description': 'Good for generating poems for creatinve questions',
            'prompt_template': poet_template
        },
        {
            'name': 'wikipedia expert',
            'description': 'Good for answering questions about general knwoledge',
            'prompt_template': wiki_template
        },
        {
            'name': 'graphical artist',
            'description': 'Good for answering questions which require an image output',
            'prompt_template': image_creator_template
        },
        {
            'name': 'legal expert',
            'description': 'Good for answering questions which are related to UK or US law',
            'prompt_template': legal_expert_template
        },
        {
            'name': word_filler_name,
            'description': 'Good at filling words in sentences with missing words',
            'prompt_template': word_filler
        }
    ]


class MyMultiPromptChain(MultiRouteChain):
    """A multi-route chain that uses an LLM router chain to choose amongst prompts."""

    router_chain: RouterChain
    """Chain for deciding a destination chain and the input to it."""
    destination_chains: Mapping[str, Union[LLMChain, SimpleSequentialChain]]
    """Map of name to candidate chains that inputs can be routed to."""
    default_chain: LLMChain
    """Default chain to use when router doesn't map input to one of the destinations."""

    @property
    def output_keys(self) -> List[str]:
        return ["text"]


def generate_destination_chains():
    """
    Creates a list of LLM chains with different prompt templates.
    Note that some of the chains are sequential chains which are supposed to generate unit tests.
    """
    prompt_factory = PromptFactory()
    destination_chains = {}
    for p_info in prompt_factory.prompt_infos:
        print("="*70)
        name = p_info['name']
        print(f"======= Prompt Name =======: \n{name}")
        prompt_template = p_info['prompt_template']
        print(f"======= Prompt Template =======: \n{prompt_template}")
        
        # callbacks = [StdOutCallbackHandler]
        
        chain = LLMChain(
            llm=cfg.llm, 
            prompt=PromptTemplate(template=prompt_template, input_variables=['input']),
            output_key='text',
            callbacks=[file_ballback_handler],
            verbose=True
        )
        
        if name not in prompt_factory.programmer_test_dict.keys() and name != prompt_factory.word_filler_name:
            print("Addition using 1st case")
            destination_chains[name] = chain
        elif name == prompt_factory.word_filler_name:
            print("Addition using 2nd case")
            transform_chain = TransformChain(
                input_variables=["input"], output_variables=["input"], transform=create_transform_func(3), callbacks=[file_ballback_handler]
            )
            destination_chains[name] = SimpleSequentialChain(
                chains=[transform_chain, chain], verbose=True, output_key='text', callbacks=[file_ballback_handler]
            )
        else:
            print("Addition using 3rd case")
            # Normal chain is used to generate code
            # Additional chain to generate unit tests
            template = prompt_factory.programmer_test_dict[name]
            prompt_template = PromptTemplate(input_variables=["input"], template=template)
            test_chain = LLMChain(llm=cfg.llm, prompt=prompt_template, output_key='text', callbacks=[file_ballback_handler])
            destination_chains[name] = SimpleSequentialChain(
                chains=[chain, test_chain], verbose=True, output_key='text', callbacks=[file_ballback_handler]
            )
        print("="*70)
        print("\n\n\n")

    print("============= Destination Chains =============")
    pprint(destination_chains)
    default_chain = ConversationChain(llm=cfg.llm, output_key="text")
    return prompt_factory.prompt_infos, destination_chains, default_chain


def generate_router_chain(prompt_infos, destination_chains, default_chain):
    """
    Generats the router chains from the prompt infos.
    :param prompt_infos The prompt informations generated above.
    :param destination_chains The LLM chains with different prompt templates
    :param default_chain A default chain
    """
    destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]
    print("====================== DESTINATIONS ======================")
    print(destinations)
    
    destinations_str = '\n'.join(destinations)
    print("====================== DESTINATIONS STRINGS ======================")
    print(destinations_str)
    
    
    router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)
    print("====================== ROUTER TEMPLATE ======================")
    print(router_template)
    
    router_prompt = PromptTemplate(
        template=router_template,
        input_variables=['input'],
        output_parser=RouterOutputParser()
    )
    print("====================== PROMPT TEMPLATE ======================")
    print(router_prompt)
    
    router_chain = LLMRouterChain.from_llm(cfg.llm, router_prompt)
    multi_route_chain = MyMultiPromptChain(
        router_chain=router_chain,
        destination_chains=destination_chains,
        default_chain=default_chain,
        verbose=True, 
        callbacks=[file_ballback_handler]
    )
    print("====================== MULTI ROUTER CHAIN ======================")
    print(multi_route_chain)
    return multi_route_chain
    


if __name__ == "__main__":

    prompt_infos, destination_chains, default_chain = generate_destination_chains()
    chain = generate_router_chain(prompt_infos, destination_chains, default_chain)
    with open('conversation.log', 'w') as f:
        while True:
            question = prompt(
                HTML("<b>Type <u>Your question</u></b>  ('q' to exit, 's' to save to html file): ")
            )
            if question == 'q':
                break
            if question in ['s', 'w'] :
                file_ballback_handler.create_html()
                continue
            result = chain.run(question)
            f.write(f"Q: {question}\n\n")
            f.write(f"A: {result}")
            f.write('\n\n ====================================================================== \n\n')
            print(result)
            print()

    
```

## Error:
```
> Entering new MyMultiPromptChain chain...
/home/gpu-titan/anaconda3/envs/Seer/lib/python3.9/site-packages/langchain/chains/llm.py:275: UserWarning: The predict_and_parse method is deprecated, instead pass an output parser directly to LLMChain.
  warnings.warn(
{
"destination": "UK",
"next_inputs": "['Inheritance Tax', 'Legal System']"}UK: {'input': "['Inheritance Tax', 'Legal System']"}Traceback (most recent call last):
  File "/home/gpu-titan/Desktop/Ramish/Seer/seer_main/complex_chain/complex_chain.py", line 360, in <module>
    result = chain.run(question)
  File "/home/gpu-titan/anaconda3/envs/Seer/lib/python3.9/site-packages/langchain/chains/base.py", line 451, in run
    return self(args[0], callbacks=callbacks, tags=tags, metadata=metadata)[
  File "/home/gpu-titan/anaconda3/envs/Seer/lib/python3.9/site-packages/langchain/chains/base.py", line 258, in __call__
    raise e
  File "/home/gpu-titan/anaconda3/envs/Seer/lib/python3.9/site-packages/langchain/chains/base.py", line 252, in __call__
    self._call(inputs, run_manager=run_manager)
  File "/home/gpu-titan/anaconda3/envs/Seer/lib/python3.9/site-packages/langchain/chains/router/base.py", line 106, in _call
    raise ValueError(
ValueError: Received invalid destination chain name 'UK'
```

### Expected behavior

Prompt is correctly routed to correct chain and returns the correct answer.

## Expected Output:
```
> Finished chain.
{'question': 'What are the main differences between the UK and US legal systems in terms of the inheritance tax?', 'text': ' The
 main difference between the UK and US legal systems regarding inheritance taxes lies in their respective approaches to
 calculating, taxing, and paying them out. In both countries, there may be variations depending on factors such as wealth or
 family structure. However, some key differences include:\n1) Inheritance Tax Rates - While the rates of inheritance tax are similar
 between the two jurisdictions (7% for assets over £325k in England/Wales; 7.5% for assets over $675k in Scotland), there may be
 variations depending on factors such as wealth or family structure, which can affect how much is owed and when it should be
 paid out to beneficiaries.\n2) Inheritance Tax Rules - The UK has a more complex inheritance tax system than the US, with
 different rules governing who qualifies for an exemption from paying inheritance tax (e.g., married couples vs unmarried
 individuals). In addition, there may be variations in how assets are taxed and when they should be transferred to beneficiaries
 before death.\n3) Inheritance Tax Planning - Both countries offer various strategies such as trusts or wills that can help reduce
 the amount of inheritance tax owed by transferring wealth out-of-estate at a lower rate than would otherwise apply, but with
 different rules governing how these plans are set up'}
```