75 lines
2.9 KiB
Python
75 lines
2.9 KiB
Python
|
"""基于TavilySearch的搜索工具"""
|
|||
|
|
|
|||
|
|
from th_agenter.core.config import get_settings
|
|||
|
|
from loguru import logger
|
|||
|
|
|
|||
|
|
from langchain.tools import BaseTool
|
|||
|
|
from langchain_community.tools.tavily_search import TavilySearchResults
|
|||
|
|
from pydantic import BaseModel, Field, PrivateAttr
|
|||
|
|
from typing import Optional, Type, ClassVar
|
|||
|
|
from langchain_tavily import TavilySearch
|
|||
|
|
|
|||
|
|
# 定义输入参数模型（替代原get_parameters()）
|
|||
|
|
class SearchInput(BaseModel):
|
|||
|
|
query: str = Field(description="搜索查询内容")
|
|||
|
|
max_results: Optional[int] = Field(
|
|||
|
|
default=5,
|
|||
|
|
description="返回结果的最大数量（默认：5）"
|
|||
|
|
)
|
|||
|
|
topic: Optional[str] = Field(
|
|||
|
|
default="general",
|
|||
|
|
description="搜索主题，可选值：general, academic, news, places"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TavilySearchTool(BaseTool):
|
|||
|
|
name:ClassVar[str] = "tavily_search_tool"
|
|||
|
|
description:ClassVar[str] = """使用Tavily搜索引擎进行网络搜索，可以获取最新信息。
|
|||
|
|
输入应该包含搜索查询(query)，可选参数包括max_results和topic。""" # 替代get_description()
|
|||
|
|
args_schema: Type[BaseModel] = SearchInput # 用Pydantic模型定义参数
|
|||
|
|
_tavily_api_key: str = PrivateAttr()
|
|||
|
|
_search_client: TavilySearchResults = PrivateAttr()
|
|||
|
|
def __init__(self, **kwargs):
|
|||
|
|
super().__init__(**kwargs)
|
|||
|
|
self._tavily_api_key = get_settings().tool.tavily_api_key
|
|||
|
|
if not self._tavily_api_key:
|
|||
|
|
raise ValueError("Tavily API key not found in settings")
|
|||
|
|
|
|||
|
|
# 初始化Tavily客户端
|
|||
|
|
self._search_client = TavilySearch(
|
|||
|
|
tavily_api_key=self._tavily_api_key
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def _run(self, query: str, max_results: int = 5, topic: str = "general"):
|
|||
|
|
try:
|
|||
|
|
logger.info(f"执行搜索：{query}")
|
|||
|
|
# 调用Tavily（LangChain已内置Tavily工具，这里直接使用）
|
|||
|
|
results = self._search_client.run({
|
|||
|
|
"query": query,
|
|||
|
|
"max_results": max_results,
|
|||
|
|
"topic": topic
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 格式化结果（根据Tavily的实际返回结构调整）
|
|||
|
|
if isinstance(results, list):
|
|||
|
|
return {
|
|||
|
|
"status": "success",
|
|||
|
|
"results": [
|
|||
|
|
{
|
|||
|
|
"title": r.get("title", ""),
|
|||
|
|
"url": r.get("url", ""),
|
|||
|
|
"content": r.get("content", "")[:200] + "..."
|
|||
|
|
} for r in results
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
else:
|
|||
|
|
return {"status": "error", "message": "Unexpected result format"}
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"搜索失败: {str(e)}")
|
|||
|
|
return {"status": "error", "message": str(e)}
|
|||
|
|
|
|||
|
|
async def _arun(self, **kwargs):
|
|||
|
|
"""异步版本"""
|
|||
|
|
"""直接调用同步版本"""
|
|||
|
|
return self._run(**kwargs) # 直接委托给同步方法
|