Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Revert "Update yaml and py file in Tavily Tool" #3464

Merged
merged 1 commit into from
Apr 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 0 additions & 7 deletions api/core/tools/provider/builtin/tavily/tavily.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,6 @@ def _validate_credentials(self, credentials: dict[str, Any]) -> None:
user_id='',
tool_parameters={
"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:
Expand Down
72 changes: 62 additions & 10 deletions api/core/tools/provider/builtin/tavily/tools/tavily_search.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any
from typing import Any, Optional

import requests

Expand All @@ -24,34 +24,87 @@ class TavilySearch:
def __init__(self, api_key: str) -> None:
self.api_key = api_key

def raw_results(self, params: dict[str, Any]) -> dict:
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:
params (Dict[str, Any]): The search parameters.
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
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, params: dict[str, Any]) -> list[dict]:
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:
params (Dict[str, Any]): The search parameters.
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(params)
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]:
Expand Down Expand Up @@ -96,14 +149,13 @@ def _invoke(
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(tool_parameters)
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)
return self.create_text_message(text=results)
116 changes: 0 additions & 116 deletions api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -25,119 +25,3 @@ parameters:
pt_BR: used for searching
llm_description: key words for searching
form: llm
- name: search_depth
type: string
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: 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: 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: 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: integer
required: false
label:
en_US: Max Results
zh_Hans: 最大结果
pt_BR: Max Results
human_description: The number of maximum search results to return
form: form
min: 1
max: 20
default: 5
- name: include_domains
type: array
required: false
label:
en_US: Include Domains
zh_Hans: 包含域
pt_BR: Include Domains
human_description: A list of domains to specifically include in the search results
form: form
- name: exclude_domains
type: array
required: false
label:
en_US: Exclude Domains
zh_Hans: 排除域
pt_BR: Exclude Domains
human_description: A list of domains to specifically exclude from the search results
Loading