{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e42733c5",
   "metadata": {},
   "source": [
    "# Adding Memory to a Multi-Input Chain\n",
    "\n",
    "Most memory objects assume a single output. In this notebook, we go over how to add memory to a chain that has multiple outputs. As an example of such a chain, we will add memory to a question/answering chain. This chain takes as inputs both related documents and a user question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "978ba52b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.embeddings.cohere import CohereEmbeddings\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.vectorstores.elastic_vector_search import ElasticVectorSearch\n",
    "from langchain.vectorstores.faiss import FAISS\n",
    "from langchain.docstore.document import Document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2ee8628b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../../state_of_the_union.txt') as f:\n",
    "    state_of_the_union = f.read()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "texts = text_splitter.split_text(state_of_the_union)\n",
    "\n",
    "embeddings = OpenAIEmbeddings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "aa70c847",
   "metadata": {},
   "outputs": [],
   "source": [
    "docsearch = FAISS.from_texts(texts, embeddings, metadatas=[{\"source\": i} for i in range(len(texts))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ea4f7d82",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What did the president say about Justice Breyer\"\n",
    "docs = docsearch.similarity_search(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d3dc4ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.question_answering import load_qa_chain\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains.conversation.memory import ConversationBufferMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9a530742",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"You are a chatbot having a conversation with a human.\n",
    "\n",
    "Given the following extracted parts of a long document and a question, create a final answer.\n",
    "\n",
    "{context}\n",
    "\n",
    "{chat_history}\n",
    "Human: {human_input}\n",
    "Chatbot:\"\"\"\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"chat_history\", \"human_input\", \"context\"], \n",
    "    template=template\n",
    ")\n",
    "memory = ConversationBufferMemory(memory_key=\"chat_history\", input_key=\"human_input\")\n",
    "chain = load_qa_chain(OpenAI(temperature=0), chain_type=\"stuff\", memory=memory, prompt=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9bb8a8b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output_text': \" President Biden honored Justice Stephen Breyer, an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. He thanked Justice Breyer for his service and said that one of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. He then announced his nomination of Circuit Court of Appeals Judge Ketanji Brown Jackson to continue Justice Breyer's legacy of excellence.\"}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"What did the president say about Justice Breyer\"\n",
    "chain({\"input_documents\": docs, \"human_input\": query}, return_only_outputs=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "82593148",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Human: What did the president say about Justice Breyer\n",
      "AI:  President Biden honored Justice Stephen Breyer, an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. He thanked Justice Breyer for his service and said that one of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. He then announced his nomination of Circuit Court of Appeals Judge Ketanji Brown Jackson to continue Justice Breyer's legacy of excellence.\n"
     ]
    }
   ],
   "source": [
    "print(chain.memory.buffer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f262b2fb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
