KeyError with MultiQueryRetriever and Custom Prompt for Fetching data from ChromaDB

[AbdelazimLokma]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

chroma_prompt = PromptTemplate(
    input_variables=["allegations", "description", "num_allegations"],
    template=(
        """You are an AI language model assistant. Your task is to analyze the following civilian complaint 
        description against a police officer, and the allegations that are raised against the officer. Identify 
        potential acts of misconduct or crimes committed by the officer, and generate {num_allegations} different queries to
        retrieve relevant sections from the Police Rulebook (one query per allegation-description combination), 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.
        
        Allegations made against officer: {allegations}
        Incident description: {description}
        """
    )
)



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

    Parameters:
    - allegations (list): a string list of allegations against an officer.
    - description (str):  a description of the event
    - ia_num (int): Internal Affairs number for logging/debugging.
    - k (int): Number of results to fetch (from chroma) per LLM generated query, 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)
    
    chain = chroma_prompt | llm

    retriever = MultiQueryRetriever.from_llm(
        retriever=db.as_retriever(search_type="similarity", search_kwargs={"k": k}),
        llm_chain=chain)
    

    # Invoke the retriever with the input dictionary
    results = retriever.invoke({
        "allegations": ", ".join(allegations), #convert list to string
    "description": description, 
    "num_allegations": str(len(allegations)) # I want one LLM generated query per allegation
    })


    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

Description

I am having trouble using MultiQueryRetriever and PromptTemplate.

My goal is to take a list of allegations against a police officer, and using the MultiQueryRetriever, have the LLM generate one query per allegation + description combination, in order to fetch the most relevant rule broken for each allegation. I have Chroma as my vector store, it contains a police department officer rule book. To do this, I am using a custom prompt that instructs the LLM to generate one Chroma query for each allegation. In order to generate this query, it must look at the allegation and try to extract potential violations (relevant to the allegation) from the description, then form a query that can be used to fetch relevant rules from Chroma. (Take a look at the actual prompt for more detail)

However, I am getting this error and have no idea why:

KeyError: "Input to PromptTemplate is missing variables {'description', 'allegations', 'num_allegations'}. Expected: ['allegations', 'description', 'num_allegations'] Received: ['question']\nNote: if you intended {description} to be part of the string and not a variable, please escape it with double curly braces like: '{{description}}'."

for some reason, it keeps saying that I only passed in a variable 'question', but when i call retriever.invoke(), I am clearly passing in the required variables.

here is an example input that is being passed in:

{'allegations': 'Conformance to Laws, Conduct Unbecoming, Respectful Treatment, Alcohol off Duty', 'description': 'Officer firstname Lastname fled from a taxicab without paying the fare. Officer Lastname was located by Officers from Area A‐7. Officer Lastname A‐7 where he was uncooperative with Sgt. Lastname and refused to talk to him. Sgt. Lastname escorted Officer Lastname back to the o his department equipment was received by Sgt. Lastname including a Glock 40 Serial # number, Radio # number, Handcuffs #number, 3 magazine Police Badge# number, 1 container of OC Spray and 1 bullet resistant vest. Department equipment to be turned over to Sgt. last name', 'num_allegations': '4'}

System Info

langchain==0.2.2
langchain-community==0.2.3
langchain-core==0.2.43
langchain-google-genai==2.0.1
langchain-openai==0.1.8
langchain-text-splitters==0.2.4

using a mac

Using Python 3.11.9

[feijoes]

@AbdelazimLokma, try upgrading LangChain to the newest version by running pip install -U langchain. I'm not sure about version 0.2.2, but in the latest version, the MultiQueryRetriever from_llm method doesn't expect the keyword argument 'llm_chain'.

try:

retriever = MultiQueryRetriever.from_llm(
        retriever=db.as_retriever(search_type="similarity", search_kwargs={"k": 2}),
        llm=llm)

[AbdelazimLokma]

ok, I've upgraded to Lanchain 0.3.7. This is how my code looks now (I'm still using the same prompt as before):

    embedding_function = OpenAIEmbeddings()
    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)

    retriever = MultiQueryRetriever.from_llm(
        retriever=db.as_retriever(search_type="similarity", search_kwargs={"k": k}),
        llm=llm, prompt = chroma_prompt)
    

    # Invoke the retriever with the input dictionary
    results = retriever.invoke({
        "allegations": ", ".join(allegations),
    "description": description,
    "num_allegations": str(len(allegations))
    })

but i am still getting the same error:

raise KeyError( KeyError: "Input to PromptTemplate is missing variables {'allegations', 'description', 'num_allegations'}. Expected: ['allegations', 'description', 'num_allegations'] Received: ['question']\nNote: if you intended {allegations} to be part of the string and not a variable, please escape it with double curly braces like: '{{allegations}}'

Alternatively I wanted to try following the tutorial on MultiQueryRetriever:

With the same prompt as before, and a new LineListOutputParser()

class LineListOutputParser(BaseOutputParser[List[str]]):
    """Output parser for a list of lines."""

    def parse(self, text: str) -> List[str]:
        lines = text.strip().split("\n")
        return list(filter(None, lines))  # Remove empty lines

This is the code that I made, following the tutorial shown here

    line_output_parser = LineListOutputParser()
    
    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"
    )
   
    # Invoke the retriever with the input dictionary
    results = retriever.invoke({
        "allegations": ", ".join(allegations),
    "description": description,
    "num_allegations": str(len(allegations))
    })

But still, the error is the same.

[feijoes]

@AbdelazimLokma I was checking the documentation and found out that the default generate_queries implementation (thats is executed from the retriever.invoke) puts all the inputs inside { 'question': question }.

(Yes, the typing is str, but the "question" is actually the dictionary you pass in the invoke method.)

I'm not entirely sure why it does that for now, but one possible solution is to create a custom MultiQueryRetriever and override the default generate_queries with your own implementation. Basically, just use self.llm_chain.invoke(question).

Here's an example:

from langchain_core.callbacks import CallbackManagerForRetrieverRun  # Only necessary for type hint
from typing import List  # Only necessary for type hint
from langchain.chains import LLMChain
from langchain.retrievers.multi_query import MultiQueryRetriever
import logging

logger = logging.getLogger(__name__)

class CustomMultiQueryRetriever(MultiQueryRetriever):
    def generate_queries(
        self, question: str, run_manager: CallbackManagerForRetrieverRun
    ) -> List[str]:
        """Generate queries based on user input.

        Args:
            question: The user query.
            run_manager: Callback manager for handling retriever runs.

        Returns:
            List of LLM-generated queries that are similar to the user input.
        """
        response = self.llm_chain.invoke(
            question, config={"callbacks": run_manager.get_child()}
        )
        if isinstance(self.llm_chain, LLMChain):
            lines = response["text"].splitlines()  # Assuming "text" is a multi-line string
        else:
            lines = response
        
        if self.verbose:
            logger.info(f"Generated queries: {lines}")
        
        return lines
``

[AbdelazimLokma]

That seems to have worked! Thanks so much for your help @feijoes. I'm also not sure why generate_queries works that way, there must be some way to pass in custom arguments rather than a single natural language query. Regardless, thanks for taking the time to help me solve this problem.