freedak
641e33b834
- 语义检索(FAISS + embedding)+ 精确查找(法规名+条号) - RAG 问答(SSE 流式,支持 thinking 折叠显示) - 法规浏览(原文阅读) - 历史记录(检索+对话持久化到 SQLite) - 设置页(系统提示词/模板/LLM 参数可配置) - 检索质量评估脚本 Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
330 lines
10 KiB
Python
330 lines
10 KiB
Python
"""索引构建脚本 — 全量/增量构建 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()
|