87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
import traceback
|
||
from typing import Annotated, Any
|
||
|
||
from src import config, graph_base
|
||
from src.agents.common.toolkits import tool
|
||
from src.agents.common.toolkits.registry import ToolExtraMetadata, _all_tool_instances, _extra_registry
|
||
from src.utils import logger
|
||
|
||
# Lazy initialization for TavilySearch (only when API key is available)
|
||
_tavily_search_instance = None
|
||
|
||
|
||
def _create_tavily_search():
|
||
"""Create and register TavilySearch tool with metadata."""
|
||
global _tavily_search_instance
|
||
if _tavily_search_instance is None:
|
||
from langchain_tavily import TavilySearch
|
||
|
||
_tavily_search_instance = TavilySearch()
|
||
|
||
return _tavily_search_instance
|
||
|
||
|
||
# 注册 TavilySearch 工具（延迟初始化）
|
||
def _register_tavily_tool():
|
||
"""Register TavilySearch tool with extra metadata."""
|
||
tavily_instance = _create_tavily_search()
|
||
# 手动注册到全局注册表
|
||
_extra_registry["tavily_search"] = ToolExtraMetadata(
|
||
category="buildin",
|
||
tags=["搜索"],
|
||
display_name="Tavily 网页搜索",
|
||
)
|
||
# 添加到工具实例列表
|
||
_all_tool_instances.append(tavily_instance)
|
||
|
||
|
||
# 模块加载时注册
|
||
if config.enable_web_search:
|
||
try:
|
||
_register_tavily_tool()
|
||
except Exception as e:
|
||
logger.warning(f"Failed to register TavilySearch tool: {e}")
|
||
|
||
|
||
@tool(category="buildin", tags=["计算"], display_name="计算器")
|
||
def calculator(a: float, b: float, operation: str) -> float:
|
||
"""计算器：对给定的2个数字进行基本数学运算"""
|
||
try:
|
||
if operation == "add":
|
||
return a + b
|
||
elif operation == "subtract":
|
||
return a - b
|
||
elif operation == "multiply":
|
||
return a * b
|
||
elif operation == "divide":
|
||
if b == 0:
|
||
raise ZeroDivisionError("除数不能为零")
|
||
return a / b
|
||
else:
|
||
raise ValueError(f"不支持的运算类型: {operation}，仅支持 add, subtract, multiply, divide")
|
||
except Exception as e:
|
||
logger.error(f"Calculator error: {e}")
|
||
raise
|
||
|
||
|
||
KG_QUERY_DESCRIPTION = """
|
||
使用这个工具可以查询知识图谱中包含的三元组信息。
|
||
关键词（query），使用可能帮助回答这个问题的关键词进行查询，不要直接使用用户的原始输入去查询。
|
||
"""
|
||
|
||
|
||
@tool(category="buildin", tags=["图谱"], display_name="查询知识图谱", description=KG_QUERY_DESCRIPTION)
|
||
def query_knowledge_graph(query: Annotated[str, "The keyword to query knowledge graph."]) -> Any:
|
||
"""使用这个工具可以查询知识图谱中包含的三元组信息。关键词（query），使用可能帮助回答这个问题的关键词进行查询，不要直接使用用户的原始输入去查询。"""
|
||
try:
|
||
logger.debug(f"Querying knowledge graph with: {query}")
|
||
result = graph_base.query_node(query, hops=2, return_format="triples")
|
||
logger.debug(
|
||
f"Knowledge graph query returned "
|
||
f"{len(result.get('triples', [])) if isinstance(result, dict) else 'N/A'} triples"
|
||
)
|
||
return result
|
||
except Exception as e:
|
||
logger.error(f"Knowledge graph query error: {e}, {traceback.format_exc()}")
|
||
return f"知识图谱查询失败: {str(e)}"
|