313 lines
14 KiB
Python
313 lines
14 KiB
Python
# utils/conversation_manager.py - 对话持久化管理器（集成Token统计）
|
||
|
||
import json
|
||
import os
|
||
import time
|
||
import tempfile
|
||
import threading
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional, Any
|
||
from dataclasses import dataclass
|
||
try:
|
||
from config import DATA_DIR, HOST_WORKSPACES_FILE
|
||
except ImportError:
|
||
import sys
|
||
from pathlib import Path
|
||
project_root = Path(__file__).resolve().parents[1]
|
||
if str(project_root) not in sys.path:
|
||
sys.path.insert(0, str(project_root))
|
||
from config import DATA_DIR, HOST_WORKSPACES_FILE
|
||
|
||
try:
|
||
from utils.perf_log import perf_log
|
||
except Exception:
|
||
perf_log = None
|
||
|
||
@dataclass
|
||
class ConversationMetadata:
|
||
"""对话元数据"""
|
||
id: str
|
||
title: str
|
||
created_at: str
|
||
updated_at: str
|
||
project_path: Optional[str]
|
||
project_relative_path: Optional[str]
|
||
thinking_mode: bool
|
||
total_messages: int
|
||
total_tools: int
|
||
run_mode: str = "fast"
|
||
model_key: Optional[str] = None
|
||
has_images: bool = False
|
||
has_videos: bool = False
|
||
status: str = "active" # active, archived, error
|
||
|
||
|
||
class IndexMixin:
|
||
"""ConversationManager index mixin 能力 mixin。"""
|
||
|
||
def _migrate_legacy_conversation_files(self):
|
||
"""把 data/conversations/conv_*.json 按 metadata.project_path 移入对应 workspace 子目录。"""
|
||
try:
|
||
legacy_files = [
|
||
p for p in self.conversations_root.glob("conv_*.json")
|
||
if p.is_file() and p.parent == self.conversations_root
|
||
]
|
||
except Exception:
|
||
return
|
||
|
||
moved_count = 0
|
||
for file_path in legacy_files:
|
||
try:
|
||
data = json.loads(file_path.read_text(encoding="utf-8"))
|
||
metadata = data.get("metadata", {}) if isinstance(data, dict) else {}
|
||
project_path = self._normalize_path(metadata.get("project_path"))
|
||
workspace_id = self._resolve_workspace_id(project_path)
|
||
if not workspace_id:
|
||
continue
|
||
target_dir = self._conversation_dir_for_workspace(workspace_id)
|
||
target_dir.mkdir(parents=True, exist_ok=True)
|
||
target = target_dir / file_path.name
|
||
if target.exists():
|
||
print(f"⚠️ legacy 对话迁移跳过，目标已存在: {target}")
|
||
continue
|
||
file_path.replace(target)
|
||
moved_count += 1
|
||
except Exception as exc:
|
||
print(f"⚠️ legacy 对话迁移失败 {file_path.name}: {exc}")
|
||
if moved_count:
|
||
print(f"🔄 已迁移 {moved_count} 个 legacy 对话到工作区目录")
|
||
|
||
def _iter_conversation_files(self, sort_by_mtime: bool = True):
|
||
"""遍历对话文件（排除索引文件），可按修改时间降序排序。"""
|
||
files = []
|
||
for p in self.conversations_dir.glob("*.json"):
|
||
if p == self.index_file:
|
||
continue
|
||
stem = p.stem or ""
|
||
# 跳过索引备份/损坏文件，避免被误当作对话文件参与重建
|
||
if stem == "index" or stem.startswith("index_corrupt_"):
|
||
continue
|
||
# 仅纳入标准对话文件，避免其他 json 干扰索引
|
||
if not stem.startswith("conv_"):
|
||
continue
|
||
files.append(p)
|
||
if sort_by_mtime:
|
||
files.sort(key=lambda p: p.stat().st_mtime, reverse=True)
|
||
return files
|
||
|
||
def _atomic_write_json(self, target_file: Path, payload: Dict[str, Any]):
|
||
"""原子写入 JSON：使用唯一临时文件，避免并发写同名 .tmp 产生竞态。"""
|
||
target_file.parent.mkdir(parents=True, exist_ok=True)
|
||
temp_path: Optional[Path] = None
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
mode="w",
|
||
encoding="utf-8",
|
||
dir=str(target_file.parent),
|
||
prefix=f".{target_file.name}.",
|
||
suffix=".tmp",
|
||
delete=False,
|
||
) as fh:
|
||
temp_path = Path(fh.name)
|
||
json.dump(payload, fh, ensure_ascii=False, indent=2)
|
||
fh.flush()
|
||
os.fsync(fh.fileno())
|
||
os.replace(str(temp_path), str(target_file))
|
||
finally:
|
||
if temp_path and temp_path.exists():
|
||
try:
|
||
temp_path.unlink()
|
||
except Exception:
|
||
pass
|
||
|
||
def _rebuild_index_from_files(self, max_count: Optional[int] = None) -> Dict:
|
||
"""
|
||
从现有对话文件重建索引。
|
||
|
||
Args:
|
||
max_count: 限制重建的条目数（按文件修改时间倒序）；None 表示全量重建。
|
||
"""
|
||
t0 = time.perf_counter()
|
||
rebuilt_index: Dict[str, Dict] = {}
|
||
files = self._iter_conversation_files(sort_by_mtime=True)
|
||
if max_count is not None:
|
||
files = files[:max(0, int(max_count))]
|
||
if perf_log:
|
||
perf_log("_rebuild_index_from_files start", extra={"file_count": len(files), "max_count": max_count})
|
||
for file_path in files:
|
||
try:
|
||
with open(file_path, "r", encoding="utf-8") as f:
|
||
raw = f.read().strip()
|
||
if not raw:
|
||
continue
|
||
data = json.loads(raw)
|
||
except Exception as exc:
|
||
print(f"⚠️ 重建索引时跳过 {file_path.name}: {exc}")
|
||
continue
|
||
|
||
conv_id = data.get("id") or file_path.stem
|
||
metadata = data.get("metadata", {}) or {}
|
||
|
||
rebuilt_index[conv_id] = {
|
||
"title": data.get("title") or "未命名对话",
|
||
"created_at": data.get("created_at"),
|
||
"updated_at": data.get("updated_at"),
|
||
"project_path": metadata.get("project_path"),
|
||
"project_relative_path": metadata.get("project_relative_path"),
|
||
"thinking_mode": metadata.get("thinking_mode", False),
|
||
"run_mode": metadata.get("run_mode") or ("thinking" if metadata.get("thinking_mode") else "fast"),
|
||
"model_key": metadata.get("model_key"),
|
||
"has_images": metadata.get("has_images", False),
|
||
"has_videos": metadata.get("has_videos", False),
|
||
"total_messages": metadata.get("total_messages", 0),
|
||
"total_tools": metadata.get("total_tools", 0),
|
||
"status": metadata.get("status", "active"),
|
||
"multi_agent_mode": bool(metadata.get("multi_agent_mode", False)),
|
||
}
|
||
elapsed_ms = (time.perf_counter() - t0) * 1000
|
||
if perf_log:
|
||
perf_log("_rebuild_index_from_files done", elapsed_ms=elapsed_ms, extra={"rebuilt_size": len(rebuilt_index), "max_count": max_count})
|
||
if rebuilt_index:
|
||
print(f"🔄 已从对话文件重建索引，共 {len(rebuilt_index)} 条记录")
|
||
return rebuilt_index
|
||
|
||
def _index_missing_conversations(self, index: Dict) -> bool:
|
||
"""检查索引是否缺失本地对话文件"""
|
||
t0 = time.perf_counter()
|
||
index_ids = set(index.keys())
|
||
files = list(self._iter_conversation_files())
|
||
missing = False
|
||
for file_path in files:
|
||
conv_id = file_path.stem
|
||
if conv_id and conv_id not in index_ids:
|
||
print(f"🔍 对话 {conv_id} 未出现在索引中，将重建索引。")
|
||
missing = True
|
||
break
|
||
elapsed_ms = (time.perf_counter() - t0) * 1000
|
||
if perf_log:
|
||
perf_log("_index_missing_conversations", elapsed_ms=elapsed_ms, extra={"file_count": len(files), "index_size": len(index), "missing": missing})
|
||
return missing
|
||
|
||
def _load_index(self, ensure_integrity: bool = False, max_rebuild: Optional[int] = None) -> Dict:
|
||
"""加载对话索引，可选地在缺失时自动重建（可限制重建条数）"""
|
||
try:
|
||
index: Dict = {}
|
||
if self.index_file.exists():
|
||
with open(self.index_file, 'r', encoding='utf-8') as f:
|
||
content = f.read().strip()
|
||
if content:
|
||
index = json.loads(content)
|
||
if index:
|
||
if ensure_integrity and not self._index_verified:
|
||
if self._index_missing_conversations(index):
|
||
rebuilt = self._rebuild_index_from_files(max_count=max_rebuild)
|
||
if rebuilt:
|
||
self._save_index(rebuilt)
|
||
index = rebuilt
|
||
self._index_verified = True
|
||
return index
|
||
# 索引为空但对话文件仍然存在时尝试重建
|
||
rebuilt = self._rebuild_index_from_files(max_count=max_rebuild)
|
||
if rebuilt:
|
||
self._save_index(rebuilt)
|
||
if ensure_integrity:
|
||
self._index_verified = True
|
||
return rebuilt
|
||
return {}
|
||
# 索引缺失但存在对话文件时重建
|
||
rebuilt = self._rebuild_index_from_files(max_count=max_rebuild)
|
||
if rebuilt:
|
||
self._save_index(rebuilt)
|
||
if ensure_integrity:
|
||
self._index_verified = True
|
||
return rebuilt
|
||
return index
|
||
except (json.JSONDecodeError, Exception) as e:
|
||
print(f"⚠️ 加载对话索引失败，将尝试重建: {e}")
|
||
backup_path = self.index_file.with_name(
|
||
f"{self.index_file.stem}_corrupt_{int(time.time() * 1000)}{self.index_file.suffix}"
|
||
)
|
||
try:
|
||
if self.index_file.exists():
|
||
self.index_file.replace(backup_path)
|
||
print(f"🗄️ 已备份损坏的索引文件到: {backup_path.name}")
|
||
except Exception as backup_exc:
|
||
print(f"⚠️ 备份损坏索引文件失败: {backup_exc}")
|
||
# 索引损坏场景优先全量重建，避免仅重建部分导致“对话丢失”错觉
|
||
rebuilt = self._rebuild_index_from_files(max_count=None)
|
||
if rebuilt:
|
||
self._save_index(rebuilt)
|
||
if ensure_integrity:
|
||
self._index_verified = True
|
||
return rebuilt
|
||
return {}
|
||
|
||
def _save_index(self, index: Dict):
|
||
"""保存对话索引"""
|
||
try:
|
||
with self._io_lock:
|
||
self._atomic_write_json(self.index_file, index)
|
||
except Exception as e:
|
||
print(f"⌘ 保存对话索引失败: {e}")
|
||
|
||
def _ensure_index_covering(self, limit: int, offset: int) -> Dict:
|
||
"""
|
||
确保索引涵盖到 offset+limit 条记录，不足时按需扩展重建（仍按 mtime 倒序，增量加载批量）。
|
||
每次先检查是否有新增对话文件未被纳入索引，避免"创建后列表不显示"的问题。
|
||
"""
|
||
t0 = time.perf_counter()
|
||
needed = max(0, int(offset) + int(limit))
|
||
index = self._load_index()
|
||
load_ms = (time.perf_counter() - t0) * 1000
|
||
if perf_log:
|
||
perf_log("_ensure_index_covering load_index", elapsed_ms=load_ms, extra={"index_size": len(index), "needed": needed})
|
||
|
||
# 主动检查新增文件：索引可能滞后于实际对话文件，导致列表不同步。
|
||
t1 = time.perf_counter()
|
||
missing = self._index_missing_conversations(index)
|
||
missing_ms = (time.perf_counter() - t1) * 1000
|
||
if perf_log:
|
||
perf_log("_ensure_index_covering check_missing", elapsed_ms=missing_ms, extra={"missing": missing, "index_size": len(index)})
|
||
if missing:
|
||
t2 = time.perf_counter()
|
||
rebuilt = self._rebuild_index_from_files(max_count=max(needed, len(index) + 10))
|
||
rebuild_ms = (time.perf_counter() - t2) * 1000
|
||
if perf_log:
|
||
perf_log("_ensure_index_covering rebuild_missing", elapsed_ms=rebuild_ms, extra={"rebuilt_size": len(rebuilt) if rebuilt else 0})
|
||
if rebuilt:
|
||
self._save_index(rebuilt)
|
||
index = rebuilt
|
||
|
||
if len(index) >= needed:
|
||
total_ms = (time.perf_counter() - t0) * 1000
|
||
if perf_log:
|
||
perf_log("_ensure_index_covering return", elapsed_ms=total_ms, extra={"index_size": len(index), "needed": needed})
|
||
return index
|
||
|
||
# 第一次尝试：扩展到需要的数量（按更新时间倒序）
|
||
t3 = time.perf_counter()
|
||
rebuilt = self._rebuild_index_from_files(max_count=needed)
|
||
rebuild_ms = (time.perf_counter() - t3) * 1000
|
||
if perf_log:
|
||
perf_log("_ensure_index_covering rebuild_needed", elapsed_ms=rebuild_ms, extra={"rebuilt_size": len(rebuilt) if rebuilt else 0})
|
||
if rebuilt:
|
||
self._save_index(rebuilt)
|
||
index = rebuilt
|
||
|
||
# 如果仍不足且存在更多文件可能未被纳入（例如首批限定过小），进行一次全量重建兜底
|
||
if len(index) < needed:
|
||
t4 = time.perf_counter()
|
||
rebuilt_full = self._rebuild_index_from_files(max_count=None)
|
||
rebuild_full_ms = (time.perf_counter() - t4) * 1000
|
||
if perf_log:
|
||
perf_log("_ensure_index_covering rebuild_full", elapsed_ms=rebuild_full_ms, extra={"rebuilt_size": len(rebuilt_full) if rebuilt_full else 0})
|
||
if rebuilt_full:
|
||
self._save_index(rebuilt_full)
|
||
index = rebuilt_full
|
||
|
||
total_ms = (time.perf_counter() - t0) * 1000
|
||
if perf_log:
|
||
perf_log("_ensure_index_covering return", elapsed_ms=total_ms, extra={"index_size": len(index), "needed": needed})
|
||
return index
|