From 955c89fccb5dc5858f880b082a109d4fd78afce3 Mon Sep 17 00:00:00 2001
From: Harrison Chase <hw.chase.17@gmail.com>
Date: Sun, 19 Feb 2023 20:47:17 -0800
Subject: [PATCH] pass in prompts to vectordbqa (#1158)

---
 .../combine_docs_examples/vector_db_qa.ipynb  | 63 ++++++++++++++++++-
 langchain/chains/vector_db_qa/base.py         | 13 +++-
 2 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/docs/modules/chains/combine_docs_examples/vector_db_qa.ipynb b/docs/modules/chains/combine_docs_examples/vector_db_qa.ipynb
index 1989394e..8510e331 100644
--- a/docs/modules/chains/combine_docs_examples/vector_db_qa.ipynb
+++ b/docs/modules/chains/combine_docs_examples/vector_db_qa.ipynb
@@ -68,7 +68,7 @@
     {
      "data": {
       "text/plain": [
-       "\" The president said that Ketanji Brown Jackson is one of the nation's top legal minds, a former top litigator in private practice, a former federal public defender, and from a family of public school educators and police officers. He also said that she is a consensus builder, and has received a broad range of support from the Fraternal Order of Police to former judges appointed by Democrats and Republicans.\""
+       "\" The president said that Ketanji Brown Jackson is one of the nation's top legal minds, a former top litigator in private practice and federal public defender, from a family of public school educators and police officers, a consensus builder, and has received a broad range of support from the Fraternal Order of Police to former judges appointed by Democrats and Republicans.\""
       ]
      },
      "execution_count": 4,
@@ -166,6 +166,67 @@
     "qa.run(query)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "90c7899a",
+   "metadata": {},
+   "source": [
+    "## Custom Prompts\n",
+    "You can pass in custom prompts to do question answering. These prompts are the same prompts as you can pass into the [base question answering chain](./question_answering.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "a45232a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.prompts import PromptTemplate\n",
+    "prompt_template = \"\"\"Use the following pieces of context to answer the question at the end. 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",
+    "\n",
+    "Question: {question}\n",
+    "Answer in Italian:\"\"\"\n",
+    "PROMPT = PromptTemplate(\n",
+    "    template=prompt_template, input_variables=[\"context\", \"question\"]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "9b5c8d1d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chain_type_kwargs = {\"prompt\": PROMPT}\n",
+    "qa = VectorDBQA.from_chain_type(llm=OpenAI(), chain_type=\"stuff\", vectorstore=docsearch, chain_type_kwargs=chain_type_kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "26ee7671",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\" Il Presidente ha detto che Ketanji Brown Jackson è uno dei pensatori legali più importanti del nostro Paese, che continuerà l'eccellente eredità di giustizia Breyer. È un ex principale litigante in pratica privata, un ex difensore federale pubblico e appartiene a una famiglia di insegnanti e poliziotti delle scuole pubbliche. È un costruttore di consenso che ha ricevuto un ampio supporto da parte di Fraternal Order of Police e giudici designati da democratici e repubblicani.\""
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
+    "qa.run(query)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "0b8c37f7",
diff --git a/langchain/chains/vector_db_qa/base.py b/langchain/chains/vector_db_qa/base.py
index 3a30a67f..882e05a1 100644
--- a/langchain/chains/vector_db_qa/base.py
+++ b/langchain/chains/vector_db_qa/base.py
@@ -1,7 +1,7 @@
 """Chain for question-answering against a vector database."""
 from __future__ import annotations
 
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel, Extra, Field, root_validator
 
@@ -120,10 +120,17 @@ class VectorDBQA(Chain, BaseModel):
 
     @classmethod
     def from_chain_type(
-        cls, llm: BaseLLM, chain_type: str = "stuff", **kwargs: Any
+        cls,
+        llm: BaseLLM,
+        chain_type: str = "stuff",
+        chain_type_kwargs: Optional[dict] = None,
+        **kwargs: Any,
     ) -> VectorDBQA:
         """Load chain from chain type."""
-        combine_documents_chain = load_qa_chain(llm, chain_type=chain_type)
+        _chain_type_kwargs = chain_type_kwargs or {}
+        combine_documents_chain = load_qa_chain(
+            llm, chain_type=chain_type, **_chain_type_kwargs
+        )
         return cls(combine_documents_chain=combine_documents_chain, **kwargs)
 
     def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]: