2025-07-21 19:25:07 +08:00
|
|
|
import os
|
2025-03-08 20:26:00 +08:00
|
|
|
from dotenv import load_dotenv
|
|
|
|
|
|
2025-07-21 13:10:15 +08:00
|
|
|
load_dotenv("src/.env", override=True)
|
2025-03-16 01:11:13 +08:00
|
|
|
|
2025-05-24 11:29:45 +08:00
|
|
|
from concurrent.futures import ThreadPoolExecutor # noqa: E402
|
2025-03-16 01:11:13 +08:00
|
|
|
executor = ThreadPoolExecutor()
|
|
|
|
|
|
2025-05-24 11:29:45 +08:00
|
|
|
from src.config import Config # noqa: E402
|
2025-03-16 01:11:13 +08:00
|
|
|
config = Config()
|
|
|
|
|
|
2025-07-21 19:25:07 +08:00
|
|
|
# 导入知识库相关模块
|
2025-07-23 19:21:45 +08:00
|
|
|
from src.knowledge.kb_factory import KnowledgeBaseFactory # noqa: E402
|
|
|
|
|
from src.knowledge.kb_manager import KnowledgeBaseManager # noqa: E402
|
|
|
|
|
from src.knowledge.lightrag_kb import LightRagKB # noqa: E402
|
|
|
|
|
from src.knowledge.chroma_kb import ChromaKB # noqa: E402
|
|
|
|
|
from src.knowledge.milvus_kb import MilvusKB # noqa: E402
|
2025-07-21 19:25:07 +08:00
|
|
|
|
|
|
|
|
# 注册知识库类型
|
|
|
|
|
KnowledgeBaseFactory.register("chroma", ChromaKB, {
|
|
|
|
|
"chunk_size": 1000,
|
|
|
|
|
"chunk_overlap": 200,
|
|
|
|
|
"description": "基于 ChromaDB 的轻量级向量知识库,适合开发和小规模部署"
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
KnowledgeBaseFactory.register("milvus", MilvusKB, {
|
|
|
|
|
"chunk_size": 1000,
|
|
|
|
|
"chunk_overlap": 200,
|
|
|
|
|
"description": "基于 Milvus 的生产级向量知识库,适合大规模高性能部署"
|
|
|
|
|
})
|
|
|
|
|
|
2025-07-26 03:36:54 +08:00
|
|
|
|
|
|
|
|
KnowledgeBaseFactory.register("lightrag", LightRagKB, {
|
|
|
|
|
"description": "基于图检索的知识库,支持实体关系构建和复杂查询"
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
|
2025-07-21 19:25:07 +08:00
|
|
|
# 创建知识库管理器
|
|
|
|
|
work_dir = os.path.join(config.save_dir, "knowledge_base_data")
|
|
|
|
|
knowledge_base = KnowledgeBaseManager(work_dir)
|
2025-03-20 19:51:46 +08:00
|
|
|
|
2025-07-23 19:21:45 +08:00
|
|
|
from src.knowledge import GraphDatabase # noqa: E402
|
2025-03-20 19:51:46 +08:00
|
|
|
graph_base = GraphDatabase()
|