"""会话与消息存取。 MySQL 是会话历史的**事实源**:图状态(checkpoint)丢了可以从这里补种, 反过来不行。所以每次对话结束都必须有一条 assistant 消息落库。 约定:``messages`` 里的 ``created_at`` 只到秒级且同秒可能重复,因此所有 「取最近 N 条」的查询都带上 ``rowid``/``id`` 的次级排序,保证分页稳定。 """ from __future__ import annotations import json from typing import Any, Sequence from db import db_conn from obs import get_logger logger = get_logger(__name__) def _dumps(value: Any) -> str | None: """把结构化字段序列化成 JSON 字符串;None 保持 None。""" if value is None: return None if isinstance(value, str): return value return json.dumps(value, ensure_ascii=False) def _loads(value: Any, default: Any) -> Any: """反序列化 JSON 字段;空值或坏数据回落默认值(不抛异常)。""" if value is None or value == "": return default if isinstance(value, (list, dict)): return value try: return json.loads(value) except (TypeError, ValueError): return default # ============================================================ # 会话 # ============================================================ def create_session(user_name: str, mode: str = "full", title: str = "") -> dict[str, Any]: """新建会话并返回其完整行。会话 ID 即 LangGraph thread_id。""" session_id = db_conn.new_id() stamp = db_conn.fmt_dt() with db_conn.transaction(): db_conn.execute( "INSERT INTO sessions (id, title, user_name, mode, pinned, created_at, updated_at)" " VALUES (?, ?, ?, ?, 0, ?, ?)", (session_id, title, user_name, mode, stamp, stamp), ) logger.info("新建会话 session_id=%s user=%s mode=%s", session_id, user_name, mode) return get_session(session_id) or {} def get_session(session_id: str) -> dict[str, Any] | None: """按 ID 取会话;不存在返回 None。""" row = db_conn.query_one("SELECT * FROM sessions WHERE id = ?", (session_id,)) return _hydrate_session(row) def get_session_for_user(session_id: str, user_name: str) -> dict[str, Any] | None: """取会话并校验归属。用于 400/404 判定前的存在性与越权检查。""" session = get_session(session_id) if session and session.get("user_name") != user_name: return None return session def list_sessions(user_name: str, limit: int = 50, offset: int = 0) -> list[dict[str, Any]]: """列出某用户的会话,置顶优先、更新时间倒序。""" rows = db_conn.query( "SELECT * FROM sessions WHERE user_name = ?" " ORDER BY pinned DESC, updated_at DESC, id DESC LIMIT ? OFFSET ?", (user_name, limit, offset), ) return [_hydrate_session(row) for row in rows] def search_sessions(user_name: str, keyword: str, limit: int = 50) -> list[dict[str, Any]]: """按标题模糊搜索会话。""" rows = db_conn.query( "SELECT * FROM sessions WHERE user_name = ? AND title LIKE ?" " ORDER BY pinned DESC, updated_at DESC, id DESC LIMIT ?", (user_name, f"%{keyword}%", limit), ) return [_hydrate_session(row) for row in rows] def update_session_title(session_id: str, title: str) -> None: """更新会话标题。""" db_conn.execute( "UPDATE sessions SET title = ?, updated_at = ? WHERE id = ?", (title, db_conn.fmt_dt(), session_id), ) def update_session_mode(session_id: str, mode: str) -> None: """更新会话的权限模式(plan | full)。""" db_conn.execute( "UPDATE sessions SET mode = ?, updated_at = ? WHERE id = ?", (mode, db_conn.fmt_dt(), session_id), ) def update_session_pin(session_id: str, pinned: bool) -> None: """设置/取消会话置顶。""" db_conn.execute( "UPDATE sessions SET pinned = ?, updated_at = ? WHERE id = ?", (1 if pinned else 0, db_conn.fmt_dt(), session_id), ) def touch_session(session_id: str) -> None: """刷新会话的 updated_at,让最近活跃的会话排在列表前面。""" db_conn.execute( "UPDATE sessions SET updated_at = ? WHERE id = ?", (db_conn.fmt_dt(), session_id) ) def delete_session(session_id: str) -> None: """删除会话及其全部消息、用量、压缩存档。 注意:调用方还须清除 LangGraph checkpoint 线程,否则图状态会残留并 在下次复用同 ID 时「诈尸」。 """ with db_conn.transaction(): db_conn.execute("DELETE FROM business_session_agents WHERE session_id = ?", (session_id,)) db_conn.execute("DELETE FROM messages WHERE session_id = ?", (session_id,)) db_conn.execute("DELETE FROM session_context_usage WHERE session_id = ?", (session_id,)) db_conn.execute("DELETE FROM context_archives WHERE session_id = ?", (session_id,)) db_conn.execute("DELETE FROM execution_runs WHERE session_id = ?", (session_id,)) db_conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,)) logger.info("删除会话 session_id=%s", session_id) # ============================================================ # 消息 # ============================================================ def save_message( session_id: str, role: str, content: str, tool_calls: Any = None, citations: Any = None, ) -> str: """落一条消息,返回消息 ID。user 与 assistant 都走这里。""" message_id = db_conn.new_id() stamp = db_conn.fmt_dt() with db_conn.transaction(): db_conn.execute( "INSERT INTO messages (id, session_id, msg_role, content, tool_calls, citations, created_at)" " VALUES (?, ?, ?, ?, ?, ?, ?)", ( message_id, session_id, role, content or "", _dumps(tool_calls), _dumps(citations), stamp, ), ) db_conn.execute( "UPDATE sessions SET updated_at = ? WHERE id = ?", (stamp, session_id) ) return message_id def get_messages( session_id: str, limit: int | None = None, offset: int = 0 ) -> list[dict[str, Any]]: """取会话消息,按时间正序返回。 ``limit`` 给定时取**最近** limit 条(先倒序截断再正序还原),这正是 冷启动补种图状态需要的语义。 """ if limit is None: rows = db_conn.query( "SELECT * FROM messages WHERE session_id = ? ORDER BY created_at ASC, id ASC" " LIMIT ? OFFSET ?", (session_id, 100000, offset), ) elif offset == 0: rows = db_conn.query( "SELECT * FROM (SELECT * FROM messages WHERE session_id = ?" " ORDER BY created_at DESC, id DESC LIMIT ?) AS recent" " ORDER BY created_at ASC, id ASC", (session_id, limit), ) else: # 带偏移的分页:先按正序定位窗口 rows = db_conn.query( "SELECT * FROM messages WHERE session_id = ?" " ORDER BY created_at ASC, id ASC LIMIT ? OFFSET ?", (session_id, limit, offset), ) return [_hydrate_message(row) for row in rows] def count_messages(session_id: str) -> int: """会话消息条数。""" return int(db_conn.scalar("SELECT COUNT(*) FROM messages WHERE session_id = ?", (session_id,), 0) or 0) def delete_message(message_id: str) -> bool: """删除单条消息,返回是否命中。""" return db_conn.execute("DELETE FROM messages WHERE id = ?", (message_id,)) > 0 # ============================================================ # 上下文用量与压缩存档 # ============================================================ def save_context_usage( session_id: str, turn: int, input_tokens: int, output_tokens: int, cache_read_tokens: int = 0, usage_source: str = "provider", compacted: bool = False, compact_count: int = 0, snapshot: Any = None, ) -> str: """记录一轮对话的上下文用量。""" row_id = db_conn.new_id() total = int(input_tokens) + int(output_tokens) db_conn.execute( "INSERT INTO session_context_usage (id, session_id, turn, input_tokens, output_tokens," " cache_read_tokens, total_tokens, usage_source, compacted, compact_count, snapshot, created_at)" " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( row_id, session_id, int(turn), int(input_tokens), int(output_tokens), int(cache_read_tokens), total, usage_source, 1 if compacted else 0, int(compact_count), _dumps(snapshot), db_conn.fmt_dt(), ), ) return row_id def get_context_usage(session_id: str, limit: int = 50) -> dict[str, Any]: """取会话的用量明细(倒序)与会话累计压缩次数。""" rows = db_conn.query( "SELECT * FROM session_context_usage WHERE session_id = ?" " ORDER BY created_at DESC, id DESC LIMIT ?", (session_id, limit), ) history = [] for row in rows: item = dict(row) item["snapshot"] = _loads(item.get("snapshot"), {}) item["compacted"] = bool(item.get("compacted")) history.append(item) compact_count = int( db_conn.scalar( "SELECT MAX(compact_count) FROM session_context_usage WHERE session_id = ?", (session_id,), 0, ) or 0 ) latest = history[0] if history else None return { "session_id": session_id, "compact_count": compact_count, "latest": latest, "history": history, } def save_context_archive( session_id: str, turn: int, file_path: str, tokens: int = 0 ) -> str: """登记一次「压缩前原文」的落盘位置,供事后追溯。""" row_id = db_conn.new_id() db_conn.execute( "INSERT INTO context_archives (id, session_id, turn, file_path, tokens, created_at)" " VALUES (?, ?, ?, ?, ?, ?)", (row_id, session_id, int(turn), file_path, int(tokens), db_conn.fmt_dt()), ) return row_id def list_context_archives(session_id: str) -> list[dict[str, Any]]: """列出会话的压缩存档索引。""" return db_conn.query( "SELECT * FROM context_archives WHERE session_id = ? ORDER BY created_at ASC", (session_id,), ) # ============================================================ # 用户偏好 # ============================================================ def save_preference(user_name: str, key: str, value: str) -> None: """写入/覆盖一条用户偏好(先查后写,避免方言相关的 upsert 语法)。""" existing = db_conn.query_one( "SELECT id FROM user_preferences WHERE user_name = ? AND pref_key = ?", (user_name, key), ) stamp = db_conn.fmt_dt() if existing: db_conn.execute( "UPDATE user_preferences SET pref_value = ?, updated_at = ? WHERE id = ?", (value, stamp, existing["id"]), ) else: db_conn.execute( "INSERT INTO user_preferences (id, user_name, pref_key, pref_value, updated_at)" " VALUES (?, ?, ?, ?, ?)", (db_conn.new_id(), user_name, key, value, stamp), ) def list_preferences(user_name: str) -> list[dict[str, Any]]: """列出某用户的全部偏好。""" return db_conn.query( "SELECT pref_key, pref_value, updated_at FROM user_preferences" " WHERE user_name = ? ORDER BY pref_key ASC", (user_name,), ) def delete_preference(user_name: str, key: str) -> bool: """删除一条偏好,返回是否命中。""" return db_conn.execute( "DELETE FROM user_preferences WHERE user_name = ? AND pref_key = ?", (user_name, key), ) > 0 def preferences_as_text(user_name: str, limit: int = 20) -> str: """把偏好渲染成可注入提示词的一行文本;无偏好返回空串。""" rows = list_preferences(user_name)[:limit] if not rows: return "" pairs = [f"{row['pref_key']}={row['pref_value']}" for row in rows] return ";".join(pairs) # ============================================================ # 行归一化 # ============================================================ def _hydrate_session(row: dict[str, Any] | None) -> dict[str, Any] | None: if not row: return None session = dict(row) session["pinned"] = bool(session.get("pinned")) return session def _hydrate_message(row: dict[str, Any]) -> dict[str, Any]: message = dict(row) message["tool_calls"] = _loads(message.get("tool_calls"), []) message["citations"] = _loads(message.get("citations"), []) return message