320 lines
11 KiB
Python
320 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
AI 聊天记录集中存储工具
|
||
|
||
命令:
|
||
python main.py # 同步一次
|
||
python main.py --watch # 持续监控
|
||
python main.py --prune # 清理孤立记录
|
||
python main.py --stats # 显示统计
|
||
python main.py --search <keyword> # 搜索
|
||
python main.py --export json|csv [--source codex|kiro|gemini|claude]
|
||
"""
|
||
import os
|
||
import sys
|
||
import subprocess
|
||
|
||
# 项目根目录
|
||
PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
VENV_DIR = os.path.join(PROJECT_DIR, '.venv')
|
||
REQUIREMENTS = os.path.join(PROJECT_DIR, 'requirements.txt')
|
||
|
||
def ensure_venv():
|
||
"""检测并创建虚拟环境，安装依赖"""
|
||
# 打包版本跳过
|
||
if getattr(sys, 'frozen', False):
|
||
return
|
||
|
||
# 已在虚拟环境中运行则跳过
|
||
if sys.prefix != sys.base_prefix:
|
||
return
|
||
|
||
# 检查 .venv 是否存在
|
||
venv_python = os.path.join(VENV_DIR, 'bin', 'python') if os.name != 'nt' else os.path.join(VENV_DIR, 'Scripts', 'python.exe')
|
||
|
||
if not os.path.exists(venv_python):
|
||
print("首次运行，创建虚拟环境...")
|
||
subprocess.run([sys.executable, '-m', 'venv', VENV_DIR], check=True)
|
||
print("安装依赖...")
|
||
pip = os.path.join(VENV_DIR, 'bin', 'pip') if os.name != 'nt' else os.path.join(VENV_DIR, 'Scripts', 'pip.exe')
|
||
subprocess.run([pip, 'install', '-r', REQUIREMENTS, '-q'], check=True)
|
||
print("环境准备完成，重新启动...\n")
|
||
|
||
# 使用虚拟环境重新执行
|
||
os.execv(venv_python, [venv_python] + sys.argv)
|
||
|
||
# 启动前检测虚拟环境
|
||
ensure_venv()
|
||
|
||
# 支持 PyInstaller 打包
|
||
if getattr(sys, 'frozen', False):
|
||
BASE_DIR = sys._MEIPASS
|
||
sys.path.insert(0, os.path.join(BASE_DIR, 'src'))
|
||
else:
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
import argparse
|
||
from config import CONFIG
|
||
from parsers import CodexParser, GeminiParser, ClaudeParser, KiroParser
|
||
from storage import ChatStorage
|
||
from logger import setup_logger, get_logger
|
||
|
||
storage: ChatStorage = None
|
||
|
||
def main():
|
||
global storage
|
||
|
||
parser = argparse.ArgumentParser(description='AI Chat Converter')
|
||
parser.add_argument('-w', '--watch', action='store_true', help='持续监控模式')
|
||
parser.add_argument('--prune', action='store_true', help='清理孤立记录')
|
||
parser.add_argument('--stats', action='store_true', help='显示统计信息')
|
||
parser.add_argument('--search', type=str, help='搜索关键词')
|
||
parser.add_argument('--export', choices=['json', 'csv'], help='导出格式')
|
||
parser.add_argument('--source', choices=['codex', 'kiro', 'gemini', 'claude'], help='指定来源')
|
||
parser.add_argument('--output', type=str, help='导出文件路径')
|
||
args = parser.parse_args()
|
||
|
||
# 初始化
|
||
setup_logger(CONFIG["log_dir"])
|
||
log = get_logger()
|
||
|
||
storage = ChatStorage(CONFIG["db_path"])
|
||
|
||
# 命令分发
|
||
if args.prune:
|
||
cmd_prune()
|
||
elif args.stats:
|
||
cmd_stats()
|
||
elif args.search:
|
||
cmd_search(args.search, args.source)
|
||
elif args.export:
|
||
cmd_export(args.export, args.source, args.output)
|
||
elif args.watch:
|
||
cmd_sync()
|
||
cmd_watch()
|
||
else:
|
||
cmd_sync()
|
||
|
||
def cmd_sync():
|
||
log = get_logger()
|
||
log.info("=" * 50)
|
||
log.info("AI 聊天记录 → 集中存储")
|
||
log.info("=" * 50)
|
||
log.info(f"数据库: {CONFIG['db_path']}")
|
||
|
||
total_added, total_updated, total_skipped, total_errors = 0, 0, 0, 0
|
||
|
||
for cli, key, parser_cls in [
|
||
('codex', 'codex_paths', lambda: CodexParser('codex')),
|
||
('kiro', 'kiro_paths', KiroParser),
|
||
('gemini', 'gemini_paths', GeminiParser),
|
||
('claude', 'claude_paths', ClaudeParser),
|
||
]:
|
||
paths = CONFIG.get(key, [])
|
||
if not paths:
|
||
continue
|
||
|
||
parser = parser_cls()
|
||
if cli in ('claude', 'kiro'):
|
||
a, u, s, e = process_multi(parser, paths, cli)
|
||
else:
|
||
a, u, s, e = process(parser, paths)
|
||
|
||
log.info(f"[{cli.capitalize()}] 新增:{a} 更新:{u} 跳过:{s} 错误:{e}")
|
||
update_cli_meta(cli)
|
||
total_added += a
|
||
total_updated += u
|
||
total_skipped += s
|
||
total_errors += e
|
||
|
||
total = storage.get_total_stats()
|
||
storage.update_total_meta(total['sessions'], total['messages'], total['tokens'])
|
||
|
||
log.info("=" * 50)
|
||
log.info(f"总计: {total['sessions']} 会话, {total['messages']} 消息")
|
||
if total_errors > 0:
|
||
log.warning(f"错误: {total_errors} 个文件解析失败")
|
||
log.info("✓ 同步完成!")
|
||
|
||
print_token_stats()
|
||
|
||
def cmd_watch():
|
||
from watcher import ChatWatcher
|
||
from datetime import datetime
|
||
|
||
log = get_logger()
|
||
log.info("")
|
||
log.info("=" * 50)
|
||
log.info("实时监听模式 (watchdog)")
|
||
log.info("=" * 50)
|
||
|
||
watch_paths = []
|
||
path_source_map = {}
|
||
|
||
for cli, key in [('codex', 'codex_paths'), ('kiro', 'kiro_paths'),
|
||
('gemini', 'gemini_paths'), ('claude', 'claude_paths')]:
|
||
for p in CONFIG.get(key, []):
|
||
if os.path.isdir(p) or os.path.isfile(p):
|
||
watch_paths.append(p)
|
||
path_source_map[p] = cli
|
||
|
||
def on_change(file_path, event_type):
|
||
now = datetime.now().strftime('%H:%M:%S')
|
||
source = None
|
||
for p, s in path_source_map.items():
|
||
if file_path.startswith(p) or file_path == p:
|
||
source = s
|
||
break
|
||
if not source:
|
||
return
|
||
|
||
try:
|
||
if source == 'kiro':
|
||
parser = KiroParser()
|
||
for sess in parser.parse_file(file_path):
|
||
storage.upsert_session(sess.session_id, sess.source, sess.file_path, sess.cwd, sess.messages, int(sess.file_mtime))
|
||
log.info(f"[{now}] kiro 更新")
|
||
elif source == 'claude':
|
||
parser = ClaudeParser()
|
||
for sess in parser.parse_file(file_path):
|
||
fp = f"claude:{sess.session_id}"
|
||
storage.upsert_session(sess.session_id, sess.source, fp, sess.cwd, sess.messages, int(sess.file_mtime))
|
||
log.info(f"[{now}] claude 更新")
|
||
else:
|
||
parser = CodexParser(source) if source == 'codex' else GeminiParser()
|
||
sess = parser.parse_file(file_path)
|
||
fp = os.path.abspath(sess.file_path)
|
||
storage.upsert_session(sess.session_id, sess.source, fp, sess.cwd, sess.messages, int(sess.file_mtime))
|
||
log.info(f"[{now}] {source} {event_type}: {os.path.basename(file_path)}")
|
||
|
||
update_cli_meta(source)
|
||
total = storage.get_total_stats()
|
||
storage.update_total_meta(total['sessions'], total['messages'], total['tokens'])
|
||
except Exception as e:
|
||
log.error(f"[{now}] 处理失败 {file_path}: {e}")
|
||
|
||
log.info(f"监听目录: {len(watch_paths)} 个")
|
||
watcher = ChatWatcher(watch_paths, on_change)
|
||
watcher.start()
|
||
|
||
def cmd_prune():
|
||
log = get_logger()
|
||
log.info("清理孤立记录...")
|
||
removed = storage.prune()
|
||
total = sum(removed.values())
|
||
if total > 0:
|
||
for cli, count in removed.items():
|
||
if count > 0:
|
||
log.info(f" {cli}: 删除 {count} 条")
|
||
log.info(f"✓ 共清理 {total} 条孤立记录")
|
||
else:
|
||
log.info("✓ 无孤立记录")
|
||
|
||
def cmd_stats():
|
||
log = get_logger()
|
||
meta = storage.get_total_meta()
|
||
tokens = storage.get_token_stats()
|
||
|
||
log.info("=" * 50)
|
||
log.info("统计信息")
|
||
log.info("=" * 50)
|
||
log.info(f"数据库: {CONFIG['db_path']}")
|
||
log.info(f"总会话: {meta['total_sessions']}")
|
||
log.info(f"总消息: {meta['total_messages']}")
|
||
log.info(f"最后同步: {meta['last_sync']}")
|
||
log.info("")
|
||
log.info("Token 统计 (tiktoken):")
|
||
total_tokens = 0
|
||
for source in ['codex', 'kiro', 'gemini', 'claude']:
|
||
t = tokens.get(source, 0)
|
||
if t > 0:
|
||
log.info(f" {source}: {t:,}")
|
||
total_tokens += t
|
||
log.info(f" 总计: {total_tokens:,}")
|
||
|
||
def cmd_search(keyword: str, source: str = None):
|
||
log = get_logger()
|
||
results = storage.search(keyword, source)
|
||
log.info(f"搜索 '{keyword}' 找到 {len(results)} 个会话:")
|
||
for r in results[:20]:
|
||
log.info(f" [{r['source']}] {r['session_id']} - {r['cwd'] or 'N/A'}")
|
||
|
||
def cmd_export(fmt: str, source: str = None, output: str = None):
|
||
log = get_logger()
|
||
if not output:
|
||
output = os.path.join(CONFIG["output_dir"], f"export.{fmt}")
|
||
|
||
if fmt == 'json':
|
||
count = storage.export_json(output, source)
|
||
else:
|
||
count = storage.export_csv(output, source)
|
||
|
||
log.info(f"✓ 导出 {count} 条到 {output}")
|
||
|
||
def print_token_stats():
|
||
log = get_logger()
|
||
tokens = storage.get_token_stats()
|
||
log.info("")
|
||
log.info("=== Token 统计 (tiktoken) ===")
|
||
total = 0
|
||
for source in ['codex', 'kiro', 'gemini', 'claude']:
|
||
t = tokens.get(source, 0)
|
||
if t > 0:
|
||
log.info(f" {source}: {t:,} tokens")
|
||
total += t
|
||
log.info(f" 总计: {total:,} tokens")
|
||
|
||
def update_cli_meta(cli: str):
|
||
stats = storage.get_cli_stats(cli)
|
||
path = CONFIG.get(f"{cli}_paths", [""])[0] if CONFIG.get(f"{cli}_paths") else ""
|
||
storage.update_cli_meta(cli, path, stats['sessions'], stats['messages'], stats['tokens'])
|
||
|
||
def process(parser, paths) -> tuple:
|
||
log = get_logger()
|
||
added, updated, skipped, errors = 0, 0, 0, 0
|
||
for f in parser.find_files(paths):
|
||
try:
|
||
s = parser.parse_file(f)
|
||
file_path = os.path.abspath(s.file_path)
|
||
db_mtime = storage.get_file_mtime(file_path)
|
||
file_mtime = int(s.file_mtime)
|
||
if db_mtime == 0:
|
||
storage.upsert_session(s.session_id, s.source, file_path, s.cwd, s.messages, file_mtime)
|
||
added += 1
|
||
elif file_mtime > db_mtime:
|
||
storage.upsert_session(s.session_id, s.source, file_path, s.cwd, s.messages, file_mtime)
|
||
updated += 1
|
||
else:
|
||
skipped += 1
|
||
except Exception as e:
|
||
log.debug(f"解析失败 {f}: {e}")
|
||
errors += 1
|
||
return added, updated, skipped, errors
|
||
|
||
def process_multi(parser, paths, source: str) -> tuple:
|
||
"""处理返回多个会话的解析器（Claude/Kiro）"""
|
||
log = get_logger()
|
||
added, updated, skipped, errors = 0, 0, 0, 0
|
||
for f in parser.find_files(paths):
|
||
try:
|
||
for s in parser.parse_file(f):
|
||
file_path = s.file_path # kiro:xxx 或 claude:xxx
|
||
db_mtime = storage.get_file_mtime(file_path)
|
||
file_mtime = int(s.file_mtime)
|
||
if db_mtime == 0:
|
||
storage.upsert_session(s.session_id, s.source, file_path, s.cwd, s.messages, file_mtime)
|
||
added += 1
|
||
elif file_mtime > db_mtime:
|
||
storage.upsert_session(s.session_id, s.source, file_path, s.cwd, s.messages, file_mtime)
|
||
updated += 1
|
||
else:
|
||
skipped += 1
|
||
except Exception as e:
|
||
log.debug(f"解析失败 {f}: {e}")
|
||
errors += 1
|
||
return added, updated, skipped, errors
|
||
|
||
if __name__ == '__main__':
|
||
main()
|