716 lines
34 KiB
Python
716 lines
34 KiB
Python
import sqlite3
|
||
import os
|
||
import sys
|
||
import re
|
||
import json
|
||
import uuid
|
||
import time
|
||
import shutil
|
||
from typing import List, Dict, Optional
|
||
|
||
# 🝙 P0 fix: DB path based on the file's location (no longer depends on the process's working directory)
|
||
# 🌟 打包（PyInstaller onedir）时：优先与源码树共用 data/chat_history.db（历史不丢失）；
|
||
# 若 exe 被复制到源码树之外 → 回退为 exe 旁的 data/chat_history.db
|
||
if getattr(sys, "frozen", False):
|
||
_exe_dir = os.path.dirname(os.path.abspath(sys.executable))
|
||
_shared = os.path.abspath(os.path.join(_exe_dir, "..", "..", "data", "chat_history.db"))
|
||
if os.path.isdir(os.path.dirname(_shared)):
|
||
_DEFAULT_DB = _shared
|
||
else:
|
||
_DEFAULT_DB = os.path.join(_exe_dir, "data", "chat_history.db")
|
||
else:
|
||
_DEFAULT_DB = os.path.join(
|
||
os.path.dirname(os.path.abspath(__file__)), "..", "data", "chat_history.db")
|
||
|
||
|
||
class _AutoCloseConn:
|
||
"""🝙 P0 fix: wraps sqlite3.Connection.
|
||
|
||
sqlite3.Connection's with only handles commit/rollback and does NOT close.
|
||
All 19 call sites use `with self.get_connection() as conn:`, so we auto-close on with exit.
|
||
"""
|
||
|
||
def __init__(self, conn):
|
||
self._conn = conn
|
||
|
||
def __getattr__(self, name):
|
||
return getattr(self._conn, name)
|
||
|
||
def __enter__(self):
|
||
self._conn.__enter__()
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc, tb):
|
||
try:
|
||
return self._conn.__exit__(exc_type, exc, tb)
|
||
finally:
|
||
try:
|
||
self._conn.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
class DBManager:
|
||
def __init__(self, db_path: str = None):
|
||
# 🝙 Default path is anchored to the core/ directory, not affected by launch CWD
|
||
self.db_path = db_path or os.path.abspath(_DEFAULT_DB)
|
||
|
||
# 🆕 附件/媒体文件根目录（copy_session 深拷贝磁盘文件用）。默认=项目根；测试可覆盖。
|
||
self.files_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
self.is_first_run = not os.path.exists(self.db_path)
|
||
d = os.path.dirname(self.db_path)
|
||
if d:
|
||
os.makedirs(d, exist_ok=True)
|
||
|
||
self._init_db()
|
||
|
||
def get_connection(self):
|
||
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA foreign_keys = ON")
|
||
return _AutoCloseConn(conn)
|
||
|
||
def _init_db(self):
|
||
"""初始化表结构并注入默认数据"""
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
|
||
# 1. 创建 sessions 表 (新增 current_leaf_msg_id)
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS sessions (
|
||
id TEXT PRIMARY KEY,
|
||
title TEXT,
|
||
created_at INTEGER,
|
||
updated_at INTEGER,
|
||
has_messages BOOLEAN DEFAULT 0,
|
||
sort_order INTEGER DEFAULT 0,
|
||
is_starred BOOLEAN DEFAULT 0,
|
||
current_leaf_msg_id TEXT
|
||
)
|
||
""")
|
||
|
||
# 2. 创建 messages 表 (新增 parent_id)
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS messages (
|
||
id TEXT PRIMARY KEY,
|
||
session_id TEXT,
|
||
role TEXT,
|
||
content TEXT,
|
||
reasoning TEXT,
|
||
is_ignored BOOLEAN,
|
||
created_at INTEGER,
|
||
attachment_metadata TEXT,
|
||
parent_id TEXT,
|
||
timeline TEXT,
|
||
usage TEXT,
|
||
stop_reason TEXT,
|
||
error_message TEXT,
|
||
FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE
|
||
)
|
||
""")
|
||
|
||
# ==========================================
|
||
# 🌟 核心性能优化：为高频查询的字段建立索引
|
||
# ==========================================
|
||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_id ON messages(session_id)")
|
||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_parent_id ON messages(parent_id)")
|
||
|
||
# 3. 🌟 自动化热升级：检测并兼容旧数据库
|
||
self._upgrade_schema(cursor)
|
||
|
||
# 4. 检查是否需要插入初始默认对话
|
||
cursor.execute("SELECT COUNT(*) FROM sessions")
|
||
if cursor.fetchone()[0] == 0:
|
||
self._seed_default_chat(cursor)
|
||
|
||
conn.commit()
|
||
|
||
|
||
def _upgrade_schema(self, cursor):
|
||
"""检测缺少的新字段并自动补齐，如果是刚升级，则自动将旧线性数据串联成链表"""
|
||
upgraded = False
|
||
|
||
# 兼容 sessions 字段
|
||
for col in ["has_messages", "sort_order", "is_starred"]:
|
||
try:
|
||
cursor.execute(f"SELECT {col} FROM sessions LIMIT 1")
|
||
except sqlite3.OperationalError:
|
||
if col == "has_messages": cursor.execute("ALTER TABLE sessions ADD COLUMN has_messages BOOLEAN DEFAULT 0")
|
||
if col == "sort_order": cursor.execute("ALTER TABLE sessions ADD COLUMN sort_order INTEGER DEFAULT 0")
|
||
if col == "is_starred": cursor.execute("ALTER TABLE sessions ADD COLUMN is_starred BOOLEAN DEFAULT 0")
|
||
|
||
# 🌟 核心：兼容链表树架构
|
||
try:
|
||
# 兼容 sessions.mode 列（chat/worker 模式锁定，NULL=未发送过）
|
||
try:
|
||
cursor.execute("SELECT mode FROM sessions LIMIT 1")
|
||
except Exception:
|
||
cursor.execute("ALTER TABLE sessions ADD COLUMN mode TEXT")
|
||
|
||
cursor.execute("SELECT current_leaf_msg_id FROM sessions LIMIT 1")
|
||
except sqlite3.OperationalError:
|
||
print("[DB]: 正在升级 sessions 表 (添加 current_leaf_msg_id)...")
|
||
cursor.execute("ALTER TABLE sessions ADD COLUMN current_leaf_msg_id TEXT")
|
||
upgraded = True
|
||
|
||
try:
|
||
cursor.execute("SELECT attachment_metadata, parent_id FROM messages LIMIT 1")
|
||
except sqlite3.OperationalError:
|
||
print("[DB]: 正在升级 messages 表 (添加 attachment_metadata, parent_id)...")
|
||
try: cursor.execute("ALTER TABLE messages ADD COLUMN attachment_metadata TEXT")
|
||
except: pass
|
||
try: cursor.execute("ALTER TABLE messages ADD COLUMN parent_id TEXT")
|
||
except: pass
|
||
upgraded = True
|
||
|
||
# 🌟 messages.timeline 列（agent 时间线 JSON: 思考/文本/工具 按事件顺序）
|
||
try:
|
||
cursor.execute("SELECT timeline FROM messages LIMIT 1")
|
||
except sqlite3.OperationalError:
|
||
print("[DB]: 正在升级 messages 表 (添加 timeline)...")
|
||
try: cursor.execute("ALTER TABLE messages ADD COLUMN timeline TEXT")
|
||
except: pass
|
||
upgraded = True
|
||
|
||
# 🆕 P1: messages.usage 列（assistant 回复的精确 usage JSON，
|
||
# 供显示/压缩估算做 usage 锚定，对照 pi 内存态 usage 回放）
|
||
try:
|
||
cursor.execute("SELECT usage FROM messages LIMIT 1")
|
||
except sqlite3.OperationalError:
|
||
print("[DB]: 正在升级 messages 表 (添加 usage)...")
|
||
try: cursor.execute("ALTER TABLE messages ADD COLUMN usage TEXT")
|
||
except: pass
|
||
upgraded = True
|
||
|
||
# 🆕 失败轮次持久化（对照 pi: message_end 无条件入库）：
|
||
# messages.stop_reason / error_message —— 区分「正常行 / 出错行」，
|
||
# 供 UI 展示与 build_api_context 回放决策（NULL = 正常）
|
||
#
|
||
# ⚠️⚠️ 绝不能置 upgraded=True：该标志会触发下方的「旧数据链表化重构」，
|
||
# 把用户的**树状分支拍平成线性链**（数据破坏）！
|
||
# 纯追加列对本迁移自身而言是安全的，与旧库结构修复无关。
|
||
try:
|
||
cursor.execute("SELECT stop_reason, error_message FROM messages LIMIT 1")
|
||
except sqlite3.OperationalError:
|
||
print("[DB]: 正在升级 messages 表 (添加 stop_reason, error_message)...")
|
||
try: cursor.execute("ALTER TABLE messages ADD COLUMN stop_reason TEXT")
|
||
except: pass
|
||
try: cursor.execute("ALTER TABLE messages ADD COLUMN error_message TEXT")
|
||
except: pass
|
||
|
||
# 如果刚才执行了树状结构升级，立即对旧数据进行“时间线串联”修复
|
||
if upgraded:
|
||
print("[DB]: 🚀 正在执行旧数据链表化重构...")
|
||
cursor.execute("SELECT id FROM sessions")
|
||
sessions = cursor.fetchall()
|
||
for s in sessions:
|
||
sid = s['id']
|
||
cursor.execute("SELECT id FROM messages WHERE session_id = ? ORDER BY created_at ASC", (sid,))
|
||
msgs = cursor.fetchall()
|
||
if not msgs: continue
|
||
|
||
# 遍历消息，将后一条的 parent_id 指向上一条
|
||
prev_id = None
|
||
for m in msgs:
|
||
mid = m['id']
|
||
if prev_id:
|
||
cursor.execute("UPDATE messages SET parent_id = ? WHERE id = ?", (prev_id, mid))
|
||
prev_id = mid
|
||
|
||
# 最后一个 msg_id 就是这棵树的末端叶子节点
|
||
cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (prev_id, sid))
|
||
print("[DB]: ✅ 数据结构完美升级为链表树！")
|
||
|
||
|
||
def _seed_default_chat(self, cursor):
|
||
session_id = f"sess_{uuid.uuid4().hex[:12]}"
|
||
now = int(time.time())
|
||
|
||
cursor.execute(
|
||
"INSERT INTO sessions (id, title, created_at, updated_at, has_messages) VALUES (?, ?, ?, ?, ?)",
|
||
(session_id, "让我们从这里开始吧", now, now, 1)
|
||
)
|
||
|
||
sys_id = f"msg_sys_init"
|
||
cursor.execute("""
|
||
INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""", (sys_id, session_id, "system", "你是一个优秀的助手！", "", 0, now, None))
|
||
|
||
default_messages = [
|
||
{"role": "user", "content": "你好呀，你是谁？"},
|
||
{"role": "assistant", "content": "嗨！我是你的 AI 助手 ✨ ..."},
|
||
{"role": "user", "content": "那你到底能帮我做什么？"},
|
||
{"role": "assistant", "content": "简单来说，能打字问的我都聊..."}
|
||
]
|
||
|
||
prev_id = sys_id
|
||
for msg in default_messages:
|
||
msg_id = f"msg_{uuid.uuid4().hex[:16]}"
|
||
now += 1
|
||
cursor.execute("""
|
||
INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""", (msg_id, session_id, msg["role"], msg["content"], "", 0, now, prev_id))
|
||
prev_id = msg_id
|
||
|
||
# 设置默认会话的叶子节点
|
||
cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (prev_id, session_id))
|
||
|
||
# ==================== 会话 (Session) 操作 ====================
|
||
def get_session_mode(self, session_id: str) -> Optional[str]:
|
||
"""读取会话锁定的模式（chat/worker），未发送过返回 None"""
|
||
with self.get_connection() as conn:
|
||
row = conn.execute("SELECT mode FROM sessions WHERE id = ?",
|
||
(session_id,)).fetchone()
|
||
return row["mode"] if row else None
|
||
|
||
def set_session_mode(self, session_id: str, mode: str):
|
||
"""锁定会话模式（首条消息发送时调用，之后不可变）"""
|
||
with self.get_connection() as conn:
|
||
conn.execute("UPDATE sessions SET mode = ? WHERE id = ?",
|
||
(mode, session_id))
|
||
|
||
def get_all_sessions(self) -> List[Dict]:
|
||
with self.get_connection() as conn:
|
||
return [dict(row) for row in conn.execute("SELECT * FROM sessions ORDER BY is_starred DESC, sort_order ASC, updated_at DESC").fetchall()]
|
||
|
||
def create_session(self, title: str = "新对话") -> Dict:
|
||
session_id = f"sess_{uuid.uuid4().hex[:12]}"
|
||
now = int(time.time())
|
||
sys_msg_id = f"msg_{uuid.uuid4().hex[:16]}"
|
||
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
min_order = self.get_min_sort_order()
|
||
|
||
# 创建 session，直接将 system 消息设为初始叶子节点
|
||
cursor.execute(
|
||
"INSERT INTO sessions (id, title, created_at, updated_at, has_messages, sort_order, current_leaf_msg_id) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||
(session_id, title, now, now, 0, min_order, sys_msg_id)
|
||
)
|
||
|
||
cursor.execute("""
|
||
INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""", (sys_msg_id, session_id, "system", "你是一个优秀的助手！", "", 0, now, None))
|
||
conn.commit()
|
||
|
||
return dict(cursor.execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone())
|
||
|
||
# ------------------------------------------------------------------
|
||
# 🆕 会话复制：深度克隆（全部分支 + 压缩标记 + 附件文件）
|
||
# ------------------------------------------------------------------
|
||
def _make_copy_title(self, base: str) -> str:
|
||
"""生成不重名的副本标题：X → X (副本) → X (副本 2) → …"""
|
||
root = re.sub(r"\s*\(副本(?:\s*\d+)?\)\s*$", "", base or "").strip() or "新对话"
|
||
with self.get_connection() as conn:
|
||
existing = {r[0] for r in conn.execute("SELECT title FROM sessions").fetchall()}
|
||
cand = f"{root} (副本)"
|
||
n = 2
|
||
while cand in existing and n < 1000:
|
||
cand = f"{root} (副本 {n})"
|
||
n += 1
|
||
return cand
|
||
|
||
def _dup_media_file(self, old_path: str, prefix: str):
|
||
"""把一个项目内媒体文件复制成新名字。
|
||
返回 (新绝对路径, 新项目相对路径)；文件不存在或异常返回 None。"""
|
||
if not old_path:
|
||
return None
|
||
old_abs = old_path if os.path.isabs(old_path) else os.path.join(self.files_root, old_path)
|
||
if not os.path.isfile(old_abs):
|
||
return None
|
||
ext = os.path.splitext(old_abs)[1] or ".bin"
|
||
new_dir = os.path.join(self.files_root, "data", "attachments")
|
||
os.makedirs(new_dir, exist_ok=True)
|
||
new_abs = ""
|
||
for _ in range(5):
|
||
new_abs = os.path.join(new_dir, f"{prefix}_{uuid.uuid4().hex[:8]}{ext}")
|
||
if not os.path.exists(new_abs):
|
||
break
|
||
shutil.copy2(old_abs, new_abs)
|
||
new_rel = os.path.relpath(new_abs, self.files_root).replace("\\", "/")
|
||
return new_abs, new_rel
|
||
|
||
def _copy_attachment_files(self, meta_json: Optional[str]) -> Optional[str]:
|
||
"""附件深拷贝：image/pdf 的磁盘文件复制改名并改写 local_path；
|
||
text 类型正文内联在 JSON 里，无需处理。
|
||
任何异常/文件缺失都原样返回，绝不让复制整体失败。"""
|
||
if not meta_json:
|
||
return meta_json
|
||
try:
|
||
meta = json.loads(meta_json)
|
||
except Exception:
|
||
return meta_json
|
||
if not isinstance(meta, dict) or not meta.get("attachments"):
|
||
return meta_json
|
||
changed = False
|
||
for att in meta.get("attachments") or []:
|
||
if not isinstance(att, dict):
|
||
continue
|
||
try:
|
||
if att.get("type") in ("image", "pdf") and att.get("local_path"):
|
||
got = self._dup_media_file(
|
||
att["local_path"], "img" if att["type"] == "image" else "pdf")
|
||
if got:
|
||
att["local_path"] = got[1]
|
||
if "abs_path" in att:
|
||
att["abs_path"] = got[0]
|
||
changed = True
|
||
if att.get("type") == "pdf":
|
||
for im in att.get("images") or []:
|
||
if not isinstance(im, dict):
|
||
continue
|
||
got = self._dup_media_file(
|
||
im.get("abs_path") or im.get("local_path"), "pdfimg")
|
||
if got:
|
||
im["abs_path"] = got[0]
|
||
im["local_path"] = got[1]
|
||
changed = True
|
||
except Exception as e:
|
||
print(f"[DB] ⚠️ 附件深拷贝失败（保留原路径）: {e}", flush=True)
|
||
return json.dumps(meta, ensure_ascii=False) if changed else meta_json
|
||
|
||
def copy_session(self, session_id: str, new_title: Optional[str] = None,
|
||
copy_attachments: bool = True) -> Optional[Dict]:
|
||
"""📋 深度复制一个会话。
|
||
|
||
- messages 全部重新生成 ID，parent_id / current_leaf_msg_id 全量重映射
|
||
→ 分支、压缩标记（role='compaction'）都原样保留
|
||
- image/pdf 附件文件物理复制成新文件 → 副本自包含，删任意一方不影响另一方
|
||
- 单事务写入；源会话零改动
|
||
返回新会话 dict；源不存在返回 None。
|
||
"""
|
||
now = int(time.time())
|
||
n_att = 0
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
src_row = cursor.execute(
|
||
"SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
|
||
if not src_row:
|
||
return None
|
||
src = dict(src_row)
|
||
rows = [dict(r) for r in cursor.execute(
|
||
"SELECT * FROM messages WHERE session_id = ? ORDER BY created_at ASC, rowid ASC",
|
||
(session_id,)).fetchall()]
|
||
|
||
# ① 新 ID 映射（保持 comp_ / msg_ 前缀约定）
|
||
idmap = {}
|
||
for m in rows:
|
||
pre = "comp_" if m.get("role") == "compaction" else "msg_"
|
||
idmap[m["id"]] = f"{pre}{uuid.uuid4().hex[:16]}"
|
||
|
||
# ② 标题（重名自动递增）
|
||
title = new_title or self._make_copy_title(src.get("title") or "新对话")
|
||
|
||
# ③ 消息 created_at 保相对间隔（同层兄弟排序不变）
|
||
t0 = min((m.get("created_at") or 0) for m in rows) if rows else now
|
||
|
||
new_sid = f"sess_{uuid.uuid4().hex[:12]}"
|
||
new_leaf = idmap.get(src.get("current_leaf_msg_id"))
|
||
if rows and not new_leaf:
|
||
# 兜底：源叶子不在链上（数据异常）→ 取副本里时间最新的一条
|
||
last = max(rows, key=lambda m: (m.get("created_at") or 0))
|
||
new_leaf = idmap.get(last["id"])
|
||
print(f"[DB] ⚠️ copy_session 源叶子异常，回退 leaf={new_leaf}", flush=True)
|
||
|
||
# ④ 新会话（置列表顶部、不带星标、模式跟随源）
|
||
cursor.execute(
|
||
"INSERT INTO sessions (id, title, created_at, updated_at, has_messages, "
|
||
"sort_order, is_starred, current_leaf_msg_id, mode) VALUES (?,?,?,?,?,?,?,?,?)",
|
||
(new_sid, title, now, now, src.get("has_messages") or 0,
|
||
self.get_min_sort_order(), 0, new_leaf, src.get("mode")))
|
||
|
||
# ⑤ 逐条复制消息（parent 重映射 + 附件深拷贝）
|
||
for m in rows:
|
||
meta = m.get("attachment_metadata")
|
||
if copy_attachments and meta:
|
||
new_meta = self._copy_attachment_files(meta)
|
||
if new_meta != meta:
|
||
n_att += 1
|
||
meta = new_meta
|
||
cursor.execute(
|
||
"INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, "
|
||
"created_at, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message) "
|
||
"VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
||
(idmap[m["id"]], new_sid, m.get("role"), m.get("content"),
|
||
m.get("reasoning"), m.get("is_ignored"),
|
||
now + ((m.get("created_at") or 0) - t0),
|
||
meta, idmap.get(m.get("parent_id")),
|
||
m.get("timeline"), m.get("usage"),
|
||
m.get("stop_reason"), m.get("error_message")))
|
||
|
||
conn.commit()
|
||
out = dict(cursor.execute(
|
||
"SELECT * FROM sessions WHERE id = ?", (new_sid,)).fetchone())
|
||
|
||
try:
|
||
print(f"[DB] copy_session {session_id[:8]} → {new_sid[:8]} "
|
||
f"消息={len(rows)} 附件深拷贝={n_att} 标题={title}", flush=True)
|
||
except Exception:
|
||
pass
|
||
return out
|
||
|
||
def update_session_title(self, session_id: str, new_title: str):
|
||
with self.get_connection() as conn:
|
||
conn.execute("UPDATE sessions SET title = ?, updated_at = ? WHERE id = ?", (new_title, int(time.time()), session_id))
|
||
conn.commit()
|
||
|
||
def delete_session(self, session_id: str):
|
||
with self.get_connection() as conn:
|
||
conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
|
||
conn.commit()
|
||
|
||
# 下方其它基本Session功能保持不变...
|
||
def mark_session_has_messages(self, session_id: str):
|
||
with self.get_connection() as conn:
|
||
conn.execute("UPDATE sessions SET has_messages = 1 WHERE id = ?", (session_id,))
|
||
conn.commit()
|
||
|
||
def check_session_needs_title(self, session_id: str) -> bool:
|
||
with self.get_connection() as conn:
|
||
row = conn.execute("SELECT title, has_messages FROM sessions WHERE id = ?", (session_id,)).fetchone()
|
||
if row: return row[0] == "新对话" and row[1] == 1
|
||
return False
|
||
|
||
def get_min_sort_order(self) -> int:
|
||
with self.get_connection() as conn:
|
||
row = conn.execute("SELECT MIN(sort_order) FROM sessions").fetchone()
|
||
return (row[0] or 0) - 1
|
||
|
||
def update_session_order(self, ordered_ids: list):
|
||
with self.get_connection() as conn:
|
||
for idx, sid in enumerate(ordered_ids):
|
||
conn.execute("UPDATE sessions SET sort_order = ? WHERE id = ?", (idx, sid))
|
||
conn.commit()
|
||
|
||
def update_session_star(self, session_id: str, is_starred: bool):
|
||
with self.get_connection() as conn:
|
||
conn.execute("UPDATE sessions SET is_starred = ?, updated_at = ? WHERE id = ?", (1 if is_starred else 0, int(time.time()), session_id))
|
||
conn.commit()
|
||
|
||
def is_session_starred(self, session_id: str) -> bool:
|
||
with self.get_connection() as conn:
|
||
row = conn.execute("SELECT is_starred FROM sessions WHERE id = ?", (session_id,)).fetchone()
|
||
return bool(row and row[0])
|
||
|
||
|
||
# ==================== 🌟 核心：消息链表树操作 ====================
|
||
|
||
def get_session_leaf(self, session_id: str) -> Optional[str]:
|
||
"""获取会话当前的叶子节点ID"""
|
||
with self.get_connection() as conn:
|
||
row = conn.execute("SELECT current_leaf_msg_id FROM sessions WHERE id = ?", (session_id,)).fetchone()
|
||
return row[0] if row else None
|
||
|
||
def update_session_leaf(self, session_id: str, leaf_msg_id: str):
|
||
"""切换时间线：手动更新当前会话的叶子节点"""
|
||
with self.get_connection() as conn:
|
||
conn.execute("UPDATE sessions SET current_leaf_msg_id = ?, updated_at = ? WHERE id = ?",
|
||
(leaf_msg_id, int(time.time()), session_id))
|
||
conn.commit()
|
||
|
||
def get_message_chain(self, session_id: str) -> List[Dict]:
|
||
"""🚀 极客级递归拉取：顺藤摸瓜，只返回当前激活时间线上的消息！彻底断绝下游污染！"""
|
||
leaf_id = self.get_session_leaf(session_id)
|
||
if not leaf_id:
|
||
return []
|
||
|
||
chain = []
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
current_id = leaf_id
|
||
seen = set() # 🐛 环检测：数据异常（如自环标记）时快速退出，防主线程死循环卡死
|
||
|
||
# 使用 Python 循环向上回溯（比 SQL CTE 更好调试，性能对于本地几千条聊天来说在 1ms 内）
|
||
while current_id and current_id not in seen:
|
||
seen.add(current_id)
|
||
cursor.execute("SELECT * FROM messages WHERE id = ?", (current_id,))
|
||
msg = cursor.fetchone()
|
||
if not msg:
|
||
break
|
||
chain.append(dict(msg))
|
||
current_id = msg['parent_id']
|
||
if current_id in seen:
|
||
try:
|
||
print(f"[DB] ⚠️ get_message_chain 检测到环（session={session_id}），已截断", flush=True)
|
||
except Exception:
|
||
pass
|
||
|
||
# 因为是向上回溯，拉出来的链条是反的，最后翻转一下恢复正序
|
||
chain.reverse()
|
||
return chain
|
||
|
||
def get_branch_info(self, parent_id: str) -> List[Dict]:
|
||
"""获取某一父节点下的所有子分支消息 (第二阶段用于UI渲染 '1/3')"""
|
||
if not parent_id: return []
|
||
with self.get_connection() as conn:
|
||
return [dict(row) for row in conn.execute(
|
||
"SELECT * FROM messages WHERE parent_id = ? ORDER BY created_at ASC", (parent_id,)
|
||
).fetchall()]
|
||
|
||
def add_message(self, session_id: str, role: str, content: str, parent_id: str,
|
||
reasoning: str = "", is_ignored: bool = False,
|
||
msg_id: Optional[str] = None, attachment_metadata: Optional[str] = None,
|
||
timeline: Optional[str] = None,
|
||
usage: Optional[str] = None,
|
||
stop_reason: Optional[str] = None,
|
||
error_message: Optional[str] = None) -> Dict:
|
||
"""添加新消息，并自动将该消息设为当前会话的最新叶子节点
|
||
🆕 P1: usage —— assistant 回复的精确 usage JSON（如 '{"input":..,"output":..}'）
|
||
🆕 失败轮次: stop_reason/error_message —— 'error' 行入库但不回退叶子
|
||
（对照 pi：错误也持久化，回放时由 build_api_context 决定取舍）"""
|
||
if not msg_id: msg_id = f"msg_{uuid.uuid4().hex[:16]}"
|
||
now = int(time.time())
|
||
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
# 1. 插入消息
|
||
cursor.execute("""
|
||
INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
""", (msg_id, session_id, role, content, reasoning, 1 if is_ignored else 0, now, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message))
|
||
|
||
# 2. 自动更新 session 的叶子节点（时间线前推）
|
||
cursor.execute("UPDATE sessions SET current_leaf_msg_id = ?, updated_at = ? WHERE id = ?",
|
||
(msg_id, now, session_id))
|
||
conn.commit()
|
||
try:
|
||
print(f"[DB] add_message role={role} 内容={len(content or '')}c "
|
||
f"思考={len(reasoning or '')}c 时间线={'有' if timeline else '无'} "
|
||
f"id={msg_id} session={session_id[:8]}", flush=True)
|
||
except Exception:
|
||
pass
|
||
|
||
return dict(cursor.execute("SELECT * FROM messages WHERE id = ?", (msg_id,)).fetchone())
|
||
|
||
# ------------------------------------------------------------------
|
||
# 🆕 压缩持久化：链上标记点（子叶子回溯到标记即停）
|
||
# 插入后链形：…→ cut_before → [MARK role="compaction"] → first_retained → …→叶子
|
||
# 历史行全部保留（前端切会话渲染照常）；API 上下文构建在标记处截断。
|
||
# 关键：不动 sessions.current_leaf_msg_id（绝不能用 add_message）；
|
||
# INSERT + UPDATE 同一事务，不留断链窗口。
|
||
# ------------------------------------------------------------------
|
||
def insert_compaction_mark(self, session_id: str, summary: str,
|
||
cut_before_id: str, first_retained_id: str,
|
||
meta_json: Optional[str] = None) -> Optional[str]:
|
||
"""在链上插入压缩标记行并把保留首条的 parent_id 改指到标记。返回 mark_id。"""
|
||
if not cut_before_id or not first_retained_id:
|
||
return None
|
||
if cut_before_id == first_retained_id:
|
||
# 🐛 防自环：同一行不能既做切点前又做保留首条（timeline 回放同 id 场景）
|
||
try:
|
||
print(f"[DB] ⚠️ insert_compaction_mark 拒绝自环 cut==retained={cut_before_id}", flush=True)
|
||
except Exception:
|
||
pass
|
||
return None
|
||
mark_id = f"comp_{uuid.uuid4().hex[:16]}"
|
||
now = int(time.time())
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
cursor.execute("""INSERT INTO messages
|
||
(id, session_id, role, content, reasoning, is_ignored,
|
||
created_at, attachment_metadata, parent_id, timeline, usage)
|
||
VALUES (?, ?, 'compaction', ?, '', 1, ?, ?, ?, NULL, NULL)""",
|
||
(mark_id, session_id, summary, now, meta_json, cut_before_id))
|
||
cursor.execute("UPDATE messages SET parent_id = ? WHERE id = ?",
|
||
(mark_id, first_retained_id))
|
||
conn.commit()
|
||
try:
|
||
print(f"[DB] insert_compaction_mark session={session_id[:8]} "
|
||
f"cut_before={cut_before_id} first_retained={first_retained_id} "
|
||
f"mark={mark_id} summary={len(summary or '')}c", flush=True)
|
||
except Exception:
|
||
pass
|
||
return mark_id
|
||
# === 在 db_manager.py 中添加这个方法 ===
|
||
def get_branch_leaf(self, msg_id: str) -> str:
|
||
"""寻找一条时间线的最末端叶子节点"""
|
||
current_id = msg_id
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
while True:
|
||
# 寻找把当前节点作为父节点的子节点，按时间倒序取最新的一条（意味着它顺着最近被聊过的那条线往下走）
|
||
cursor.execute("SELECT id FROM messages WHERE parent_id = ? ORDER BY created_at DESC LIMIT 1", (current_id,))
|
||
child = cursor.fetchone()
|
||
if child:
|
||
current_id = child[0]
|
||
else:
|
||
break # 没有子节点了，它自己就是叶子！
|
||
return current_id
|
||
def delete_message_branch(self, session_id: str, msg_id: str):
|
||
"""🚀 精准剪枝：删AI只删当前分支，删User连根拔起，并自动平滑回退时间线"""
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
|
||
# 1. 查出要删除的节点的父亲
|
||
cursor.execute("SELECT parent_id FROM messages WHERE id = ?", (msg_id,))
|
||
row = cursor.fetchone()
|
||
if not row: return
|
||
safe_parent_id = row['parent_id']
|
||
|
||
# 2. 目标就是传进来的 msg_id 本身 (不再强制上移到 parent)
|
||
target_id = msg_id
|
||
|
||
# 3. 拉取全会话，构建亲属关系树
|
||
cursor.execute("SELECT id, parent_id FROM messages WHERE session_id = ?", (session_id,))
|
||
all_msgs = cursor.fetchall()
|
||
|
||
children_map = {}
|
||
for m in all_msgs:
|
||
pid = m['parent_id']
|
||
if pid not in children_map:
|
||
children_map[pid] = []
|
||
children_map[pid].append(m['id'])
|
||
|
||
# 4. 广度优先搜索 (BFS)：找出目标消息及其所有子孙
|
||
to_delete = set([target_id])
|
||
queue = [target_id]
|
||
while queue:
|
||
curr = queue.pop(0)
|
||
if curr in children_map:
|
||
for child in children_map[curr]:
|
||
if child not in to_delete:
|
||
to_delete.add(child)
|
||
queue.append(child)
|
||
|
||
# 5. 判断当前时间线的“叶子节点”是否在被波及的名单里
|
||
cursor.execute("SELECT current_leaf_msg_id FROM sessions WHERE id = ?", (session_id,))
|
||
leaf_row = cursor.fetchone()
|
||
leaf_needs_update = leaf_row and leaf_row['current_leaf_msg_id'] in to_delete
|
||
|
||
# 6. 🌟 执行物理删除前，收集将被删除的附件元数据
|
||
deleted_metadata = []
|
||
for d_id in to_delete:
|
||
# 先查出它的 metadata
|
||
cursor.execute("SELECT attachment_metadata FROM messages WHERE id = ?", (d_id,))
|
||
row = cursor.fetchone()
|
||
if row and row['attachment_metadata']:
|
||
deleted_metadata.append(row['attachment_metadata'])
|
||
|
||
# 然后再执行物理删除
|
||
cursor.execute("DELETE FROM messages WHERE id = ?", (d_id,))
|
||
|
||
|
||
# 7. 🌟 核心：如果时间线断了，自动寻找平滑降落点
|
||
sibling_row = None
|
||
if leaf_needs_update:
|
||
# 尝试寻找被删节点的最新“兄弟姐妹” (例如删了分支2，寻找分支1)
|
||
cursor.execute("SELECT id FROM messages WHERE parent_id = ? ORDER BY created_at DESC LIMIT 1", (safe_parent_id,))
|
||
sibling_row = cursor.fetchone()
|
||
|
||
# 如果有兄弟，降落到兄弟；如果没兄弟(只有1次回答)，退回原点(提问)
|
||
new_leaf = sibling_row['id'] if sibling_row else safe_parent_id
|
||
cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (new_leaf, session_id))
|
||
|
||
conn.commit()
|
||
|
||
# 8. 如果降落到了兄弟分支，兄弟可能还有下文，需再次对齐真实叶子节点
|
||
if leaf_needs_update and sibling_row:
|
||
real_leaf = self.get_branch_leaf(new_leaf)
|
||
self.update_session_leaf(session_id, real_leaf)
|
||
|
||
return deleted_metadata # 🌟 返回被删除的元数据，交给 MainWindow 去粉碎文件
|