如何使用 LangChain 查看 MultiQueryRetriever 管道中的最终提示？

我目前正在使用 LangChain 库开展一个项目，我想从向量数据库中检索相关文档，然后使用 Ollama LLM 根据这些文档生成答案。

以下是我当前的实现：

import logging

logging.basicConfig()
logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.INFO)

# Define the prompt template for generating multiple query versions
QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template="""You are an AI language model assistant. Your task is to generate five
    different versions of the given user question to retrieve relevant documents from
    a vector database. By generating multiple perspectives on the user question, your
    goal is to help the user overcome some of the limitations of the distance-based
    similarity search. Provide these alternative questions separated by newlines.
    Original question: {question}""",
)

# Initialize the MultiQueryRetriever
retriever = MultiQueryRetriever.from_llm(
    vectordb.as_retriever(), 
    ollama,
    prompt=QUERY_PROMPT
)

# Modified RAG prompt for generating the final response
template = """Answer the question based ONLY on the following context:
{context}
Question: {question}
"""

# Create the final QA chain
prompt = ChatPromptTemplate.from_template(template)

from langchain_core.runnables import RunnableLambda


def inspect(state):
    """Print the state passed between Runnables in a langchain and pass it on"""
    print(state)
    return state


qa_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | RunnableLambda(inspect)  # Add the inspector here to print the intermediate results
    | prompt
    | ollama
    | StrOutputParser()
)

# Invoke the QA chain with a sample query
qa_chain.invoke("Give 10 quotes from this articles related to love?")

我该如何在发送到 Ollama LLM 进行处理之前查看由 qa_chain 生成的最终提示？我想查看包含上下文和用户问题的确切提示。

代码中添加的 inspect 函数可以打印出 retriever 和 prompt 之间的中间状态，但无法看到最终发送给 Ollama LLM 的提示。这是因为 prompt 本身是一个 Runnable 对象，它会在内部将上下文和问题组合成最终提示。

为了查看最终提示，需要修改 prompt 的调用方式，或者在 Ollama LLM 调用之前拦截最终提示。以下提供两种方法：

方法一：修改 prompt 调用方式

可以将 prompt 的调用从管道中分离出来，以便在调用 LLM 之前查看最终提示。

# ... (之前的代码) ...

def inspect_final_prompt(context, question):
    """打印最终提示并返回上下文和问题"""
    final_prompt = prompt.format(context=context, question=question)
    print("最终提示：\n" + final_prompt)
    return final_prompt

# 修改后的 QA 链
qa_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | RunnableLambda(inspect)  # 保留之前的 inspector
    | RunnableMap({"context": lambda x: x["context"], "question": lambda x: x["question"]}) # 将字典拆分为两个参数
    | RunnableLambda(inspect_final_prompt)  # 添加新的 inspector 来打印最终提示
    | ollama
    | StrOutputParser()
)

# ... (之后的代码) ...

在这个修改后的代码中，我们添加了一个新的 inspect_final_prompt 函数，它接收上下文和问题作为参数，然后使用 prompt.format() 方法生成最终提示并打印出来。

方法二：使用自定义 LLM Wrapper

可以创建一个自定义的 LLM wrapper，它可以在调用底层 LLM 之前拦截并打印最终提示。

from langchain_core.llms import BaseLLM

class MyOllama(BaseLLM):
    """自定义 Ollama LLM wrapper，用于打印最终提示"""

    def __init__(self, ollama_instance):
        self.ollama_instance = ollama_instance

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        print("最终提示：\n" + prompt)
        return self.ollama_instance(prompt, stop=stop)

# ... (之前的代码) ...

# 使用自定义 LLM wrapper
ollama = MyOllama(ollama) 

# ... (之后的代码) ...

在这个例子中，我们创建了一个名为 MyOllama 的自定义 LLM wrapper，它继承自 BaseLLM 。在 _call 方法中，我们首先打印最终提示，然后调用原始的 Ollama LLM 实例。

选择哪种方法取决于的偏好和具体需求。