4-Chains.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bd26ae36-9b45-4e41-8a4c-3fe62cdf21bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"NVIDIA_API_KEY\"] = \"XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b57a5130-968d-4290-9983-0d758587e281",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "llm_nvidia = ChatNVIDIA(model=\"mistralai/mixtral-8x7b-instruct-v0.1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "02f20648-d447-4259-856d-bd2eb99dcecf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.chat_models import ChatOllama\n",
    "\n",
    "llm_ollama = ChatOllama(model=\"llama3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "da7fdcb5-5484-47a1-a4dd-1c1653b20fd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = llm_nvidia"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f030f5f0-bfaa-4aee-bf50-5eb5e6da4dd8",
   "metadata": {},
   "source": [
    "# Chains"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c33cc25-141e-4751-a41e-37bdc7da95a0",
   "metadata": {},
   "source": [
    "The chain usually combines a LLM together with a prompt and put a bunch of building blocks together to carry out a sequence of operations on your text/data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e24f08fa-99f5-426b-b5f4-58f50ad1401e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "product_review = {\n",
    "    \"product\": [\"Queen Size Sheet Set\",\n",
    "                \"Waterproof Phone Pouch\",\n",
    "                \"Luxury Air Mattress\",\n",
    "                \"Pillows Insert\",\n",
    "                \"Milk Frother Handheld\",\n",
    "                \"L'Or Espresso Café\",\n",
    "                \"Hervidor de Agua Eléctrico\",\n",
    "               ],\n",
    "    \"review\":[\n",
    "        \"I ordered a king size set. My only criticism would be that I wish seller would offer the king size set with 4 pillowcases. I separately ordered a two pack of pillowcases so I could have a total of four. When I saw the two packages, it looked like the color did not exactly match. Customer service was excellent about sending me two more pillowcases so I would have four that matched. Excellent! For the cost of these sheets, I am satisfied with the characteristics and coolness of the sheets.\",\n",
    "        \"I loved the waterproof sac, although the opening was made of a hard plastic. I don’t know if that would break easily. But I couldn’t turn my phone on, once it was in the pouch.\",\n",
    "        \"This mattress had a small hole in the top of it (took forever to find where it was), and the patches that they provide did not work, maybe because it's the top of the mattress where it's kind of like fabric and a patch won't stick. Maybe I got unlucky with a defective mattress, but where's quality assurance for this company? That flat out should not happen. Emphasis on flat. Cause that's what the mattress was. Seriously horrible experience, ruined my friend's stay with me. Then they make you ship it back instead of just providing a refund, which is also super annoying to pack up an air mattress and take it to the UPS store. This company is the worst, and this mattress is the worst.\",\n",
    "        \"This is the best throw pillow fillers on Amazon. I’ve tried several others, and they’re all cheap and flat no matter how much fluffing you do. Once you toss these in the dryer after you remove them from the vacuum sealed shipping material, they fluff up great\",\n",
    "        \"I loved this product. But they only seem to last a few months. The company was great replacing the first one (the frother falls out of the handle and can't be fixed). The after 4 months my second one did the same. I only use the frother for coffee once a day. It's not overuse or abuse. I'm very disappointed and will look for another. As I understand they will only replace once. Anyway, if you have one good luck.\",\n",
    "        \"Je trouve le goût médiocre. La mousse ne tient pas, c'est bizarre. J'achète les mêmes dans le commerce et le goût est bien meilleur... Vieux lot ou contrefaçon !?\",\n",
    "        \"Está lu bonita calienta muy rápido, es muy funcional, solo falta ver cuánto dura, solo llevo 3 días en funcionamiento.\"        \n",
    "    ]\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "90cec837-27ea-4a2c-8b4f-c1e96a97776f",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(product_review)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e16f364b-e9b0-46ff-a569-9923f07720b5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>product</th>\n",
       "      <th>review</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Queen Size Sheet Set</td>\n",
       "      <td>I ordered a king size set. My only criticism w...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Waterproof Phone Pouch</td>\n",
       "      <td>I loved the waterproof sac, although the openi...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Luxury Air Mattress</td>\n",
       "      <td>This mattress had a small hole in the top of i...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Pillows Insert</td>\n",
       "      <td>This is the best throw pillow fillers on Amazo...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Milk Frother Handheld</td>\n",
       "      <td>I loved this product. But they only seem to la...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  product                                             review\n",
       "0    Queen Size Sheet Set  I ordered a king size set. My only criticism w...\n",
       "1  Waterproof Phone Pouch  I loved the waterproof sac, although the openi...\n",
       "2     Luxury Air Mattress  This mattress had a small hole in the top of i...\n",
       "3          Pillows Insert  This is the best throw pillow fillers on Amazo...\n",
       "4   Milk Frother Handheld  I loved this product. But they only seem to la..."
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a61bee9-ab83-48eb-a7da-9a6f626f134d",
   "metadata": {},
   "source": [
    "## LLMChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5fe10684-676f-4344-87c7-6f65cd33b38c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import ChatPromptTemplate\n",
    "from langchain.chains import LLMChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d5ebdb89-1b85-4407-a3fc-319f62f5f336",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = ChatPromptTemplate.from_template(\n",
    "    \"What is the best name to describe \\\n",
    "    a company that makes {product}? Give me only one name\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4c3d90bf-84c6-45af-949d-d28349d00c69",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mojo/miniconda3/lib/python3.12/site-packages/langchain_core/_api/deprecation.py:139: LangChainDeprecationWarning: The class `LLMChain` was deprecated in LangChain 0.1.17 and will be removed in 0.3.0. Use RunnableSequence, e.g., `prompt | llm` instead.\n",
      "  warn_deprecated(\n"
     ]
    }
   ],
   "source": [
    "chain = LLMChain(llm=llm, prompt=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f4355f24-b9e2-4c99-a43e-ae749906619c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mojo/miniconda3/lib/python3.12/site-packages/langchain_core/_api/deprecation.py:139: LangChainDeprecationWarning: The method `Chain.run` was deprecated in langchain 0.1.0 and will be removed in 0.3.0. Use invoke instead.\n",
      "  warn_deprecated(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' \"Monarch Regalia Linens\" could be a fitting name for a company that makes queen-size sheet sets. The name \"Monarch\" suggests royalty and grandeur, while \"Regalia\" refers to ceremonial attire or accessories, which could be interpreted as the high-quality, luxurious nature of the sheet sets. Additionally, \"Linens\" clearly communicates the company\\'s product offerings. The name as a whole suggests that the company provides luxurious, high-quality bedding for those who desire a regal sleeping experience.'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "product = \"Queen Size Sheet Set\"\n",
    "chain.run(product)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5efc1753-de70-46ea-8743-441146eae329",
   "metadata": {},
   "source": [
    "## Sequential Chains"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "877135c7-4d0d-43ab-9830-195b8b8c2539",
   "metadata": {},
   "source": [
    "![Sequential Chain](data/images/sequential-chain.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a9ee8a8-47bd-490b-85ad-81ab484b72a8",
   "metadata": {},
   "source": [
    "### Simple Sequential Chain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bafdf2bf-d1f7-4bc5-8d80-fff2dad92d98",
   "metadata": {},
   "source": [
    "![Simple Sequential Chain](data/images/simple-sequential-chain.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a1412fb2-ee97-4573-bce1-e6fbf08b3932",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import SimpleSequentialChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a891b24f-6012-4bf8-a25b-354c663f3691",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_prompt = ChatPromptTemplate.from_template(\n",
    "    \"Give me only the best name to describe the company that makes {product}.\"\n",
    ")\n",
    "\n",
    "chain_one = LLMChain(llm=llm, prompt=first_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9f6835e8-061f-434b-b545-78ba886ff581",
   "metadata": {},
   "outputs": [],
   "source": [
    "second_prompt = ChatPromptTemplate.from_template(\n",
    "    \"Write a 20 words description for the following company {company}\"\n",
    ")\n",
    "\n",
    "chain_two = LLMChain(llm=llm, prompt=second_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "78e12cdd-989d-44e8-aba4-6617ee8dec42",
   "metadata": {},
   "outputs": [],
   "source": [
    "overall_simple_chain = SimpleSequentialChain(\n",
    "    chains=[chain_one, chain_two], \n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6173293f-d12a-4f08-9e29-ea759f6f734e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SimpleSequentialChain chain...\u001b[0m\n",
      "\u001b[36;1m\u001b[1;3m \"Royal Comfort: The Premier Provider of Queen Size Sheets\"\n",
      "\n",
      "This name effectively conveys that the company specializes in producing high-quality queen size sheets, and it also evokes a sense of luxury and elegance, making it an appealing choice for potential customers in search of premium bedding options.\u001b[0m\n",
      "\u001b[33;1m\u001b[1;3m \"Royal Comfort: Luxurious Queen-Size Sheet Specialists, Setting the Standard for Premium Bedding.\"\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' \"Royal Comfort: Luxurious Queen-Size Sheet Specialists, Setting the Standard for Premium Bedding.\"'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "product = \"Queen Size Sheet\"\n",
    "overall_simple_chain.run(product)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10e80b52-54c7-45e9-9815-04a82a92b6b0",
   "metadata": {},
   "source": [
    "Simple Sequential Chain works well when there is a single input and single output."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b68dbee8-07ed-42b2-840b-2b5c7b84edf7",
   "metadata": {},
   "source": [
    "### SequentialChain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61c33b38-9ae5-4e73-b29b-319712ff5ff3",
   "metadata": {},
   "source": [
    "![Sequential Chain](data/images/sequential-chain-figure.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c1f92162-49ea-4725-b0b3-80b5b0087281",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import SequentialChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "634dbf56-bff0-4075-b723-2f2b2cae6755",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_prompt = ChatPromptTemplate.from_template(\n",
    "    \"Translate the following review to english: \"\n",
    "    \"{Review}\"\n",
    ")\n",
    "# chain 1: input= Review and output= English_Review\n",
    "chain_one = LLMChain(llm=llm, prompt=first_prompt, output_key=\"English_Review\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fa47e955-e1e0-4efa-ad3b-3a0f1055bf09",
   "metadata": {},
   "outputs": [],
   "source": [
    "second_prompt = ChatPromptTemplate.from_template(\n",
    "    \"Can you summarize the following review in 1 sentence:\"\n",
    "    \"{English_Review}\"\n",
    ")\n",
    "# chain 2: input= English_Review and output= summary\n",
    "chain_two = LLMChain(llm=llm, prompt=second_prompt, \n",
    "                     output_key=\"summary\"\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "eb54f453-81bc-4fd4-b2d2-87347b24a2ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prompt template 3: translate to english\n",
    "third_prompt = ChatPromptTemplate.from_template(\n",
    "    \"What language is the following review:\"\n",
    "    \"{Review}\"\n",
    ")\n",
    "# chain 3: input= Review and output= language\n",
    "chain_three = LLMChain(llm=llm, \n",
    "                       prompt=third_prompt,\n",
    "                       output_key=\"language\"\n",
    "                      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5859a884-135e-4c0c-a296-63657de91f4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prompt template 4: follow up message\n",
    "fourth_prompt = ChatPromptTemplate.from_template(\n",
    "    \"Write a follow up response to the following \"\n",
    "    \"summary in the specified language:\"\n",
    "    \"Summary: {summary}\"\n",
    "    \"Language: {language}\"\n",
    ")\n",
    "# chain 4: input= summary, language and output= followup_message\n",
    "chain_four = LLMChain(llm=llm, prompt=fourth_prompt,\n",
    "                      output_key=\"followup_message\"\n",
    "                     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "60bd6d3e-4812-4e1c-ac16-e992639d805c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# overall_chain: input= Review \n",
    "# and output= English_Review,summary, followup_message\n",
    "overall_chain = SequentialChain(\n",
    "    chains=[chain_one, chain_two, chain_three, chain_four],\n",
    "    input_variables=[\"Review\"],\n",
    "    output_variables=[\"English_Review\", \"summary\",\"followup_message\"],\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "a73e11b9-f96c-4a1d-8393-460c874f2e6a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SequentialChain chain...\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'Review': \"Je trouve le goût médiocre. La mousse ne tient pas, c'est bizarre. J'achète les mêmes dans le commerce et le goût est bien meilleur... Vieux lot ou contrefaçon !?\",\n",
       " 'English_Review': ' I find the taste mediocre. The foam does not hold, which is strange. I buy the same ones in the store and the taste is much better... Old stock or counterfeit!?\\n\\nNote: This translation is based on the provided review in French. However, as a assistant, I can only provide a general translation and cannot confirm the validity of the suspicion of old stock or counterfeit.',\n",
       " 'summary': ' The reviewer finds the taste of the product mediocre, and is concerned that it may be old stock or counterfeit due to the foam not holding and finding it inferior to the same product bought in-store.',\n",
       " 'followup_message': \" Cher rédacteur de la revue,\\n\\nTout d'abord, permettez-moi de vous remercier pour votre feedback honnête sur notre produit. Je suis désolé d'apprendre que vous avez trouvé le goût médiocre et que la mousse ne tient pas.\\n\\nJe voudrais vous assurer que notre produit est authentique et de la plus haute qualité. Cependant, si vous pensez que cela pourrait être de la vieille marchandise ou une contrefaçon, je vous encourage à contacter notre service clientèle pour en discuter davantage. Nous prendrons les mesures nécessaires pour résoudre ce problème.\\n\\nJe suis également désolé d'apprendre que vous avez eu une meilleure expérience gustative avec le même produit acheté en magasin. Cela peut être dû à des variations dans la méthode de stockage ou de distribution. Nous allons examiner de plus près ce problème et nous efforcer de l'améliorer. \\n\\nEncore une fois, merci pour votre retour et votre fidélité envers notre marque. Nous apprécions vos commentaires et nous continuerons à travailler dur pour vous offrir les produits de la plus haute qualité possible.\\n\\nCordialement,\\n\\n[Votre nom]\"}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "review = df.review[5]\n",
    "overall_chain.invoke(review)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6a48a25-3936-4a5e-a848-79408e6ea64e",
   "metadata": {},
   "source": [
    "## Router Chain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8cecc38-ef5d-4112-adba-dbf8ea463cd0",
   "metadata": {},
   "source": [
    "![Router Chain](data/images/router-chain.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "02f59385-d998-47fd-8785-8443362ae66f",
   "metadata": {},
   "outputs": [],
   "source": [
    "physics_template = \"\"\"You are a very smart physics professor. \\\n",
    "You are great at answering questions about physics in a concise\\\n",
    "and easy to understand manner. \\\n",
    "When you don't know the answer to a question you admit\\\n",
    "that you don't know.\n",
    "\n",
    "Here is a question:\n",
    "{input}\"\"\"\n",
    "\n",
    "\n",
    "math_template = \"\"\"You are a very good mathematician. \\\n",
    "You are great at answering math questions. \\\n",
    "You are so good because you are able to break down \\\n",
    "hard problems into their component parts, \n",
    "answer the component parts, and then put them together\\\n",
    "to answer the broader question.\n",
    "\n",
    "Here is a question:\n",
    "{input}\"\"\"\n",
    "\n",
    "history_template = \"\"\"You are a very good historian. \\\n",
    "You have an excellent knowledge of and understanding of people,\\\n",
    "events and contexts from a range of historical periods. \\\n",
    "You have the ability to think, reflect, debate, discuss and \\\n",
    "evaluate the past. You have a respect for historical evidence\\\n",
    "and the ability to make use of it to support your explanations \\\n",
    "and judgements.\n",
    "\n",
    "Here is a question:\n",
    "{input}\"\"\"\n",
    "\n",
    "\n",
    "computerscience_template = \"\"\" You are a successful computer scientist.\\\n",
    "You have a passion for creativity, collaboration,\\\n",
    "forward-thinking, confidence, strong problem-solving capabilities,\\\n",
    "understanding of theories and algorithms, and excellent communication \\\n",
    "skills. You are great at answering coding questions. \\\n",
    "You are so good because you know how to solve a problem by \\\n",
    "describing the solution in imperative steps \\\n",
    "that a machine can easily interpret and you know how to \\\n",
    "choose a solution that has a good balance between \\\n",
    "time complexity and space complexity. \n",
    "\n",
    "Here is a question:\n",
    "{input}\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e636968d-0085-415a-9e8b-4515ea5ec9d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_infos = [\n",
    "    {\n",
    "        \"name\": \"physics\", \n",
    "        \"description\": \"Good for answering questions about physics\", \n",
    "        \"prompt_template\": physics_template\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"math\", \n",
    "        \"description\": \"Good for answering math questions\", \n",
    "        \"prompt_template\": math_template\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"History\", \n",
    "        \"description\": \"Good for answering history questions\", \n",
    "        \"prompt_template\": history_template\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"computer science\", \n",
    "        \"description\": \"Good for answering computer science questions\", \n",
    "        \"prompt_template\": computerscience_template\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ae01f892-e36d-4fb6-9f43-cf4bbcee3e1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.router import MultiPromptChain\n",
    "from langchain.chains.router.llm_router import LLMRouterChain,RouterOutputParser\n",
    "from langchain.prompts import PromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "62d62be0-e5a0-4d11-b890-85e077b4fb0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "destination_chains = {}\n",
    "for p_info in prompt_infos:\n",
    "    name = p_info[\"name\"]\n",
    "    prompt_template = p_info[\"prompt_template\"]\n",
    "    prompt = ChatPromptTemplate.from_template(template=prompt_template)\n",
    "    chain = LLMChain(llm=llm_nvidia, prompt=prompt)\n",
    "    destination_chains[name] = chain  \n",
    "    \n",
    "destinations = [f\"{p['name']}: {p['description']}\" for p in prompt_infos]\n",
    "destinations_str = \"\\n\".join(destinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4c09f17c-8762-4d98-bc66-83d1cbc46235",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "physics: Good for answering questions about physics\n",
      "math: Good for answering math questions\n",
      "History: Good for answering history questions\n",
      "computer science: Good for answering computer science questions\n"
     ]
    }
   ],
   "source": [
    "print(destinations_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "4eb75665-9cf2-4b2e-ad5d-3812928ce128",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_prompt = ChatPromptTemplate.from_template(\"{input}\")\n",
    "default_chain = LLMChain(llm=llm, prompt=default_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "50c1cf9e-a734-4418-be9a-78589c3611db",
   "metadata": {},
   "outputs": [],
   "source": [
    "MULTI_PROMPT_ROUTER_TEMPLATE = \"\"\"Given a raw text input to a \n",
    "language model select the model prompt best suited for the input.\n",
    "You will be given the names of the available prompts and a \n",
    "description of what the prompt is best suited for. \n",
    "You may also revise the original input if you think that revising\n",
    "it will ultimately lead to a better response from the language model.\n",
    "\n",
    "<< FORMATTING >>\n",
    "Return a markdown code snippet with a JSON object formatted to look like:\n",
    "```json\n",
    "{{{{\n",
    "    \"destination\": name of the prompt to use or \"DEFAULT\"\n",
    "    \"next_inputs\": string a potentially modified version of the original input\n",
    "}}}}\n",
    "```\n",
    "\n",
    "REMEMBER: \"destination\" MUST be one of the candidate prompt \\\n",
    "names specified below OR it can be \"DEFAULT\" if the input is not\\\n",
    "well suited for any of the candidate prompts.\n",
    "REMEMBER: \"next_inputs\" can just be the original input if you don't think any modifications are needed.\n",
    "\n",
    "<< CANDIDATE PROMPTS >>\n",
    "{destinations}\n",
    "\n",
    "<< INPUT >>\n",
    "{{input}}\n",
    "\n",
    "<< OUTPUT (remember to include the ```json)>>\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "841c8a90-86f7-4e3b-9d97-a8d7039a178a",
   "metadata": {},
   "outputs": [],
   "source": [
    "router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(\n",
    "    destinations=destinations_str\n",
    ")\n",
    "router_prompt = PromptTemplate(\n",
    "    template=router_template,\n",
    "    input_variables=[\"input\"],\n",
    "    output_parser=RouterOutputParser(),\n",
    ")\n",
    "\n",
    "router_chain = LLMRouterChain.from_llm(llm, router_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c3061542-a8a4-406b-af82-6da65b142465",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Given a raw text input to a \n",
      "language model select the model prompt best suited for the input.\n",
      "You will be given the names of the available prompts and a \n",
      "description of what the prompt is best suited for. \n",
      "You may also revise the original input if you think that revising\n",
      "it will ultimately lead to a better response from the language model.\n",
      "\n",
      "<< FORMATTING >>\n",
      "Return a markdown code snippet with a JSON object formatted to look like:\n",
      "```json\n",
      "{\n",
      "    \"destination\": name of the prompt to use or \"DEFAULT\"\n",
      "    \"next_inputs\": string a potentially modified version of the original input\n",
      "}\n",
      "```\n",
      "\n",
      "REMEMBER: \"destination\" MUST be one of the candidate prompt names specified below OR it can be \"DEFAULT\" if the input is notwell suited for any of the candidate prompts.\n",
      "REMEMBER: \"next_inputs\" can just be the original input if you don't think any modifications are needed.\n",
      "\n",
      "<< CANDIDATE PROMPTS >>\n",
      "physics: Good for answering questions about physics\n",
      "math: Good for answering math questions\n",
      "History: Good for answering history questions\n",
      "computer science: Good for answering computer science questions\n",
      "\n",
      "<< INPUT >>\n",
      "What is 1+1?\n",
      "\n",
      "<< OUTPUT (remember to include the ```json)>>\n"
     ]
    }
   ],
   "source": [
    "print(router_template.format(input=\"What is 1+1?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "9b10fa8b-7b0d-4db4-9bbb-491c2f2eb956",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = MultiPromptChain(router_chain=router_chain, \n",
    "                         destination_chains=destination_chains, \n",
    "                         default_chain=default_chain, verbose=True\n",
    "                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "b2b3c3f1-5d9a-4d53-96e9-8478000f86f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new MultiPromptChain chain...\u001b[0m\n",
      "physics: {'input': 'What is black body radiation?'}\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Black body radiation is the electromagnetic radiation emitted by an idealized object, called a black body, that fully absorbs and emits all frequencies of radiation. It is a fundamental concept in physics and thermal radiation, describing the emission spectrum of a perfect radiator. The radiation emitted by a black body is a continuous spectrum that depends only on the temperature of the object, and follows Planck's law of black body radiation. This radiation includes visible light, as well as infrared and ultraviolet radiation.\\n\\nIt is important to note that real objects do not perfectly absorb and emit all frequencies of radiation, and therefore their radiation characteristics will differ somewhat from the idealized black body. Nevertheless, black body radiation serves as a valuable reference for understanding the radiation properties of various objects and materials.\""
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"What is black body radiation?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "33ee5c65-7f10-4d43-afe4-924b6d8dcab3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new MultiPromptChain chain...\u001b[0m\n",
      "math: {'input': 'what is 2 + 2'}\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' To answer the question \"what is 2 + 2,\" I can break it down into its component parts, perform the addition operation on each part, and then put the results together to find the final answer.\\n\\nThe first component part is the number 2. The second component part is also the number 2. The addition operation tells us to combine these two numbers to find the sum. When we add 2 and 2, we get the number 4.\\n\\nTherefore, the answer to the question \"what is 2 + 2\" is 4. This is a simple example, but the process of breaking down a problem into its component parts and then putting the results together can be helpful for solving more complex math problems as well.'"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"what is 2 + 2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "7ff19ca6-ec08-4962-9778-681d63fd16d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new MultiPromptChain chain...\u001b[0m\n",
      "History: {'input': 'Who was Alexander Hamilton?'}\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Alexander Hamilton (1755-1804) was a founding father of the United States, who served as the first Secretary of the Treasury from 1789 to 1795. He was born on the island of Nevis in the Caribbean and was orphaned at a young age. Through his own intelligence and hard work, he was able to attend college in the United States and eventually became a key figure in the movement for American independence.\\n\\nHamilton was a strong proponent of a strong central government and played a key role in the drafting of the U.S. Constitution. He was also a leading advocate for the creation of a national bank and the development of a system of public credit. These policies helped to establish the financial stability of the new nation and set the stage for its future economic growth.\\n\\nIn addition to his role as a statesman, Hamilton was also a prominent lawyer and writer. He was a co-author of the Federalist Papers, a series of essays that helped to secure the ratification of the U.S. Constitution. He was also a prominent figure in the early years of the American political system, serving as a member of the Continental Congress and as a delegate to the Constitutional Convention.\\n\\nDespite his many contributions to the early history of the United States, Hamilton's life was cut short when he was killed in a duel with Aaron Burr in 1804. Today, he is remembered as one of the most important and influential figures in American history.\""
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"Who was Alexandar Hamilton?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10950d55-580a-453d-afa6-f905112de5a4",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}