71 lines
2.0 KiB
Python
71 lines
2.0 KiB
Python
|
"""
|
|||
|
|
RAG 检索工具
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from langchain_core.tools import tool
|
|||
|
|
from backend.app.logger import info
|
|||
|
|
|
|||
|
|
|
|||
|
|
_rag_pipeline = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _get_rag_pipeline():
|
|||
|
|
"""获取或创建 RAG pipeline 单例"""
|
|||
|
|
global _rag_pipeline
|
|||
|
|
if _rag_pipeline is None:
|
|||
|
|
from backend.app.rag.pipeline import RAGPipeline
|
|||
|
|
_rag_pipeline = RAGPipeline(
|
|||
|
|
num_queries=3,
|
|||
|
|
rerank_top_n=5,
|
|||
|
|
use_rerank=True,
|
|||
|
|
return_parent_docs=True,
|
|||
|
|
)
|
|||
|
|
return _rag_pipeline
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _format_confidence(result) -> str:
|
|||
|
|
"""格式化置信度描述"""
|
|||
|
|
if result.confidence < 0.4:
|
|||
|
|
return "低"
|
|||
|
|
elif result.confidence < 0.6:
|
|||
|
|
return "中"
|
|||
|
|
return "高"
|
|||
|
|
|
|||
|
|
|
|||
|
|
@tool
|
|||
|
|
async def rag_search(query: str) -> str:
|
|||
|
|
"""
|
|||
|
|
检索知识库获取相关信息
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
包含检索结果和置信度的结构化回复
|
|||
|
|
"""
|
|||
|
|
info(f"[Tool] rag_search: {query[:30]}...")
|
|||
|
|
try:
|
|||
|
|
pipeline = _get_rag_pipeline()
|
|||
|
|
result = await pipeline.aretrieve_with_confidence(query, original_query=query)
|
|||
|
|
|
|||
|
|
if not result.content:
|
|||
|
|
return "【RAG检索结果】\n未在知识库中找到相关内容。\n置信度：0.0\n建议：可尝试联网搜索获取信息。"
|
|||
|
|
|
|||
|
|
confidence_desc = _format_confidence(result)
|
|||
|
|
is_useful_note = "✅ 检索结果可信，可直接使用" if result.is_useful else "⚠️ 检索结果置信度较低，可能需要联网搜索补充"
|
|||
|
|
|
|||
|
|
response = f"""【RAG检索结果】
|
|||
|
|
{result.content}
|
|||
|
|
|
|||
|
|
【置信度评估】
|
|||
|
|
- 综合置信度：{result.confidence:.2f}（{confidence_desc}）
|
|||
|
|
- 向量相似度：{result.scores['embedding']:.2f}
|
|||
|
|
- 重排分数：{result.scores['rerank']:.2f}
|
|||
|
|
- LLM评估：{result.scores['llm']:.2f}
|
|||
|
|
|
|||
|
|
{is_useful_note}"""
|
|||
|
|
|
|||
|
|
info(f"[Tool] rag_search 完成: confidence={result.confidence:.3f}, is_useful={result.is_useful}")
|
|||
|
|
return response
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
info(f"[Tool] rag_search 失败: {e}")
|
|||
|
|
return f"【RAG检索失败】\n错误：{str(e)}\n建议：请稍后重试或使用联网搜索"
|