LangChain检索链运行时提示词超出上下文窗口问题排查

运行LangChain的RetrievalQA链时，调用chain.run()出现「Prompt size exceeds maximum context window」错误，即使模板和问题的token数远小于模型2048的上下文限制。以下是实现代码：

from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.chains import RetrievalQA
from langchain_community.llms import GPT4All
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain import PromptTemplate

local_path = (
   "/home/user/coding/facts/models/llama-2-7b-chat.Q5_K_M.gguf"  
)

callbacks = [StreamingStdOutCallbackHandler()]

llm = GPT4All(
    model=local_path, 
    callbacks=callbacks, 
    verbose=True,
    temp=0,
    streaming=True)

prompt_template = """
You are a chatbot which retrieves data from a vector database.
Your task is to answer user questions by only using the data provided by the retriever.
If the retriever does not find any source documents say: "I don't know the answer".

    -----------
    Context: {context}
    -----------
    Question: {question}
"""

qa_prompt=PromptTemplate(template=prompt_template, input_variables=['context', 'question'],max_length = 2048)

print("loading embeddings model")
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': False}
embeddings = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)
db = Chroma(
    persist_directory="emb",
    embedding_function=embeddings
)
retriever = db.as_retriever()

chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    chain_type="stuff",
    chain_type_kwargs={'prompt': qa_prompt}
)

print(len(qa_prompt.template))

result = chain.run("When was Mark Twain born?")

print(result)

原因分析 #

• 你只统计了模板和问题的token数，但忽略了**{context}变量实际填充的检索文档内容**。stuff链类型会将检索到的所有文档内容直接拼接进context位置，如果这些文档的总token数（加上模板、问题）超过模型2048的上下文窗口，就会触发错误。
• PromptTemplate的max_length参数仅用于限制模板本身的长度，无法控制最终填充后prompt的总长度，因此设置该参数无效。

解决方法 #

1. 查看context的实际内容 #

可以手动获取检索结果，查看文档内容并计算总token数：

# 手动获取检索到的文档
docs = retriever.get_relevant_documents("When was Mark Twain born?")
print("检索到的文档：")
for doc in docs:
    print(doc.page_content)
    print("---")

# 计算填充后prompt的总token数
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-mpnet-base-v2")
filled_prompt = qa_prompt.template.format(
    context="\n".join([d.page_content for d in docs]),
    question="When was Mark Twain born?"
)
total_tokens = len(tokenizer.encode(filled_prompt))
print(f"填充后prompt总token数：{total_tokens}")

2. 减少单次检索的文档数量 #

修改retriever的search_kwargs参数，限制返回的文档数量：

retriever = db.as_retriever(search_kwargs={"k": 2})  # 仅返回前2个最相关的文档

3. 改用其他链类型 #

如果文档较长，替换stuff为map_reduce、refine或map_rerank类型，这些类型会拆分文档分批处理：

chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    chain_type="map_reduce",
    chain_type_kwargs={'prompt': qa_prompt}
)

4. 优化文档分块 #

如果原始文档本身过长，在生成向量库时需将文档拆分为更小的片段，降低单篇文档的token数，避免拼接后超出窗口限制。

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