23-day5-memory-agent_with_SQLite
✅ 虚拟环境
# 创建目录
mkdir -p day5_memory && cd day5_memory
# 创建虚拟环境(Python 3.10+)
python3 -m venv day5-memory
source day5-memory/bin/activate
# 升级 pip
pip install --upgrade pip -i https://mirrors.aliyun.com/pypi/simple/
# 创建 .env 文件(替换为您的实际API Key)
echo "DASHSCOPE_API_KEY=sk-6adfbafe2c854374b32c720982666ce5" > .env
# 安装依赖
pip install \
langchain \
langchain-community \
langchain-openai \
python-dotenv \
-i https://mirrors.aliyun.com/pypi/simple/
✅ 程序:agent_with_sqlite.py
tee agent_with_sqlite.py <<'EOF'
# agent_with_sqlite.py
import os
import sqlite3
from contextlib import contextmanager
from typing import List, Optional
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
# ==============================
# 🔽 自定义 SQLite 聊天历史类
# ==============================
DB_FILE = "chat_memory.db"
@contextmanager
def get_db_connection():
"""数据库连接上下文管理器,确保自动关闭"""
conn = sqlite3.connect(DB_FILE, check_same_thread=False)
try:
yield conn
finally:
conn.close()
def init_db():
"""初始化数据库表"""
with get_db_connection() as conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS chat_history (
session_id TEXT NOT NULL,
message_index INTEGER NOT NULL,
role TEXT NOT NULL CHECK(role IN ('human', 'ai')),
content TEXT NOT NULL,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (session_id, message_index)
)
""")
# 可选:创建索引加速查询
conn.execute("CREATE INDEX IF NOT EXISTS idx_session ON chat_history(session_id);")
conn.commit()
class SQLiteChatMessageHistory(BaseChatMessageHistory):
"""基于 SQLite 的聊天历史存储,兼容 LangChain"""
def __init__(self, session_id: str):
if not session_id:
raise ValueError("session_id 不能为空")
self.session_id = session_id
init_db() # 确保表存在
@property
def messages(self) -> List[BaseMessage]:
"""从数据库加载消息(供 LangChain 读取)"""
messages = []
with get_db_connection() as conn:
cur = conn.execute(
"""
SELECT role, content
FROM chat_history
WHERE session_id = ?
ORDER BY message_index
""",
(self.session_id,)
)
for row in cur.fetchall():
role, content = row
if role == "human":
messages.append(HumanMessage(content=content))
elif role == "ai":
messages.append(AIMessage(content=content))
return messages
def add_message(self, message: BaseMessage) -> None:
"""保存单条消息到数据库"""
if isinstance(message, HumanMessage):
role = "human"
elif isinstance(message, AIMessage):
role = "ai"
else:
raise ValueError(f"不支持的消息类型: {type(message)}")
with get_db_connection() as conn:
# 获取当前会话最大 index
cur = conn.execute(
"SELECT COALESCE(MAX(message_index), -1) FROM chat_history WHERE session_id = ?",
(self.session_id,)
)
next_index = cur.fetchone()[0] + 1
conn.execute(
"INSERT INTO chat_history (session_id, message_index, role, content) VALUES (?, ?, ?, ?)",
(self.session_id, next_index, role, message.content)
)
conn.commit()
def clear(self) -> None:
"""清空当前会话历史"""
with get_db_connection() as conn:
conn.execute("DELETE FROM chat_history WHERE session_id = ?", (self.session_id,))
conn.commit()
# ==============================
# 🔽 LangChain 配置
# ==============================
load_dotenv()
llm = ChatOpenAI(
model="qwen-max",
openai_api_key=os.getenv("DASHSCOPE_API_KEY"),
openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
temperature=0.7
)
prompt = ChatPromptTemplate.from_messages([
("system", "你是一个有记忆的 AI 助手。请记住用户之前说过的话。"),
MessagesPlaceholder(variable_name="history"),
("human", "{input}")
])
chain = prompt | llm
# 使用自定义历史类
with_message_history = RunnableWithMessageHistory(
chain,
lambda session_id: SQLiteChatMessageHistory(session_id),
input_messages_key="input",
history_messages_key="history",
)
# ==============================
# 🔽 主程序
# ==============================
def show_memory(session_id: str):
"""显示当前会话记忆"""
history = SQLiteChatMessageHistory(session_id)
msgs = history.messages
print("🧠 当前记忆内容:")
if not msgs:
print(" (无记忆)")
else:
for i, msg in enumerate(msgs, 1):
role = "👤 用户" if isinstance(msg, HumanMessage) else "🤖 Agent"
print(f" {i}. {role}: {msg.content}")
print()
if __name__ == "__main__":
print("🤖 多会话记忆型 Agent 启动(SQLite 持久化版)!")
print("指令:")
print(" - 输入 'quit' 退出")
print(" - 输入 'show_memory' 查看当前会话记忆")
print(" - 输入 'switch <session_id>' 切换会话(如 switch alice)")
print(" - 默认会话 ID: default\n")
current_session = "default"
while True:
try:
user_input = input(f"👤 [{current_session}] 你: ").strip()
except (KeyboardInterrupt, EOFError):
print("\n👋 再见!")
break
if not user_input:
continue
if user_input.lower() == "quit":
break
elif user_input.lower() == "show_memory":
show_memory(current_session)
continue
elif user_input.lower().startswith("switch "):
parts = user_input.split(" ", 1)
new_session = parts[1].strip() if len(parts) > 1 else ""
if not new_session:
print("⚠️ 用法: switch <session_id>")
continue
current_session = new_session
print(f"🔄 已切换到会话: '{current_session}'\n")
continue
# 调用带记忆的链(自动使用 SQLite 存储)
response = with_message_history.invoke(
{"input": user_input},
config={"configurable": {"session_id": current_session}}
)
print(f"🤖 Agent: {response.content}\n")
EOF
运行:
python agent_with_sqlite.py
数据文件:
- 所有聊天记录自动保存,数据库文件名
chat_memory.db - 数据表名:
chat_history
🔗 LangChain 调用方法说明
核心机制:
with_message_history = RunnableWithMessageHistory(
chain,
lambda session_id: SQLiteChatMessageHistory(session_id), # 👈 关键:返回历史对象
input_messages_key="input",
history_messages_key="history",
)
- 每次调用
.invoke(..., config={"configurable": {"session_id": "xxx"}})时:- LangChain 会调用
lambda session_id: ...创建SQLiteChatMessageHistory实例 - 自动从 SQLite 加载该
session_id的历史消息 → 注入history占位符 - 模型回复后,自动调用
.add_message()保存新消息到 SQLite
- LangChain 会调用
✅ 优势总结
| 特性 | 说明 |
|---|---|
| 持久化 | 重启后记忆不丢失 |
| 多会话 | session_id 隔离不同用户/对话 |
| 无缝集成 | 完全兼容 LangChain 接口 |
| 轻量可靠 | 仅依赖标准库 sqlite3 |
| 可扩展 | 未来可加搜索、过期清理、Web API 等 |
浙公网安备 33010602011771号