query-pipelines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Query Pipelines Llamaindex\n",
    "\n",
    "LlamaIndex provides a declarative query API that allows you to chain together different modules in order to orchestrate simple-to-advanced workflows over your data.\n",
    "\n",
    "This is centered around our QueryPipeline abstraction. Load in a variety of modules (from LLMs to prompts to retrievers to other pipelines), connect them all together into a sequential chain or DAG, and run it end2end.\n",
    "\n",
    "So what are the advantages of QueryPipeline?\n",
    "\n",
    "- Express common workflows with fewer lines of code/boilerplate\n",
    "- Greater readability\n",
    "- Greater parity / better integration points with common low-code / no-code solutions (e.g. LangFlow)\n",
    "- [In the future] A declarative interface allows easy serializability of pipeline components, providing portability of pipelines/easier deployment to different systems."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing the libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_pipeline.query import QueryPipeline\n",
    "from llama_index.llms import OpenAI\n",
    "from llama_index.prompts import PromptTemplate\n",
    "from llama_index.storage import StorageContext\n",
    "from llama_index import (\n",
    "    VectorStoreIndex,\n",
    "    ServiceContext,\n",
    "    SimpleDirectoryReader,\n",
    "    load_index_from_storage,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the API Keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "# Load the enviroment variables\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 1 docs\n"
     ]
    }
   ],
   "source": [
    "reader = SimpleDirectoryReader(\n",
    "    input_files=[\"./data/paul_graham_essay.txt\"]\n",
    ")\n",
    "docs = reader.load_data()\n",
    "print(f\"Loaded {len(docs)} docs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create or rebuild storage and the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if not os.path.exists(\"storage\"):\n",
    "    index = VectorStoreIndex.from_documents(docs)\n",
    "    # save index to disk\n",
    "    index.set_index_id(\"vector_index\")\n",
    "    index.storage_context.persist(\"./storage\")\n",
    "else:\n",
    "    # rebuild storage context\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"storage\")\n",
    "    # load index\n",
    "    index = load_index_from_storage(storage_context, index_id=\"vector_index\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chain Together Prompt and LLM\n",
    "\n",
    "In this section we show a super simple workflow of chaining together a prompt with LLM.\n",
    "\n",
    "We simply define chain on initialization. This is a special case of a query pipeline where the components are purely sequential, and we automatically convert outputs into the right format for the next inputs.\n",
    "\n",
    "#### Defining a Sequential Chain\n",
    "\n",
    "Some simple pipelines are purely linear in nature - the output of the previous module directly goes into the input of the next module.\n",
    "\n",
    "Some examples:\n",
    "\n",
    "- prompt -> LLM -> output parsing\n",
    "- prompt -> LLM -> prompt -> LLM\n",
    "- retriever -> response synthesizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# try chaining basic prompts\n",
    "prompt_str = \"Please generate related movies to {movie_name}\"\n",
    "prompt_tmpl = PromptTemplate(prompt_str)\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "# Define the query pipeline\n",
    "p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module a3aea295-140f-4182-99e1-d68a8beee6f0 with input: \n",
      "movie_name: The Departed\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module f785a2ed-cbda-4bc6-85a9-51f04986f0b8 with input: \n",
      "messages: Please generate related movies to The Departed\n",
      "\n",
      "\u001b[0massistant: 1. Infernal Affairs (2002) - The Departed is actually a remake of this Hong Kong crime thriller, which follows a similar storyline of undercover cops infiltrating a criminal organization.\n",
      "\n",
      "2. The Town (2010) - Directed by Ben Affleck, this crime drama revolves around a group of bank robbers in Boston and the FBI agent determined to bring them down.\n",
      "\n",
      "3. Heat (1995) - Directed by Michael Mann, this classic crime film features an intense cat-and-mouse game between a skilled detective and a professional thief in Los Angeles.\n",
      "\n",
      "4. American Gangster (2007) - Based on a true story, this crime drama stars Denzel Washington as a Harlem drug lord and Russell Crowe as the detective determined to bring him to justice.\n",
      "\n",
      "5. Training Day (2001) - Denzel Washington won an Academy Award for his role as a corrupt narcotics detective who takes a rookie cop (Ethan Hawke) on a dangerous ride-along through the streets of Los Angeles.\n",
      "\n",
      "6. The Departed (2006) - Although it's the movie you mentioned, it's worth including it in the list for those who haven't seen it. Directed by Martin Scorsese, this crime thriller follows an undercover cop who infiltrates an Irish gang in Boston, while a mole within the police force works to uncover his true identity.\n",
      "\n",
      "7. Donnie Brasco (1997) - Based on a true story, this crime drama stars Johnny Depp as an undercover FBI agent who infiltrates the Mafia, forming a close bond with a mobster played by Al Pacino.\n",
      "\n",
      "8. The Untouchables (1987) - Directed by Brian De Palma, this crime drama tells the story of Eliot Ness (Kevin Costner) and his team of law enforcement agents as they try to bring down Al Capone (Robert De Niro) during the Prohibition era.\n",
      "\n",
      "9. The French Connection (1971) - This gritty crime thriller follows two New York City detectives as they investigate a massive heroin smuggling operation. It won five Academy Awards, including Best Picture.\n",
      "\n",
      "10. Internal Affairs (1990) - Starring Richard Gere and Andy Garcia, this crime thriller explores the corrupt practices within the Los Angeles Police Department as an Internal Affairs officer investigates a fellow officer's suspicious activities.\n"
     ]
    }
   ],
   "source": [
    "output = p.run(movie_name=\"The Departed\")\n",
    "print(str(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chain multiple Prompts with Streaming\n",
    "\n",
    "The query pipelines have LLM streaming support (simply do as_query_component(streaming=True)). Intermediate outputs will get autoconverted, and the final output can be a streaming output. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_str = \"Please generate related movies to {movie_name}\"\n",
    "prompt_tmpl = PromptTemplate(prompt_str)\n",
    "# let's add some subsequent prompts for fun\n",
    "prompt_str2 = \"\"\"\\\n",
    "Here's some text:\n",
    "\n",
    "{text}\n",
    "\n",
    "Can you rewrite this with a summary of each movie?\n",
    "\"\"\"\n",
    "prompt_tmpl2 = PromptTemplate(prompt_str2)\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "llm_c = llm.as_query_component(streaming=True)\n",
    "\n",
    "p = QueryPipeline(\n",
    "    chain=[prompt_tmpl, llm_c, prompt_tmpl2, llm_c], verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module 75059c2a-0656-48ab-a6f1-9c6b21705828 with input: \n",
      "movie_name: The Dark Knight\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 68f4a780-4fe7-485c-8650-48878708a24d with input: \n",
      "messages: Please generate related movies to The Dark Knight\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 4df40a5f-b8f2-4c5c-b2bc-05e21242cfec with input: \n",
      "text: <generator object llm_chat_callback.<locals>.wrap.<locals>.wrapped_llm_chat.<locals>.wrapped_gen at 0x0000025C53B82020>\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 30e10c95-c374-4d9e-b893-f7c1ce3d003d with input: \n",
      "messages: Here's some text:\n",
      "\n",
      "1. Batman Begins (2005)\n",
      "2. The Dark Knight Rises (2012)\n",
      "3. Batman v Superman: Dawn of Justice (2016)\n",
      "4. Man of Steel (2013)\n",
      "5. The Avengers (2012)\n",
      "6. Iron Man (2008)\n",
      "7. Captain Amer...\n",
      "\n",
      "\u001b[0m1. Batman Begins (2005): A young Bruce Wayne becomes Batman to protect Gotham City from the League of Shadows and their leader, Ra's al Ghul.\n",
      "2. The Dark Knight Rises (2012): Batman returns to save Gotham City from the ruthless terrorist Bane, who plans to destroy it.\n",
      "3. Batman v Superman: Dawn of Justice (2016): Batman and Superman clash as they face off against each other, but eventually join forces to stop the villainous Lex Luthor.\n",
      "4. Man of Steel (2013): The origin story of Superman, as he embraces his powers and battles against General Zod, a fellow Kryptonian seeking to destroy Earth.\n",
      "5. The Avengers (2012): Earth's mightiest heroes, including Iron Man, Captain America, and Thor, unite to stop Loki and his alien army from conquering the world.\n",
      "6. Iron Man (2008): Billionaire Tony Stark becomes Iron Man after being kidnapped and builds a high-tech suit to fight against terrorists and protect the innocent.\n",
      "7. Captain America: The Winter Soldier (2014): Captain America teams up with Black Widow and Falcon to uncover a conspiracy within S.H.I.E.L.D. and face the deadly Winter Soldier.\n",
      "8. The Amazing Spider-Man (2012): Peter Parker gains spider-like abilities and becomes Spider-Man, using his powers to battle the Lizard and protect New York City.\n",
      "9. Watchmen (2009): Set in an alternate reality, a group of retired superheroes investigates the murder of one of their own, uncovering a dark conspiracy threatening humanity.\n",
      "10. Sin City (2005): A collection of interconnected stories set in the crime-ridden city of Basin City, featuring corrupt cops, femme fatales, and vigilante justice.\n",
      "11. V for Vendetta (2005): In a dystopian future, a masked vigilante known as V fights against a totalitarian regime and inspires the people to rise up against oppression.\n",
      "12. Blade Runner 2049 (2017): A young blade runner uncovers a long-buried secret that leads him to seek out former blade runner Rick Deckard, while unraveling the mysteries of a future society.\n",
      "13. Inception (2010): A skilled thief enters people's dreams to steal information, but is tasked with planting an idea instead, leading to a complex and mind-bending heist.\n",
      "14. The Matrix (1999): A computer hacker discovers the truth about reality, joining a group of rebels fighting against sentient machines that have enslaved humanity.\n",
      "15. The Crow (1994): A murdered musician is resurrected by a supernatural crow, seeking vengeance against those who killed him and his fiancée."
     ]
    }
   ],
   "source": [
    "output = p.run(movie_name=\"The Dark Knight\")\n",
    "for o in output:\n",
    "    print(o.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chain Together Query Rewriting Workflow (prompts + LLM) with Retrieval\n",
    "\n",
    "Here we try a slightly more complex workflow where we send the input through two prompts before initiating retrieval.\n",
    "\n",
    "- Generate question about given topic.\n",
    "\n",
    "- Hallucinate answer given question, for better retrieval.\n",
    "\n",
    "Since each prompt only takes in one input, note that the QueryPipeline will automatically chain LLM outputs into the prompt and then into the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor import CohereRerank\n",
    "\n",
    "# First prompt: generate question regarding topic\n",
    "prompt_str1 = \"Please generate a concise question about Paul Graham's life regarding the following topic {topic}\"\n",
    "prompt_tmpl1 = PromptTemplate(prompt_str1)\n",
    "# Second prompt: use HyDE to hallucinate answer.\n",
    "prompt_str2 = (\n",
    "    \"Please write a passage to answer the question\\n\"\n",
    "    \"Try to include as many key details as possible.\\n\"\n",
    "    \"\\n\"\n",
    "    \"\\n\"\n",
    "    \"{query_str}\\n\"\n",
    "    \"\\n\"\n",
    "    \"\\n\"\n",
    "    'Passage:\"\"\"\\n'\n",
    ")\n",
    "prompt_tmpl2 = PromptTemplate(prompt_str2)\n",
    "# Define the LLM and the retriever\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "retriever = index.as_retriever(similarity_top_k=5)\n",
    "# Build the pipeline\n",
    "p = QueryPipeline(\n",
    "    chain=[prompt_tmpl1, llm, prompt_tmpl2, llm, retriever], verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module 6adb7b56-54f3-4564-a6a7-8d512ee8ae34 with input: \n",
      "topic: college\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 57bd8005-ea8e-40c6-9ea5-e714cc166401 with input: \n",
      "messages: Please generate a concise question about Paul Graham's life regarding the following topic college\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module a901183d-16ab-4786-aace-15bb59d49844 with input: \n",
      "query_str: assistant: How did Paul Graham's college experience shape his career and entrepreneurial mindset?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module af775f22-0342-4930-89fc-6a14dc37dbc0 with input: \n",
      "messages: Please write a passage to answer the question\n",
      "Try to include as many key details as possible.\n",
      "\n",
      "\n",
      "How did Paul Graham's college experience shape his career and entrepreneurial mindset?\n",
      "\n",
      "\n",
      "Passage:\"\"\"\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module e68186fe-29c2-4fc4-929b-731d4125fc49 with input: \n",
      "input: assistant: Paul Graham's college experience played a pivotal role in shaping his career and entrepreneurial mindset. As a student at Cornell University, Graham immersed himself in the world of compute...\n",
      "\n",
      "\u001b[0m"
     ]
    },
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nodes = p.run(topic=\"college\")\n",
    "len(nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Full RAG Pipeline as a DAG\n",
    "\n",
    "Here we chain together a full RAG pipeline consisting of query rewriting, retrieval, reranking, and response synthesis.\n",
    "\n",
    "Here we can’t use chain syntax because certain modules depend on multiple inputs (for instance, response synthesis expects both the retrieved nodes and the original question). Instead we’ll construct a DAG explicitly, through add_modules and then add_link."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RAG Pipeline with Query Rewriting\n",
    "\n",
    "We use an LLM to rewrite the query first before passing it to our downstream modules - retrieval/reranking/synthesis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor import CohereRerank\n",
    "from llama_index.response_synthesizers import TreeSummarize\n",
    "from llama_index import ServiceContext"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from llama_index import download_loader\n",
    "\n",
    "PDFReader = download_loader(\"PDFReader\")\n",
    "\n",
    "loader = PDFReader()\n",
    "docs = loader.load_data(file=Path('./data/Attention is all you need.pdf'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if not os.path.exists(\"storage\"):\n",
    "    index = VectorStoreIndex.from_documents(docs)\n",
    "    # save index to disk\n",
    "    index.set_index_id(\"vector_index\")\n",
    "    index.storage_context.persist(\"./storage\")\n",
    "else:\n",
    "    # rebuild storage context\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"storage\")\n",
    "    # load index\n",
    "    index = load_index_from_storage(storage_context, index_id=\"vector_index\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define modules\n",
    "prompt_str = \"Please generate a question about the Transformer model regarding the following topic {topic}\"\n",
    "prompt_tmpl = PromptTemplate(prompt_str)\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "retriever = index.as_retriever(similarity_top_k=3)\n",
    "reranker = CohereRerank()\n",
    "summarizer = TreeSummarize(\n",
    "    service_context=ServiceContext.from_defaults(llm=llm)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the query pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define query pipeline\n",
    "p = QueryPipeline(verbose=True)\n",
    "# Adding the modules to the pipeline\n",
    "p.add_modules(\n",
    "    {\n",
    "        \"llm\": llm,\n",
    "        \"prompt_tmpl\": prompt_tmpl,\n",
    "        \"retriever\": retriever,\n",
    "        \"summarizer\": summarizer,\n",
    "        \"reranker\": reranker,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we draw links between modules with add_link. add_link takes in the source/destination module ids, and optionally the source_key and dest_key. Specify the source_key or dest_key if there are multiple outputs/inputs respectively.\n",
    "\n",
    "You can view the set of input/output keys for each module through module.as_query_component().input_keys and module.as_query_component().output_keys.\n",
    "\n",
    "Here we explicitly specify dest_key for the reranker and summarizer modules because they take in two inputs (query_str and nodes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "required_keys={'query_str', 'nodes'} optional_keys=set()\n"
     ]
    }
   ],
   "source": [
    "#Define the links between modules\n",
    "p.add_link(\"prompt_tmpl\", \"llm\")\n",
    "p.add_link(\"llm\", \"retriever\")\n",
    "p.add_link(\"retriever\", \"reranker\", dest_key=\"nodes\")\n",
    "p.add_link(\"llm\", \"reranker\", dest_key=\"query_str\")\n",
    "p.add_link(\"reranker\", \"summarizer\", dest_key=\"nodes\")\n",
    "p.add_link(\"llm\", \"summarizer\", dest_key=\"query_str\")\n",
    "\n",
    "# look at summarizer input keys\n",
    "print(summarizer.as_query_component().input_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module prompt_tmpl with input: \n",
      "topic: Positional encoding\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
      "messages: Please generate a question about the Transformer model regarding the following topic Positional encoding\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
      "input: assistant: How does the positional encoding in the Transformer model help capture the sequential information in the input data?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module reranker with input: \n",
      "query_str: assistant: How does the positional encoding in the Transformer model help capture the sequential information in the input data?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='547e0749-a842-4d59-89df-850726120f2c', embedding=None, metadata={'page_label': '2', 'file_name': 'Attention is all you need.pdf'}, excluded_embed_metadata_keys=[], ex...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
      "query_str: assistant: How does the positional encoding in the Transformer model help capture the sequential information in the input data?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='c9685c43-7c69-4da5-bd4a-e8314bc52d32', embedding=None, metadata={'page_label': '5', 'file_name': 'Attention is all you need.pdf'}, excluded_embed_metadata_keys=[], ex...\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "# Run the pipeline\n",
    "response = p.run(topic=\"Positional encoding\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The positional encoding in the Transformer model helps capture the sequential information in the input data by injecting information about the relative or absolute position of each token in the sequence. Since the Transformer model does not have any recurrence or convolution, the positional encoding is necessary for the model to understand the order of the sequence. This allows the model to differentiate between tokens at different positions and capture the sequential relationships between them.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can do async too"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running modules and inputs in parallel: \n",
      "Module key: prompt_tmpl. Input: \n",
      "topic: PÑositional encoding\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running modules and inputs in parallel: \n",
      "Module key: llm. Input: \n",
      "messages: Please generate a question about the Transformer model regarding the following topic PÑositional encoding\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running modules and inputs in parallel: \n",
      "Module key: retriever. Input: \n",
      "input: assistant: What is the purpose of positional encoding in the Transformer model and how does it help in capturing the sequential information of input tokens?\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running modules and inputs in parallel: \n",
      "Module key: reranker. Input: \n",
      "query_str: assistant: What is the purpose of positional encoding in the Transformer model and how does it help in capturing the sequential information of input tokens?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='547e0749-a842-4d59-89df-850726120f2c', embedding=None, metadata={'page_label': '2', 'file_name': 'Attention is all you need.pdf'}, excluded_embed_metadata_keys=[], ex...\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running modules and inputs in parallel: \n",
      "Module key: summarizer. Input: \n",
      "query_str: assistant: What is the purpose of positional encoding in the Transformer model and how does it help in capturing the sequential information of input tokens?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='c9685c43-7c69-4da5-bd4a-e8314bc52d32', embedding=None, metadata={'page_label': '5', 'file_name': 'Attention is all you need.pdf'}, excluded_embed_metadata_keys=[], ex...\n",
      "\n",
      "\n",
      "\u001b[0mThe purpose of positional encoding in the Transformer model is to inject information about the relative or absolute position of the input tokens. Since the Transformer model does not have any recurrence or convolution, positional encoding is necessary for the model to make use of the order of the sequence. \n",
      "\n",
      "Positional encoding helps in capturing the sequential information of input tokens by adding positional information to the input embeddings. This allows the model to differentiate between tokens based on their position in the sequence. By encoding the position information, the model can learn to attend to different positions in the input sequence and capture the dependencies between tokens that are important for understanding the sequential nature of the data.\n"
     ]
    }
   ],
   "source": [
    "response = await p.arun(topic=\"Positional encoding\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RAG Pipeline without Query Rewriting\n",
    "\n",
    "Here we setup a RAG pipeline without the query rewriting step.\n",
    "\n",
    "Here we need a way to link the input query to both the retriever, reranker, and summarizer. We can do this by defining a special InputComponent, allowing us to link the inputs to multiple downstream modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor import CohereRerank\n",
    "from llama_index.response_synthesizers import TreeSummarize\n",
    "from llama_index import ServiceContext\n",
    "from llama_index.query_pipeline import InputComponent\n",
    "\n",
    "retriever = index.as_retriever(similarity_top_k=5)\n",
    "summarizer = TreeSummarize(\n",
    "    service_context=ServiceContext.from_defaults(\n",
    "        llm=OpenAI(model=\"gpt-3.5-turbo\")\n",
    "    )\n",
    ")\n",
    "reranker = CohereRerank()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = QueryPipeline(verbose=True)\n",
    "p.add_modules(\n",
    "    {\n",
    "        \"input\": InputComponent(),\n",
    "        \"retriever\": retriever,\n",
    "        \"reranker\": reranker,\n",
    "        \"summarizer\": summarizer,\n",
    "    }\n",
    ")\n",
    "p.add_link(\"input\", \"retriever\")\n",
    "p.add_link(\"retriever\", \"reranker\", dest_key=\"nodes\")\n",
    "p.add_link(\"input\", \"reranker\", dest_key=\"query_str\")\n",
    "p.add_link(\"input\", \"summarizer\", dest_key=\"query_str\")\n",
    "p.add_link(\"reranker\", \"summarizer\", dest_key=\"nodes\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module input with input: \n",
      "input: what is the purpose of positional encoding in the Transformer architecture?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
      "input: what is the purpose of positional encoding in the Transformer architecture?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module reranker with input: \n",
      "query_str: what is the purpose of positional encoding in the Transformer architecture?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='f14f266b-2905-49d8-94ef-f2f94e7f421f', embedding=None, metadata={'page_label': '3', 'file_name': 'Attention is all you need.pdf'}, excluded_embed_metadata_keys=[], ex...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
      "query_str: what is the purpose of positional encoding in the Transformer architecture?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='f14f266b-2905-49d8-94ef-f2f94e7f421f', embedding=None, metadata={'page_label': '3', 'file_name': 'Attention is all you need.pdf'}, excluded_embed_metadata_keys=[], ex...\n",
      "\n",
      "\u001b[0mThe purpose of positional encoding in the Transformer architecture is to provide the model with information about the relative or absolute position of the tokens in the sequence. Since the Transformer model does not have any recurrence or convolution, positional encoding is necessary for the model to understand the order of the sequence. This allows the model to effectively process and generate sequences of variable length.\n"
     ]
    }
   ],
   "source": [
    "output = p.run(input=\"what is the purpose of positional encoding in the Transformer architecture?\")\n",
    "print(str(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining a RAG pipeline with Sentence Window Retrieval"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we'll apply sentence window retrieval to our RAG pipeline. The sentence window is a post-processing technique we have to include in the right position of the pipeline. \n",
    "\n",
    "This technique parses documents into single sentences per node. Each node also contains a \"window\", a richer context, with the sentences on either side of the node sentence. During retrieval, before passing the retrieved sentences to the LLM, the single sentences are replaced with a window containing the surrounding sentences, the richer context. This is most useful for large documents/indexes, as it helps to retrieve more fine-grained details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we define a new Service Context to apply the Sentence Window parser and the index creation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.node_parser import SentenceWindowNodeParser\n",
    "from typing import List\n",
    "import pathlib\n",
    "\n",
    "def sentence_window_context(embed_model, window_size: int = 3, chunk_size: int=256, chunk_overlap: int = 0):\n",
    "        # create the sentence window node parser w/ default settings\n",
    "        node_parser = SentenceWindowNodeParser.from_defaults(\n",
    "            #sentence_splitter= text_splitter,\n",
    "            window_size=window_size,\n",
    "            window_metadata_key=\"window\",\n",
    "            original_text_metadata_key=\"original_text\",\n",
    "            include_metadata= False\n",
    "        )        \n",
    "        # Define a context\n",
    "        print(\"Defining the Service Context\")\n",
    "        service_context = ServiceContext.from_defaults(\n",
    "                embed_model=embed_model,\n",
    "                # embed_model=OpenAIEmbedding(model=\"text-embedding-3-small\",)\n",
    "                node_parser=node_parser)\n",
    "    \n",
    "        return service_context\n",
    "\n",
    "def create_or_rebuild_index(docs: List, service_context, storage_path: str = \"storage\"):\n",
    "\n",
    "    pathlib.Path().resolve()\n",
    "    if not os.path.exists(storage_path):\n",
    "        # Build the index acording to the service\n",
    "        print(\"Creating the index\")\n",
    "        index = VectorStoreIndex.from_documents(docs, service_context=service_context)\n",
    "        # save index to disk\n",
    "        index.set_index_id(\"vector_index\")\n",
    "        print(\"Saving the index\")\n",
    "        index.storage_context.persist(os.path.join(pathlib.Path().resolve(), storage_path))\n",
    "    else:\n",
    "        print(\"Restoring the storage\")\n",
    "        # rebuild storage context\n",
    "        storage_context = StorageContext.from_defaults(persist_dir=storage_path)\n",
    "        # load index\n",
    "        print(\"Loading the index\")\n",
    "        index = load_index_from_storage(storage_context, index_id=\"vector_index\")    \n",
    "    \n",
    "    return index\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can start defining our embeddings and creating the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defining the Service Context\n",
      "Creating the index\n",
      "Saving the index\n"
     ]
    }
   ],
   "source": [
    "from llama_index.embeddings import OpenAIEmbedding\n",
    "import os\n",
    "\n",
    "embed_model_name=\"text-embedding-3-small\"\n",
    "embed_model= OpenAIEmbedding(model=embed_model_name)\n",
    "# Create the service context\n",
    "service_context = sentence_window_context(embed_model, window_size=3)\n",
    "# Create or rebuild the index\n",
    "index = create_or_rebuild_index(docs, service_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we need a way to link the output of the retriever to our post-processor to extract the whole window and then inject it to the summarizer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.postprocessor import MetadataReplacementPostProcessor\n",
    "from llama_index.query_pipeline import InputComponent\n",
    "from llama_index.response_synthesizers import TreeSummarize\n",
    "\n",
    "# Set the retriever\n",
    "retriever = index.as_retriever(similarity_top_k=3)\n",
    "# Define the summarizer\n",
    "summarizer = TreeSummarize(\n",
    "        service_context=ServiceContext.from_defaults(\n",
    "            llm=OpenAI(model=\"gpt-3.5-turbo\",\n",
    "                       temperature=0.2,\n",
    "                       max_tokens=512)\n",
    "        )\n",
    ")\n",
    "# Define the reranker\n",
    "postprocessor = MetadataReplacementPostProcessor(\n",
    "                                target_metadata_key=\"window\"\n",
    "                    )\n",
    "# Define the query pipeline\n",
    "p = QueryPipeline(verbose=True)\n",
    "p.add_modules(\n",
    "        {\n",
    "            \"input\": InputComponent(),\n",
    "            \"retriever\": retriever,\n",
    "            \"postprocessor\": postprocessor,\n",
    "            \"summarizer\": summarizer,\n",
    "        }\n",
    ")\n",
    "    # Set the links between components\n",
    "p.add_link(\"input\", \"retriever\")\n",
    "p.add_link(\"input\", \"postprocessor\", dest_key=\"query_str\")\n",
    "p.add_link(\"retriever\", \"postprocessor\", dest_key=\"nodes\")\n",
    "p.add_link(\"input\", \"summarizer\", dest_key=\"query_str\")\n",
    "p.add_link(\"postprocessor\", \"summarizer\", dest_key=\"nodes\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module input with input: \n",
      "input: what is the purpose of positional encoding in the Transformer architecture?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
      "input: what is the purpose of positional encoding in the Transformer architecture?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module postprocessor with input: \n",
      "query_str: what is the purpose of positional encoding in the Transformer architecture?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='d6a7afad-0fa4-434f-a842-24fd306ab9d0', embedding=None, metadata={'window': 'We also use the usual learned linear transfor-\\nmation and softmax function to convert the...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
      "query_str: what is the purpose of positional encoding in the Transformer architecture?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='d6a7afad-0fa4-434f-a842-24fd306ab9d0', embedding=None, metadata={'window': 'We also use the usual learned linear transfor-\\nmation and softmax function to convert the...\n",
      "\n",
      "\u001b[0mThe purpose of positional encoding in the Transformer architecture is to inject information about the relative or absolute position of tokens in the input sequence. This is necessary because the Transformer model does not have any recurrence or convolution, so positional encoding helps the model make use of the order of the sequence. The positional encodings, which have the same dimension as the embeddings, are added to the input embeddings at the bottoms of the encoder and decoder stacks.\n"
     ]
    }
   ],
   "source": [
    "# Run the pipeline\n",
    "output = p.run(input=\"what is the purpose of positional encoding in the Transformer architecture?\")\n",
    "print(str(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAG pipeline with Auto Merging Retrieval and Rerank"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The idea here is pretty much similar to Sentence Window Retriever — to search for more granular pieces of information and then to extend the context window before feeding said context to an LLM for reasoning. Documents are split into smaller child chunks referring to larger parent chunks.\n",
    "Fetch smaller chunks during retrieval first, then if more than n chunks in top k retrieved chunks are linked to the same parent node (larger chunk), we replace the context fed to the LLM by this parent node — works like auto merging a few retrieved chunks into a larger parent chunk, hence the method name\n",
    "\n",
    "In this notebook, we showcase our `AutoMergingRetriever`, which looks at a set of leaf nodes and recursively “merges” subsets of leaf nodes that reference a parent node beyond a given threshold. This allows us to consolidate potentially disparate, smaller contexts into a larger context that might help synthesis.\n",
    "\n",
    "You can define this hierarchy yourself over a set of documents, or you can make use of our brand-new text parser: a `HierarchicalNodeParser` that takes in a candidate set of documents and outputs an entire hierarchy of nodes, from “coarse-to-fine”.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we need to build the index using the `HierarchicalNodeParser`;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import ServiceContext, VectorStoreIndex, StorageContext, load_index_from_storage\n",
    "from llama_index.query_pipeline.query import QueryPipeline\n",
    "from llama_index.llms import OpenAI\n",
    "from llama_index.embeddings import OpenAIEmbedding\n",
    "\n",
    "from llama_index.node_parser import HierarchicalNodeParser, get_leaf_nodes, get_root_nodes\n",
    "#from llama_index.indices.postprocessor import SentenceTransformerRerank\n",
    "from llama_index.response_synthesizers import Refine\n",
    "from llama_index.retrievers import AutoMergingRetriever\n",
    "from llama_index.query_engine import RetrieverQueryEngine\n",
    "from llama_index.postprocessor import CohereRerank\n",
    "from llama_index import ServiceContext\n",
    "from llama_index.query_pipeline import InputComponent\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "In this section we make use of the HierarchicalNodeParser. This will output a hierarchy of nodes, from top-level nodes with bigger chunk sizes to child nodes with smaller chunk sizes, where each child node has a parent node with a bigger chunk size. By default, the hierarchy is:\n",
    "\n",
    "- 1st level: chunk size 2048\n",
    "- 2nd level: chunk size 512\n",
    "- 3rd level: chunk size 128\n",
    "\n",
    "We then load these nodes into storage. The leaf nodes are indexed and retrieved via a vector store - these are the nodes that will first be directly retrieved via similarity search. The other nodes will be retrieved from a docstore.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "168\n",
      "Recent work has achieved\n",
      "signiﬁcant improvements in computational efﬁciency through factorization tricks [ 21] and conditional\n",
      "computation [ 32], while also improving model performance in case of the latter. The fundamental\n",
      "constraint of sequential computation, however, remains.\n",
      "Attention mechanisms have become an integral part of compelling sequence modeling and transduc-\n",
      "tion models in various tasks, allowing modeling of dependencies without regard to their distance in\n",
      "the input or output sequences [ 2,19].\n",
      "126\n"
     ]
    }
   ],
   "source": [
    "# create the hierarchical node parser w/ default settings\n",
    "node_parser = HierarchicalNodeParser.from_defaults(\n",
    "    chunk_sizes=[2048, 512, 128]\n",
    ")\n",
    "# Get the nodes\n",
    "nodes = node_parser.get_nodes_from_documents(docs)\n",
    "print(len(nodes))\n",
    "# Get the leaf nodes\n",
    "leaf_nodes = get_leaf_nodes(nodes)\n",
    "print(leaf_nodes[10].text)\n",
    "print(len(leaf_nodes))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the embddings model\n",
    "embed_model_name=\"text-embedding-3-small\"\n",
    "embed_model= OpenAIEmbedding(model=embed_model_name)\n",
    "# DEfine the LLM\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)\n",
    "# Create the context definition\n",
    "auto_merging_context = ServiceContext.from_defaults(\n",
    "    llm=llm,\n",
    "    embed_model=embed_model,\n",
    "    node_parser=node_parser,\n",
    ")\n",
    "# Create the storage context\n",
    "storage_context = StorageContext.from_defaults()\n",
    "storage_context.docstore.add_documents(nodes)\n",
    "# Build the Vector index\n",
    "automerging_index = VectorStoreIndex(\n",
    "    leaf_nodes, storage_context=storage_context, service_context=auto_merging_context\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Pipeline\n",
    "\n",
    "Here, we define the retriever and a post processor to rerank the retrieved documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "automerging_retriever = automerging_index.as_retriever(\n",
    "    similarity_top_k=6\n",
    ")\n",
    "\n",
    "retriever = AutoMergingRetriever(\n",
    "    automerging_retriever, \n",
    "    automerging_index.storage_context, \n",
    "    verbose=True\n",
    ")\n",
    "# Reranker\n",
    "reranker = CohereRerank()\n",
    "\n",
    "# Define the summarizer\n",
    "summarizer = Refine(\n",
    "        service_context=ServiceContext.from_defaults(\n",
    "            llm=OpenAI(model=\"gpt-3.5-turbo\",\n",
    "                       temperature=0.2,\n",
    "                       max_tokens=512)\n",
    "        )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the query pipeline\n",
    "p = QueryPipeline(verbose=True)\n",
    "p.add_modules(\n",
    "        {\n",
    "            \"input\": InputComponent(),\n",
    "            \"retriever\": retriever,\n",
    "            \"reranker\": reranker,\n",
    "            \"summarizer\": summarizer,\n",
    "        }\n",
    ")\n",
    "# Set the links between components\n",
    "p.add_link(\"input\", \"retriever\")\n",
    "p.add_link(\"retriever\", \"reranker\", dest_key=\"nodes\")\n",
    "p.add_link(\"input\", \"reranker\", dest_key=\"query_str\")\n",
    "p.add_link(\"input\", \"summarizer\", dest_key=\"query_str\")\n",
    "p.add_link(\"reranker\", \"summarizer\", dest_key=\"nodes\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module input with input: \n",
      "input: What is the purpose of positional encoding? Explain it with details\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
      "input: What is the purpose of positional encoding? Explain it with details\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module reranker with input: \n",
      "query_str: What is the purpose of positional encoding? Explain it with details\n",
      "nodes: [NodeWithScore(node=TextNode(id_='c10a01ef-3fa8-4a78-82b5-4d94d20641fc', embedding=None, metadata={'page_label': '6', 'file_name': 'Attention is all you need.pdf'}, excluded_embed_metadata_keys=[], ex...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
      "query_str: What is the purpose of positional encoding? Explain it with details\n",
      "nodes: [NodeWithScore(node=TextNode(id_='c10a01ef-3fa8-4a78-82b5-4d94d20641fc', embedding=None, metadata={'page_label': '6', 'file_name': 'Attention is all you need.pdf'}, excluded_embed_metadata_keys=[], ex...\n",
      "\n",
      "\u001b[0mPositional encoding is used to provide information about the position of each word or token in a sequence. It is added to the input embeddings at the bottom of the encoder and decoder stacks. The positional encodings have the same dimension as the embeddings, allowing them to be summed together. This helps the model differentiate between words that have the same embedding but appear in different positions within the sequence. By incorporating positional information, the model can better understand the order and structure of the input sequence, which is important for tasks such as machine translation or language understanding.\n"
     ]
    }
   ],
   "source": [
    "# Run the pipeline\n",
    "output = p.run(input=\"What is the purpose of positional encoding? Explain it with details\")\n",
    "print(str(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining a Custom Component in a Query Pipeline\n",
    "\n",
    "You can easily define a custom component. Simply subclass a QueryComponent, implement validation/run functions + some helpers, and plug it in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_pipeline import (\n",
    "    CustomQueryComponent,\n",
    "    InputKeys,\n",
    "    OutputKeys,\n",
    ")\n",
    "from typing import Dict, Any\n",
    "from llama_index.llms.llm import BaseLLM\n",
    "from pydantic import Field\n",
    "\n",
    "\n",
    "class RelatedMovieComponent(CustomQueryComponent):\n",
    "    \"\"\"Related movie component.\"\"\"\n",
    "\n",
    "    llm: BaseLLM = Field(..., description=\"OpenAI LLM\")\n",
    "\n",
    "    def _validate_component_inputs(\n",
    "        self, input: Dict[str, Any]\n",
    "    ) -> Dict[str, Any]:\n",
    "        \"\"\"Validate component inputs during run_component.\"\"\"\n",
    "        # NOTE: this is OPTIONAL but we show you here how to do validation as an example\n",
    "        return input\n",
    "\n",
    "    @property\n",
    "    def _input_keys(self) -> set:\n",
    "        \"\"\"Input keys dict.\"\"\"\n",
    "        # NOTE: These are required inputs. If you have optional inputs please override\n",
    "        # `optional_input_keys_dict`\n",
    "        return {\"movie\"}\n",
    "\n",
    "    @property\n",
    "    def _output_keys(self) -> set:\n",
    "        return {\"output\"}\n",
    "\n",
    "    def _run_component(self, **kwargs) -> Dict[str, Any]:\n",
    "        \"\"\"Run the component.\"\"\"\n",
    "        # use QueryPipeline itself here for convenience\n",
    "        prompt_str = \"Please generate related movies to {movie_name}\"\n",
    "        prompt_tmpl = PromptTemplate(prompt_str)\n",
    "        p = QueryPipeline(chain=[prompt_tmpl, llm])\n",
    "        return {\"output\": p.run(movie_name=kwargs[\"movie\"])}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s try the custom component out! We’ll also add a step to convert the output to Shakespeare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "component = RelatedMovieComponent(llm=llm)\n",
    "\n",
    "# let's add some subsequent prompts for fun\n",
    "prompt_str = \"\"\"\\\n",
    "Here's some text:\n",
    "\n",
    "{text}\n",
    "\n",
    "Can you rewrite this in the voice of Shakespeare?\n",
    "\"\"\"\n",
    "prompt_tmpl = PromptTemplate(prompt_str)\n",
    "\n",
    "p = QueryPipeline(chain=[component, prompt_tmpl, llm], verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module 1ebfa41f-5ac9-441c-b009-b18cf45edec0 with input: \n",
      "movie: Love Actually\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module e9ce341a-1ae8-4379-9d04-790a0ed024f9 with input: \n",
      "text: assistant: 1. \"Valentine's Day\" (2010) - This romantic comedy follows the lives of several interconnected couples and singles in Los Angeles as they navigate love and relationships on Valentine's Day....\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module dd728f26-fcaa-4cad-bd57-fa7ee052936f with input: \n",
      "messages: Here's some text:\n",
      "\n",
      "1. \"Valentine's Day\" (2010) - This romantic comedy follows the lives of several interconnected couples and singles in Los Angeles as they navigate love and relationships on Valentin...\n",
      "\n",
      "\u001b[0massistant: 1. \"Valentine's Daye\" (2010) - Thise romantic comedy doth follow the lives of several interconnected couples and singles in Los Angeles as they doth navigate love and relationships on Valentine's Daye.\n",
      "\n",
      "2. \"New Year's Eve\" (2011) - Similar to \"Love Actually,\" thise film doth tell the story of multiple characters whose lives doth intertwine on New Year's Eve in New York City, exploring themes of love, hope, and second chances.\n",
      "\n",
      "3. \"Crazy, Stupid, Love\" (2011) - Thise movie doth revolve around a middle-aged man who, after his wife doth ask him for a divorce, seeketh the help of a young womanizer to rediscover his confidence and find love again.\n",
      "\n",
      "4. \"The Holiday\" (2006) - In thise heartwarming film, two women from different countries doth swap homes for the holidays and unexpectedly find love in the process. It doth explore themes of self-discovery, second chances, and the magic of the holiday season.\n",
      "\n",
      "5. \"Notting Hill\" (1999) - Starring Hugh Grant and Julia Roberts, thise romantic comedy doth follow the unlikely love story between a British bookstore owner and a famous American actress who doth meet by chance in the charming neighborhood of Notting Hill, London.\n",
      "\n",
      "6. \"Bridget Jones's Diary\" (2001) - Based on the popular novel, thise film doth follow the life of Bridget Jones, a quirky and relatable woman who doth navigate her love life, career, and self-image whilst juggling two potential suitors.\n",
      "\n",
      "7. \"Four Weddings and a Funeral\" (1994) - Thise British romantic comedy doth center around a group of friends who doth attend various weddings and a funeral, exploring themes of love, friendship, and the complexities of relationships.\n",
      "\n",
      "8. \"About Time\" (2013) - From the creators of \"Love Actually,\" thise film doth tell the story of a young man who doth discover he can time travel and useth thise ability to find love and make the most of every moment in his life.\n",
      "\n",
      "9. \"Serendipity\" (2001) - Thise romantic comedy doth follow the story of two strangers who doth meet by chance and feel an instant connection. Despite going their separate ways, they doth believe in fate and spendeth years trying to find each other again.\n",
      "\n",
      "10. \"The Best Exotic Marigold Hotel\" (2011) - A group of British retirees doth decide to spend their golden years in a seemingly luxurious hotel in India. As they doth navigate new friendships, love interests, and cultural differences, they doth discover that life can still hold surprises and love can bloom at any age.\n"
     ]
    }
   ],
   "source": [
    "output = p.run(movie=\"Love Actually\")\n",
    "print(str(output))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}