使用MultiQueryRetriever与自定义Prompt从ChromaDB取数时遇KeyError

问题背景 #

在使用MultiQueryRetriever结合自定义PromptTemplate从Chroma向量库检索警队规则手册时，明明传入了allegations、description、num_allegations三个变量，却触发KeyError，提示缺少这些变量，只收到了question变量。

问题原因 #

MultiQueryRetriever的默认逻辑会将输入统一包装为question变量传递给LLMChain，不管你传入的是字符串还是多变量字典。而自定义PromptTemplate声明需要三个特定变量，两者不匹配导致报错。

解决办法 #

调整PromptTemplate，改用单一question变量，将所有需要的信息整合到question文本中，让MultiQueryRetriever可以正常传递输入。

修改后的代码示例 #

1. 调整PromptTemplate

chroma_prompt = PromptTemplate(
    input_variables=["question"],
    template=(
        """You are an AI language model assistant. Your task is to analyze the following civilian complaint 
        details against a police officer, identify potential acts of misconduct or crimes, and generate the required number of queries to
        retrieve relevant sections from the Police Rulebook (one query per allegation), stored in a vector database.
        By generating multiple perspectives on the analysis, your goal is to help the user overcome some of the limitations of the 
        distance-based similarity search. Provide these alternative analyses as distinct queries, separated by newlines.
        
        {question}
        """
    )
)

2. 修改fetch_from_chroma函数的调用逻辑

def fetch_from_chroma(allegations, description, ia_num, llm, k=2):
    """
    Fetches relevant documents from Chroma using Maximal Marginal Relevance (MMR).

    Parameters:
    - allegations (list): List of allegations against the officer.
    - description (str): Incident description.
    - ia_num (int): Internal Affairs number for logging/debugging.
    - llm: Language model instance.
    - k (int): Number of results to fetch, set to 2 by default.

    Returns:
    - context_text (str): Combined context text from retrieved documents.
    - sources (list): List of source metadata.
    """
    embedding_function = OpenAIEmbeddings()
    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
    
    line_output_parser = LineListOutputParser()
    
    # 构造包含所有信息的question文本
    num_queries = len(allegations)
    question_text = (
        f"Allegations made against officer: {', '.join(allegations)}\n"
        f"Incident description: {description}\n"
        f"Generate {num_queries} distinct queries, one for each allegation."
    )
    
    llm_chain = chroma_prompt | llm | line_output_parser
    
    retriever = MultiQueryRetriever(
        retriever=db.as_retriever(search_type="similarity", search_kwargs={"k": k}), 
        llm_chain=llm_chain, 
        parser_key="lines"
    )
    
    # 传递整合后的question变量
    results = retriever.invoke({"question": question_text})

    if len(results) == 0:
        print(f"{ia_num} - Unable to find matching results.")
        return "No Context Available", "No Sources Available"

    context_text = "\n\n---\n\n".join([doc.page_content for doc in results])
    sources = [doc.metadata.get("source", None) for doc in results]
    print(f"{ia_num} - Found matching results.")
    return context_text, sources

其他细节说明 #

• 若坚持使用多变量输入，可通过重写MultiQueryRetriever的generate_queries方法，或使用RunnableMap映射输入变量，但复杂度较高，不如单变量整合直接。
• 确保LineListOutputParser逻辑正确，能正确拆分LLM生成的多行查询。

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