46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
import os
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv("src/.env", override=True)
|
|
|
|
from concurrent.futures import ThreadPoolExecutor # noqa: E402
|
|
|
|
executor = ThreadPoolExecutor()
|
|
|
|
from src.config import Config # noqa: E402
|
|
|
|
config = Config()
|
|
|
|
# 导入知识库相关模块
|
|
from src.knowledge.chroma_kb import ChromaKB # noqa: E402
|
|
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.milvus_kb import MilvusKB # noqa: E402
|
|
|
|
# 注册知识库类型
|
|
KnowledgeBaseFactory.register(
|
|
"chroma",
|
|
ChromaKB,
|
|
{"chunk_size": 1000, "chunk_overlap": 200, "description": "基于 ChromaDB 的轻量级向量知识库，适合开发和小规模部署"},
|
|
)
|
|
|
|
KnowledgeBaseFactory.register(
|
|
"milvus",
|
|
MilvusKB,
|
|
{"chunk_size": 1000, "chunk_overlap": 200, "description": "基于 Milvus 的生产级向量知识库，适合大规模高性能部署"},
|
|
)
|
|
|
|
|
|
KnowledgeBaseFactory.register("lightrag", LightRagKB, {"description": "基于图检索的知识库，支持实体关系构建和复杂查询"})
|
|
|
|
|
|
# 创建知识库管理器
|
|
work_dir = os.path.join(config.save_dir, "knowledge_base_data")
|
|
knowledge_base = KnowledgeBaseManager(work_dir)
|
|
|
|
from src.knowledge import GraphDatabase # noqa: E402
|
|
|
|
graph_base = GraphDatabase()
|