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"}}) 时:
    1. LangChain 会调用 lambda session_id: ... 创建 SQLiteChatMessageHistory 实例
    2. 自动从 SQLite 加载该 session_id 的历史消息 → 注入 history 占位符
    3. 模型回复后,自动调用 .add_message() 保存新消息到 SQLite

✅ 优势总结

特性 说明
持久化 重启后记忆不丢失
多会话 session_id 隔离不同用户/对话
无缝集成 完全兼容 LangChain 接口
轻量可靠 仅依赖标准库 sqlite3
可扩展 未来可加搜索、过期清理、Web API 等
posted @ 2026-01-30 00:09  船山薪火  阅读(18)  评论(0)    收藏  举报
![image](https://img2024.cnblogs.com/blog/3174785/202601/3174785-20260125205854513-941832118.jpg)