ForcePilot/src/knowledge/factory.py

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()