rag/fix_chromadb.py at main · chenncy/rag · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
# -*- coding: utf-8 -*-
# fix_chromadb.py
"""
修复 ChromaDB 数据库损坏问题
当遇到 "Error in compaction" 或 "hnsw" 相关错误时，可以使用此脚本修复
"""
import os
import shutil
from app.config import DB_DIR

def fix_chromadb():
    """修复或重建 ChromaDB 数据库"""
    print("=" * 60)
    print("🔧 ChromaDB 数据库修复工具")
    print("=" * 60)

    if not os.path.exists(DB_DIR):
        print(f"❌ 数据库目录不存在: {DB_DIR}")
        return False

    print(f"📂 数据库目录: {DB_DIR}")

    # 检查目录内容
    files = os.listdir(DB_DIR)
    print(f"📋 目录中的文件: {len(files)} 个")

    # 备份选项
    backup_dir = DB_DIR + "_backup"
    if os.path.exists(backup_dir):
        shutil.rmtree(backup_dir)

    print(f"\n📦 正在备份数据库到: {backup_dir}")
    try:
        shutil.copytree(DB_DIR, backup_dir)
        print("✅ 备份完成")
    except Exception as e:
        print(f"⚠️  备份失败: {e}")
        response = input("是否继续修复？(y/n): ")
        if response.lower() != 'y':
            print("❌ 已取消")
            return False

    # 修复选项
    print("\n请选择修复方式:")
    print("1. 删除损坏的索引文件（保留数据，重建索引）")
    print("2. 完全重建数据库（删除所有数据，需要重新上传文件）")

    choice = input("请选择 (1/2): ").strip()

    if choice == "1":
        # 只删除索引相关文件，保留 SQLite 数据库
        print("\n🔧 正在删除损坏的索引文件...")
        deleted_count = 0
        for filename in files:
            # 删除 HNSW 索引相关文件
            if filename.startswith("chroma-") or filename.endswith(".index") or "hnsw" in filename.lower():
                file_path = os.path.join(DB_DIR, filename)
                try:
                    if os.path.isfile(file_path):
                        os.remove(file_path)
                        deleted_count += 1
                        print(f"  ✓ 已删除: {filename}")
                    elif os.path.isdir(file_path):
                        shutil.rmtree(file_path)
                        deleted_count += 1
                        print(f"  ✓ 已删除目录: {filename}")
                except Exception as e:
                    print(f"  ✗ 删除失败 {filename}: {e}")

        print(f"\n✅ 已删除 {deleted_count} 个索引文件")
        print("💡 提示: 重启服务器后，索引会在下次查询时自动重建")

    elif choice == "2":
        # 完全删除数据库
        print("\n⚠️  警告: 这将删除所有向量数据！")
        confirm = input("确认删除？(yes/no): ").strip().lower()

        if confirm == "yes":
            print("\n🗑️  正在删除数据库...")
            try:
                shutil.rmtree(DB_DIR)
                os.makedirs(DB_DIR, exist_ok=True)
                print("✅ 数据库已清空")
                print("💡 提示: 需要重新上传文件以重建知识库")
            except Exception as e:
                print(f"❌ 删除失败: {e}")
                return False
        else:
            print("❌ 已取消")
            return False
    else:
        print("❌ 无效选择")
        return False

    print("\n" + "=" * 60)
    print("✅ 修复完成！")
    print("=" * 60)
    print("\n💡 下一步:")
    print("1. 重启服务器")
    print("2. 如果选择方式1，索引会在首次查询时自动重建")
    print("3. 如果选择方式2，需要重新上传文件")

    return True

if __name__ == "__main__":
    try:
        fix_chromadb()
    except KeyboardInterrupt:
        print("\n\n❌ 用户中断")
    except Exception as e:
        print(f"\n❌ 修复过程出错: {e}")
        import traceback
        traceback.print_exc()