"""索引构建脚本 — 全量/增量构建 FAISS 索引 + SQLite metadata 用法: python scripts/build_index.py --mode full # 全量构建 python scripts/build_index.py --mode incremental # 增量构建 流程: 1. 扫描 law-pack 目录,按类别切片 2. 写入 SQLite metadata(laws + clauses 表) 3. 批量调用 embedding 服务向量化所有条文 4. 构建 FAISS HNSW 索引并持久化 5. 输出统计信息 安全: - 先写临时文件,成功后原子替换 - 中断不损坏已有索引 """ import argparse import asyncio import json import logging import os import sqlite3 import sys import time from pathlib import Path from typing import List, Dict, Any, Optional # 添加项目根目录到 path sys.path.insert(0, str(Path(__file__).parent.parent)) import numpy as np from app.config import ( LAW_PACK_DIR, LAW_KB_DATA_DIR, FAISS_INDEX_PATH, SQLITE_PATH, LOGS_DIR, REGION_MAPPING_PATH, CATEGORY_DIRS, EMBEDDING_BATCH_SIZE, HNSW_M, HNSW_EF_CONSTRUCTION, ) from scripts.parse_clause import ( Clause, load_region_mapping, scan_category_dir, ) logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", ) logger = logging.getLogger(__name__) # ===== SQLite Schema ===== SCHEMA_SQL = """ CREATE TABLE IF NOT EXISTS laws ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, category TEXT NOT NULL, publish_date TEXT, province TEXT, city TEXT, region_level TEXT, file_path TEXT NOT NULL, clause_count INTEGER DEFAULT 0, file_mtime REAL DEFAULT 0, indexed_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS clauses ( id INTEGER PRIMARY KEY AUTOINCREMENT, law_id INTEGER NOT NULL, chapter TEXT, clause_no TEXT, content TEXT NOT NULL, faiss_idx INTEGER, FOREIGN KEY (law_id) REFERENCES laws(id) ); CREATE INDEX IF NOT EXISTS idx_clauses_law_id ON clauses(law_id); CREATE INDEX IF NOT EXISTS idx_clauses_faiss_idx ON clauses(faiss_idx); CREATE INDEX IF NOT EXISTS idx_laws_category ON laws(category); CREATE INDEX IF NOT EXISTS idx_laws_province ON laws(province); CREATE TABLE IF NOT EXISTS prompt_configs ( id INTEGER PRIMARY KEY AUTOINCREMENT, config_key TEXT NOT NULL UNIQUE, config_value TEXT, updated_at TEXT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_prompt_configs_key ON prompt_configs(config_key); """ def init_sqlite(db_path: Path): """初始化 SQLite(创建表)""" db_path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(str(db_path)) conn.executescript(SCHEMA_SQL) conn.commit() return conn def reset_sqlite(db_path: Path): """重置 SQLite(全量构建时用)""" if db_path.exists(): db_path.unlink() return init_sqlite(db_path) async def build_index(mode: str = "full"): """构建索引""" start_time = time.time() timestamp = time.strftime("%Y-%m-%d %H:%M:%S") logger.info("=" * 60) logger.info(f"QYLAW 索引构建 — {mode} 模式") logger.info(f"时间: {timestamp}") logger.info(f"法规目录: {LAW_PACK_DIR}") logger.info(f"数据目录: {LAW_KB_DATA_DIR}") logger.info("=" * 60) # 确保目录存在 LAW_KB_DATA_DIR.mkdir(parents=True, exist_ok=True) (LAW_KB_DATA_DIR / "faiss").mkdir(parents=True, exist_ok=True) LOGS_DIR.mkdir(parents=True, exist_ok=True) # 加载区域映射 region_mapping = load_region_mapping(REGION_MAPPING_PATH) logger.info(f"加载区域映射: {len(region_mapping)} 条") # 扫描所有类别 all_clauses: List[Clause] = [] all_skipped: List[str] = [] category_stats: Dict[str, Dict] = {} for category in CATEGORY_DIRS: logger.info(f"扫描 [{category}]...") clauses, skipped = scan_category_dir(LAW_PACK_DIR, category, region_mapping) all_clauses.extend(clauses) all_skipped.extend(skipped) law_count = len(set(c.law_name for c in clauses)) category_stats[category] = { "laws": law_count, "clauses": len(clauses), "skipped": len(skipped), } logger.info(f" [{category}] {law_count} 法规, {len(clauses)} 条文, {len(skipped)} 跳过") logger.info(f"总计: {len(all_clauses)} 条文, {len(all_skipped)} 跳过文件") if not all_clauses: logger.error("无有效条文,终止构建") return # 写入 SQLite if mode == "full": conn = reset_sqlite(SQLITE_PATH) else: conn = init_sqlite(SQLITE_PATH) # 按法规分组 laws_map: Dict[str, Clause] = {} # law_name -> 首个 clause(取 metadata) law_clauses: Dict[str, List[Clause]] = {} for c in all_clauses: if c.law_name not in laws_map: laws_map[c.law_name] = c law_clauses[c.law_name] = [] law_clauses[c.law_name].append(c) # 写 laws 表 law_id_map: Dict[str, int] = {} for law_name, first_clause in laws_map.items(): file_path = Path(first_clause.file_path) file_mtime = file_path.stat().st_mtime if file_path.exists() else 0 cur = conn.execute( """ INSERT INTO laws (name, category, publish_date, province, city, region_level, file_path, clause_count, file_mtime, indexed_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( law_name, first_clause.category, first_clause.publish_date, first_clause.province, first_clause.city, first_clause.region_level, first_clause.file_path, len(law_clauses[law_name]), file_mtime, timestamp, ), ) law_id_map[law_name] = cur.lastrowid # 写 clauses 表(暂不写 faiss_idx,向量化后更新) clause_rows: List[tuple] = [] for law_name, clauses in law_clauses.items(): law_id = law_id_map[law_name] for c in clauses: clause_rows.append((law_id, c.chapter, c.clause_no, c.content)) conn.executemany( "INSERT INTO clauses (law_id, chapter, clause_no, content) VALUES (?, ?, ?, ?)", clause_rows, ) conn.commit() # 获取 clause id 顺序(与写入顺序一致) clause_ids = [row[0] for row in conn.execute("SELECT id FROM clauses ORDER BY id").fetchall()] clause_texts = [row[0] for row in conn.execute("SELECT content FROM clauses ORDER BY id").fetchall()] logger.info(f"SQLite 写入完成: {len(clause_ids)} 条文, {len(law_id_map)} 法规") # 批量向量化 logger.info(f"开始向量化(批量大小 {EMBEDDING_BATCH_SIZE})...") # 延迟导入 embedding 服务(避免循环依赖) from app.services.embedding import embed_batch embed_start = time.time() try: all_vecs = await embed_batch(clause_texts, batch_size=EMBEDDING_BATCH_SIZE) except Exception as e: logger.error(f"向量化失败: {e}") conn.close() return embed_time = time.time() - embed_start logger.info(f"向量化完成: {len(all_vecs)} 向量, 耗时 {embed_time:.0f}s") if len(all_vecs) != len(clause_ids): logger.error(f"向量数 {len(all_vecs)} != 条文数 {len(clause_ids)}") conn.close() return vec_dim = len(all_vecs[0]) if all_vecs else 0 logger.info(f"向量维度: {vec_dim}") # 更新 clauses 表的 faiss_idx for i, cid in enumerate(clause_ids): conn.execute("UPDATE clauses SET faiss_idx = ? WHERE id = ?", (i, cid)) conn.commit() # 构建 FAISS 索引 # 注:HNSW 构建内存峰值高(68万向量约 30GB),114 内存紧张时 OOM # 改用 IndexFlatL2(暴力检索,零额外内存,68万向量检索 <100ms) # 若内存充裕可改回 HNSW:faiss.IndexHNSWFlat(vec_dim, HNSW_M) logger.info("构建 FAISS Flat 索引(暴力检索,省内存)...") import faiss vecs_array = np.array(all_vecs, dtype=np.float32) index = faiss.IndexFlatL2(vec_dim) index.add(vecs_array) logger.info(f"FAISS 索引构建完成: {index.ntotal} 向量") # 原子写入(先写临时文件,成功后替换) tmp_faiss = FAISS_INDEX_PATH.with_suffix(".faiss.tmp") faiss.write_index(index, str(tmp_faiss)) # 替换 if FAISS_INDEX_PATH.exists(): FAISS_INDEX_PATH.unlink() tmp_faiss.rename(FAISS_INDEX_PATH) logger.info(f"FAISS 索引持久化: {FAISS_INDEX_PATH}") conn.close() # 输出统计 total_time = time.time() - start_time index_size = FAISS_INDEX_PATH.stat().st_size / 1024 / 1024 stats = { "mode": mode, "timestamp": timestamp, "total_clauses": len(all_clauses), "total_laws": len(law_id_map), "category_stats": category_stats, "skipped_files": all_skipped, "vector_dim": vec_dim, "faiss_index_size_mb": round(index_size, 1), "embed_time_seconds": round(embed_time, 0), "total_time_seconds": round(total_time, 0), } # 写统计到日志文件 stats_path = LOGS_DIR / f"build_{time.strftime('%Y%m%d_%H%M%S')}.json" with open(stats_path, "w", encoding="utf-8") as f: json.dump(stats, f, ensure_ascii=False, indent=2) # 控制台输出 print("\n" + "=" * 60) print(f" QYLAW 索引构建完成 — {mode} 模式") print(f" 耗时: {total_time:.0f}s ({total_time/3600:.1f}h)") print("=" * 60) for cat, s in category_stats.items(): print(f" [{cat}] {s['laws']} 法规, {s['clauses']} 条文, {s['skipped']} 跳过") print(f" 总切片: {len(all_clauses)} 条") print(f" 向量维度: {vec_dim}") print(f" FAISS 索引: {index_size:.1f} MB") print(f" 向量化耗时: {embed_time:.0f}s") print(f" 跳过文件: {len(all_skipped)} 篇") print(f" 索引文件: {FAISS_INDEX_PATH}") print(f" SQLite: {SQLITE_PATH}") print(f" 统计日志: {stats_path}") print("=" * 60) def main(): parser = argparse.ArgumentParser(description="QYLAW 索引构建") parser.add_argument( "--mode", choices=["full", "incremental"], default="full", help="构建模式: full=全量重建, incremental=增量更新", ) args = parser.parse_args() asyncio.run(build_index(args.mode)) if __name__ == "__main__": main()