mirror of
https://github.com/langgenius/dify.git
synced 2024-11-15 19:22:36 +08:00
FEAT: add tavily tool for searching... A search engine for LLM (#2681)
This commit is contained in:
parent
69a5ce1e31
commit
7f891939f1
BIN
api/core/tools/provider/builtin/tavily/_assets/icon.png
Normal file
BIN
api/core/tools/provider/builtin/tavily/_assets/icon.png
Normal file
Binary file not shown.
After Width: | Height: | Size: 2.0 KiB |
22
api/core/tools/provider/builtin/tavily/tavily.py
Normal file
22
api/core/tools/provider/builtin/tavily/tavily.py
Normal file
|
@ -0,0 +1,22 @@
|
|||
from typing import Any
|
||||
|
||||
from core.tools.errors import ToolProviderCredentialValidationError
|
||||
from core.tools.provider.builtin.tavily.tools.tavily_search import TavilySearchTool
|
||||
from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
|
||||
|
||||
|
||||
class TavilyProvider(BuiltinToolProviderController):
|
||||
def _validate_credentials(self, credentials: dict[str, Any]) -> None:
|
||||
try:
|
||||
TavilySearchTool().fork_tool_runtime(
|
||||
meta={
|
||||
"credentials": credentials,
|
||||
}
|
||||
).invoke(
|
||||
user_id='',
|
||||
tool_parameters={
|
||||
"query": "Sachin Tendulkar",
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
raise ToolProviderCredentialValidationError(str(e))
|
29
api/core/tools/provider/builtin/tavily/tavily.yaml
Normal file
29
api/core/tools/provider/builtin/tavily/tavily.yaml
Normal file
|
@ -0,0 +1,29 @@
|
|||
identity:
|
||||
author: Yash Parmar
|
||||
name: tavily
|
||||
label:
|
||||
en_US: Tavily
|
||||
zh_Hans: Tavily
|
||||
pt_BR: Tavily
|
||||
description:
|
||||
en_US: Tavily
|
||||
zh_Hans: Tavily
|
||||
pt_BR: Tavily
|
||||
icon: icon.png
|
||||
credentials_for_provider:
|
||||
tavily_api_key:
|
||||
type: secret-input
|
||||
required: true
|
||||
label:
|
||||
en_US: Tavily API key
|
||||
zh_Hans: Tavily API key
|
||||
pt_BR: Tavily API key
|
||||
placeholder:
|
||||
en_US: Please input your Tavily API key
|
||||
zh_Hans: 请输入你的 Tavily API key
|
||||
pt_BR: Please input your Tavily API key
|
||||
help:
|
||||
en_US: Get your Tavily API key from Tavily
|
||||
zh_Hans: 从 TavilyApi 获取您的 Tavily API key
|
||||
pt_BR: Get your Tavily API key from Tavily
|
||||
url: https://docs.tavily.com/docs/tavily-api/introduction
|
161
api/core/tools/provider/builtin/tavily/tools/tavily_search.py
Normal file
161
api/core/tools/provider/builtin/tavily/tools/tavily_search.py
Normal file
|
@ -0,0 +1,161 @@
|
|||
from typing import Any, Optional
|
||||
|
||||
import requests
|
||||
|
||||
from core.tools.entities.tool_entities import ToolInvokeMessage
|
||||
from core.tools.tool.builtin_tool import BuiltinTool
|
||||
|
||||
TAVILY_API_URL = "https://api.tavily.com"
|
||||
|
||||
|
||||
class TavilySearch:
|
||||
"""
|
||||
A class for performing search operations using the Tavily Search API.
|
||||
|
||||
Args:
|
||||
api_key (str): The API key for accessing the Tavily Search API.
|
||||
|
||||
Methods:
|
||||
raw_results: Retrieves raw search results from the Tavily Search API.
|
||||
results: Retrieves cleaned search results from the Tavily Search API.
|
||||
clean_results: Cleans the raw search results.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: str) -> None:
|
||||
self.api_key = api_key
|
||||
|
||||
def raw_results(
|
||||
self,
|
||||
query: str,
|
||||
max_results: Optional[int] = 3,
|
||||
search_depth: Optional[str] = "advanced",
|
||||
include_domains: Optional[list[str]] = [],
|
||||
exclude_domains: Optional[list[str]] = [],
|
||||
include_answer: Optional[bool] = False,
|
||||
include_raw_content: Optional[bool] = False,
|
||||
include_images: Optional[bool] = False,
|
||||
) -> dict:
|
||||
"""
|
||||
Retrieves raw search results from the Tavily Search API.
|
||||
|
||||
Args:
|
||||
query (str): The search query.
|
||||
max_results (int, optional): The maximum number of results to retrieve. Defaults to 3.
|
||||
search_depth (str, optional): The search depth. Defaults to "advanced".
|
||||
include_domains (List[str], optional): The domains to include in the search. Defaults to [].
|
||||
exclude_domains (List[str], optional): The domains to exclude from the search. Defaults to [].
|
||||
include_answer (bool, optional): Whether to include answer in the search results. Defaults to False.
|
||||
include_raw_content (bool, optional): Whether to include raw content in the search results. Defaults to False.
|
||||
include_images (bool, optional): Whether to include images in the search results. Defaults to False.
|
||||
|
||||
Returns:
|
||||
dict: The raw search results.
|
||||
|
||||
"""
|
||||
params = {
|
||||
"api_key": self.api_key,
|
||||
"query": query,
|
||||
"max_results": max_results,
|
||||
"search_depth": search_depth,
|
||||
"include_domains": include_domains,
|
||||
"exclude_domains": exclude_domains,
|
||||
"include_answer": include_answer,
|
||||
"include_raw_content": include_raw_content,
|
||||
"include_images": include_images,
|
||||
}
|
||||
response = requests.post(f"{TAVILY_API_URL}/search", json=params)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
def results(
|
||||
self,
|
||||
query: str,
|
||||
max_results: Optional[int] = 3,
|
||||
search_depth: Optional[str] = "advanced",
|
||||
include_domains: Optional[list[str]] = [],
|
||||
exclude_domains: Optional[list[str]] = [],
|
||||
include_answer: Optional[bool] = False,
|
||||
include_raw_content: Optional[bool] = False,
|
||||
include_images: Optional[bool] = False,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Retrieves cleaned search results from the Tavily Search API.
|
||||
|
||||
Args:
|
||||
query (str): The search query.
|
||||
max_results (int, optional): The maximum number of results to retrieve. Defaults to 3.
|
||||
search_depth (str, optional): The search depth. Defaults to "advanced".
|
||||
include_domains (List[str], optional): The domains to include in the search. Defaults to [].
|
||||
exclude_domains (List[str], optional): The domains to exclude from the search. Defaults to [].
|
||||
include_answer (bool, optional): Whether to include answer in the search results. Defaults to False.
|
||||
include_raw_content (bool, optional): Whether to include raw content in the search results. Defaults to False.
|
||||
include_images (bool, optional): Whether to include images in the search results. Defaults to False.
|
||||
|
||||
Returns:
|
||||
list: The cleaned search results.
|
||||
|
||||
"""
|
||||
raw_search_results = self.raw_results(
|
||||
query,
|
||||
max_results=max_results,
|
||||
search_depth=search_depth,
|
||||
include_domains=include_domains,
|
||||
exclude_domains=exclude_domains,
|
||||
include_answer=include_answer,
|
||||
include_raw_content=include_raw_content,
|
||||
include_images=include_images,
|
||||
)
|
||||
return self.clean_results(raw_search_results["results"])
|
||||
|
||||
def clean_results(self, results: list[dict]) -> list[dict]:
|
||||
"""
|
||||
Cleans the raw search results.
|
||||
|
||||
Args:
|
||||
results (list): The raw search results.
|
||||
|
||||
Returns:
|
||||
list: The cleaned search results.
|
||||
|
||||
"""
|
||||
clean_results = []
|
||||
for result in results:
|
||||
clean_results.append(
|
||||
{
|
||||
"url": result["url"],
|
||||
"content": result["content"],
|
||||
}
|
||||
)
|
||||
# return clean results as a string
|
||||
return "\n".join([f"{res['url']}\n{res['content']}" for res in clean_results])
|
||||
|
||||
|
||||
class TavilySearchTool(BuiltinTool):
|
||||
"""
|
||||
A tool for searching Tavily using a given query.
|
||||
"""
|
||||
|
||||
def _invoke(
|
||||
self, user_id: str, tool_parameters: dict[str, Any]
|
||||
) -> ToolInvokeMessage | list[ToolInvokeMessage]:
|
||||
"""
|
||||
Invokes the Tavily search tool with the given user ID and tool parameters.
|
||||
|
||||
Args:
|
||||
user_id (str): The ID of the user invoking the tool.
|
||||
tool_parameters (Dict[str, Any]): The parameters for the Tavily search tool.
|
||||
|
||||
Returns:
|
||||
ToolInvokeMessage | list[ToolInvokeMessage]: The result of the Tavily search tool invocation.
|
||||
"""
|
||||
query = tool_parameters.get("query", "")
|
||||
api_key = self.runtime.credentials["tavily_api_key"]
|
||||
if not query:
|
||||
return self.create_text_message("Please input query")
|
||||
tavily_search = TavilySearch(api_key)
|
||||
results = tavily_search.results(query)
|
||||
print(results)
|
||||
if not results:
|
||||
return self.create_text_message(f"No results found for '{query}' in Tavily")
|
||||
else:
|
||||
return self.create_text_message(text=results)
|
|
@ -0,0 +1,27 @@
|
|||
identity:
|
||||
name: tavily_search
|
||||
author: Yash Parmar
|
||||
label:
|
||||
en_US: TavilySearch
|
||||
zh_Hans: TavilySearch
|
||||
pt_BR: TavilySearch
|
||||
description:
|
||||
human:
|
||||
en_US: A tool for search engine built specifically for AI agents (LLMs), delivering real-time, accurate, and factual results at speed.
|
||||
zh_Hans: 专为人工智能代理 (LLM) 构建的搜索引擎工具,可快速提供实时、准确和真实的结果。
|
||||
pt_BR: A tool for search engine built specifically for AI agents (LLMs), delivering real-time, accurate, and factual results at speed.
|
||||
llm: A tool for search engine built specifically for AI agents (LLMs), delivering real-time, accurate, and factual results at speed.
|
||||
parameters:
|
||||
- name: query
|
||||
type: string
|
||||
required: true
|
||||
label:
|
||||
en_US: Query string
|
||||
zh_Hans: 查询语句
|
||||
pt_BR: Query string
|
||||
human_description:
|
||||
en_US: used for searching
|
||||
zh_Hans: 用于搜索网页内容
|
||||
pt_BR: used for searching
|
||||
llm_description: key words for searching
|
||||
form: llm
|
Loading…
Reference in New Issue
Block a user