langchain/docs/modules/chat/examples/vector_db_qa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "07c1e3b9",
   "metadata": {},
   "source": [
    "# Retrieval Question/Answering\n",
    "\n",
    "This example showcases using a chat model to do question answering over a vector database.\n",
    "\n",
    "This notebook is very similar to the example of using an LLM in the RetrievalQA. The only differences here are (1) using a ChatModel, and (2) passing in a ChatPromptTemplate (optimized for chat models)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "82525493",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.chains import RetrievalQA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5c7049db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Chroma using direct local API.\n",
      "Using DuckDB in-memory for database. Data will be transient.\n"
     ]
    }
   ],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "loader = TextLoader('../../state_of_the_union.txt')\n",
    "documents = loader.load()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "\n",
    "embeddings = OpenAIEmbeddings()\n",
    "docsearch = Chroma.from_documents(texts, embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35f99145",
   "metadata": {},
   "source": [
    "We can now set up the chat model and chat model specific prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "32a49412",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.prompts.chat import (\n",
    "    ChatPromptTemplate,\n",
    "    SystemMessagePromptTemplate,\n",
    "    AIMessagePromptTemplate,\n",
    "    HumanMessagePromptTemplate,\n",
    ")\n",
    "from langchain.schema import (\n",
    "    AIMessage,\n",
    "    HumanMessage,\n",
    "    SystemMessage\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f231fb9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_template=\"\"\"Use the following pieces of context to answer the users question. \n",
    "If you don't know the answer, just say that you don't know, don't try to make up an answer.\n",
    "----------------\n",
    "{context}\"\"\"\n",
    "messages = [\n",
    "    SystemMessagePromptTemplate.from_template(system_template),\n",
    "    HumanMessagePromptTemplate.from_template(\"{question}\")\n",
    "]\n",
    "prompt = ChatPromptTemplate.from_messages(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3018f865",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain_type_kwargs = {\"prompt\": prompt}\n",
    "qa = RetrievalQA.from_chain_type(llm=ChatOpenAI(), chain_type=\"stuff\", retriever=docsearch.as_retriever(), chain_type_kwargs=chain_type_kwargs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "032a47f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"The president nominated Ketanji Brown Jackson to serve on the United States Supreme Court. He referred to her as one of our nation's top legal minds, a former federal public defender, a consensus builder, and from a family of public school educators and police officers. Since she's been nominated, she has received a broad range of support from the Fraternal Order of Police to former judges appointed by Democrats and Republicans.\""
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "qa.run(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b403637",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "b1677b440931f40d89ef8be7bf03acb108ce003de0ac9b18e8d43753ea2e7103"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}