How do i add memory to RetrievalQA.from_chain_type? or, how do I add a custom prompt to ConversationalRetrievalChain?

Question

How do i add memory to RetrievalQA.from_chain_type? or, how do I add a custom prompt to ConversationalRetrievalChain?

For the past 2 weeks ive been trying to make a chatbot that can chat over documents (so not in just a semantic search/qa so with memory) but also with a custom prompt. I've tried every combination of all the chains and so far the closest I've gotten is ConversationalRetrievalChain, but without custom prompts, and RetrievalQA.from_chain_type but without memory

Answer 1

Update: This post answers the first part of OP's question:

how do i add memory to RetrievalQA.from_chain_type?

For the second part, see @andrew_reece's answer

or, how do I add a custom prompt to ConversationalRetrievalChain?

Original:

Have you tried passing in chain_type_kwargs (at the bottom is a screenshot from the source code for quick references)?

The documentation hasn't make it very easy to unserstand what's under the hood, but here is something that could achieve your goal.

You could find the notebook at this GitHub Link setup

from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain import PromptTemplate
from langchain.retrievers import TFIDFRetriever


retriever = TFIDFRetriever.from_texts(
    ["Our client, a gentleman named Jason, has a dog whose name is Dobby",
     "Jason has a good friend called Emma",
     "Emma has a cat whose name is Sullivan"])

Then define your customized prompt:

template = """
Use the following context (delimited by <ctx></ctx>) and the chat history (delimited by <hs></hs>) to answer the question:
------
<ctx>
{context}
</ctx>
------
<hs>
{history}
</hs>
------
{question}
Answer:
"""
prompt = PromptTemplate(
    input_variables=["history", "context", "question"],
    template=template,
)

Take the note of what you used for your input variables, especially 'history' and 'question', since you will need to match these when setting up the memory:

qa = RetrievalQA.from_chain_type(
    llm=ChatOpenAI(),
    chain_type='stuff',
    retriever=retriever,
    verbose=True,
    chain_type_kwargs={
        "verbose": True,
        "prompt": prompt,
        "memory": ConversationBufferMemory(
            memory_key="history",
            input_key="question"),
    }
)

Now you can call qa.run({"query": "who's the client's friend?"})

"The client's friend is Emma."

and then qa.run("and her pet's name is?")

"Emma's pet's name is Sullivan."

To check and verify the memory/chat history: qa.combine_documents_chain.memory

ConversationBufferMemory(chat_memory=ChatMessageHistory(messages=[HumanMessage(content="who's the client's friend?", additional_kwargs={}), AIMessage(content="The client's friend is Emma.", additional_kwargs={}), HumanMessage(content="and her pet's name is?", additional_kwargs={}), AIMessage(content="Emma's pet's name is Sullivan.", additional_kwargs={})]), output_key=None, input_key='question', return_messages=False, human_prefix='Human', ai_prefix='AI', memory_key='history')