ailine/rag_indexer/index_builder.py

"""
离线 RAG 索引构建核心流水线。

自定义实现父子块策略，支持 Qdrant 混合检索（Dense + Sparse）。
"""

import asyncio
import logging
import sys
from pathlib import Path
from dataclasses import dataclass, field
from typing import List, Union, Optional, Any, Dict

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.stores import BaseStore
from langchain_text_splitters import RecursiveCharacterTextSplitter, TextSplitter

from .loaders import DocumentLoader
from .splitters import SplitterType, get_splitter

from backend.rag_core import get_embeddings, QdrantHybridStore, create_docstore

logger = logging.getLogger(__name__)

# ---------- 配置数据类 ----------
@dataclass
class DocstoreConfig:
    """文档存储配置（用于父文档存储）。"""
    pool_config: Dict[str, Any] | None = None
    max_concurrency: int | None = None
    # 若要从外部注入已创建好的 docstore，可直接设置此字段
    instance: BaseStore | None = None

@dataclass
class IndexBuilderConfig:
    """索引构建器配置。"""
    collection_name: str = "rag_documents"
    splitter_type: SplitterType = SplitterType.PARENT_CHILD

    # 父块切分参数（仅当 splitter_type 为 PARENT_CHILD 时生效）
    parent_chunk_size: int = 1000
    parent_chunk_overlap: int = 100
    # 子块切分参数
    child_chunk_size: int = 200
    child_chunk_overlap: int = 20
    child_splitter_type: SplitterType = SplitterType.SEMANTIC   # 子块默认语义切分
    # 子块语义切分参数
    child_buffer_size: int = 1
    child_breakpoint_threshold_type: str = "percentile"
    child_breakpoint_threshold_amount: float = 90  # 降低阈值，让切分更激进
    child_min_chunk_size: int = 50  # 降低最小块大小

    # 检索参数
    search_k: int = 5

    # 文档存储配置（仅父子块模式需要）
    docstore: DocstoreConfig = field(default_factory=DocstoreConfig)

    # 其他切分器参数（当 splitter_type 非父子块时使用）
    extra_splitter_kwargs: Dict[str, Any] = field(default_factory=dict)

# ---------- 索引构建器 ----------
class IndexBuilder:
    """RAG 索引构建主流水线，支持单块切分与父子块切分，支持混合检索。"""

    def __init__(self, config: Optional[IndexBuilderConfig] = None, **kwargs):
        """
        Args:
            config: 索引构建器配置对象，优先级高于 kwargs
            **kwargs: 可直接传入配置参数，会合并到 config 中（为方便使用保留）
        """
        if config is None:
            config = IndexBuilderConfig(**kwargs)
        elif kwargs:
            # 合并 kwargs 到 config 的字段（仅更新已有字段）
            for key, value in kwargs.items():
                if hasattr(config, key):
                    setattr(config, key, value)

        self.config = config
        self._docstore_conn: Optional[str] = None  # 用于记录由 create_docstore 创建的连接信息

        # 初始化基础组件
        self.loader = DocumentLoader()

        # 设置嵌入模型 - 完全使用服务内部提供
        self.embeddings = get_embeddings()
        logger.info("使用统一嵌入服务")

        # 初始化向量存储（自动支持稠密+稀疏混合检索）
        self.vector_store = QdrantHybridStore(
            collection_name=config.collection_name,
        )
        logger.info("✅ 混合检索向量存储初始化成功（稠密+BM25稀疏）")

        # 根据切分类型初始化相关组件
        self._init_splitters_and_retriever()

    # ---------- 私有初始化方法 ----------
    def _init_splitters_and_retriever(self) -> None:
        """根据配置初始化切分器和检索器。"""
        if self.config.splitter_type == SplitterType.PARENT_CHILD:
            self._init_parent_child_mode()
        else:
            self._init_single_splitter_mode()

    def _init_single_splitter_mode(self) -> None:
        """单一切分模式（递归或语义）。"""
        splitter_kwargs = self.config.extra_splitter_kwargs.copy()
        if self.config.splitter_type == SplitterType.SEMANTIC:
            splitter_kwargs["embeddings"] = self.embeddings
        self.splitter = get_splitter(self.config.splitter_type, **splitter_kwargs)
        self.retriever = None
        self.docstore = None
        logger.info("使用单一 %s 切分器", self.config.splitter_type.value)

    def _init_parent_child_mode(self) -> None:
        cfg = self.config

        # 父块切分器
        self.parent_splitter = RecursiveCharacterTextSplitter(
            chunk_size=cfg.parent_chunk_size,
            chunk_overlap=cfg.parent_chunk_overlap,
        )

        # 子块切分器
        if cfg.child_splitter_type == SplitterType.SEMANTIC:
            self.child_splitter = get_splitter(
                SplitterType.SEMANTIC,
                embeddings=self.embeddings,
                buffer_size=cfg.child_buffer_size,
                breakpoint_threshold_type=cfg.child_breakpoint_threshold_type,
                breakpoint_threshold_amount=cfg.child_breakpoint_threshold_amount,
                min_chunk_size=cfg.child_min_chunk_size,
                **cfg.extra_splitter_kwargs
            )
        else:
            self.child_splitter = RecursiveCharacterTextSplitter(
                chunk_size=cfg.child_chunk_size,
                chunk_overlap=cfg.child_chunk_overlap,
            )

        # 文档存储
        self.docstore = self._create_or_use_docstore()

        # 注意：不再使用 LangChain 的 ParentDocumentRetriever
        # 改为自定义实现，以支持稀疏向量
        self.retriever = None
        logger.info("父子文档模式初始化完成（使用自定义索引逻辑）")

    def _create_or_use_docstore(self) -> BaseStore:
        """创建或获取文档存储实例。"""
        cfg = self.config.docstore
        if cfg.instance is not None:
            logger.debug("使用外部注入的文档存储")
            return cfg.instance

        # 使用 create_docstore 创建 PostgreSQL 存储
        docstore, conn_info = create_docstore(
            pool_config=cfg.pool_config,
            max_concurrency=cfg.max_concurrency,
        )
        self._docstore_conn = conn_info
        logger.info("文档存储已创建（PostgreSQL）")
        return docstore

    # ---------- 公共构建方法 ----------
    async def build_from_file(self, file_path: Union[str, Path]) -> int:
        """从单个文件构建索引。"""
        logger.info("加载文件: %s", file_path)
        documents = self.loader.load_file(file_path)
        logger.info("已加载 %d 个文档", len(documents))
        return await self._process_documents(documents)

    async def build_from_directory(self, directory_path: Union[str, Path], recursive: bool = True) -> int:
        """从目录递归构建索引。"""
        logger.info("加载目录: %s (递归=%s)", directory_path, recursive)
        documents = self.loader.load_directory(directory_path, recursive=recursive)
        logger.info("已从目录加载 %d 个文档", len(documents))
        return await self._process_documents(documents)

    async def _process_documents(self, documents: List[Document]) -> int:
        """处理文档列表，分发给相应的索引逻辑。"""
        if not documents:
            logger.warning("没有文档需要处理")
            return 0

        if self.config.splitter_type == SplitterType.PARENT_CHILD:
            return await self._index_with_parent_child(documents)
        else:
            return await self._index_with_single_splitter(documents)

    async def _index_with_single_splitter(self, documents: List[Document]) -> int:
        """单一切分模式：切分后直接写入向量库（异步）。"""
        chunks = self.splitter.split_documents(documents)
        logger.info("已切分为 %d 个块", len(chunks))

        self.vector_store.create_collection()
        await self.vector_store.aadd_documents(chunks)
        return len(chunks)

    async def _index_with_parent_child(self, documents: List[Document]) -> int:
        """父子块模式：自定义实现，支持稠密+稀疏双向量。"""
        self.vector_store.create_collection()
        assert self.docstore is not None

        import uuid
        total_chunks = 0

        # 1. 切分父块
        parent_chunks = self.parent_splitter.split_documents(documents)
        logger.info("切分出 %d 个父块", len(parent_chunks))

        # 2. 为每个父块生成 UUID 并存储
        parent_docs_with_ids = []
        for parent_chunk in parent_chunks:
            parent_id = str(uuid.uuid4())
            parent_chunk.metadata["id"] = parent_id
            parent_chunk.metadata["is_parent"] = True
            parent_docs_with_ids.append((parent_id, parent_chunk))

        # 3. 父文档批量存入 PostgreSQL
        await self.docstore.amset(parent_docs_with_ids)
        logger.info("已存入 %d 个父文档到 PostgreSQL", len(parent_docs_with_ids))

        # 4. 切分子块并添加 parent_id
        all_child_chunks = []
        for parent_id, parent_chunk in parent_docs_with_ids:
            child_chunks = self.child_splitter.split_documents([parent_chunk])
            for child_chunk in child_chunks:
                child_chunk.metadata["parent_id"] = parent_id
                child_chunk.metadata["is_parent"] = False
                # 继承父文档的重要元数据
                child_chunk.metadata["source"] = parent_chunk.metadata.get("source")
                child_chunk.metadata["page"] = parent_chunk.metadata.get("page")
                child_chunk.metadata["file_path"] = parent_chunk.metadata.get("file_path")
                all_child_chunks.append(child_chunk)

        total_chunks = len(all_child_chunks)
        logger.info("切分出 %d 个子块", total_chunks)

        # 5. 子文档分批存入 Qdrant（双向量，异步）
        batch_size = 100
        for i in range(0, total_chunks, batch_size):
            batch = all_child_chunks[i:i+batch_size]
            await self.vector_store.aadd_documents(batch)
            logger.info("已向 Qdrant 存入子文档批次 %d/%d", 
                       i // batch_size + 1, 
                       (total_chunks + batch_size - 1) // batch_size)

        logger.info("父子文档索引完成：%d 父文档，%d 子文档", 
                   len(parent_docs_with_ids), total_chunks)
        return total_chunks

    async def _add_batch_with_retry(self, batch: List[Document], batch_no: int) -> None:
        """这个方法不再使用，保留只是为了兼容（不再被调用）"""
        # 这个方法现在不需要了，因为我们重写了 _index_with_parent_child
        pass

    # ---------- 信息获取方法 ----------
    def get_collection_info(self) -> Any:
        """获取向量库集合信息。"""
        return self.vector_store.get_collection_info()

    def get_child_splitter(self) -> TextSplitter:
        """获取当前使用的子块切分器。"""
        if self.config.splitter_type == SplitterType.PARENT_CHILD:
            return self.child_splitter
        return self.splitter

    def get_parent_splitter(self) -> RecursiveCharacterTextSplitter:
        """获取父块切分器（仅父子块模式可用）。"""
        if self.config.splitter_type != SplitterType.PARENT_CHILD:
            raise RuntimeError("父块切分器仅在父子块模式下可用")
        return self.parent_splitter

    def get_docstore(self) -> BaseStore:
        """获取文档存储实例（仅父子块模式可用）。"""
        if self.config.splitter_type != SplitterType.PARENT_CHILD:
            raise RuntimeError("文档存储仅在父子块模式下可用")
        assert self.docstore is not None
        return self.docstore

    # ---------- 资源管理 ----------
    def close(self) -> None:
        """关闭资源（同步版本，供上下文管理器使用）。"""
        if self.docstore is not None and hasattr(self.docstore, "aclose"):
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                # 无运行中的事件循环，创建临时循环
                loop = asyncio.new_event_loop()
                loop.run_until_complete(self.docstore.aclose())
                loop.close()
            else:
                # 已有运行中的循环，创建任务（用户自行等待）
                loop.create_task(self.docstore.aclose())
        logger.info("IndexBuilder 资源已关闭")

    async def aclose(self) -> None:
        """异步关闭资源。"""
        if self.docstore is not None and hasattr(self.docstore, "aclose"):
            await self.docstore.aclose()
        logger.info("IndexBuilder 资源已异步关闭")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.aclose()
        return False