refactor!: 完全异步化 RAG 系统，移除 LangChain ParentDocumentRetriever 依赖
Some checks failed
构建并部署 AI Agent 服务 / deploy (push) Failing after 6m34s
Some checks failed
构建并部署 AI Agent 服务 / deploy (push) Failing after 6m34s
- 重写 rag_core/vector_store.py：完全异步实现 aadd_documents、asimilarity_search - 重写 app/rag/retriever.py：异步混合检索，移除同步兼容代码 - 修改 rag_indexer/index_builder.py：全链路异步调用 - 删除 rag_core/retriever_factory.py：不再使用 LangChain ParentDocumentRetriever - 清理冗余导入和代码：移除 model_services 兼容、不需要的异常导入 - 更新 rag_indexer/README.md：反映新架构 核心改进： - 完全异步化：索引构建和检索全链路 async/await - 自定义实现：不再依赖 LangChain 的 ParentDocumentRetriever - 双向量支持：子文档同时存储 dense + sparse 向量到 Qdrant - 架构清晰：rag_core 公共组件、rag_indexer 索引、app/rag 检索
This commit is contained in:
@@ -1,8 +1,7 @@
|
||||
"""
|
||||
离线 RAG 索引构建核心流水线。
|
||||
|
||||
使用 LangChain 的 ParentDocumentRetriever 实现父子块策略。
|
||||
支持 Qdrant 混合检索（Dense + Sparse）。
|
||||
自定义实现父子块策略，支持 Qdrant 混合检索（Dense + Sparse）。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -12,33 +11,22 @@ from pathlib import Path
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Union, Optional, Any, Dict
|
||||
|
||||
from httpx import RemoteProtocolError
|
||||
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 qdrant_client.http.exceptions import ResponseHandlingException
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http.models import SparseVectorParams
|
||||
|
||||
from .loaders import DocumentLoader
|
||||
from .splitters import SplitterType, get_splitter
|
||||
|
||||
from backend.rag_core import LlamaCppEmbedder, QdrantVectorStore, create_docstore, create_parent_retriever
|
||||
|
||||
# 尝试导入新的 model_services（如果可用）
|
||||
try:
|
||||
from backend.app.model_services import get_embedding_service
|
||||
HAS_MODEL_SERVICES = True
|
||||
except ImportError:
|
||||
HAS_MODEL_SERVICES = False
|
||||
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，可直接设置此字段
|
||||
@@ -71,11 +59,10 @@ class IndexBuilderConfig:
|
||||
class IndexBuilder:
|
||||
"""RAG 索引构建主流水线，支持单块切分与父子块切分，支持混合检索。"""
|
||||
|
||||
def __init__(self, config: Optional[IndexBuilderConfig] = None, embeddings: Optional[Embeddings] = None, **kwargs):
|
||||
def __init__(self, config: Optional[IndexBuilderConfig] = None, **kwargs):
|
||||
"""
|
||||
Args:
|
||||
config: 索引构建器配置对象，优先级高于 kwargs
|
||||
embeddings: 可选的外部嵌入模型实例，如果提供则使用它
|
||||
**kwargs: 可直接传入配置参数，会合并到 config 中（为方便使用保留）
|
||||
"""
|
||||
if config is None:
|
||||
@@ -91,29 +78,15 @@ class IndexBuilder:
|
||||
|
||||
# 初始化基础组件
|
||||
self.loader = DocumentLoader()
|
||||
|
||||
# 设置嵌入模型 - 优先使用外部提供的，然后尝试使用新服务，最后回退到原来的方式
|
||||
if embeddings is not None:
|
||||
self.embeddings = embeddings
|
||||
self._embedder = None
|
||||
logger.info("使用外部提供的嵌入模型")
|
||||
elif HAS_MODEL_SERVICES:
|
||||
try:
|
||||
self.embeddings = get_embedding_service()
|
||||
self._embedder = None
|
||||
logger.info("使用 model_services 提供的嵌入服务")
|
||||
except Exception as e:
|
||||
logger.warning(f"获取嵌入服务失败，回退到 LlamaCppEmbedder: {e}")
|
||||
self._embedder = LlamaCppEmbedder()
|
||||
self.embeddings = self._embedder.as_langchain_embeddings()
|
||||
else:
|
||||
self._embedder = LlamaCppEmbedder()
|
||||
self.embeddings = self._embedder.as_langchain_embeddings()
|
||||
|
||||
# 设置嵌入模型 - 完全使用服务内部提供
|
||||
self.embeddings = get_embeddings()
|
||||
logger.info("使用统一嵌入服务")
|
||||
|
||||
# 初始化向量存储（自动支持稠密+稀疏混合检索）
|
||||
self.vector_store = QdrantVectorStore(
|
||||
self.vector_store = QdrantHybridStore(
|
||||
collection_name=config.collection_name,
|
||||
embeddings=self.embeddings if self._embedder is None else None
|
||||
embeddings=self.embeddings,
|
||||
)
|
||||
logger.info("✅ 混合检索向量存储初始化成功（稠密+BM25稀疏）")
|
||||
|
||||
@@ -141,13 +114,13 @@ class IndexBuilder:
|
||||
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,
|
||||
@@ -163,16 +136,10 @@ class IndexBuilder:
|
||||
# 文档存储
|
||||
self.docstore = self._create_or_use_docstore()
|
||||
|
||||
# 使用工厂函数创建检索器，避免重复代码
|
||||
self.retriever = create_parent_retriever(
|
||||
collection_name=cfg.collection_name,
|
||||
parent_splitter=self.parent_splitter,
|
||||
child_splitter=self.child_splitter,
|
||||
docstore=self.docstore,
|
||||
search_k=cfg.search_k,
|
||||
embeddings=self.embeddings if self._embedder is None else None,
|
||||
)
|
||||
logger.info("ParentDocumentRetriever 初始化完成")
|
||||
# 注意：不再使用 LangChain 的 ParentDocumentRetriever
|
||||
# 改为自定义实现，以支持稀疏向量
|
||||
self.retriever = None
|
||||
logger.info("父子文档模式初始化完成（使用自定义索引逻辑）")
|
||||
|
||||
def _create_or_use_docstore(self) -> BaseStore:
|
||||
"""创建或获取文档存储实例。"""
|
||||
@@ -217,54 +184,71 @@ class IndexBuilder:
|
||||
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()
|
||||
self.vector_store.add_documents(chunks)
|
||||
await self.vector_store.aadd_documents(chunks)
|
||||
return len(chunks)
|
||||
|
||||
async def _index_with_parent_child(self, documents: List[Document]) -> int:
|
||||
"""父子块模式：使用 ParentDocumentRetriever 批量添加。"""
|
||||
"""父子块模式：自定义实现，支持稠密+稀疏双向量。"""
|
||||
self.vector_store.create_collection()
|
||||
assert self.retriever is not None
|
||||
assert self.docstore is not None
|
||||
|
||||
batch_size = 10
|
||||
total = len(documents)
|
||||
processed = 0
|
||||
import uuid
|
||||
total_chunks = 0
|
||||
|
||||
for i in range(0, total, batch_size):
|
||||
batch = documents[i:i+batch_size]
|
||||
await self._add_batch_with_retry(batch, i // batch_size + 1)
|
||||
processed += len(batch)
|
||||
logger.info("批次 %d: 已处理 %d/%d", i // batch_size + 1, processed, total)
|
||||
# 1. 切分父块
|
||||
parent_chunks = self.parent_splitter.split_documents(documents)
|
||||
logger.info("切分出 %d 个父块", len(parent_chunks))
|
||||
|
||||
logger.info("ParentDocumentRetriever 索引完成，共处理 %d 个文档", processed)
|
||||
return processed
|
||||
# 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:
|
||||
"""添加批次，失败时自动重试（处理网络波动）。"""
|
||||
max_retries = 5
|
||||
base_delay = 2
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
await self.retriever.aadd_documents(batch)
|
||||
logger.info("批次 %d 成功添加 %d 个文档", batch_no, len(batch))
|
||||
return
|
||||
except (RemoteProtocolError, ConnectionError, OSError, ResponseHandlingException) as e:
|
||||
if attempt == max_retries - 1:
|
||||
logger.error("批次 %d 重试 %d 次后仍然失败: %s", batch_no, max_retries, e)
|
||||
raise
|
||||
wait_time = base_delay * (2 ** attempt)
|
||||
error_type = type(e).__name__
|
||||
logger.warning(
|
||||
"批次 %d 遇到网络异常 [%s]，%d秒后重试 (%d/%d): %s",
|
||||
batch_no, error_type, wait_time, attempt + 1, max_retries, e
|
||||
)
|
||||
self.vector_store.refresh_client()
|
||||
logger.debug("批次 %d 已刷新 Qdrant 客户端连接", batch_no)
|
||||
await asyncio.sleep(wait_time)
|
||||
"""这个方法不再使用，保留只是为了兼容（不再被调用）"""
|
||||
# 这个方法现在不需要了，因为我们重写了 _index_with_parent_child
|
||||
pass
|
||||
|
||||
# ---------- 信息获取方法 ----------
|
||||
def get_collection_info(self) -> Any:
|
||||
|
||||
Reference in New Issue
Block a user