基于 Streamlit 与 LangChain 构建具备对话记忆的 AI Agent

一、引言

在构建智能对话系统时，对话记忆是实现连贯多轮对话的关键。本文将以 basic_memory.py 为例，详细讲解如何基于 Streamlit 和 LangChain 构建具备对话记忆的 AI Agent。

二、代码剖析

1. 初始化对话记忆

msgs = StreamlitChatMessageHistory(key="langchain_messages")
if len(msgs.messages) == 0:msgs.add_ai_message("How can I help you?")

原理剖析：

• 使用 StreamlitChatMessageHistory 类创建对话记忆对象，指定 key 参数为 "langchain_messages"。
• 如果对话消息为空，则添加一条 AI 消息 "How can I help you?" 作为对话起点。
• 内存存储机制：StreamlitChatMessageHistory 会将消息存储在 Streamlit 的 Session State 中，利用 Streamlit 的会话机制实现记忆的持久化。

流程图：

2. 构建对话链

prompt = ChatPromptTemplate.from_messages([("system", "You are an AI chatbot having a conversation with a human."),MessagesPlaceholder(variable_name="history"),("human", "{question}"),]
)
chain = prompt | ChatOpenAI(api_key=openai_api_key)

原理剖析：

• 使用 ChatPromptTemplate 构建对话提示模板，包含系统消息、历史消息占位符和人类消息。
• 历史消息传递：通过 MessagesPlaceholder 将历史消息融入提示模板，确保模型能获取完整对话上下文。
• 将提示模板与 OpenAI 模型组合成对话链，实现对话逻辑的封装。

架构图：

3. 整合对话历史

chain_with_history = RunnableWithMessageHistory(chain,lambda session_id: msgs,input_messages_key="question",history_messages_key="history",
)

原理剖析：

• 使用 RunnableWithMessageHistory 将对话链与对话记忆整合。
• 记忆绑定：通过 lambda 函数将 msgs 对象绑定到对话链，实现记忆的自动加载与更新。
• 指定输入消息键和历史消息键，确保消息的正确传递。

流程图：

4. 渲染对话消息

for msg in msgs.messages:st.chat_message(msg.type).write(msg.content)

原理剖析：

• 遍历对话消息列表，使用 Streamlit 的 chat_message 方法渲染对话消息。
• 消息类型区分：根据消息类型（人类或 AI）进行不同样式的消息展示。

界面图：

5. 条件设置

if "openai_api_key" in st.secrets:openai_api_key = st.secrets.openai_api_key
else:openai_api_key = st.sidebar.text_input("OpenAI API Key", type="password")
if not openai_api_key:st.info("Enter an OpenAI API Key to continue")st.stop()

原理剖析：

• 检查 Streamlit 秘密配置中是否存在 OpenAI API 密钥，若不存在则通过侧边栏输入获取。
• 权限控制：若未获取到 API 密钥，则阻止后续操作，防止未授权访问。

流程图：

6. 处理新消息

if prompt := st.chat_input():st.chat_message("human").write(prompt)response = chain_with_history.invoke({"question": prompt}, config)st.chat_message("ai").write(response.content)

原理剖析：

• 捕获用户输入的新消息，通过 RunnableWithMessageHistory 调用对话链生成响应。
• 自动记忆更新：新消息和响应会自动存储到对话记忆中，实现多轮对话的记忆延续。

流程图：

三、代码整体架构图与解析

整体架构图

整体架构解析

1. Streamlit 应用入口：

   • 使用 st.set_page_config 和 st.title 设置应用的基本信息。
   • 提供应用的简介和源码链接。

2. 初始化对话记忆：

   • 创建 StreamlitChatMessageHistory 对象，用于存储对话消息。
   • 如果消息列表为空，添加一条初始 AI 消息。

3. 设置 API 密钥：

   • 检查 Streamlit 秘密配置中是否存在 OpenAI API 密钥。
   • 若不存在，通过侧边栏输入获取密钥。

4. 构建对话链：

   • 使用 ChatPromptTemplate 定义对话提示模板。
   • 将提示模板与 OpenAI 模型组合成对话链。

5. 整合对话历史：

   • 使用 RunnableWithMessageHistory 将对话链与对话记忆整合。
   • 绑定消息历史，确保对话链能获取完整上下文。

6. 渲染对话消息：

   • 遍历消息列表，使用 Streamlit 的 chat_message 方法渲染对话消息。

7. 条件设置：

   • 检查是否获取到有效的 API 密钥。
   • 如果未获取到密钥，阻止后续操作。

8. 处理新消息：

   • 捕获用户输入的新消息。
   • 通过整合后的对话链生成响应。

9. 更新对话记忆：

   • 自动将新消息和响应存储到对话记忆中。
   • 确保多轮对话的连贯性。

通过以上架构，代码实现了对话记忆的存储、调用和更新，确保了多轮对话的连贯性和上下文一致性。

四、完整代码

from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_openai import ChatOpenAIimport streamlit as stst.set_page_config(page_title="StreamlitChatMessageHistory", page_icon="📖")
st.title("📖 StreamlitChatMessageHistory")"""
A basic example of using StreamlitChatMessageHistory to help LLMChain remember messages in a conversation.
The messages are stored in Session State across re-runs automatically. You can view the contents of Session State
in the expander below. View the
[source code for this app](https://github.com/langchain-ai/streamlit-agent/blob/main/streamlit_agent/basic_memory.py).
"""# Set up memory
msgs = StreamlitChatMessageHistory(key="langchain_messages")
if len(msgs.messages) == 0:msgs.add_ai_message("How can I help you?")view_messages = st.expander("View the message contents in session state")# Get an OpenAI API Key before continuing
if "openai_api_key" in st.secrets:openai_api_key = st.secrets.openai_api_key
else:openai_api_key = st.sidebar.text_input("OpenAI API Key", type="password")
if not openai_api_key:st.info("Enter an OpenAI API Key to continue")st.stop()# Set up the LangChain, passing in Message Historyprompt = ChatPromptTemplate.from_messages([("system", "You are an AI chatbot having a conversation with a human."),MessagesPlaceholder(variable_name="history"),("human", "{question}"),]
)chain = prompt | ChatOpenAI(api_key=openai_api_key)
chain_with_history = RunnableWithMessageHistory(chain,lambda session_id: msgs,input_messages_key="question",history_messages_key="history",
)# Render current messages from StreamlitChatMessageHistory
for msg in msgs.messages:st.chat_message(msg.type).write(msg.content)# If user inputs a new prompt, generate and draw a new response
if prompt := st.chat_input():st.chat_message("human").write(prompt)# Note: new messages are saved to history automatically by Langchain during runconfig = {"configurable": {"session_id": "any"}}response = chain_with_history.invoke({"question": prompt}, config)st.chat_message("ai").write(response.content)# Draw the messages at the end, so newly generated ones show up immediately
with view_messages:"""Message History initialized with:```pythonmsgs = StreamlitChatMessageHistory(key="langchain_messages")```Contents of `st.session_state.langchain_messages`:"""view_messages.json(st.session_state.langchain_messages)

本教程中用作参考剖析的代码来源于 langchain-ai 开源项目。