-
Notifications
You must be signed in to change notification settings - Fork 8.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
FEAT: add tavily tool for searching... A search engine for LLM (#2681)
- Loading branch information
Showing
5 changed files
with
239 additions
and
0 deletions.
There are no files selected for viewing
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 changes: 161 additions & 0 deletions
161
api/core/tools/provider/builtin/tavily/tools/tavily_search.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
27 changes: 27 additions & 0 deletions
27
api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |