zhenxun_bot/zhenxun/services/ai/context/rag/builder.py

from typing import Any

from zhenxun.services.ai.context.rag.backends import StorageBackend
from zhenxun.services.ai.context.rag.configs import RAGConfig
from zhenxun.services.ai.context.rag.engine import ScopedRAGClient
from zhenxun.services.ai.context.rag.ingestion import (
    ChunkingStrategy,
    DedupNode,
    DocumentChunking,
    DynamicChunkingNode,
    EmbeddingNode,
    IndexPipeline,
    StorageCommitNode,
)
from zhenxun.services.ai.context.rag.retrieval import (
    BaseRetriever,
    DatabaseSparseRetriever,
    HybridRetriever,
    LifecyclePostProcessor,
    PipelineRetriever,
    PostProcessor,
    PreProcessor,
    RerankRetriever,
    VectorDBRetriever,
)
from zhenxun.services.log import logger


class RAGBuilder:
    """
    RAG 管线组装构建器 (Fluent Builder Pattern)。
    使用内部状态驱动模式 (The Memory Pattern)，内部维护私有的 RAGConfig 实例。
    """

    def __init__(
        self, storage: StorageBackend | None = None, config: RAGConfig | None = None
    ):
        self._config = config or RAGConfig()
        if storage is not None:
            self._config.storage = storage

    def with_embedder(self, embedder: Any) -> "RAGBuilder":
        """
        设置向量化引擎。

        参数:
            embedder: 实现了向量化协议的引擎实例 (如 DefaultEmbedder, FastEmbedder)。
        """
        self._config.embedder = embedder
        return self

    def with_retriever(self, retriever: BaseRetriever) -> "RAGBuilder":
        """
        替换底层的向量库查表算法，注入自定义召回器。

        参数:
            retriever: 实现了 BaseRetriever 协议的自定义召回器实例。
        """
        self._config.custom_retriever = retriever
        return self

    def with_scope(self, scopes: str | list[str]) -> "RAGBuilder":
        """
        设置数据隔离作用域。

        参数:
            scopes: 支持单作用域前缀(字符串)或联合检索多作用域(列表)。
            如 "/" 或 ["/group_1", "/user_2"]。
        """
        self._config.scopes = scopes
        return self

    def with_chunking(self, strategy: ChunkingStrategy) -> "RAGBuilder":
        """
        设置文档切块策略。

        参数:
            strategy: 切块策略实例 (如 DocumentChunking, RecursiveCharacterChunking)。
        """
        self._config.chunking.strategy = strategy
        return self

    def enable_dedup(self, threshold: float = 0.98) -> "RAGBuilder":
        """
        开启批处理入库去重。
        在文本分块入库前，通过对比向量相似度拦截高度重复的内容。

        参数:
            threshold: 去重相似度阈值，大于该值的块将被判定为重复并丢弃。
        """
        self._config.dedup.enable = True
        self._config.dedup.threshold = threshold
        return self

    def disable_dedup(self) -> "RAGBuilder":
        """
        关闭批处理入库去重。
        """
        self._config.dedup.enable = False
        return self

    def enable_rerank(self, model_name: str, top_n: int = 5) -> "RAGBuilder":
        """
        开启大模型交叉注意力重排。
        对初筛召回的结果使用专用的 Rerank 模型进行二次排序，极大提升召回准确率。

        参数:
            model_name: 用于重排序的模型名称 (如 'BAAI/bge-reranker-v2-m3')。
            top_n: 重排后最终保留并返回的文档数量。
        """
        self._config.rerank.enable = True
        self._config.rerank.model_name = model_name
        self._config.rerank.top_n = top_n
        return self

    def enable_hybrid_search(
        self, dense_weight: float = 0.7, sparse_weight: float = 0.3
    ) -> "RAGBuilder":
        """
        开启双轨混合检索 (Hybrid Search) 及 RRF 融合。
        并发调用向量检索 (Dense) 和关键词检索 (Sparse)，结合两者优势。

        参数:
            dense_weight: 稠密向量检索的分数计算权重。
            sparse_weight: 稀疏关键词(BM25等)检索的分数计算权重。
        """
        self._config.hybrid.enable = True
        self._config.hybrid.dense_weight = dense_weight
        self._config.hybrid.sparse_weight = sparse_weight
        return self

    def enable_lifecycle_scoring(
        self,
        half_life_days: int = 30,
        decay_weight: float = 0.3,
        semantic_weight: float = 0.7,
        importance_weight: float = 0.0,
        reinforcement_weight: float = 0.2,
    ) -> "RAGBuilder":
        """
        开启生命周期打分后处理。
        综合考虑记忆的时间新鲜度、基础语义相似度、重要性以及访问频次，模拟人类记忆遗忘曲线。

        参数:
            half_life_days: 时间衰减的半衰期(天)。经过这么多天后，时间得分衰减一半。
            decay_weight: 时间衰减得分所占权重。
            semantic_weight: 基础语义相似度所占权重。
            importance_weight: 客观重要性所占权重。
            reinforcement_weight: 访问强化(被回想次数越多越容易想起)所占权重。
        """
        self._config.lifecycle.enable = True
        self._config.lifecycle.half_life_days = half_life_days
        self._config.lifecycle.decay_weight = decay_weight
        self._config.lifecycle.semantic_weight = semantic_weight
        self._config.lifecycle.importance_weight = importance_weight
        self._config.lifecycle.reinforcement_weight = reinforcement_weight
        return self

    def add_pre_processor(self, processor: PreProcessor) -> "RAGBuilder":
        """
        挂载自定义查询前处理器。

        参数:
            processor: 实现了 PreProcessor 协议的处理器实例。
        """
        self._config.pre_processors.append(processor)
        return self

    def add_post_processor(self, processor: PostProcessor) -> "RAGBuilder":
        """
        挂载自定义检索后处理器。

        参数:
            processor: 实现了 PostProcessor 协议的处理器实例。
        """
        self._config.post_processors.append(processor)
        return self

    @classmethod
    def resolve(cls, config: RAGConfig | dict | Any | None) -> RAGConfig:
        """
        解析多种格式的配置，统一转换为 RAGConfig 对象。

        参数:
            config: RAGConfig 实例、字典或其他类型配置。
        """
        if isinstance(config, RAGConfig):
            return config
        if isinstance(config, dict):
            return RAGConfig(**config)
        return RAGConfig()

    def build(self) -> ScopedRAGClient:
        """
        完成所有积木组装，输出最终的 RAG 客户端实体。
        """
        cfg = self._config
        storage = cfg.storage
        if not storage:
            raise ValueError("RAGBuilder 必须配置 storage 才能 build。")

        embedder = cfg.embedder
        if not embedder:
            from zhenxun.services.ai.context.rag.backends import DefaultEmbedder
            from zhenxun.services.ai.llm.manager import get_default_model

            embedder = DefaultEmbedder(model_name=get_default_model("embedding"))
            logger.debug("RAGBuilder: 未指定 Embedder，已使用系统默认 Embedder。")

        chunking_strategy = cfg.chunking.strategy or DocumentChunking()

        scopes = cfg.scopes
        nodes = [
            DynamicChunkingNode(chunking_strategy),
            EmbeddingNode(embedder),
        ]
        if cfg.dedup.enable:
            nodes.append(DedupNode(cfg.dedup.threshold))

        nodes.append(StorageCommitNode(storage))
        pipeline = IndexPipeline(nodes)

        base_retriever: BaseRetriever = cfg.custom_retriever or VectorDBRetriever(
            storage,
            embedder,
            scopes[0] if isinstance(scopes, list) else scopes,
        )

        if cfg.hybrid.enable:
            database_sparse_retriever = DatabaseSparseRetriever(
                storage, scopes[0] if isinstance(scopes, list) else scopes
            )
            base_retriever = HybridRetriever(
                dense_retriever=base_retriever,
                sparse_retriever=database_sparse_retriever,
                dense_weight=cfg.hybrid.dense_weight,
                sparse_weight=cfg.hybrid.sparse_weight,
            )

        if cfg.rerank.enable and cfg.rerank.model_name:
            base_retriever = RerankRetriever(
                base_retriever, cfg.rerank.model_name, cfg.rerank.top_n
            )

        pre_processors = list(cfg.pre_processors)

        post_processors = list(cfg.post_processors)
        if cfg.lifecycle.enable:
            post_processors.append(
                LifecyclePostProcessor(
                    half_life_days=cfg.lifecycle.half_life_days,
                    decay_weight=cfg.lifecycle.decay_weight,
                    semantic_weight=cfg.lifecycle.semantic_weight,
                    importance_weight=cfg.lifecycle.importance_weight,
                    reinforcement_weight=cfg.lifecycle.reinforcement_weight,
                )
            )

        if pre_processors or post_processors:
            base_retriever = PipelineRetriever(
                base_retriever,
                pre_processors=pre_processors,
                post_processors=post_processors,
            )

        return ScopedRAGClient(
            storage=storage,
            retriever=base_retriever,
            pipeline=pipeline,
            scopes=scopes,
            config=cfg,
        )