This commit is contained in:
273
app/rag/tools.py
273
app/rag/tools.py
@@ -1,230 +1,89 @@
|
||||
"""
|
||||
RAG 工具包装
|
||||
RAG 工具模块
|
||||
|
||||
将 RAG 流水线包装成 LangChain Tool，供 Agent 调用。
|
||||
将检索功能封装为 LangChain Tool，供 Agent 调用。
|
||||
"""
|
||||
|
||||
from typing import Optional, Dict, Any
|
||||
from langchain.tools import tool
|
||||
from langchain_core.tools import Tool
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langchain_core.language_models import BaseLanguageModel
|
||||
|
||||
from .pipeline import RAGPipeline, RAGConfig, RAGLevel
|
||||
from langchain_core.tools import tool
|
||||
from rag_core import LlamaCppEmbedder, QDRANT_URL, QDRANT_API_KEY
|
||||
from .pipeline import RAGPipeline, RAGLevel
|
||||
|
||||
|
||||
class RAGTool:
|
||||
"""
|
||||
RAG 工具包装器
|
||||
@tool
|
||||
async def search_knowledge_base(query: str, rag_level: str = "rerank") -> str:
|
||||
"""在知识库中搜索与查询相关的文档片段。
|
||||
|
||||
将 RAG 流水线包装成 Agent 可调用的工具。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pipeline: RAGPipeline,
|
||||
tool_name: str = "search_knowledge_base",
|
||||
tool_description: str = None,
|
||||
):
|
||||
"""
|
||||
初始化 RAG 工具
|
||||
|
||||
Args:
|
||||
pipeline: RAG 流水线实例
|
||||
tool_name: 工具名称
|
||||
tool_description: 工具描述
|
||||
"""
|
||||
self.pipeline = pipeline
|
||||
self.tool_name = tool_name
|
||||
|
||||
# 默认工具描述
|
||||
self.tool_description = tool_description or (
|
||||
"在知识库中搜索相关信息。当用户询问需要查阅文档、规章制度、"
|
||||
"专业知识或需要基于已知信息回答的问题时使用此工具。"
|
||||
"输入应为要搜索的查询文本。"
|
||||
)
|
||||
|
||||
# 创建 LangChain 工具
|
||||
self._tool = self._create_tool()
|
||||
|
||||
def _create_tool(self) -> Tool:
|
||||
"""创建 LangChain 工具"""
|
||||
|
||||
@tool(self.tool_name, args_schema=None)
|
||||
def search_knowledge_base(query: str) -> str:
|
||||
"""
|
||||
在知识库中搜索相关信息
|
||||
|
||||
Args:
|
||||
query: 搜索查询
|
||||
|
||||
Returns:
|
||||
格式化后的搜索结果
|
||||
"""
|
||||
try:
|
||||
# 执行检索
|
||||
result = self.pipeline.retrieve(query)
|
||||
|
||||
if not result.documents:
|
||||
return "在知识库中未找到相关信息。"
|
||||
|
||||
# 格式化上下文
|
||||
context = self.pipeline.format_context(
|
||||
result.documents,
|
||||
max_length=4000, # 限制上下文长度
|
||||
)
|
||||
|
||||
# 构建响应
|
||||
response = (
|
||||
f"🔍 在知识库中找到了 {len(result.documents)} 条相关信息：\n\n"
|
||||
f"{context}\n\n"
|
||||
f"⏱️ 检索耗时: {result.query_time:.2f}秒"
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"检索过程中发生错误: {str(e)}"
|
||||
if self.pipeline.config.verbose:
|
||||
print(f"RAG 工具错误: {error_msg}")
|
||||
return error_msg
|
||||
|
||||
# 设置工具描述
|
||||
search_knowledge_base.description = self.tool_description
|
||||
|
||||
return search_knowledge_base
|
||||
|
||||
def get_tool(self) -> Tool:
|
||||
"""获取 LangChain 工具"""
|
||||
return self._tool
|
||||
|
||||
def __call__(self, query: str) -> str:
|
||||
"""直接调用工具"""
|
||||
return self._tool.invoke({"query": query})
|
||||
|
||||
|
||||
def create_rag_tool(
|
||||
embeddings: Embeddings,
|
||||
llm: Optional[BaseLanguageModel] = None,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
tool_name: str = "search_knowledge_base",
|
||||
tool_description: Optional[str] = None,
|
||||
) -> Tool:
|
||||
"""
|
||||
创建 RAG 工具（便捷函数）
|
||||
适用于事实性问题、背景知识查询。
|
||||
|
||||
Args:
|
||||
embeddings: 嵌入模型
|
||||
llm: 语言模型（用于高级 RAG 功能）
|
||||
config: RAG 配置字典
|
||||
tool_name: 工具名称
|
||||
tool_description: 工具描述
|
||||
|
||||
query: 查询字符串
|
||||
rag_level: 检索级别，可选值：basic（基础向量检索）、rerank（基础检索+重排序）、fusion（RAG-Fusion）
|
||||
|
||||
Returns:
|
||||
LangChain Tool 实例
|
||||
检索到的相关文档内容
|
||||
"""
|
||||
# 初始化嵌入模型
|
||||
embedder = LlamaCppEmbedder()
|
||||
embeddings = embedder.as_langchain_embeddings()
|
||||
|
||||
# 创建 RAG 流水线
|
||||
pipeline = RAGPipeline.create_from_config(
|
||||
pipeline = RAGPipeline(
|
||||
embeddings=embeddings,
|
||||
llm=llm,
|
||||
config_dict=config,
|
||||
config={
|
||||
"rag_level": rag_level,
|
||||
"collection_name": "rag_documents",
|
||||
"rerank_top_n": 5,
|
||||
}
|
||||
)
|
||||
|
||||
# 创建工具包装器
|
||||
rag_tool = RAGTool(
|
||||
pipeline=pipeline,
|
||||
tool_name=tool_name,
|
||||
tool_description=tool_description,
|
||||
)
|
||||
# 执行检索
|
||||
try:
|
||||
documents = await pipeline.aretrieve(query)
|
||||
if not documents:
|
||||
return "未找到相关信息。"
|
||||
|
||||
# 格式化结果
|
||||
context = pipeline.format_context(documents)
|
||||
return context
|
||||
except Exception as e:
|
||||
return f"检索过程中发生错误: {str(e)}"
|
||||
|
||||
|
||||
@tool
|
||||
def search_knowledge_base_sync(query: str, rag_level: str = "rerank") -> str:
|
||||
"""同步版本的知识库搜索工具。
|
||||
|
||||
return rag_tool.get_tool()
|
||||
|
||||
|
||||
# 导出便捷函数
|
||||
search_knowledge_base_tool = create_rag_tool
|
||||
|
||||
|
||||
def bind_rag_to_agent(
|
||||
agent_llm: BaseLanguageModel,
|
||||
embeddings: Embeddings,
|
||||
rag_llm: Optional[BaseLanguageModel] = None,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
tool_name: str = "search_knowledge_base",
|
||||
) -> BaseLanguageModel:
|
||||
"""
|
||||
将 RAG 工具绑定到 Agent 模型
|
||||
适用于事实性问题、背景知识查询。
|
||||
|
||||
Args:
|
||||
agent_llm: Agent 使用的语言模型
|
||||
embeddings: 嵌入模型
|
||||
rag_llm: RAG 流水线使用的语言模型（如果与 agent_llm 不同）
|
||||
config: RAG 配置
|
||||
tool_name: 工具名称
|
||||
|
||||
Returns:
|
||||
绑定工具后的模型
|
||||
"""
|
||||
# 如果未指定 RAG LLM，使用 Agent LLM
|
||||
if rag_llm is None:
|
||||
rag_llm = agent_llm
|
||||
query: 查询字符串
|
||||
rag_level: 检索级别，可选值：basic（基础向量检索）、rerank（基础检索+重排序）、fusion（RAG-Fusion）
|
||||
|
||||
# 创建 RAG 工具
|
||||
rag_tool = create_rag_tool(
|
||||
Returns:
|
||||
检索到的相关文档内容
|
||||
"""
|
||||
# 初始化嵌入模型
|
||||
embedder = LlamaCppEmbedder()
|
||||
embeddings = embedder.as_langchain_embeddings()
|
||||
|
||||
# 创建 RAG 流水线
|
||||
pipeline = RAGPipeline(
|
||||
embeddings=embeddings,
|
||||
llm=rag_llm,
|
||||
config=config,
|
||||
tool_name=tool_name,
|
||||
config={
|
||||
"rag_level": rag_level,
|
||||
"collection_name": "rag_documents",
|
||||
"rerank_top_n": 5,
|
||||
}
|
||||
)
|
||||
|
||||
# 绑定工具到模型
|
||||
return agent_llm.bind_tools([rag_tool])
|
||||
|
||||
|
||||
def create_agentic_rag_pipeline(
|
||||
embeddings: Embeddings,
|
||||
agent_llm: BaseLanguageModel,
|
||||
rag_llm: Optional[BaseLanguageModel] = None,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
创建完整的 Agentic RAG 流水线（Level 4）
|
||||
|
||||
Args:
|
||||
embeddings: 嵌入模型
|
||||
agent_llm: Agent 模型
|
||||
rag_llm: RAG 专用模型
|
||||
config: 配置
|
||||
# 执行检索
|
||||
try:
|
||||
documents = pipeline.retrieve(query)
|
||||
if not documents:
|
||||
return "未找到相关信息。"
|
||||
|
||||
Returns:
|
||||
包含模型和工具的字典
|
||||
"""
|
||||
# 配置 Agentic RAG 级别
|
||||
if config is None:
|
||||
config = {}
|
||||
config["rag_level"] = RAGLevel.AGENTIC.value
|
||||
|
||||
# 创建 RAG 工具
|
||||
rag_tool = create_rag_tool(
|
||||
embeddings=embeddings,
|
||||
llm=rag_llm or agent_llm,
|
||||
config=config,
|
||||
tool_name="search_knowledge_base",
|
||||
tool_description=(
|
||||
"在知识库中搜索相关信息。当用户询问需要查阅文档、规章制度、"
|
||||
"专业知识或需要基于已知信息回答的问题时使用此工具。"
|
||||
"Agent 应该先判断是否需要使用此工具，然后调用它获取上下文。"
|
||||
),
|
||||
)
|
||||
|
||||
# 绑定工具到模型
|
||||
bound_llm = agent_llm.bind_tools([rag_tool])
|
||||
|
||||
return {
|
||||
"llm": bound_llm,
|
||||
"tool": rag_tool,
|
||||
"pipeline": RAGPipeline.create_from_config(
|
||||
embeddings=embeddings,
|
||||
llm=rag_llm or agent_llm,
|
||||
config_dict=config,
|
||||
),
|
||||
}
|
||||
# 格式化结果
|
||||
context = pipeline.format_context(documents)
|
||||
return context
|
||||
except Exception as e:
|
||||
return f"检索过程中发生错误: {str(e)}"
|
||||
|
||||
Reference in New Issue
Block a user