SmartBizTrip/memory.py

# 记忆是被 ChatService 调用的辅助模块，不是独立服务。
# 启动 S4.2：创建空记忆 → 注入数据库连接 → 恢复偏好、查询历史和短期消息。
# 主流程 3/11：add_message → 截取最近消息 → 尝试写入数据库。
# 主流程 4/6：读取历史和偏好；查询类意图另外记录整段查询和当前任务。
# 旁路：网页保存偏好或清空记忆时才调用对应方法；下面各方法不是依次全执行。
# 当前对象和数据库表均未使用 user_id/session_id 区分用户；current_task 仅保存在内存。

"""
需求：实现SmartBizTrip的记忆管理模块，包括短期对话记忆、用户偏好记忆和当前任务上下文
支持 MySQL 数据库持久化，服务重启后自动恢复记忆数据

架构图：

┌────────────────────────────────────────────────────────────────────────────┐
│                         ConversationMemory  类                             │
│                                                                            │
│  ┌─────────────────── 内存层 (Memory) ───────────────────┐                │
│  │                                                        │                │
│  │  short_term_messages[]  ── 短期对话 (最多10条)         │                │
│  │  user_profile{}         ── 用户偏好 (座位/舱位等)      │                │
│  │  current_task{}         ── 当前任务上下文 (类型/城市)   │                │
│  │  entity_history[]       ── 提取的实体历史 (最多20条)    │                │
│  └────────────────────┬──────────────────────────────────┘                │
│                       │ add / update / get                                 │
│                       ▼                                                    │
│  ┌─────────────────── 持久层 (MySQL) ────────────────────┐                │
│  │                                                        │                │
│  │  _db_conn ── mysql.connector 连接                      │                │
│  │  set_db_connection() / _ensure_db()  连接管理          │                │
│  │                                                        │                │
│  │  ┌──────────────────────────────────────────────┐     │                │
│  │  │  user_profiles 表                             │     │                │
│  │  │  (profile_key PK, profile_value)              │     │                │
│  │  │  save_profile_to_db()  ← INSERT ... UPSERT   │     │                │
│  │  │  load_profile_from_db()  ← SELECT            │     │                │
│  │  └──────────────────────────────────────────────┘     │                │
│  │  ┌──────────────────────────────────────────────┐     │                │
│  │  │  query_history 表                             │     │                │
│  │  │  (id PK, intent_type, query_content,          │     │                │
│  │  │   query_time)                                 │     │                │
│  │  │  save_entity_to_db()   ← INSERT               │     │                │
│  │  │  load_entities_from_db() ← SELECT ORDER BY    │     │                │
│  │  └──────────────────────────────────────────────┘     │                │
│  │  ┌──────────────────────────────────────────────┐     │                │
│  │  │  short_term_messages 表                       │     │                │
│  │  │  (id PK, role, content, message_time,         │     │                │
│  │  │   message_order)                              │     │                │
│  │  │  save_messages_to_db()   ← DELETE + INSERT    │     │                │
│  │  │  load_messages_from_db() ← SELECT ORDER BY    │     │                │
│  │  └──────────────────────────────────────────────┘     │                │
│  │                                                        │                │
│  │  clear_all_from_db()  ← 清空三张表                     │                │
│  └────────────────────────────────────────────────────────┘                │
│                                                                            │
│  ┌─────────────────── 序列化层 ──────────────────────────┐                │
│  │                                                        │                │
│  │  to_dict()   → 导出记忆为 dict (实体仅取最近5条)       │                │
│  │  from_dict() ← 从 dict 恢复记忆 (@classmethod)         │                │
│  └────────────────────────────────────────────────────────┘                │
└────────────────────────────────────────────────────────────────────────────┘

数据流向：

    用户消息 → add_message() → short_term_messages[] → save_messages_to_db()
                                                      → MySQL: short_term_messages 表
                                                      → (DELETE + 批量 INSERT)

    偏好更新 → update_profile() → user_profile{}     → save_profile_to_db()
                                                      → MySQL: user_profiles 表
                                                      → (INSERT ... ON DUPLICATE KEY UPDATE)

    实体提取 → extract_entities() → entity_history[]  → save_entity_to_db()
                                                      → MySQL: query_history 表
                                                      → (单条 INSERT)

    服务重启 → load_profile_from_db()   → 恢复 user_profile{}
              load_messages_from_db()   → 恢复 short_term_messages[]
              load_entities_from_db()   → 恢复 entity_history[]

方法分类：

    写操作 (Write)          读操作 (Read)           序列化 (Serialize)
    ──────────────────     ──────────────────      ──────────────────
    add_message()          get_short_term_text()   to_dict()
    update_profile()       get_profile_text()      from_dict()
    update_task_context()
    extract_entities()
    clear()

    持久化 (DB Write)       加载 (DB Read)          连接管理
    ──────────────────     ──────────────────      ──────────────────
    save_messages_to_db()  load_messages_from_db() set_db_connection()
    save_profile_to_db()   load_profile_from_db()  _ensure_db()
    save_entity_to_db()    load_entities_from_db()
    clear_all_from_db()
"""
from datetime import datetime
import json
import mysql.connector
import pytz


class ConversationMemory:
    """管理对话记忆的类，包括短期对话、用户偏好和任务上下文"""

    def __init__(self, short_term_limit: int = 10):
        # 【记忆 M1：初始化】先创建空容器，数据库连接由 ChatService 稍后注入。
        # 由 ChatService 创建。short_term_messages 按时间顺序保存消息，user_profile 按键保存偏好。
        # current_task 保存最近正在处理的任务字段，entity_history 保存查询记录；它们用途不同，不会自动相互同步。
        # _db_conn 初始为 None，表示尚未启用持久化；内存操作仍可以先进行。
        self.short_term_messages = []  # 短期对话消息列表，最多保留short_term_limit条
        self.user_profile = {}  # 用户偏好，如 {"座位喜好": "二等座", "cabin_type": "经济舱"}
        self.current_task = {}  # 当前任务上下文，如 {"type": "train", "departure_city": "北京", "arrival_city": "上海"}
        self.short_term_limit = short_term_limit  # 短期记忆最大长度
        self.entity_history = []  # 历史提取的关键实体
        self._db_conn = None  # 数据库连接

    def set_db_connection(self, db_conn):
        """设置数据库连接（由 ChatService 注入）"""
        # 【记忆 M2：连接注入】保存已有连接，不在此方法里创建数据库连接。
        # 调用方传入已经建立的 MySQL 连接，也可以传 None 表示只使用内存。
        # 这是保存引用：后面的读写方法使用同一个连接对象，而不是每次复制一条连接。
        self._db_conn = db_conn

    def _ensure_db(self):
        """确保数据库连接有效，断开则自动重连"""
        # 【持久化内部步骤 1】每次读写前检查连接；有连接对象时通过 ping 尝试重连。
        # 由各个数据库读写方法在执行 SQL 前调用；正常完成不需要返回数据。
        # 没有连接就抛 RuntimeError；已有连接则 ping 检查并允许驱动重连，失败时转成统一错误。
        if self._db_conn is None:
            raise RuntimeError("数据库连接未初始化")
        try:
            self._db_conn.ping(reconnect=True)
        except Exception:
            raise RuntimeError("数据库连接已断开")

    def add_message(self, role: str, content: str):
        """添加消息到短期记忆，并持久化到数据库"""
        # 【主流程 3/11，记忆内部步骤 1】追加角色、文本和时间；用户消息和助手消息分别调用一次。
        # 输入 role 是 user/assistant，content 是要保存的全文；本方法通过修改列表产生作用，没有业务返回值。
        # 用户输入时调用一次，助手最终回复时再调用一次；时间戳用于显示消息记录时间。
        self.short_term_messages.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now(pytz.timezone('Asia/Shanghai')).strftime('%H:%M:%S')
        })
        # 超过限制则移除最旧的消息
        # 【记忆内部步骤 2】只保留最近 short_term_limit 条消息；“条”按消息计，不按对话轮数计。
        if len(self.short_term_messages) > self.short_term_limit:
            # 负数切片从末尾向前数，保留最后 N 条；不是按角色各保留 N 条，也不按自然日划分。
            self.short_term_messages = self.short_term_messages[-self.short_term_limit:]
        # 持久化到数据库
        self.save_messages_to_db()

    def get_short_term_text(self) -> str:
        """获取短期对话的文本格式，用于意图识别和代理调用"""
        # 【主流程 4/6】按消息列表顺序拼成 User/Assistant 文本，供模型和子助手理解上下文。
        # 输入来自实例里的短期消息列表，输出一个用换行连接的字符串。
        # 例如两条消息变成“User: 北京明天天气
        # Assistant: ...”；模型收到的是这段文本，不是数据库游标。
        lines = []
        for msg in self.short_term_messages:
            role_label = "User" if msg["role"] == "user" else "Assistant"
            lines.append(f"{role_label}: {msg['content']}")
        return '\n'.join(lines)

    def update_profile(self, profile_update: dict, replace: bool = False):
        """更新用户偏好，并持久化到数据库"""
        # 【偏好内部步骤 1】先留存旧值，再按 replace 选择整体替换或增量合并。
        # 调用方给出 profile_update 字典；replace=True 是整体替换，False 是按键合并。
        # 例如旧值 {"座席": "二等座", "舱位": "经济舱"}，只提交新座席时，合并保留舱位，替换会移除舱位。
        # copy 保存旧字典的浅拷贝，用于写入过程抛出异常时恢复；不是再次向数据库读取旧值。
        previous = self.user_profile.copy()
        if replace:
            self.user_profile = dict(profile_update)
        else:
            self.user_profile.update(profile_update)
        try:
            # 【偏好内部步骤 2】尝试写库；若异常向外抛出，就恢复内存中的旧偏好。
            # 当前保存方法只有 replace=True 时会继续抛出写库异常；合并模式失败只打印信息。
            self.save_profile_to_db(replace=replace)
        except Exception:
            # 恢复内存后用单独的 raise 把原异常继续向上抛，让 API 调用方知道保存失败。
            self.user_profile = previous
            raise

    def update_task_context(self, task_update: dict):
        """更新当前任务上下文"""
        # 【主流程 6，记忆旁路】合并当前任务字段；此方法不写数据库，也不调用模型提取字段。
        # dict.update 将本次字段覆盖到 current_task；未提交的旧键仍保留。
        # 当前 ChatService 通常只提交 type/query，不能把它理解为已自动提取出完整行程字段。
        self.current_task.update(task_update)

    def extract_entities(self, intent_type: str, query: str):
        """记录意图类型和查询文本到历史，并持久化到数据库。"""
        # 【主流程 6，查询历史内部步骤 1】实际保存的是意图类型和整段 query。
        # 方法名虽然叫 extract_entities，这里没有再解析城市、日期，也没有额外模型调用。
        # 输入 intent_type 记录是哪类查询，query 记录该次查询文本。
        # 先追加内存历史，再裁剪为最多 20 条，最后单独向 query_history 表追加记录。
        self.entity_history.append({
            "type": intent_type,
            "query": query,
            "timestamp": datetime.now(pytz.timezone('Asia/Shanghai')).strftime('%Y-%m-%d %H:%M:%S')
        })
        # 只保留最近20条实体记录
        if len(self.entity_history) > 20:
            self.entity_history = self.entity_history[-20:]
        # 持久化单条到数据库
        self.save_entity_to_db(intent_type, query)

    def get_profile_text(self) -> str:
        """获取用户偏好的文本描述，用于注入到prompt中"""
        # 【主流程 4】把偏好键值对转成可放入提示词的文本；没有偏好时返回固定说明。
        # 字典非空时逐对转换为“键: 值”，再用中文逗号连接；返回给 intent_agent 放进提示模板。
        if not self.user_profile:
            return "无已知的用户偏好"
        items = [f"{k}: {v}" for k, v in self.user_profile.items()]
        return "，".join(items)

    def clear(self):
        """清空所有记忆，同时清除数据库数据"""
        # 【清空内部步骤 1】先清空四类内存状态，再尝试删除数据库中的三张记忆表数据。
        # 本方法先清内存再调用 clear_all_from_db；数据库清理失败时该方法不会自动恢复原内存。
        # 因此这是代码当前执行顺序，不是跨内存和数据库的一笔原子事务。
        self.short_term_messages = []
        self.user_profile = {}
        self.current_task = {}
        self.entity_history = []
        self.clear_all_from_db()

    def to_dict(self) -> dict:
        """导出记忆为字典，用于序列化"""
        # 【辅助导出】把当前内存整理成字典，查询历史仅取最近 5 条；不是完整数据库备份。
        # 返回便于传递的普通字典；列表和字典多数仍引用原对象，没有在此处做完整深拷贝。
        return {
            "short_term_messages": self.short_term_messages,
            "user_profile": self.user_profile,
            "current_task": self.current_task,
            "entity_history": self.entity_history[-5:]  # 只导出最近5条
        }

    @classmethod
    def from_dict(cls, data: dict) -> 'ConversationMemory':
        """从字典导入记忆"""
        # 【辅助导入 1 → 2】创建新对象并填入字典中的数据，不建立连接，也不自动写回数据库。
        # @classmethod 使第一个参数 cls 指向类；cls() 创建一个新记忆实例。
        # 与普通方法的 self 不同，这里无需先有实例；缺失字段使用空列表或空字典补齐。
        memory = cls()
        memory.short_term_messages = data.get("short_term_messages", [])
        memory.user_profile = data.get("user_profile", {})
        memory.current_task = data.get("current_task", {})
        memory.entity_history = data.get("entity_history", [])
        return memory

    # ==================== 数据库持久化 ====================

    def save_profile_to_db(self, replace: bool = False):
        """将用户偏好持久化到数据库（UPSERT）"""
        # 【偏好内部步骤 2】无连接则仅保留内存；有连接时检查连接并开始写偏好。
        # 调用者 update_profile 已修改内存，此处负责同步偏好表。
        # 每个键对应一条记录；数据库用唯一键判断 INSERT 还是 UPDATE，最后统一 commit。
        if self._db_conn is None:
            return
        # 先给 cursor 一个明确初值；即使创建游标之前就失败，finally 也能判断是否需要关闭。
        cursor = None
        try:
            self._ensure_db()
            cursor = self._db_conn.cursor()
            # 替换模式先删除旧偏好，保证网页删除的键不会残留；合并模式只做 UPSERT。
            if replace:
                cursor.execute("DELETE FROM user_profiles")
            for key, value in self.user_profile.items():
                cursor.execute(
                    "INSERT INTO user_profiles (profile_key, profile_value) "
                    "VALUES (%s, %s) ON DUPLICATE KEY UPDATE profile_value = %s",
                    # 三个参数依次对应 INSERT 的键、INSERT 的值、冲突时 UPDATE 的值。
                    # value 转成字符串写库，因此重新加载后不会自动恢复为原来的数字或列表类型。
                    (key, str(value), str(value))
                )
            # 【偏好内部步骤 3】本轮写入提交后才完成；异常分支 rollback，finally 关闭游标。
            self._db_conn.commit()
        except Exception as e:
            # 回滚本次尚未提交的数据库修改；后面的 if replace 决定是否继续向调用方抛出异常。
            self._db_conn.rollback()
            if replace:
                raise
            print(f"保存用户偏好到数据库失败: {e}")
        # finally 在正常和异常路径都会执行；只负责释放游标，不改变本轮成功或失败的结果。
        finally:
            if cursor is not None:
                cursor.close()

    def load_profile_from_db(self):
        """从数据库加载用户偏好"""
        # 【启动恢复 1】读取数据库偏好，重建内存中的 user_profile 字典。
        # 服务初始化时调用。返回的 rows 是字典列表，最后的字典推导式把它改成“偏好键→偏好值”。
        # 这里只 SELECT，不需要提交写事务；加载结果会整体替换当前 user_profile。
        if self._db_conn is None:
            return
        try:
            self._ensure_db()
            cursor = self._db_conn.cursor(dictionary=True)
            cursor.execute("SELECT profile_key, profile_value FROM user_profiles")
            rows = cursor.fetchall()
            cursor.close()
            self.user_profile = {row["profile_key"]: row["profile_value"] for row in rows}
        except Exception as e:
            print(f"从数据库加载用户偏好失败: {e}")

    def save_entity_to_db(self, intent_type: str, query: str):
        """将单条查询实体持久化到数据库"""
        # 【查询历史内部步骤 2】每次新增一条历史记录；内存只保留 20 条不代表表中也只存 20 条。
        # 输入是本次查询类型和文本；query_content 保存成 JSON，query_time 单独保存完整日期时间。
        # 此处只追加一条记录，不会因为内存裁剪而删除旧数据库历史。
        if self._db_conn is None:
            return
        try:
            self._ensure_db()
            cursor = self._db_conn.cursor()
            now = datetime.now(pytz.timezone('Asia/Shanghai')).strftime('%Y-%m-%d %H:%M:%S')
            cursor.execute(
                "INSERT INTO query_history (intent_type, query_content, query_time) "
                "VALUES (%s, %s, %s)",
                (intent_type, json.dumps({"query": query}, ensure_ascii=False), now)
            )
            self._db_conn.commit()
            cursor.close()
        except Exception as e:
            print(f"保存查询实体到数据库失败: {e}")

    def load_entities_from_db(self, limit: int = 20):
        """从数据库加载查询历史，按时间倒序取最近N条"""
        # 【启动恢复 2】数据库先取最新的 limit 条，再反转为从旧到新的阅读顺序。
        # limit 限制加载条数，默认 20。先取最新记录，再恢复为最旧在前、最新在后的列表。
        # 每条 query_content 单独尝试 json.loads；某条解析失败时，该条 query 使用空串，继续处理其余记录。
        if self._db_conn is None:
            return
        try:
            self._ensure_db()
            cursor = self._db_conn.cursor(dictionary=True)
            cursor.execute(
                "SELECT intent_type, query_content, query_time FROM query_history "
                "ORDER BY query_time DESC LIMIT %s",
                # 逗号表示只有一个元素的元组，用作 SQL 参数容器；(limit) 仍只是普通数值。
                (limit,)
            )
            rows = cursor.fetchall()
            cursor.close()
            # 按时间正序排列（倒序取出来后翻转）
            # SQL 用 DESC 取最近记录，reverse 原地反转列表，恢复聊天中习惯的时间正序。
            rows.reverse()
            self.entity_history = []
            for row in rows:
                try:
                    query_data = json.loads(row["query_content"])
                    query_text = query_data.get("query", "")
                except Exception:
                    query_text = ""
                self.entity_history.append({
                    "type": row["intent_type"],
                    "query": query_text,
                    "timestamp": row["query_time"].strftime('%Y-%m-%d %H:%M:%S') if row["query_time"] else ""
                })
        except Exception as e:
            print(f"从数据库加载查询历史失败: {e}")

    def save_messages_to_db(self):
        """将短期对话覆盖写入数据库（先删除旧数据，再写入当前列表）"""
        # 【主流程 3/11，记忆内部步骤 3】无连接时跳过；有连接时用当前短期列表覆盖表内数据。
        # 保存的是截取后的最近消息，不是完整聊天记录；message_order 用于恢复原顺序。
        # 保存策略是当前短期消息的整表快照：DELETE 旧快照 → 逐条 INSERT → commit。
        # 与 query_history 的“每次追加”不同，short_term_messages 表不会在此路径保留无限历史。
        if self._db_conn is None:
            return
        try:
            self._ensure_db()
            cursor = self._db_conn.cursor()
            # 【记忆内部步骤 4】先删旧列表，再按序插入当前列表，最后 commit。
            cursor.execute("DELETE FROM short_term_messages")
            # enumerate 同时给出下标 i 和消息 msg；i 从 0 开始，写入 message_order 供加载时排序。
            for i, msg in enumerate(self.short_term_messages):
                cursor.execute(
                    "INSERT INTO short_term_messages (role, content, message_time, message_order) "
                    "VALUES (%s, %s, %s, %s)",
                    (msg["role"], msg["content"], msg["timestamp"], i)
                )
            self._db_conn.commit()
            cursor.close()
        except Exception as e:
            print(f"保存短期对话到数据库失败: {e}")

    def load_messages_from_db(self):
        """从数据库加载短期对话"""
        # 【启动恢复 3】按 message_order 恢复短期消息；随后 ChatService 同步到自己的 messages。
        # 按 message_order 升序还原保存时的列表次序，而不是按格式化后的时间字符串排序。
        # 若没有查到记录，列表推导式得到 []，短期记忆恢复为空列表。
        if self._db_conn is None:
            return
        try:
            self._ensure_db()
            cursor = self._db_conn.cursor(dictionary=True)
            cursor.execute(
                "SELECT role, content, message_time FROM short_term_messages ORDER BY message_order ASC"
            )
            rows = cursor.fetchall()
            cursor.close()
            self.short_term_messages = [
                {"role": row["role"], "content": row["content"], "timestamp": row["message_time"]}
                for row in rows
            ]
        except Exception as e:
            print(f"从数据库加载短期对话失败: {e}")

    def clear_all_from_db(self):
        """清空数据库中所有记忆数据"""
        # 【清空内部步骤 2】删除三张记忆表的记录并提交；不删除天气、票务等业务数据。
        # 三条 DELETE 依次清除短期消息、查询历史、用户偏好，最后统一提交。
        # 这三条语句没有 WHERE 用户条件，影响的是这些表中的所有记录。
        if self._db_conn is None:
            return
        try:
            self._ensure_db()
            cursor = self._db_conn.cursor()
            cursor.execute("DELETE FROM short_term_messages")
            cursor.execute("DELETE FROM query_history")
            cursor.execute("DELETE FROM user_profiles")
            self._db_conn.commit()
            cursor.close()
        except Exception as e:
            print(f"清空数据库记忆失败: {e}")