langchain/docs/examples/demos/custom_routing_chains.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0af33207",
   "metadata": {},
   "source": [
    "# Custom Routing Chains\n",
    "\n",
    "This covers how to implement a custom routing chain. That problem really reduces to how to implement a custom router. This also acts as a design doc of sorts for routers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16773dc8",
   "metadata": {},
   "source": [
    "## Terminology\n",
    "\n",
    "Before going through any code, let's align on some terminology.\n",
    "- Tool: A function that performs a specific duty. This can be things like: Google Search, Database lookup, Python REPL. The interface for a tool is currently a function that is expected to have a string as an input, with a string as an output.\n",
    "- Tool Input: The input string to a tool.\n",
    "- Observation: The output from calling a tool on a particular input.\n",
    "- Router: The object responsible for deciding which tools to call and when. Exposes a `route` method, which takes in a string and returns a Router Output.\n",
    "- Router Output: The object returned from calling `Router.route` on a string. Consists of:\n",
    "    - The tool to use\n",
    "    - The input to that tool\n",
    "    - A log of the router's thinking.\n",
    "- Routing Chain: A chain which is made up of a router and suite of tools. When passed a string, the Routing Chain will iterative call tools as needed until it arrives at a Final Answer.\n",
    "- Final Answer: The final output of a Routing Chain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eaca15e",
   "metadata": {},
   "source": [
    "## Router\n",
    "A central piece of this chain is the router. The router is responsible for taking user input and deciding which tools, if any, to use. Although it doesn't necessarily have to be backed by a language model (LLM), for pretty much all current use cases it is. LLMs make great routers because they are really good at understanding human intent, which makes them perfect for choosing which tools to use (and for interpreting the output of those tools).\n",
    "\n",
    "Below is the interface we expect routers to expose, along with the RouterOutput definition.\n",
    "\n",
    "```python\n",
    "\n",
    "class RouterOutput(NamedTuple):\n",
    "    \"\"\"Output of a router.\"\"\"\n",
    "\n",
    "    tool: str\n",
    "    tool_input: str\n",
    "    log: str\n",
    "        \n",
    "\n",
    "class Router(ABC):\n",
    "    \"\"\"Chain responsible for deciding the action to take.\"\"\"\n",
    "\n",
    "    @abstractmethod\n",
    "    def route(self, text: str) -> RouterOutput:\n",
    "        \"\"\"Given input, decided how to route it.\n",
    "\n",
    "        Args:\n",
    "            text: input string\n",
    "\n",
    "        Returns:\n",
    "            RouterOutput specifying what tool to use.\n",
    "        \"\"\"\n",
    "\n",
    "    @property\n",
    "    @abstractmethod\n",
    "    def observation_prefix(self) -> str:\n",
    "        \"\"\"Prefix to append the observation with before calling the router again.\"\"\"\n",
    "\n",
    "    @property\n",
    "    @abstractmethod\n",
    "    def router_prefix(self) -> str:\n",
    "        \"\"\"Prefix to prepend the router call with.\"\"\"\n",
    "\n",
    "    @property\n",
    "    def finish_tool_name(self) -> str:\n",
    "        \"\"\"Name of the tool to use to finish the chain.\"\"\"\n",
    "        return \"Final Answer\"\n",
    "\n",
    "    @property\n",
    "    def starter_string(self) -> str:\n",
    "        \"\"\"Put this string after user input but before first router call.\"\"\"\n",
    "        return \"\\n\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "471389be",
   "metadata": {},
   "source": [
    "In order to understand why the router interface is what it is, let's take a look at how it is used in the RoutingChain class:\n",
    "\n",
    "```python\n",
    "def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:\n",
    "    # Construct a mapping of tool name to tool for easy lookup\n",
    "    name_to_tool_map = {tc.tool_name: tc.tool for tc in self.tool_configs}\n",
    "    # Construct the initial string to pass into the router. This is made up\n",
    "    # of the user input, the special starter string, and then the router prefix.\n",
    "    # The starter string is a special string that may be used by a router to\n",
    "    # immediately follow the user input. The router prefix is a string that\n",
    "    # prompts the router to start routing.\n",
    "    starter_string = (\n",
    "        inputs[self.input_key]\n",
    "        + self.router.starter_string\n",
    "        + self.router.router_prefix\n",
    "    )\n",
    "    # We use the ChainedInput class to iteratively add to the input over time.\n",
    "    chained_input = ChainedInput(starter_string, verbose=self.verbose)\n",
    "    # We construct a mapping from each tool to a color, used for logging.\n",
    "    color_mapping = get_color_mapping(\n",
    "        [c.tool_name for c in self.tool_configs], excluded_colors=[\"green\"]\n",
    "    )\n",
    "    # We now enter the router loop (until it returns something).\n",
    "    while True:\n",
    "        # Call the router to see what to do.\n",
    "        output = self.router.route(chained_input.input)\n",
    "        # Add the log to the Chained Input.\n",
    "        chained_input.add(output.log, color=\"green\")\n",
    "        # If the tool chosen is the finishing tool, then we end and return.\n",
    "        if output.tool == self.router.finish_tool_name:\n",
    "            return {self.output_key: output.tool_input}\n",
    "        # Otherwise we lookup the tool\n",
    "        chain = name_to_tool_map[output.tool]\n",
    "        # We then call the tool on the tool input to get an observation\n",
    "        observation = chain(output.tool_input)\n",
    "        # We then log the observation\n",
    "        chained_input.add(f\"\\n{self.router.observation_prefix}\")\n",
    "        chained_input.add(observation, color=color_mapping[output.tool])\n",
    "        # We then add the router prefix into the prompt to get the router to start\n",
    "        # thinking, and start the loop all over.\n",
    "        chained_input.add(f\"\\n{self.router.router_prefix}\")\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9f6ca91",
   "metadata": {},
   "source": [
    "Once we have the custom router written, it is pretty easy to construct the routing chain:\n",
    "\n",
    "```python\n",
    "tools: List[ToolConfig] = ...\n",
    "router = CustomRouter(....)\n",
    "routing_chain = RoutingChain(tools=tools, router=router, verbose=True)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d0c7662",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}