69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
|
"""独立重建 KB3，避免交互命令超时中止向量索引的持久化。"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import gc
|
||
|
|
import json
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
SERVER_DIR = Path(__file__).resolve().parents[1]
|
||
|
|
if str(SERVER_DIR) not in sys.path:
|
||
|
|
sys.path.insert(0, str(SERVER_DIR))
|
||
|
|
|
||
|
|
from sqlmodel import Session
|
||
|
|
|
||
|
|
from app.db import engine, init_db
|
||
|
|
from app.models import KnowledgeBase, now
|
||
|
|
from app.rag import store
|
||
|
|
from app.rag.pipeline import rebuild_kb
|
||
|
|
|
||
|
|
|
||
|
|
LOG_PATH = SERVER_DIR / "data" / "kb3-rebuild.log"
|
||
|
|
|
||
|
|
|
||
|
|
def log(payload: dict) -> None:
|
||
|
|
payload["at"] = now()
|
||
|
|
with LOG_PATH.open("a", encoding="utf-8") as stream:
|
||
|
|
stream.write(json.dumps(payload, ensure_ascii=False) + "\n")
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
init_db()
|
||
|
|
with Session(engine) as session:
|
||
|
|
kb = session.get(KnowledgeBase, 3)
|
||
|
|
if not kb:
|
||
|
|
raise RuntimeError("KB3 不存在")
|
||
|
|
kb.index_status = "needs_reindex"
|
||
|
|
kb.chunk_count = 0
|
||
|
|
session.add(kb); session.commit()
|
||
|
|
log({"status": "started", "files": kb.file_count})
|
||
|
|
completed = 0
|
||
|
|
def progress(file_name: str, chunks: int) -> None:
|
||
|
|
nonlocal completed
|
||
|
|
completed += 1
|
||
|
|
log({"status": "progress", "completed_files": completed, "file": file_name, "chunks": chunks})
|
||
|
|
result = rebuild_kb(session, 3, progress_cb=progress)
|
||
|
|
# 给持久化 HNSW writer 一个明确的落盘窗口，再以同一独立进程验证计数。
|
||
|
|
gc.collect()
|
||
|
|
time.sleep(5)
|
||
|
|
vector_count = store.count(3)
|
||
|
|
kb = session.get(KnowledgeBase, 3)
|
||
|
|
if vector_count <= 0 or vector_count != kb.chunk_count:
|
||
|
|
kb.index_status = "needs_reindex"
|
||
|
|
session.add(kb); session.commit()
|
||
|
|
log({"status": "verification_failed", "result": result,
|
||
|
|
"database_chunks": kb.chunk_count, "vector_chunks": vector_count})
|
||
|
|
raise RuntimeError("KB3 向量计数与数据库不一致")
|
||
|
|
kb.index_status = "ready"
|
||
|
|
session.add(kb); session.commit()
|
||
|
|
log({"status": "completed", "result": result, "vector_chunks": vector_count})
|
||
|
|
# Chroma 1.5 的 Rust compactor 在后台写 HNSW；显式停止系统确保索引文件
|
||
|
|
# 在进程退出前完整落盘，避免新进程读取到只有 index_metadata 的半成品。
|
||
|
|
store.get_client()._system.stop()
|
||
|
|
log({"status": "persisted", "vector_chunks": vector_count})
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|