Langchain中RetrievalQA无法传入自定义Prompt模板的问题求助

问题场景 #

作为LangChain新手，按照官方《Retrieval QA - Langchain》文档开发时，自定义Prompt模板后通过chain_type_kwargs传入RetrievalQA，触发了Pydantic的StuffDocumentsChain验证错误；移除该参数程序可正常运行，需解决如何正确传入自定义Prompt模板的问题。

错误信息 #

File /usr/local/lib/python3.11/site-packages/pydantic/main.py:341, in pydantic.main.BaseModel.__init__()

ValidationError: 1 validation error for StuffDocumentsChain
__root__
  document_variable_name context was not found in llm_chain input_variables: ['question'] (type=value_error)

问题原因 #

使用stuff类型的chain时，StuffDocumentsChain需要将检索到的文档内容注入到Prompt模板的指定变量中，默认变量名为context。但自定义的Prompt模板里仅声明了question变量，缺少context变量，导致Chain无法找到填充文档内容的位置，触发验证错误。

解决方法 #

有两种修复方式：

方式一：在Prompt模板中添加默认的context变量 #

修改Prompt模板，加入{context}占位符，让Chain能把检索到的文档内容填充进去：

template = """
使用以下上下文信息回答问题：
{context}

Question: {question}
Answer: 
"""
PROMPT = PromptTemplate(template=template, input_variables=['context', 'question'])

方式二：自定义文档变量名并在参数中指定 #

如果不想用context作为变量名，可以自定义变量名（比如docs_content），同时在chain_type_kwargs中通过document_variable_name指定这个名称：

template = """
使用以下文档内容回答问题：
{docs_content}

Question: {question}
Answer: 
"""
PROMPT = PromptTemplate(template=template, input_variables=['docs_content', 'question'])

qa = RetrievalQA.from_chain_type(
    llm=ChatOpenAI(model_name='gpt-3.5-turbo-16k'),
    chain_type="stuff",
    chain_type_kwargs={
        "prompt": PROMPT,
        "document_variable_name": "docs_content"  # 指定对应模板中的变量名
    },
    retriever=docsearch.as_retriever(),
)

完整修改后代码示例 #

import json, os

from langchain.chains import RetrievalQA
from langchain.document_loaders import JSONLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain import PromptTemplate

from pathlib import Path
from pprint import pprint

os.environ["OPENAI_API_KEY"] = "my-key"

def metadata_func(record: dict, metadata: dict) -> dict:
    metadata["drug_name"] = record["drug_name"]
    return metadata

loader = JSONLoader(
    file_path='./drugs_data_v2.json', 
    jq_schema='.drugs[]',
    content_key="data",
    metadata_func=metadata_func)

docs = loader.load()

text_splitter = CharacterTextSplitter(chunk_size=5000, chunk_overlap=200)
texts = text_splitter.split_documents(docs)

embeddings = OpenAIEmbeddings()
docsearch = Chroma.from_documents(texts, embeddings)

# 修复后的Prompt模板
template = """
使用以下上下文信息回答问题：
{context}

Question: {question}
Answer: 
"""
PROMPT = PromptTemplate(template=template, input_variables=['context', 'question'])

qa = RetrievalQA.from_chain_type(
    llm=ChatOpenAI(model_name='gpt-3.5-turbo-16k'),
    chain_type="stuff",
    chain_type_kwargs={"prompt": PROMPT},
    retriever=docsearch.as_retriever(),
)

query = "What did the president say about Ketanji Brown Jackson"
qa.run(query)

内容的提问来源于stack exchange，提问作者umair mehmood