612bd9f7b0
- 修改查询方法以支持不同返回格式(graph/triples) - 移除不再使用的格式化方法和废弃查询方法 - 优化节点和关系的查询性能及结果处理 - 更新相关API端点以直接返回原始查询结果
153 lines
5.2 KiB
Python
153 lines
5.2 KiB
Python
import json
|
||
import asyncio
|
||
import inspect
|
||
import types
|
||
from collections.abc import Callable
|
||
from typing import Annotated, Any
|
||
|
||
from pydantic import BaseModel, Field
|
||
from langchain_core.tools import StructuredTool, tool
|
||
from langchain_tavily import TavilySearch
|
||
|
||
from src import config, graph_base, knowledge_base
|
||
from src.utils import logger
|
||
|
||
|
||
# 工具注册表 - 移到前面以避免NameError
|
||
_TOOLS_REGISTRY = {}
|
||
|
||
|
||
class KnowledgeRetrieverModel(BaseModel):
|
||
query_text: str = Field(
|
||
description=(
|
||
"查询的关键词，查询的时候，应该尽量以可能帮助回答这个问题的关键词进行查询，"
|
||
"不要直接使用用户的原始输入去查询。"
|
||
)
|
||
)
|
||
|
||
def get_runnable_tools():
|
||
"""获取所有可运行的工具（给大模型使用）"""
|
||
tools = _TOOLS_REGISTRY.copy()
|
||
|
||
# 获取所有知识库
|
||
for db_Id, retrieve_info in knowledge_base.get_retrievers().items():
|
||
_id = f"retrieve_{db_Id[:8]}" # Deepseek does not support non-alphanumeric characters in tool names
|
||
description = (
|
||
f"使用 {retrieve_info['name']} 知识库进行检索。\n"
|
||
f"下面是这个知识库的描述：\n{retrieve_info['description']}"
|
||
)
|
||
|
||
# 创建异步工具，确保正确处理异步检索器
|
||
async def async_retriever_wrapper(query_text: str, db_id=db_Id, retriever_info=retrieve_info):
|
||
"""异步检索器包装函数"""
|
||
retriever = retriever_info["retriever"]
|
||
try:
|
||
if asyncio.iscoroutinefunction(retriever):
|
||
result = await retriever(query_text)
|
||
else:
|
||
result = retriever(query_text)
|
||
return result
|
||
except Exception as e:
|
||
logger.error(f"Error in retriever {db_id}: {e}")
|
||
return f"检索失败: {str(e)}"
|
||
|
||
# 使用 StructuredTool.from_function 创建异步工具
|
||
tools[_id] = StructuredTool.from_function(
|
||
coroutine=async_retriever_wrapper, # 指定为协程
|
||
name=_id,
|
||
description=description,
|
||
args_schema=KnowledgeRetrieverModel,
|
||
metadata=retrieve_info
|
||
)
|
||
|
||
return tools
|
||
|
||
def get_all_tools_info():
|
||
"""获取所有工具的信息（用于前端展示）"""
|
||
tools_info = {}
|
||
|
||
tools = get_runnable_tools()
|
||
|
||
# 获取注册的工具信息
|
||
for _id, tool_obj in tools.items():
|
||
|
||
metadata = getattr(tool_obj, 'metadata', {}) or {}
|
||
info = {
|
||
"id": _id,
|
||
"name": metadata.get('name', _id),
|
||
"description": metadata.get('description') or getattr(tool_obj, 'description', ''),
|
||
'metadata': metadata,
|
||
"args": []
|
||
}
|
||
|
||
# 获取工具参数信息
|
||
if hasattr(tool_obj, 'args_schema') and tool_obj.args_schema:
|
||
schema = tool_obj.args_schema.schema()
|
||
if 'properties' in schema:
|
||
for arg_name, arg_info in schema['properties'].items():
|
||
info["args"].append({
|
||
"name": arg_name,
|
||
"type": arg_info.get('type', ''),
|
||
"description": arg_info.get('description', '')
|
||
})
|
||
tools_info[info['id']] = info
|
||
|
||
return tools_info
|
||
|
||
class BaseToolOutput:
|
||
"""
|
||
LLM 要求 Tool 的输出为 str，但 Tool 用在别处时希望它正常返回结构化数据。
|
||
只需要将 Tool 返回值用该类封装，能同时满足两者的需要。
|
||
基类简单的将返回值字符串化，或指定 format="json" 将其转为 json。
|
||
用户也可以继承该类定义自己的转换方法。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
data: Any,
|
||
format: str | Callable | None = None,
|
||
data_alias: str = "",
|
||
**extras: Any,
|
||
) -> None:
|
||
self.data = data
|
||
self.format = format
|
||
self.extras = extras
|
||
if data_alias:
|
||
setattr(self, data_alias, property(lambda obj: obj.data))
|
||
|
||
def __str__(self) -> str:
|
||
if self.format == "json":
|
||
return json.dumps(self.data, ensure_ascii=False, indent=2)
|
||
elif callable(self.format):
|
||
return self.format(self)
|
||
else:
|
||
return str(self.data)
|
||
|
||
@tool
|
||
def calculator(a: float, b: float, operation: str) -> float:
|
||
"""Calculate two numbers. operation: add, subtract, multiply, divide"""
|
||
if operation == "add":
|
||
return a + b
|
||
elif operation == "subtract":
|
||
return a - b
|
||
elif operation == "multiply":
|
||
return a * b
|
||
elif operation == "divide":
|
||
return a / b
|
||
else:
|
||
raise ValueError(f"Invalid operation: {operation}, only support add, subtract, multiply, divide")
|
||
|
||
@tool
|
||
def query_knowledge_graph(query: Annotated[str, "The keyword to query knowledge graph."]):
|
||
"""Use this to query knowledge graph, which include some food domain knowledge."""
|
||
return graph_base.query_node(query, hops=2, return_format='triples')
|
||
|
||
# 更新工具注册表
|
||
_TOOLS_REGISTRY.update({
|
||
"Calculator": calculator,
|
||
"QueryKnowledgeGraph": query_knowledge_graph,
|
||
})
|
||
|
||
if config.enable_web_search:
|
||
_TOOLS_REGISTRY["WebSearchWithTavily"] = TavilySearch(max_results=10)
|