{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "839f3c76",
   "metadata": {},
   "source": [
    "# Conversational Retrieval Agent\n",
    "\n",
    "This is an agent specifically optimized for doing retrieval when necessary and also holding a conversation.\n",
    "\n",
    "To start, we will set up the retriever we want to use, and then turn it into a retriever tool. Next, we will use the high level constructor for this type of agent. Finally, we will walk through how to construct a conversational retrieval agent from components."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc66a262",
   "metadata": {},
   "source": [
    "## The Retriever\n",
    "\n",
    "To start, we need a retriever to use! The code here is mostly just example code. Feel free to use your own retriever and skip to the section on creating a retriever tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "22533f46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "loader = TextLoader('../../../../../docs/extras/modules/state_of_the_union.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c252c7e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "\n",
    "documents = loader.load()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "embeddings = OpenAIEmbeddings()\n",
    "db = FAISS.from_documents(texts, embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5b8a404f",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = db.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cd528f5",
   "metadata": {},
   "source": [
    "## Retriever Tool\n",
    "\n",
    "Now we need to create a tool for our retriever. The main things we need to pass in are a name for the retriever as well as a description. These will both be used by the language model, so they should be informative."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9a82f72a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.agent_toolkits import create_retriever_tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dfbd92a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tool = create_retriever_tool(\n",
    "    retriever, \n",
    "    \"search_state_of_union\",\n",
    "    \"Searches and returns documents regarding the state-of-the-union.\"\n",
    ")\n",
    "tools = [tool]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d9b7210",
   "metadata": {},
   "source": [
    "## Agent Constructor\n",
    "\n",
    "Here, we will use the high level `create_conversational_retrieval_agent` API to construct the agent.\n",
    "\n",
    "Notice that beside the list of tools, the only thing we need to pass in is a language model to use.\n",
    "Under the hood, this agent is using the OpenAIFunctionsAgent, so we need to use an ChatOpenAI model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0cd147eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.agent_toolkits import create_conversational_retrieval_agent "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9fa4661b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "llm = ChatOpenAI(temperature = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "da2ad764",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_executor = create_conversational_retrieval_agent(llm, tools, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85ee7677",
   "metadata": {},
   "source": [
    "We can now try it out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "03059322",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mHello Bob! How can I assist you today?\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "result = agent_executor({\"input\": \"hi, im bob\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "33073ff0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Hello Bob! How can I assist you today?'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"output\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f1f0b79",
   "metadata": {},
   "source": [
    "Notice that it remembers your name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4ad92bc7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mYour name is Bob.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "result = agent_executor({\"input\": \"whats my name?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7ae62ecd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Your name is Bob.'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"output\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e258ade",
   "metadata": {},
   "source": [
    "Notice that it now does retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6cd17d67",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `search_state_of_union` with `{'query': 'Kentaji Brown Jackson'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m[Document(page_content='Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \\n\\nTonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \\n\\nOne of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \\n\\nAnd I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.', metadata={'source': '../../../../../docs/extras/modules/state_of_the_union.txt'}), Document(page_content='One was stationed at bases and breathing in toxic smoke from “burn pits” that incinerated wastes of war—medical and hazard material, jet fuel, and more. \\n\\nWhen they came home, many of the world’s fittest and best trained warriors were never the same. \\n\\nHeadaches. Numbness. Dizziness. \\n\\nA cancer that would put them in a flag-draped coffin. \\n\\nI know. \\n\\nOne of those soldiers was my son Major Beau Biden. \\n\\nWe don’t know for sure if a burn pit was the cause of his brain cancer, or the diseases of so many of our troops. \\n\\nBut I’m committed to finding out everything we can. \\n\\nCommitted to military families like Danielle Robinson from Ohio. \\n\\nThe widow of Sergeant First Class Heath Robinson.  \\n\\nHe was born a soldier. Army National Guard. Combat medic in Kosovo and Iraq. \\n\\nStationed near Baghdad, just yards from burn pits the size of football fields. \\n\\nHeath’s widow Danielle is here with us tonight. They loved going to Ohio State football games. He loved building Legos with their daughter.', metadata={'source': '../../../../../docs/extras/modules/state_of_the_union.txt'}), Document(page_content='A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \\n\\nAnd if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \\n\\nWe can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling.  \\n\\nWe’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers.  \\n\\nWe’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \\n\\nWe’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.', metadata={'source': '../../../../../docs/extras/modules/state_of_the_union.txt'}), Document(page_content='We can’t change how divided we’ve been. But we can change how we move forward—on COVID-19 and other issues we must face together. \\n\\nI recently visited the New York City Police Department days after the funerals of Officer Wilbert Mora and his partner, Officer Jason Rivera. \\n\\nThey were responding to a 9-1-1 call when a man shot and killed them with a stolen gun. \\n\\nOfficer Mora was 27 years old. \\n\\nOfficer Rivera was 22. \\n\\nBoth Dominican Americans who’d grown up on the same streets they later chose to patrol as police officers. \\n\\nI spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. \\n\\nI’ve worked on these issues a long time. \\n\\nI know what works: Investing in crime preventionand community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety.', metadata={'source': '../../../../../docs/extras/modules/state_of_the_union.txt'})]\u001b[0m\u001b[32;1m\u001b[1;3mIn the most recent state of the union, the President mentioned Kentaji Brown Jackson. The President nominated Circuit Court of Appeals Judge Ketanji Brown Jackson to serve on the United States Supreme Court. The President described Judge Ketanji Brown Jackson as one of our nation's top legal minds who will continue Justice Breyer's legacy of excellence.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "result = agent_executor({\"input\": \"what did the president say about kentaji brown jackson in the most recent state of the union?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c51de037",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"In the most recent state of the union, the President mentioned Kentaji Brown Jackson. The President nominated Circuit Court of Appeals Judge Ketanji Brown Jackson to serve on the United States Supreme Court. The President described Judge Ketanji Brown Jackson as one of our nation's top legal minds who will continue Justice Breyer's legacy of excellence.\""
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"output\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3bad540",
   "metadata": {},
   "source": [
    "Notice that the follow up question asks about information previously retrieved, so no need to do another retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "527ea5fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mThe President nominated Judge Ketanji Brown Jackson four days ago.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "result = agent_executor({\"input\": \"how long ago did he nominate her?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a34d20fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The President nominated Judge Ketanji Brown Jackson four days ago.'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"output\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e599dbd3",
   "metadata": {},
   "source": [
    "## Creating from components\n",
    "\n",
    "What actually is going on underneath the hood? Let's take a look so we can understand how to modify going forward.\n",
    "\n",
    "There are a few components:\n",
    "\n",
    "- The memory\n",
    "- The prompt template\n",
    "- The agent\n",
    "- The agent executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1b21be1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is needed for both the memory and the prompt\n",
    "memory_key = \"history\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f827f95f",
   "metadata": {},
   "source": [
    "### The Memory\n",
    "\n",
    "In this example, we want the agent to remember not only previous conversations, but also previous intermediate steps. For that, we can use `AgentTokenBufferMemory`. Note that if you want to change whether the agent remembers intermediate steps, or how the long the buffer is, or anything like that you should change this part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "138b0675",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.openai_functions_agent.agent_token_buffer_memory import AgentTokenBufferMemory\n",
    "\n",
    "memory = AgentTokenBufferMemory(memory_key=memory_key, llm=llm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4827993f",
   "metadata": {},
   "source": [
    "## The Prompt Template\n",
    "\n",
    "For the prompt template, we will use the `OpenAIFunctionsAgent` default way of creating one, but pass in a system prompt and a placeholder for memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "779272dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent\n",
    "from langchain.schema.messages import SystemMessage\n",
    "from langchain.prompts import MessagesPlaceholder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "824bc74e",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = SystemMessage(\n",
    "        content=(\n",
    "            \"Do your best to answer the questions. \"\n",
    "            \"Feel free to use any tools available to look up \"\n",
    "            \"relevant information, only if neccessary\"\n",
    "        )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "07e41722",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = OpenAIFunctionsAgent.create_prompt(\n",
    "        system_message=system_message,\n",
    "        extra_prompt_messages=[MessagesPlaceholder(variable_name=memory_key)]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "001e455e",
   "metadata": {},
   "source": [
    "## The Agent\n",
    "\n",
    "We will use the OpenAIFunctionsAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "adf4b5b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c5c321e",
   "metadata": {},
   "source": [
    "## The Agent Executor\n",
    "\n",
    "Importantly, we pass in `return_intermediate_steps=True` since we are recording that with our memory object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "2e7ffe96",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e39a095f",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_executor = AgentExecutor(agent=agent, tools=tools, memory=memory, verbose=True,\n",
    "                                   return_intermediate_steps=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "96136958",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mHello Bob! How can I assist you today?\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "result = agent_executor({\"input\": \"hi, im bob\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "8de674cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mYour name is Bob.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "result = agent_executor({\"input\": \"whats my name\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf655a48",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
