Fix error in [Update yaml and py file in Tavily Tool] (#3465)

Co-authored-by: Yeuoly <admin@srmxy.cn>
This commit is contained in:
Richards Tu 2024-04-19 16:51:51 +08:00 committed by GitHub
parent 9b7b133cbc
commit d7f0056e2d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 161 additions and 62 deletions

View File

@ -16,6 +16,13 @@ class TavilyProvider(BuiltinToolProviderController):
user_id='', user_id='',
tool_parameters={ tool_parameters={
"query": "Sachin Tendulkar", "query": "Sachin Tendulkar",
"search_depth": "basic",
"include_answer": True,
"include_images": False,
"include_raw_content": False,
"max_results": 5,
"include_domains": "",
"exclude_domains": ""
}, },
) )
except Exception as e: except Exception as e:

View File

@ -1,4 +1,4 @@
from typing import Any, Optional from typing import Any
import requests import requests
@ -24,87 +24,43 @@ class TavilySearch:
def __init__(self, api_key: str) -> None: def __init__(self, api_key: str) -> None:
self.api_key = api_key self.api_key = api_key
def raw_results( def raw_results(self, params: dict[str, Any]) -> dict:
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. Retrieves raw search results from the Tavily Search API.
Args: Args:
query (str): The search query. params (Dict[str, Any]): The search parameters.
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: Returns:
dict: The raw search results. dict: The raw search results.
""" """
params = { params["api_key"] = self.api_key
"api_key": self.api_key, if 'exclude_domains' in params and isinstance(params['exclude_domains'], str) and params['exclude_domains'] != 'None':
"query": query, params['exclude_domains'] = params['exclude_domains'].split()
"max_results": max_results, else:
"search_depth": search_depth, params['exclude_domains'] = []
"include_domains": include_domains, if 'include_domains' in params and isinstance(params['include_domains'], str) and params['include_domains'] != 'None':
"exclude_domains": exclude_domains, params['include_domains'] = params['include_domains'].split()
"include_answer": include_answer, else:
"include_raw_content": include_raw_content, params['include_domains'] = []
"include_images": include_images,
}
response = requests.post(f"{TAVILY_API_URL}/search", json=params) response = requests.post(f"{TAVILY_API_URL}/search", json=params)
response.raise_for_status() response.raise_for_status()
return response.json() return response.json()
def results( def results(self, params: dict[str, Any]) -> list[dict]:
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. Retrieves cleaned search results from the Tavily Search API.
Args: Args:
query (str): The search query. params (Dict[str, Any]): The search parameters.
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: Returns:
list: The cleaned search results. list: The cleaned search results.
""" """
raw_search_results = self.raw_results( raw_search_results = self.raw_results(params)
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"]) return self.clean_results(raw_search_results["results"])
def clean_results(self, results: list[dict]) -> list[dict]: def clean_results(self, results: list[dict]) -> list[dict]:
@ -149,11 +105,12 @@ class TavilySearchTool(BuiltinTool):
ToolInvokeMessage | list[ToolInvokeMessage]: The result of the Tavily search tool invocation. ToolInvokeMessage | list[ToolInvokeMessage]: The result of the Tavily search tool invocation.
""" """
query = tool_parameters.get("query", "") query = tool_parameters.get("query", "")
api_key = self.runtime.credentials["tavily_api_key"] api_key = self.runtime.credentials["tavily_api_key"]
if not query: if not query:
return self.create_text_message("Please input query") return self.create_text_message("Please input query")
tavily_search = TavilySearch(api_key) tavily_search = TavilySearch(api_key)
results = tavily_search.results(query) results = tavily_search.results(tool_parameters)
print(results) print(results)
if not results: if not results:
return self.create_text_message(f"No results found for '{query}' in Tavily") return self.create_text_message(f"No results found for '{query}' in Tavily")

View File

@ -25,3 +25,138 @@ parameters:
pt_BR: used for searching pt_BR: used for searching
llm_description: key words for searching llm_description: key words for searching
form: llm form: llm
- name: search_depth
type: select
required: false
label:
en_US: Search Depth
zh_Hans: 搜索深度
pt_BR: Search Depth
human_description:
en_US: The depth of search results
zh_Hans: 搜索结果的深度
pt_BR: The depth of search results
form: form
options:
- value: basic
label:
en_US: Basic
zh_Hans: 基本
pt_BR: Basic
- value: advanced
label:
en_US: Advanced
zh_Hans: 高级
pt_BR: Advanced
default: basic
- name: include_images
type: boolean
required: false
label:
en_US: Include Images
zh_Hans: 包含图片
pt_BR: Include Images
human_description:
en_US: Include images in the search results
zh_Hans: 在搜索结果中包含图片
pt_BR: Include images in the search results
form: form
options:
- value: true
label:
en_US: Yes
zh_Hans:
pt_BR: Yes
- value: false
label:
en_US: No
zh_Hans:
pt_BR: No
default: false
- name: include_answer
type: boolean
required: false
label:
en_US: Include Answer
zh_Hans: 包含答案
pt_BR: Include Answer
human_description:
en_US: Include answers in the search results
zh_Hans: 在搜索结果中包含答案
pt_BR: Include answers in the search results
form: form
options:
- value: true
label:
en_US: Yes
zh_Hans:
pt_BR: Yes
- value: false
label:
en_US: No
zh_Hans:
pt_BR: No
default: false
- name: include_raw_content
type: boolean
required: false
label:
en_US: Include Raw Content
zh_Hans: 包含原始内容
pt_BR: Include Raw Content
human_description:
en_US: Include raw content in the search results
zh_Hans: 在搜索结果中包含原始内容
pt_BR: Include raw content in the search results
form: form
options:
- value: true
label:
en_US: Yes
zh_Hans:
pt_BR: Yes
- value: false
label:
en_US: No
zh_Hans:
pt_BR: No
default: false
- name: max_results
type: number
required: false
label:
en_US: Max Results
zh_Hans: 最大结果
pt_BR: Max Results
human_description:
en_US: The number of maximum search results to return
zh_Hans: 返回的最大搜索结果数
pt_BR: The number of maximum search results to return
form: form
min: 1
max: 20
default: 5
- name: include_domains
type: string
required: false
label:
en_US: Include Domains
zh_Hans: 包含域
pt_BR: Include Domains
human_description:
en_US: A list of domains to specifically include in the search results
zh_Hans: 在搜索结果中特别包含的域名列表
pt_BR: A list of domains to specifically include in the search results
form: form
- name: exclude_domains
type: string
required: false
label:
en_US: Exclude Domains
zh_Hans: 排除域
pt_BR: Exclude Domains
human_description:
en_US: A list of domains to specifically exclude from the search results
zh_Hans: 从搜索结果中特别排除的域名列表
pt_BR: A list of domains to specifically exclude from the search results
form: form