基于LCEL增强Neo4j高级RAG实现返回模型输入源功能

要让你的RAG链同时返回模型生成的回答和对应的源上下文，需要从链结构、数据格式和前端展示三个部分修改，具体步骤如下：

1. 定义输出数据结构 #

创建一个Pydantic模型来规范返回结果，确保同时包含回答和源信息：

from pydantic import BaseModel

class AnswerWithSources(BaseModel):
    answer: str
    sources: list  # 存储检索到的上下文文档

2. 修改链结构，保留上下文并组合结果 #

原链仅输出回答，现在需要调整链逻辑，保留检索到的上下文并与回答组合：

from langchain_core.runnables import RunnablePassthrough

# 定义生成回答的子链
generate_answer = (
    prompt
    | model
    | StrOutputParser()
)

# 重构主链：获取上下文→生成回答→组合结果
chain = (
    RunnableParallel(
        {
            "context": itemgetter("question") | retriever,
            "question": itemgetter("question"),
            "history": itemgetter("history")
        }
    )
    | RunnablePassthrough.assign(answer=generate_answer)  # 保留上下文同时生成回答
    | lambda x: {"answer": x["answer"], "sources": x["context"]}  # 整理成目标格式
    | AnswerWithSources.model_validate  # 用Pydantic验证输出结构
)

# 更新输入输出类型绑定
chain = chain.with_types(input_type=Question, output_type=AnswerWithSources)

3. 修改前端展示逻辑 #

在Streamlit部分添加源上下文的展示，让用户可以查看模型参考的原始内容：

if user_question := st.chat_input():
    st.chat_message("human").write(user_question)
    config = {"configurable": {"session_id": "any"}}

    response = chain_with_history.invoke({"question": user_question}, config)
    
    print("Response:", response)
    # 显示AI回答
    st.chat_message("ai").write(response.answer)
    # 展开显示源上下文
    with st.expander("查看参考源"):
        for idx, source in enumerate(response.sources, 1):
            st.markdown(f"**源 {idx}:**")
            # 处理Document类型的源（带page_content和metadata）
            if hasattr(source, 'page_content'):
                st.write(source.page_content)
                if source.metadata:
                    st.write("元数据:", source.metadata)
            else:
                st.write(str(source))

可选：优化提示词（关联回答与源） #

如果需要让AI在回答中标注引用的源位置，可以修改提示词：

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个AI聊天机器人，仅基于给定上下文回答问题。无法回答时请直接说「我不知道」。引用上下文内容时，请用[1][2]等标记对应源的序号。"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "给定对话历史：{history}，以及上下文：\n{context}\n，请回答问题：{question}。"),
    ]
)

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