diff --git a/api/README.md b/api/README.md index 0862cb6..98e125c 100644 --- a/api/README.md +++ b/api/README.md @@ -21,6 +21,22 @@ python init_db.py - 创建 `data/app.db` SQLite 数据库 - 初始化默认声音数据 +可选参数(按需重建): + +```bash +# 仅重建数据库(drop + create)并初始化默认数据 +python init_db.py --rebuild-db + +# 仅重建向量库集合(不动 DB 表结构);会重置文档索引状态为 pending +python init_db.py --rebuild-vector-store + +# 同时重建 DB 和向量库 +python init_db.py --rebuild-db --rebuild-vector-store + +# 仅执行重建,不写入默认数据 +python init_db.py --rebuild-db --skip-seed +``` + ### 3. 启动服务 ```bash diff --git a/api/init_db.py b/api/init_db.py index ff4d8fe..7a6d571 100644 --- a/api/init_db.py +++ b/api/init_db.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 """初始化数据库""" +import argparse import os import sys @@ -7,7 +8,7 @@ import sys sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from app.db import Base, engine, DATABASE_URL -from app.models import Voice, Assistant, KnowledgeBase, Workflow, LLMModel, ASRModel +from app.models import Voice, Assistant, KnowledgeBase, Workflow, LLMModel, ASRModel, KnowledgeDocument def init_db(): @@ -22,6 +23,52 @@ def init_db(): print("✅ 数据库表创建完成") +def rebuild_vector_store(reset_doc_status: bool = True): + """重建知识库向量集合(按 DB 中的 KB 列表重建 collection 壳)。""" + from app.db import SessionLocal + from app.vector_store import vector_store + + db = SessionLocal() + try: + print("🧹 重建向量库集合...") + kb_list = db.query(KnowledgeBase).all() + + # 删除现有 KB 集合 + try: + collections = vector_store.client.list_collections() + except Exception as exc: + raise RuntimeError(f"无法读取向量集合列表: {exc}") from exc + + for col in collections: + name = getattr(col, "name", None) or str(col) + if name.startswith("kb_"): + try: + vector_store.client.delete_collection(name=name) + print(f" - removed {name}") + except Exception as exc: + print(f" - skip remove {name}: {exc}") + + # 按 DB 重建 KB 集合 + for kb in kb_list: + vector_store.create_collection(kb.id, kb.embedding_model) + print(f" + created kb_{kb.id} ({kb.embedding_model})") + + if reset_doc_status: + kb.chunk_count = 0 + docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.kb_id == kb.id).all() + kb.doc_count = 0 + for doc in docs: + doc.chunk_count = 0 + doc.status = "pending" + doc.error_message = None + doc.processed_at = None + + db.commit() + print("✅ 向量库重建完成(仅重建集合壳,文档需重新索引)") + finally: + db.close() + + def init_default_data(): from sqlalchemy.orm import Session from app.db import SessionLocal @@ -355,15 +402,45 @@ def init_default_asr_models(): if __name__ == "__main__": + parser = argparse.ArgumentParser(description="初始化/重建 AI VideoAssistant 数据与向量库") + parser.add_argument( + "--rebuild-db", + action="store_true", + help="重建数据库(drop + create tables)", + ) + parser.add_argument( + "--rebuild-vector-store", + action="store_true", + help="重建向量库 KB 集合(清空后按 DB 的 knowledge_bases 重建 collection)", + ) + parser.add_argument( + "--skip-seed", + action="store_true", + help="跳过默认数据初始化", + ) + args = parser.parse_args() + + # 无参数时保持旧行为:重建 DB + 初始化默认数据 + if not args.rebuild_db and not args.rebuild_vector_store and not args.skip_seed: + args.rebuild_db = True + # 确保 data 目录存在 data_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data") os.makedirs(data_dir, exist_ok=True) - init_db() - init_default_data() - init_default_assistants() - init_default_workflows() - init_default_knowledge_bases() - init_default_llm_models() - init_default_asr_models() - print("🎉 数据库初始化完成!") + if args.rebuild_db: + init_db() + + if not args.skip_seed: + init_default_data() + init_default_assistants() + init_default_workflows() + init_default_knowledge_bases() + init_default_llm_models() + init_default_asr_models() + print("✅ 默认数据初始化完成") + + if args.rebuild_vector_store: + rebuild_vector_store(reset_doc_status=True) + + print("🎉 初始化脚本执行完成!")