2112d4833d
- 优化知识库文件夹结构，增强可维护性 - 新增知识库管理器和工厂类，统一管理不同类型的知识库实例 - 实现知识库的基本操作接口，包括创建、更新和删除数据库 - 引入新的知识库实现类（ChromaKB、MilvusKB、LightRagKB），支持多种存储后端 - 增强文件路径验证，确保安全性 - 更新文档，提供新的知识库接口说明
108 lines
3.1 KiB
Python
108 lines
3.1 KiB
Python
from src.knowledge.base import KBNotFoundError, KnowledgeBase
|
|
from src.utils import logger
|
|
|
|
|
|
class KnowledgeBaseFactory:
|
|
"""知识库工厂类，负责创建不同类型的知识库实例"""
|
|
|
|
# 注册的知识库类型映射 {kb_type: kb_class}
|
|
_kb_types: dict[str, type[KnowledgeBase]] = {}
|
|
|
|
# 每种类型的默认配置
|
|
_default_configs: dict[str, dict] = {}
|
|
|
|
@classmethod
|
|
def register(cls, kb_type: str, kb_class: type[KnowledgeBase], default_config: dict = None):
|
|
"""
|
|
注册知识库类型
|
|
|
|
Args:
|
|
kb_type: 知识库类型标识
|
|
kb_class: 知识库类
|
|
default_config: 默认配置
|
|
"""
|
|
if not issubclass(kb_class, KnowledgeBase):
|
|
raise ValueError("Knowledge base class must inherit from KnowledgeBase")
|
|
|
|
cls._kb_types[kb_type] = kb_class
|
|
cls._default_configs[kb_type] = default_config or {}
|
|
# logger.info(f"Registered knowledge base type: {kb_type}")
|
|
|
|
@classmethod
|
|
def create(cls, kb_type: str, work_dir: str, **kwargs) -> KnowledgeBase:
|
|
"""
|
|
创建知识库实例
|
|
|
|
Args:
|
|
kb_type: 知识库类型
|
|
work_dir: 工作目录
|
|
**kwargs: 其他初始化参数
|
|
|
|
Returns:
|
|
知识库实例
|
|
|
|
Raises:
|
|
KBNotFoundError: 未知的知识库类型
|
|
"""
|
|
if kb_type not in cls._kb_types:
|
|
available_types = list(cls._kb_types.keys())
|
|
raise KBNotFoundError(f"Unknown knowledge base type: {kb_type}. Available types: {available_types}")
|
|
|
|
kb_class = cls._kb_types[kb_type]
|
|
|
|
# 合并默认配置和用户配置
|
|
config = cls._default_configs[kb_type].copy()
|
|
config.update(kwargs)
|
|
|
|
try:
|
|
# 创建实例
|
|
instance = kb_class(work_dir, **config)
|
|
logger.info(f"Created {kb_type} knowledge base instance at {work_dir}")
|
|
return instance
|
|
except Exception as e:
|
|
logger.error(f"Failed to create {kb_type} knowledge base: {e}")
|
|
raise
|
|
|
|
@classmethod
|
|
def get_available_types(cls) -> dict[str, dict]:
|
|
"""
|
|
获取所有可用的知识库类型
|
|
|
|
Returns:
|
|
知识库类型信息字典
|
|
"""
|
|
result = {}
|
|
for kb_type, kb_class in cls._kb_types.items():
|
|
result[kb_type] = {
|
|
"class_name": kb_class.__name__,
|
|
"description": kb_class.__doc__ or "",
|
|
"default_config": cls._default_configs[kb_type],
|
|
}
|
|
return result
|
|
|
|
@classmethod
|
|
def is_type_supported(cls, kb_type: str) -> bool:
|
|
"""
|
|
检查是否支持指定的知识库类型
|
|
|
|
Args:
|
|
kb_type: 知识库类型
|
|
|
|
Returns:
|
|
是否支持
|
|
"""
|
|
return kb_type in cls._kb_types
|
|
|
|
@classmethod
|
|
def get_default_config(cls, kb_type: str) -> dict:
|
|
"""
|
|
获取指定类型的默认配置
|
|
|
|
Args:
|
|
kb_type: 知识库类型
|
|
|
|
Returns:
|
|
默认配置字典
|
|
"""
|
|
return cls._default_configs.get(kb_type, {}).copy()
|