{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "sidebar_position: 3\n",
    "title: \"RunnableBranch: Dynamically route logic based on input\"\n",
    "keywords: [RunnableBranch, LCEL]\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b47436a",
   "metadata": {},
   "source": [
    "# Dynamically route logic based on input\n",
    "\n",
    "This notebook covers how to do routing in the LangChain Expression Language.\n",
    "\n",
    "Routing allows you to create non-deterministic chains where the output of a previous step defines the next step. Routing helps provide structure and consistency around interactions with LLMs.\n",
    "\n",
    "There are two ways to perform routing:\n",
    "\n",
    "1. Using a `RunnableBranch`.\n",
    "2. Writing custom factory function that takes the input of a previous step and returns a **runnable**. Importantly, this should return a **runnable** and NOT actually execute.\n",
    "\n",
    "We'll illustrate both methods using a two step sequence where the first step classifies an input question as being about `LangChain`, `Anthropic`, or `Other`, then routes to a corresponding prompt chain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f885113d",
   "metadata": {},
   "source": [
    "## Using a RunnableBranch\n",
    "\n",
    "A `RunnableBranch` is initialized with a list of (condition, runnable) pairs and a default runnable. It selects which branch by passing each condition the input it's invoked with. It selects the first condition to evaluate to True, and runs the corresponding runnable to that condition with the input. \n",
    "\n",
    "If no provided conditions match, it runs the default runnable.\n",
    "\n",
    "Here's an example of what it looks like in action:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1aa13c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatAnthropic\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.schema.output_parser import StrOutputParser"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed84c59a",
   "metadata": {},
   "source": [
    "First, let's create a chain that will identify incoming questions as being about `LangChain`, `Anthropic`, or `Other`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3ec03886",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = (\n",
    "    PromptTemplate.from_template(\n",
    "        \"\"\"Given the user question below, classify it as either being about `LangChain`, `Anthropic`, or `Other`.\n",
    "\n",
    "Do not respond with more than one word.\n",
    "\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\n",
    "Classification:\"\"\"\n",
    "    )\n",
    "    | ChatAnthropic()\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "87ae7c1c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' Anthropic'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\"question\": \"how do I call Anthropic?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aa0a365",
   "metadata": {},
   "source": [
    "Now, let's create three sub chains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d479962a",
   "metadata": {},
   "outputs": [],
   "source": [
    "langchain_chain = (\n",
    "    PromptTemplate.from_template(\n",
    "        \"\"\"You are an expert in langchain. \\\n",
    "Always answer questions starting with \"As Harrison Chase told me\". \\\n",
    "Respond to the following question:\n",
    "\n",
    "Question: {question}\n",
    "Answer:\"\"\"\n",
    "    )\n",
    "    | ChatAnthropic()\n",
    ")\n",
    "anthropic_chain = (\n",
    "    PromptTemplate.from_template(\n",
    "        \"\"\"You are an expert in anthropic. \\\n",
    "Always answer questions starting with \"As Dario Amodei told me\". \\\n",
    "Respond to the following question:\n",
    "\n",
    "Question: {question}\n",
    "Answer:\"\"\"\n",
    "    )\n",
    "    | ChatAnthropic()\n",
    ")\n",
    "general_chain = (\n",
    "    PromptTemplate.from_template(\n",
    "        \"\"\"Respond to the following question:\n",
    "\n",
    "Question: {question}\n",
    "Answer:\"\"\"\n",
    "    )\n",
    "    | ChatAnthropic()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "593eab06",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema.runnable import RunnableBranch\n",
    "\n",
    "branch = RunnableBranch(\n",
    "    (lambda x: \"anthropic\" in x[\"topic\"].lower(), anthropic_chain),\n",
    "    (lambda x: \"langchain\" in x[\"topic\"].lower(), langchain_chain),\n",
    "    general_chain,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "752c732e",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_chain = {\"topic\": chain, \"question\": lambda x: x[\"question\"]} | branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "29231bb8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\" As Dario Amodei told me, here are some ways to use Anthropic:\\n\\n- Sign up for an account on Anthropic's website to access tools like Claude, Constitutional AI, and Writer. \\n\\n- Use Claude for tasks like email generation, customer service chat, and QA. Claude can understand natural language prompts and provide helpful responses.\\n\\n- Use Constitutional AI if you need an AI assistant that is harmless, honest, and helpful. It is designed to be safe and aligned with human values.\\n\\n- Use Writer to generate natural language content for things like marketing copy, stories, reports, and more. Give it a topic and prompt and it will create high-quality written content.\\n\\n- Check out Anthropic's documentation and blog for tips, tutorials, examples, and announcements about new capabilities as they continue to develop their AI technology.\\n\\n- Follow Anthropic on social media or subscribe to their newsletter to stay up to date on new features and releases.\\n\\n- For most people, the easiest way to leverage Anthropic's technology is through their website - just create an account to get started!\", additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"how do I use Anthropic?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c67d8733",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' As Harrison Chase told me, here is how you use LangChain:\\n\\nLangChain is an AI assistant that can have conversations, answer questions, and generate text. To use LangChain, you simply type or speak your input and LangChain will respond. \\n\\nYou can ask LangChain questions, have discussions, get summaries or explanations about topics, and request it to generate text on a subject. Some examples of interactions:\\n\\n- Ask general knowledge questions and LangChain will try to answer factually. For example \"What is the capital of France?\"\\n\\n- Have conversations on topics by taking turns speaking. You can prompt the start of a conversation by saying something like \"Let\\'s discuss machine learning\"\\n\\n- Ask for summaries or high-level explanations on subjects. For example \"Can you summarize the main themes in Shakespeare\\'s Hamlet?\" \\n\\n- Give creative writing prompts or requests to have LangChain generate text in different styles. For example \"Write a short children\\'s story about a mouse\" or \"Generate a poem in the style of Robert Frost about nature\"\\n\\n- Correct LangChain if it makes an inaccurate statement and provide the right information. This helps train it.\\n\\nThe key is interacting naturally and giving it clear prompts and requests', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"how do I use LangChain?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "935ad949",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' 2 + 2 = 4', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"whats 2 + 2\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d8d042c",
   "metadata": {},
   "source": [
    "## Using a custom function\n",
    "\n",
    "You can also use a custom function to route between different outputs. Here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "687492da",
   "metadata": {},
   "outputs": [],
   "source": [
    "def route(info):\n",
    "    if \"anthropic\" in info[\"topic\"].lower():\n",
    "        return anthropic_chain\n",
    "    elif \"langchain\" in info[\"topic\"].lower():\n",
    "        return langchain_chain\n",
    "    else:\n",
    "        return general_chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "02a33c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema.runnable import RunnableLambda\n",
    "\n",
    "full_chain = {\"topic\": chain, \"question\": lambda x: x[\"question\"]} | RunnableLambda(\n",
    "    route\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c2e977a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' As Dario Amodei told me, to use Anthropic IPC you first need to import it:\\n\\n```python\\nfrom anthroipc import ic\\n```\\n\\nThen you can create a client and connect to the server:\\n\\n```python \\nclient = ic.connect()\\n```\\n\\nAfter that, you can call methods on the client and get responses:\\n\\n```python\\nresponse = client.ask(\"What is the meaning of life?\")\\nprint(response)\\n```\\n\\nYou can also register callbacks to handle events: \\n\\n```python\\ndef on_poke(event):\\n  print(\"Got poked!\")\\n\\nclient.on(\\'poke\\', on_poke)\\n```\\n\\nAnd that\\'s the basics of using the Anthropic IPC client library for Python! Let me know if you have any other questions!', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"how do I use Anthropic?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "48913dc6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' As Harrison Chase told me, to use LangChain you first need to sign up for an API key at platform.langchain.com. Once you have your API key, you can install the Python library and write a simple Python script to call the LangChain API. Here is some sample code to get started:\\n\\n```python\\nimport langchain\\n\\napi_key = \"YOUR_API_KEY\"\\n\\nlangchain.set_key(api_key)\\n\\nresponse = langchain.ask(\"What is the capital of France?\")\\n\\nprint(response.response)\\n```\\n\\nThis will send the question \"What is the capital of France?\" to the LangChain API and print the response. You can customize the request by providing parameters like max_tokens, temperature, etc. The LangChain Python library documentation has more details on the available options. The key things are getting an API key and calling langchain.ask() with your question text. Let me know if you have any other questions!', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"how do I use LangChain?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a14d0dca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' 4', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"whats 2 + 2\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46802d04",
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
