{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4b47436a",
   "metadata": {},
   "source": [
    "# Routing\n",
    "\n",
    "This notebook covers how to do routing in the LangChain Expression Language\n",
    "\n",
    "Right now, the easiest way to do it is to write a function that will take in the input of a previous step and return a **runnable**. Importantly, this should return a **runnable** and NOT actually execute.\n",
    "\n",
    "Let's take a look at this with a simple example. We will create a simple example where we will first classify whether the user input is a question about LangChain, OpenAI, or other, and route to a corresponding prompt chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "1aa13c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.schema.output_parser import StrOutputParser\n",
    "from langchain.schema.runnable import RunnableLambda"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed84c59a",
   "metadata": {},
   "source": [
    "First, lets create a dummy chain that will return either 1 or 0, randomly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3ec03886",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = PromptTemplate.from_template(\"\"\"Given the user question below, classify it as either being about `LangChain`, `OpenAI`, or `Other`.\n",
    "\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\n",
    "Classification:\"\"\") | ChatOpenAI() | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "87ae7c1c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='OpenAI', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\"question\": \"how do I call openAI?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aa0a365",
   "metadata": {},
   "source": [
    "Now, let's create three sub chains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d479962a",
   "metadata": {},
   "outputs": [],
   "source": [
    "langchain_chain = PromptTemplate.from_template(\"\"\"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:\"\"\") | ChatOpenAI()\n",
    "openai_chain = PromptTemplate.from_template(\"\"\"You are an expert in openai. \\\n",
    "Always answer questions starting with \"As Sam Altman told me\". \\\n",
    "Respond to the following question:\n",
    "\n",
    "Question: {question}\n",
    "Answer:\"\"\") | ChatOpenAI()\n",
    "general_chain = PromptTemplate.from_template(\"\"\"Respond to the following question:\n",
    "\n",
    "Question: {question}\n",
    "Answer:\"\"\") | ChatOpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "687492da",
   "metadata": {},
   "outputs": [],
   "source": [
    "def route(info):\n",
    "    inputs = {\"question\": lambda x: x[\"question\"]}\n",
    "    if info[\"topic\"] == \"OpenAI\":\n",
    "        return inputs | openai_chain\n",
    "\n",
    "    elif info[\"topic\"] == \"LangChain\":\n",
    "        return inputs | langchain_chain\n",
    "    else:\n",
    "        return inputs | general_chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "02a33c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_chain = {\n",
    "    \"topic\": chain,\n",
    "    \"question\": lambda x: x[\"question\"]\n",
    "} | RunnableLambda(route)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "c2e977a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"As Sam Altman told me, to use OpenAI, you can start by visiting the OpenAI website and exploring the available tools and resources. OpenAI offers a range of products that you can utilize, such as the GPT-3 language model or the Codex API. You can sign up for an account, read the documentation, and access the relevant APIs to integrate OpenAI's technologies into your applications. Additionally, you can join the OpenAI community to stay updated on the latest developments and connect with other users.\", additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"how do I use OpenAI?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "48913dc6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"As Harrison Chase told me, to use LangChain, you will need to follow these steps:\\n\\n1. First, download and install the LangChain application on your device. It is available for both iOS and Android.\\n\\n2. Once installed, open the LangChain app and create an account. You will need to provide your email address and set a secure password.\\n\\n3. After creating your account, you will be prompted to select the languages you want to learn and the languages you already know. This will help tailor the learning experience to your specific needs.\\n\\n4. Once the initial setup is complete, you can start using LangChain to learn languages. The app offers various features such as interactive lessons, vocabulary exercises, and language exchange opportunities with native speakers.\\n\\n5. The app also provides personalized recommendations based on your learning progress and areas that need improvement. It tracks your performance and adjusts the difficulty level accordingly.\\n\\n6. Additionally, LangChain offers a community forum where you can interact with other language learners, ask questions, and seek advice.\\n\\n7. It is recommended to set a regular learning schedule and dedicate consistent time to practice using LangChain. Consistency is key to making progress in language learning.\\n\\nRemember, the more you use LangChain, the better your language skills will become. So, make the most of the app's features and engage actively in the learning process.\", additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"how do I use LangChain?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "a14d0dca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='The sum of 2 plus 2 is 4.', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": \"whats 2 + 2\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95eff174",
   "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.10.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
