Added a non-agent chain version for comparison

explore.ipynb

@@ -39,7 +39,8 @@
     "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
     "from langchain.chains import create_history_aware_retriever, create_retrieval_chain\n",
     "from langchain.tools.retriever import create_retriever_tool\n",
-    "from langchain_core.runnables import RunnablePassthrough, RunnableParallel\n",
+    "from langchain_core.runnables import RunnablePassthrough, RunnableParallel, chain\n",
+    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
     "\n",
     "# Agents\n",
     "from langchain import hub\n",
@@ -167,23 +168,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Deleting the previous database and creating a new one (because otherwise content is duplicated in the db every time this block is run)\n",
-      "The Wookieepedia database contains 10838 chunks, with mean length of 645 characters\n",
+      "The Wookieepedia database contains 10927 chunks, with mean length of 645 characters\n",
       "Current source pages in Wookieepedia db:\n",
-      "393\n"
+      "395\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "{'wookieepedia\\\\1995_Topps_Star_Wars_Widevision.txt',\n",
+       "{'https://starwars.fandom.com/wiki/Darth_Plagueis',\n",
+       " 'https://starwars.fandom.com/wiki/Voss',\n",
+       " 'wookieepedia\\\\1995_Topps_Star_Wars_Widevision.txt',\n",
        " 'wookieepedia\\\\1995_Topps_Star_Wars___The_Empire_Strikes_Back_Widevision.txt',\n",
        " 'wookieepedia\\\\1996_Topps_Star_Wars___Return_of_the_Jedi_Widevision.txt',\n",
        " 'wookieepedia\\\\2_systems_control_droid.txt',\n",
@@ -578,7 +580,7 @@
        " 'wookieepedia\\\\Zev.txt'}"
       ]
      },
-     "execution_count": 4,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -644,7 +646,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -662,7 +664,7 @@
        " Document(page_content='\\nAt an opera show, Palpatine and Skywalker discussed the situation with Palpatine revealing to know that the Council wanted Skywalker to spy on him. At the same time, Skywalker, who was dealing with the fear of his wife\\'s loss was beginning to distrust the Jedi Council. This allowed Sidious to orchestrate Skywalker\\'s downfall; he claimed that the Jedi and the Sith were alike in nearly every way, setting the stage for the young Jedi to switch sides while playing to his devotion to the Jedi ideology. Next, Palpatine told him the story of his master, \"Darth Plagueis the Wise.\" According to the story, Plagueis possessed such strength in the Force that he could use it to prevent the ones closest to him from dying. Skywalker asked what happened to him, to which Palpatine said that he was afraid of losing his power, but had taught his apprentice everything, and then his apprentice killed him in his sleep. Skywalker, enthralled by the story, wished to know more about this power, asking him if he himself could learn it with Palpatine stating that, though possible, such things could not be taught by members of the Order. Skywalker, however, was unware that Palpatine was really talking about how he killed his mentor Plagueis, in order to lure him to the dark side.', metadata={'source': 'wookieepedia\\\\Darth_Sidious.txt'})]"
       ]
      },
-     "execution_count": 20,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -718,7 +720,7 @@
     "\n",
     "    return wdb.similarity_search(original_query, k = 10)\n",
     "\n",
-    "get_wookieepedia_context('Do you know the Tragedy of Darth Plagueis the Wise?', 'Darth Plagueis', woo_db)"
+    "# get_wookieepedia_context('Do you know the Tragedy of Darth Plagueis the Wise?', 'Darth Plagueis', woo_db)"
    ]
   },
   {
@@ -731,7 +733,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -744,12 +746,12 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Non-agent version"
+    "### Base version (only one retriever)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -782,7 +784,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -809,7 +811,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -835,7 +837,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
@@ -843,13 +845,13 @@
       "text/plain": [
        "{'input': \"Who participates in Han's rescue from Jabba? And where is the palace?\",\n",
        " 'context': [Document(page_content=\"After the smuggler Han Solo failed to repay him for lost cargo, Jabba placed a high price on his head. Solo was eventually delivered to him by one of his bounty hunters, Boba Fett, as a gift from Darth Vader. However, this capture brought him to the attention of Jedi Knight Luke Skywalker, who sought to rescue his friend from Jabba's imprisonment. As he attempted to execute the Jedi and his allies in the Great Pit of Carkoon, Jabba was choked to death by Leia Organa. With the Hutts unable to decide who would inherit Jabba's criminal ventures, many of his slaves, including the Niktos, were free, and his palace was occupied by his former Majordomo Bib Fortuna, who took his place as Daimyo of Tatooine until Fett killed and usurped him. Fett sought to rebuild Jabba's criminal empire in his own image, intending to rule with respect rather than the fear that the Hutt instilled in his followers.\", metadata={'source': 'wookieepedia\\\\Jabba_Desilijic_Tiure.txt'}),\n",
-       "  Document(page_content=\"Jabba sends Han Solo, Chewbacca, and Greedo to travel to Corellia to recover the ashes of archrival Krestrel D'Naran.\\nHan Solo meets Corbus Tyra, who claims to be his father, Ovan.\\nChewbacca rescues Han Solo from bounty hunter Krrsantan and they later conduct a heist on Coruscant. Although the urn they have been after does not contain the ashes, it contains the neural core of Ajax Sigma, which Han Solo and Chewbacca bury away.\\n8/7 BBY\\nJabba Desilijic Tiure sends Krrsantan to hunt down Obi-Wan Kenobi.\\n7 BBY\\nBodhi Rook enters the Terrabe Sector Service Academy to train to be a pilot.\\nThe mission to Horuz occurs.\\nUn'hee is born.\\nObi-Wan Kenobi rescues 12-year old Luke Skywalker and his friend Windy after they are left unconscious during a sandstorm.\\n6 BBY\\nThe mission to Inusagi occurs.[source?]\\nThe Malkhani Insurrections occurs.\\nThe Leonis family migrates to Lothal.\\nThe Westhills massacre occurs.\", metadata={'source': 'wookieepedia\\\\Timeline_of_galactic_history.txt'}),\n",
+       "  Document(page_content=\"On a distant spaceport, Mira manages to escape from the stormtroopers who were guarding her and sneak into the Millennium Falcon, which was being loading with spice. Han Solo and Chewbacca are going to make another Kessel Run for Jabba the Hutt, but first he was going to pick up Luke Skywalker. On the hidden Rebel base on the planet Randa, Han picks up Luke, Artoo and Threepio and finds his unexpected guest. Mira tells Luke and Han about Bzorn's plan and asks Han to deliver her to the professor, knowing that he will stop the destruction after he learns that Mira is safe.\", metadata={'source': 'wookieepedia\\\\The_Second_Kessel_Run.txt'}),\n",
        "  Document(page_content=\"Djarin takes the Darksaber while Gideon is taken into New Republic custody, though he is broken out en route to trial.\\nAs Bo-Katan Kryze returned to her forces without the Darksaber, her Mandalorians abandon her and become privateers under Axe Woves.\\nThe Jedi Luke Skywalker beckons Grogu to train as a Jedi with him. Removing his Mandalorian helmet, Djarin gives his blessing to Grogu before letting the latter go.\\nLuke Skywalker builds a Jedi Temple on Ossus. After a brief time with Skywalker there, Grogu decides to return to Din Djarin.      Jabba's Palace under new management  \\nBoba Fett and Fennec Shand attack Jabba's Palace, replacing Bib Fortuna's place as master of the palace.\\nA gang war erupts on Tatooine, with Daimyo Boba Fett's Fett gotra defeating the Pyke Syndicate presence on the world and emerging as its leading criminal syndicate.\\nc. 9 ABY\", metadata={'source': 'wookieepedia\\\\Timeline_of_galactic_history.txt'}),\n",
        "  Document(page_content=\"At the chancellor's request, Windu dispatched Kenobi and Skywalker on a mission to save Jabba's son.  \\nSoon afterward Windu received a message from the Advanced Recon Force Scout Troopers who tracked Rotta's location to an old monastery on the planet Teth. Due to the monastery's heavy fortification, Windu instructed the troopers to stand by and await Republic reinforcements. He then reassigned three of Admiral Wurtz's cruisers to the rescue operation. In spite of Count Dooku's plot to form an alliance between the Confederacy and the Hutt Clan, Rotta was ultimately redeemed from captivity by Anakin Skywalker and his new Padawan, Ahsoka Tano. The Separatists were defeated on Teth and Jabba was made aware of Dooku's conspiracy, which in turn allowed the Republic to secure passage through the Hyperspace lanes in Hutt territory.\\n\\nLiberation of Ryloth\", metadata={'source': 'wookieepedia\\\\Mace_Windu.txt'})],\n",
-       " 'answer': \"Han Solo is rescued from Jabba the Hutt by a group consisting of Jedi Knight Luke Skywalker, Princess Leia Organa (disguised as the bounty hunter Boushh), Chewbacca, Lando Calrissian, R2-D2, and C-3PO. The rescue takes place at Jabba's Palace on the desert planet of Tatooine.\"}"
+       " 'answer': \"Han Solo's rescue from Jabba the Hutt's palace involves a team consisting of Jedi Knight Luke Skywalker, Princess Leia Organa (disguised as the bounty hunter Boushh), Chewbacca, Lando Calrissian, R2-D2, and C-3PO. The palace where Han is held captive is located on the desert planet of Tatooine. Jabba's palace is a large and imposing structure situated in the Dune Sea of Tatooine, serving as the crime lord's headquarters and residence.\"}"
       ]
      },
-     "execution_count": 25,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -858,10 +860,10 @@
     "# full_chain = create_retrieval_chain(script_retriever_chain, document_chain)\n",
     "full_chain = create_retrieval_chain(woo_retriever_chain, document_chain)\n",
     "\n",
-    "full_chain.invoke(dict(\n",
-    "    # chat_history = [HumanMessage(content = 'Are you knowledgeable about Star Wars?'), AIMessage(content = 'Very')],\n",
-    "    input = \"Who participates in Han's rescue from Jabba? And where is the palace?\"\n",
-    "))"
+    "# full_chain.invoke(dict(\n",
+    "#     # chat_history = [HumanMessage(content = 'Are you knowledgeable about Star Wars?'), AIMessage(content = 'Very')],\n",
+    "#     input = \"Who participates in Han's rescue from Jabba? And where is the palace?\"\n",
+    "# ))"
    ]
   },
   {
@@ -870,15 +872,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "simplify_query_prompt = ChatPromptTemplate.from_messages([\n",
-    "    ('system', 'Given the above conversation, generate a search query to find a relevant page in the Star Wars fandom wiki; the query should be something simple, at most 4 words, such as the name of a character, place, event, item, etc.'),\n",
-    "    MessagesPlaceholder('chat_history', optional = True), # Using this form since not clear how to have optional = True in the tuple form\n",
-    "    ('human', '{query}')\n",
-    "])\n",
+    "# simplify_query_prompt = ChatPromptTemplate.from_messages([\n",
+    "#     ('system', 'Given the above conversation, generate a search query to find a relevant page in the Star Wars fandom wiki; the query should be something simple, at most 4 words, such as the name of a character, place, event, item, etc.'),\n",
+    "#     MessagesPlaceholder('chat_history', optional = True), # Using this form since not clear how to have optional = True in the tuple form\n",
+    "#     ('human', '{query}')\n",
+    "# ])\n",
     "\n",
-    "simplify_query_chain = simplify_query_prompt | llm | StrOutputParser() # To extract just the message\n",
+    "# simplify_query_chain = simplify_query_prompt | llm | StrOutputParser() # To extract just the message\n",
     "\n",
-    "# simplify_query_chain.invoke(dict(context = 'You are an expert of Star Wars lore', query = 'Do you know the tragedy of Darth Plagueis the Wise?'))"
+    "# # simplify_query_chain.invoke(dict(context = 'You are an expert of Star Wars lore', query = 'Do you know the tragedy of Darth Plagueis the Wise?'))"
    ]
   },
   {
@@ -890,7 +892,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -916,7 +918,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
@@ -927,36 +929,10 @@
       "\n",
       "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
       "\u001b[32;1m\u001b[1;3m\n",
-      "Invoking: `search_film_scripts` with `{'query': 'tragedy of Darth Plagueis the Wise'}`\n",
+      "Invoking: `search_film_scripts` with `{'query': 'Darth Plagueis the Wise'}`\n",
       "\n",
       "\n",
-      "\u001b[0m\u001b[36;1m\u001b[1;3mIn a few minutes the entire passageway is ablaze with \n",
-      "laserfire. The deadly bolts ricochet in wild random patterns \n",
-      "creating huge explosions. Stormtroopers scatter and duck \n",
-      "behind storage lockers. Laserbolts hit several Rebel soldiers \n",
-      "who scream and stagger through the smoke, holding shattered \n",
-      "arms and faces.\n",
-      "\n",
-      "An explosion hits near the robots.\n",
-      "\n",
-      "THREEPIO\n",
-      "I should have known better than to \n",
-      "trust the logic of a half-sized \n",
-      "thermocapsulary dehousing assister...\n",
-      "\n",
-      "Artoo counters with an angry rebuttal as the battle rages \n",
-      "around the two hapless robots.\n",
-      "\n",
-      "EXT. TATOOINE - DESERT WASTELAND - DAY\n",
-      "\n",
-      "A death-white wasteland stretches from horizon to horizon. \n",
-      "The tremendous heat of two huge twin suns settle on a lone \n",
-      "figure, Luke Skywalker, a farm boy with heroic aspirations \n",
-      "who looks much younger than his eighteen years. His shaggy \n",
-      "hair and baggy tunic give him the air of a simple but lovable \n",
-      "lad with a prize-winning smile.\n",
-      "\n",
-      "OBI-WAN: (continuing) It can't be . . . It can't be . . .\n",
+      "\u001b[0m\u001b[36;1m\u001b[1;3mOBI-WAN: (continuing) It can't be . . . It can't be . . .\n",
       "\n",
       "As ANAKIN surveys the carnage, a DARK-ROBED SITH LORD enters. ANAKIN turns to DARTH SIDIOUS and kneels before him.\n",
       "\n",
@@ -1016,25 +992,55 @@
       "his sword from Vader's then watches his opponent with a serene \n",
       "look on his face.\n",
       "\n",
-      "The towering creature brings down his curved, double-pointed \n",
-      "gaderffii -- the dreaded axe blade that has struck terror in \n",
-      "the heart of the local settlers. But Luke manages to block \n",
-      "the blow with his laser rifle, which is smashed to pieces. \n",
-      "The terrified farm boy scrambles backward until he is forced \n",
-      "to the edge of a deep crevice. The sinister Raider stands \n",
-      "over him with his weapon raised and lets out a horrible \n",
-      "shrieking laugh.\n",
+      "LUKE\n",
+      "Hurry up! Come with me! What are you \n",
+      "waiting for?! Get in gear!\n",
+      "\n",
+      "The robot scoots around in a tight circle, stops short, and \n",
+      "smoke begins to pour out of every joint. Luke throws his \n",
+      "arms up in disgust. Exasperated, the young farm boy jumps \n",
+      "into his Landspeeder leaving the smoldering robot to hum \n",
+      "madly.\n",
       "\n",
-      "EXT. TATOOINE - ROCK CANYON - FLOOR - DAY\n",
+      "INT. REBEL BLOCKADE RUNNER - MAIN HALLWAY\n",
       "\n",
-      "Artoo forces himself into the shadows of a small alcove in \n",
-      "the rocks as the vicious Sandpeople walk past carrying the \n",
-      "inert Luke Skywalker, who is dropped in a heap before the \n",
-      "speeder.The Sandpeople ransack the speeder, throwing parts \n",
-      "and supplies in all directions. Suddenly they stop. Then \n",
-      "everything is quiet for a few moments. A great howling moan \n",
-      "is heard echoing throughout the canyon which sends the \n",
-      "Sandpeople fleeing in terror.\u001b[0m\u001b[32;1m\u001b[1;3mThe dialogue from the Star Wars film scripts does not directly mention the tragedy of Darth Plagueis the Wise. However, in \"Star Wars: Episode III - Revenge of the Sith,\" Chancellor Palpatine tells Anakin Skywalker the story of Darth Plagueis the Wise. According to Palpatine, Darth Plagueis was a Dark Lord of the Sith who was so powerful and wise that he could influence the midi-chlorians to create life and prevent death. The tragedy of Darth Plagueis the Wise is that he was eventually betrayed and killed by his own apprentice, who was seeking to obtain his power. This story plays a significant role in Anakin's fall to the dark side as he becomes intrigued by the idea of cheating death, which ultimately leads him to become Darth Vader.\u001b[0m\n",
+      "The awesome, seven-foot-tall Dark Lord of the Sith makes his \n",
+      "way into the blinding light of the main passageway. This is \n",
+      "Darth Vader, right hand of the Emperor. His face is obscured \n",
+      "by his flowing black robes and grotesque breath mask, which \n",
+      "stands out next to the fascist white armored suits of the \n",
+      "Imperial stormtroopers. Everyone instinctively backs away \n",
+      "from the imposing warrior and a deathly quiet sweeps through \n",
+      "the Rebel troops. Several of the Rebel troops break and run \n",
+      "in a frenzied panic.\n",
+      "\n",
+      "INT. REBEL BLOCKADE RUNNER\n",
+      "\n",
+      "A woman's hand puts a card into an opening in Artoo's dome.\n",
+      "Artoo makes beeping sounds.\n",
+      "\n",
+      "47INT DEATH STAR - CORRIDOR TO DOCKING BAY\n",
+      "\n",
+      "Lord Vader strides down the hallway, accompanied by a very nervous \n",
+      "Death Star commander.\n",
+      "\n",
+      "48INT DOCKING BAY - DEATH STAR\n",
+      "\n",
+      "Thousands of Imperial troops in tight formation fill the mammoth \n",
+      "docking bay. Vader and the officer walk to the landing platform, where \n",
+      "the shuttle is coming to rest.\n",
+      "\n",
+      "The shuttle's ramp lowers and the Emperor's Royal Guards come out and \n",
+      "create a lethal perimeter. The assembled troops move to rigid attention \n",
+      "with a momentous SNAP.\n",
+      "\n",
+      "Then, in the huge SILENCE which follows, the EMPEROR appears. He is a \n",
+      "rather small, shriveled old man. His bent frame slowly makes its way \n",
+      "down the ramp with the aid of a gnarled cane. He wears a hooded cloak \n",
+      "similar to the one Ben wears, except that it is black. The Emperor's \n",
+      "face is shrouded and difficult to see, except for his piercing yellow \n",
+      "eyes. Commander Jerjerrod and Darth Vader kneel to him. The Supreme \n",
+      "Ruler of the galaxy beckons to the Dark Lord.\u001b[0m\u001b[32;1m\u001b[1;3mThe dialogue from the Star Wars film scripts does not directly mention the tragedy of Darth Plagueis the Wise. However, in \"Star Wars: Episode III - Revenge of the Sith,\" Chancellor Palpatine tells Anakin Skywalker the story of Darth Plagueis the Wise. According to Palpatine, Darth Plagueis was a Dark Lord of the Sith who was so powerful and wise that he could influence the midi-chlorians to create life and prevent death. Unfortunately, Darth Plagueis was betrayed and killed by his own apprentice, who was seeking to obtain his power. This story plays a significant role in Anakin's fall to the dark side as he becomes intrigued by the idea of cheating death, ultimately leading him to become Darth Vader under the influence of Palpatine.\u001b[0m\n",
       "\n",
       "\u001b[1m> Finished chain.\u001b[0m\n"
      ]
@@ -1045,10 +1051,10 @@
        "{'chat_history': [HumanMessage(content='Are you knowledgeable about Star Wars?'),\n",
        "  AIMessage(content='Very')],\n",
        " 'input': 'Do you know the tragedy of Darth Plagueis the Wise?',\n",
-       " 'output': 'The dialogue from the Star Wars film scripts does not directly mention the tragedy of Darth Plagueis the Wise. However, in \"Star Wars: Episode III - Revenge of the Sith,\" Chancellor Palpatine tells Anakin Skywalker the story of Darth Plagueis the Wise. According to Palpatine, Darth Plagueis was a Dark Lord of the Sith who was so powerful and wise that he could influence the midi-chlorians to create life and prevent death. The tragedy of Darth Plagueis the Wise is that he was eventually betrayed and killed by his own apprentice, who was seeking to obtain his power. This story plays a significant role in Anakin\\'s fall to the dark side as he becomes intrigued by the idea of cheating death, which ultimately leads him to become Darth Vader.'}"
+       " 'output': 'The dialogue from the Star Wars film scripts does not directly mention the tragedy of Darth Plagueis the Wise. However, in \"Star Wars: Episode III - Revenge of the Sith,\" Chancellor Palpatine tells Anakin Skywalker the story of Darth Plagueis the Wise. According to Palpatine, Darth Plagueis was a Dark Lord of the Sith who was so powerful and wise that he could influence the midi-chlorians to create life and prevent death. Unfortunately, Darth Plagueis was betrayed and killed by his own apprentice, who was seeking to obtain his power. This story plays a significant role in Anakin\\'s fall to the dark side as he becomes intrigued by the idea of cheating death, ultimately leading him to become Darth Vader under the influence of Palpatine.'}"
       ]
      },
-     "execution_count": 28,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1079,10 +1085,97 @@
     "#     input = 'Do you know the tragedy of Darth Plagueis the Wise?'\n",
     "# )\n",
     "\n",
-    "agent_executor.invoke(dict(\n",
-    "    chat_history = [HumanMessage(content = 'Are you knowledgeable about Star Wars?'), AIMessage(content = 'Very')],\n",
-    "    input = 'Do you know the tragedy of Darth Plagueis the Wise?'\n",
-    "))"
+    "# agent_executor.invoke(dict(\n",
+    "#     chat_history = [HumanMessage(content = 'Are you knowledgeable about Star Wars?'), AIMessage(content = 'Very')],\n",
+    "#     input = 'Do you know the tragedy of Darth Plagueis the Wise?'\n",
+    "# ))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Non-agent chain-logic version"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Determine which retriever is best and generate an appropriate query for it\n",
+    "\n",
+    "# Again, many LangChain objects seem to use the class-scope fields instead of instance ones in __init__...\n",
+    "# Separately, the advantage of using a class over a dictionary is simply the descriptions\n",
+    "class DirectedQuery(BaseModel):\n",
+    "    '''Determine whether a query is best answered by looking at scripts rather than articles'''\n",
+    "\n",
+    "    query: str = Field(\n",
+    "        ...,\n",
+    "        description = '''The query to either search film scripts or wiki articles.\n",
+    "        A film script query should include character names and relevant keywords of what they are saying in the a scene which is likely to contain the required information.\n",
+    "        A wiki articles search should instead be at most 4 words, simply being the name of a character or location or event whose page is likely to contain the required information.''',\n",
+    "    )\n",
+    "    source: str = Field(\n",
+    "        ...,\n",
+    "        description = 'Either \"wiki\" or \"scripts\", indicating which source the query should be passed to.',\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "# output_parser = PydanticToolsParser(tools = [DirectedQuery])\n",
+    "\n",
+    "query_analyser_prompt = ChatPromptTemplate.from_messages([\n",
+    "        ('system', 'You have the ability to issue search queries of one of two kinds to get information to help answer questions.'),\n",
+    "        ('human', '{question}'),\n",
+    "])\n",
+    "structured_llm = llm.with_structured_output(DirectedQuery)\n",
+    "query_generator = dict(question = RunnablePassthrough()) | query_analyser_prompt | structured_llm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "retrievers = dict(wiki = woo_db.as_retriever(search_kwargs = dict(k = 4)), scripts = script_db.as_retriever(search_kwargs = dict(k = 4)))\n",
+    "\n",
+    "@chain\n",
+    "def compound_retriever(question):\n",
+    "    response = query_generator.invoke(question)\n",
+    "    retriever = retrievers[response.source]\n",
+    "    return retriever.invoke(response.query)\n",
+    "\n",
+    "\n",
+    "# compound_retriever.invoke('Do you know the tragedy of Darth Plagueis the Wise?')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'input': 'Do you know the tragedy of Darth Plagueis the Wise?',\n",
+       " 'context': [Document(page_content='The Tragedy of Darth Plagueis the Wise', metadata={'source': 'https://starwars.fandom.com/wiki/Darth_Plagueis'}),\n",
+       "  Document(page_content='Biography\\nEarly life\\nDarth Plagueis was a legendary Dark Lord of the Sith trained by the Sith Master Darth Tenebrous. At some point during his life, Plagueis acquired the protocol droid 11-4D. During his time as a Sith Lord and studying the Force, Plagueis acquired a vast amount of knowledge about the dark side and its teachings.\\n\\nMaster of Darth Sidious', metadata={'source': 'https://starwars.fandom.com/wiki/Darth_Plagueis'}),\n",
+       "  Document(page_content='\"Darth Plagueis was a Dark Lord of the Sith so powerful and so wise, he could use the Force to influence the midi-chlorians to create life. He had such a knowledge of the dark side, he could even keep the ones he cared about from dying.\"\\n―Sheev Palpatine', metadata={'source': 'https://starwars.fandom.com/wiki/Darth_Plagueis'}),\n",
+       "  Document(page_content='Behind the scenes\\n     Darth Plagueis as depicted in Star Wars Legends  \\nDarth Plagueis was first mentioned on-screen in the 2005 film Star Wars: Episode III Revenge of the Sith, the third and final installment of the Star Wars prequel trilogy. Although Palpatine was never confirmed to be Plagueis\\' apprentice in the movie itself, a link to the official encyclopedia on StarWars.com did refer to Sidious as having been \"trained by Darth Plagueis.\"\\nThe character was created by George Lucas as early as the first draft of Revenge of the Sith—dated April 2003—and possibly earlier. His story was massively expanded upon in the 2012 Star Wars Legends novel Darth Plagueis, written by James Luceno. It notably established Plagueis to be a Muun just as Lucas proposed.', metadata={'source': 'https://starwars.fandom.com/wiki/Darth_Plagueis'})],\n",
+       " 'answer': 'Yes, I am familiar with the tragedy of Darth Plagueis the Wise. Darth Plagueis was a legendary Dark Lord of the Sith who was so powerful and knowledgeable about the dark side of the Force that he could manipulate midi-chlorians to create life and prevent death. However, despite his immense power and wisdom, he was ultimately betrayed and killed by his own apprentice, Darth Sidious (Emperor Palpatine), who later became the Emperor of the Galactic Empire. This tragic story was famously recounted by Palpatine to Anakin Skywalker in \"Star Wars: Episode III Revenge of the Sith.\"'}"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "compound_chain = create_retrieval_chain(compound_retriever, document_chain)\n",
+    "\n",
+    "compound_chain.invoke(dict(input = 'Do you know the tragedy of Darth Plagueis the Wise?'))"
    ]
   }
  ],

serve.py

@@ -13,7 +13,8 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain.chains import create_history_aware_retriever, create_retrieval_chain
 from langchain.tools.retriever import create_retriever_tool
-from langchain_core.runnables import RunnablePassthrough, RunnableParallel
+from langchain_core.runnables import RunnablePassthrough, RunnableParallel, chain
+from langchain_core.pydantic_v1 import BaseModel, Field
 
 # Agents
 from langchain import hub
@@ -25,7 +26,6 @@ from langchain_core.documents import Document
 
 # To serve the app
 from fastapi import FastAPI
-from langchain.pydantic_v1 import BaseModel, Field
 from langchain_core.messages import BaseMessage
 from langserve import add_routes, CustomUserType
 
@@ -90,6 +90,8 @@ def get_wookieepedia_context(original_query: str, simple_query: str, wdb: Chroma
 llm = ChatOpenAI(model = 'gpt-3.5-turbo-0125', temperature = 0)
 
 
+## Base version (only one retriever)
+
 document_prompt_system_text = '''
 You are very knowledgeable about Star Wars and your job is to answer questions about its plot, characters, etc.
 Use the context below to produce your answers with as much detail as possible.
@@ -129,13 +131,13 @@ full_chain = create_retrieval_chain(woo_retriever_chain, document_chain)
 
 
 
-simplify_query_prompt = ChatPromptTemplate.from_messages([
-    ('system', 'Given the above conversation, generate a search query to find a relevant page in the Star Wars fandom wiki; the query should be something simple, at most 4 words, such as the name of a character, place, event, item, etc.'),
-    MessagesPlaceholder('chat_history', optional = True), # Using this form since not clear how to have optional = True in the tuple form
-    ('human', '{query}')
-])
+# simplify_query_prompt = ChatPromptTemplate.from_messages([
+#     ('system', 'Given the above conversation, generate a search query to find a relevant page in the Star Wars fandom wiki; the query should be something simple, at most 4 words, such as the name of a character, place, event, item, etc.'),
+#     MessagesPlaceholder('chat_history', optional = True), # Using this form since not clear how to have optional = True in the tuple form
+#     ('human', '{query}')
+# ])
 
-simplify_query_chain = simplify_query_prompt | llm | StrOutputParser() # To extract just the message
+# simplify_query_chain = simplify_query_prompt | llm | StrOutputParser() # To extract just the message
 
 
 
@@ -171,6 +173,41 @@ agent_executor = AgentExecutor(agent = agent, tools = tools, verbose = True)
 
 
 
+## Non-agent chain-logic version
+
+# Determine which retriever is best and generate an appropriate query for it
+class DirectedQuery(BaseModel):
+    '''Determine whether a query is best answered by looking at scripts rather than articles'''
+
+    query: str = Field(
+        ...,
+        description = '''The query to either search film scripts or wiki articles.
+        A film script query should include character names and relevant keywords of what they are saying in the a scene which is likely to contain the required information.
+        A wiki articles search should instead be at most 4 words, simply being the name of a character or location or event whose page is likely to contain the required information.''',
+    )
+    source: str = Field(
+        ...,
+        description = 'Either "wiki" or "scripts", indicating which source the query should be passed to.',
+    )
+query_analyser_prompt = ChatPromptTemplate.from_messages([
+        ('system', 'You have the ability to issue search queries of one of two kinds to get information to help answer questions.'),
+        ('human', '{question}'),
+])
+structured_llm = llm.with_structured_output(DirectedQuery)
+query_generator = dict(question = RunnablePassthrough()) | query_analyser_prompt | structured_llm
+
+retrievers = dict(wiki = woo_db.as_retriever(search_kwargs = dict(k = 4)), scripts = script_db.as_retriever(search_kwargs = dict(k = 4)))
+
+@chain
+def compound_retriever(question):
+    response = query_generator.invoke(question)
+    retriever = retrievers[response.source]
+    return retriever.invoke(response.query)
+
+compound_chain = create_retrieval_chain(compound_retriever, document_chain)
+
+
+
 ## Type specifications (with unusual class-scope fields)
 
 class StrInput(BaseModel):
@@ -198,13 +235,15 @@ app = FastAPI(
 
 
 # add_routes(app, script_db.as_retriever())
-add_routes(app, full_chain.with_types(input_type = StrInput, output_type = Output), playground_type = 'default')
+# add_routes(app, full_chain.with_types(input_type = StrInput, output_type = Output), playground_type = 'default')
 
 # NOTE: The chat playground type has a web page issue (flashes and becomes white, hence non-interactable; this was supposedly solved in an issue late last year)
 
 # add_routes(app, agent_executor, playground_type = 'chat')
 # add_routes(app, agent_executor.with_types(input_type = StrInput, output_type = Output))
 
+add_routes(app, compound_chain.with_types(input_type = StrInput))
+
 
 
 if __name__ == '__main__':