54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
from typing import Any, Dict
|
||
|
||
from tavily import TavilyClient
|
||
|
||
from ai.tool import Tool
|
||
|
||
|
||
class TavilySearchTool(Tool):
|
||
def __init__(self, tavily_token: str):
|
||
self._client = TavilyClient(api_key=tavily_token)
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "tavily_search"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return "Поиск информации в интернете"
|
||
|
||
@property
|
||
def parameters(self) -> Dict[str, Any]:
|
||
return {
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {
|
||
"type": "string",
|
||
"description": "Запрос для поиска (на русском или английском языке)"
|
||
}
|
||
},
|
||
"required": ["query"]
|
||
}
|
||
|
||
async def execute(self, args: Dict[str, Any], _artifacts: Dict[str, Any]) -> str:
|
||
query = args.get("query", "")
|
||
print(f"Веб-поиск: {query}")
|
||
|
||
try:
|
||
results = self._client.search(query=query, max_results=5)
|
||
|
||
if not results or "results" not in results:
|
||
return "Не удалось получить результаты поиска."
|
||
|
||
answer_parts = []
|
||
for i, result in enumerate(results["results"], 1):
|
||
title = result.get("title", "Без названия")
|
||
url = result.get("url", "")
|
||
content = result.get("content", "")
|
||
answer_parts.append(f"{i}. {title}\n {url}\n {content}\n")
|
||
|
||
answer = "\n".join(answer_parts)
|
||
return f"По запросу \"{query}\" найдено:\n\n{answer}"
|
||
except Exception as e:
|
||
print(f"Ошибка веб-поиска: {e}")
|
||
return f"Не удалось выполнить веб-поиск: {e}"
|