""" 离线 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