KeyError с MultiQueryRetriever и пользовательским приглашением для получения данных из ChromaDB ⇐ Python

[Anonymous]

У меня возникли проблемы с использованием MultiQueryRetriever и PromptTemplate.
Моя цель — взять список обвинений против полицейского и с помощью MultiQueryRetriever заставить LLM генерировать один запрос на каждое обвинение. + комбинация описания, чтобы получить наиболее подходящее нарушенное правило для каждого обвинения. В качестве векторного хранилища у меня есть Chroma, в нем есть свод правил для сотрудников полицейского управления. Для этого я использую специальную подсказку, которая инструктирует LLM генерировать один запрос Chroma для каждого утверждения. Чтобы сгенерировать этот запрос, он должен просмотреть обвинение и попытаться извлечь потенциальные нарушения (имеющие отношение к обвинению) из описания, а затем сформировать запрос, который можно использовать для получения соответствующих правил из Chroma. (Для получения более подробной информации ознакомьтесь с фактическим приглашением)
Это 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

Это специальное приглашение, которое я разработал:

Код: Выделить всё

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}
"""
)
)

Это фрагмент кода, который извлекается из Chroma:

Код: Выделить всё

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

Parameters:
- query (str): The query string.
- ia_num (int): Internal Affairs number for logging/debugging.
- k (int): Number of results to fetch, set to 3 by default.
- lambda_mult (float): MMR diversity parameter. Values closer to 1 prioritize diversity, closer to 0 prioritize relevance.

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()

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))
})

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

Однако я получаю эту ошибку и понятия не имею, почему:

Код: Выделить всё

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}}'."

по какой-то причине он продолжает говорить, что я передал только переменную «вопрос», но когда я вызываю метод restart.invoke(), я явно передаю необходимые переменные.
p>
вот пример передаваемых входных данных:

Код: Выделить всё

{'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.3.7
langchain-community==0.2.3
langchain-core==0.3.19
langchain-google-genai==2.0.1
langchain-openai==0.2.8
langchain-text-splitters==0.3.2

с помощью Mac
с использованием Python 3.11.9

Подробнее здесь: https://stackoverflow.com/questions/792 ... from-chrom