27 lines
1021 B
Python
27 lines
1021 B
Python
# app/rag_initializer.py
|
|
from app.rag.tools import create_rag_tool_sync, create_rag_tool_async
|
|
from rag_core import create_parent_retriever
|
|
from app.model_services import get_embedding_service
|
|
from app.logger import info, warning
|
|
|
|
async def init_rag_tool(local_llm_creator):
|
|
"""初始化 RAG 工具，失败返回 None"""
|
|
try:
|
|
info("🔄 正在初始化 RAG 检索系统...")
|
|
# 使用统一的嵌入服务获取接口
|
|
embeddings = get_embedding_service()
|
|
retriever = create_parent_retriever(
|
|
collection_name="rag_documents",
|
|
search_k=5,
|
|
embeddings=embeddings
|
|
)
|
|
rewrite_llm = local_llm_creator()
|
|
rag_tool = create_rag_tool_async(
|
|
retriever, rewrite_llm,
|
|
num_queries=3, rerank_top_n=5
|
|
)
|
|
info("✅ RAG 检索工具初始化成功（异步版本）")
|
|
return rag_tool
|
|
except Exception as e:
|
|
warning(f"⚠️ RAG 检索工具初始化失败: {e}")
|
|
return None |