root
8db63e7a8d
新增功能： - 创建 app/model_services 模块，提供统一的模型服务获取接口 - 实现 BaseServiceProvider 基类和 FallbackServiceChain 降级链 - 实现 get_embedding_service()：优先本地 llama.cpp，降级到智谱 API - 实现 get_rerank_service()：优先本地 llama.cpp，降级到智谱 API - 支持单例管理，确保全局只有一个服务实例 修改内容： - 更新 app/config.py，添加智谱 API 相关配置 - 修改 rag_core/vector_store.py：支持接受外部传入的 embeddings - 修改 rag_core/retriever_factory.py：支持接受外部传入的 embeddings - 修改 app/agent/rag_initializer.py：使用 get_embedding_service() - 修改 app/rag/pipeline.py：使用 get_rerank_service() - 修改 app/memory/mem0_client.py：智能判断可用服务配置 mem0 - 修改 rag_indexer/index_builder.py：支持使用新服务，保持向后兼容 - 修改 rag_indexer/config.py：添加智谱配置 环境变量： - ZHIPUAI_API_KEY：智谱 API 密钥（必选） - ZHIPU_EMBEDDING_MODEL：可选，默认 embedding-3 - ZHIPU_RERANK_MODEL：可选，默认 rerank-2 - ZHIPU_API_BASE：可选，默认 https://open.bigmodel.cn/api/paas/v4
86 lines
3.1 KiB
Python
86 lines
3.1 KiB
Python
# rag/pipeline.py
|
|
|
|
import asyncio
|
|
import os
|
|
from typing import List
|
|
from langchain_core.documents import Document
|
|
from langchain_core.language_models import BaseLanguageModel
|
|
|
|
from ..model_services import get_rerank_service
|
|
from .query_transform import MultiQueryGenerator
|
|
from .fusion import reciprocal_rank_fusion
|
|
|
|
class RAGPipeline:
|
|
"""
|
|
固定流程的 RAG 检索流水线：
|
|
多路改写 → 并行检索 → RRF融合 → 重排序 → 返回父文档
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
retriever, # 基础检索器（应返回父文档，例如 ParentDocumentRetriever 实例）
|
|
llm: BaseLanguageModel,
|
|
num_queries: int = 3,
|
|
rerank_top_n: int = 5,
|
|
):
|
|
"""
|
|
Args:
|
|
retriever: 基础检索器对象，需实现 ainvoke(query) 异步方法
|
|
llm: 用于生成多路查询的语言模型
|
|
num_queries: 生成的查询变体数量
|
|
rerank_top_n: 最终返回的文档数量
|
|
rerank_model: 重排序模型名称
|
|
"""
|
|
self.retriever = retriever
|
|
self.llm = llm
|
|
self.num_queries = num_queries
|
|
self.rerank_top_n = rerank_top_n
|
|
|
|
# 初始化组件 - 使用统一的重排服务获取接口
|
|
self.query_generator = MultiQueryGenerator(llm=llm, num_queries=num_queries)
|
|
self.reranker = get_rerank_service()
|
|
|
|
async def aretrieve(self, query: str) -> List[Document]:
|
|
"""
|
|
异步执行完整检索流程
|
|
"""
|
|
# Step 1: 生成多路查询
|
|
queries = await self.query_generator.agenerate(query)
|
|
# 包含原始查询，确保至少有一条
|
|
if query not in queries:
|
|
queries.insert(0, query)
|
|
else:
|
|
# 如果原始查询已在列表中，将其移至首位
|
|
queries.remove(query)
|
|
queries.insert(0, query)
|
|
|
|
# Step 2: 并行检索（每个查询获取文档列表）
|
|
tasks = [self.retriever.ainvoke(q) for q in queries]
|
|
doc_lists = await asyncio.gather(*tasks)
|
|
|
|
# Step 3: RRF 融合
|
|
fused_docs = reciprocal_rank_fusion(doc_lists)
|
|
|
|
# Step 4: 重排序
|
|
try:
|
|
final_docs = self.reranker.compress_documents(fused_docs, query, top_n=self.rerank_top_n)
|
|
except Exception:
|
|
# 若重排序器不可用，直接返回融合后的前 N 个结果
|
|
final_docs = fused_docs[:self.rerank_top_n]
|
|
|
|
return final_docs
|
|
|
|
def retrieve(self, query: str) -> List[Document]:
|
|
"""同步检索入口（内部调用异步方法）"""
|
|
return asyncio.run(self.aretrieve(query))
|
|
|
|
def format_context(self, documents: List[Document]) -> str:
|
|
"""将文档列表格式化为上下文字符串"""
|
|
if not documents:
|
|
return ""
|
|
|
|
parts = []
|
|
for i, doc in enumerate(documents, 1):
|
|
source = doc.metadata.get("source", "未知来源")
|
|
parts.append(f"【资料 {i}】来源：{source}\n{doc.page_content}\n---\n")
|
|
return "\n".join(parts) |