diff --git a/api/core/tools/provider/builtin/tavily/tavily.py b/api/core/tools/provider/builtin/tavily/tavily.py index a013d41fcf8691..3a4a0854562f6a 100644 --- a/api/core/tools/provider/builtin/tavily/tavily.py +++ b/api/core/tools/provider/builtin/tavily/tavily.py @@ -16,6 +16,13 @@ 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: diff --git a/api/core/tools/provider/builtin/tavily/tools/tavily_search.py b/api/core/tools/provider/builtin/tavily/tools/tavily_search.py index 9a4d27376b6d7a..a4a9003f17c4c9 100644 --- a/api/core/tools/provider/builtin/tavily/tools/tavily_search.py +++ b/api/core/tools/provider/builtin/tavily/tools/tavily_search.py @@ -1,4 +1,4 @@ -from typing import Any, Optional +from typing import Any import requests @@ -24,87 +24,34 @@ class TavilySearch: 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: + def raw_results(self, params: dict[str, Any]) -> 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. + params (Dict[str, Any]): The search parameters. 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, - } + params["api_key"] = self.api_key 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]: + def results(self, params: dict[str, Any]) -> 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. + params (Dict[str, Any]): The search parameters. 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, - ) + raw_search_results = self.raw_results(params) return self.clean_results(raw_search_results["results"]) def clean_results(self, results: list[dict]) -> list[dict]: @@ -149,13 +96,14 @@ 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(query) + results = tavily_search.results(tool_parameters) 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) \ No newline at end of file diff --git a/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml b/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml index ccdb9408fcbc89..aed1c93eabb260 100644 --- a/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml +++ b/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml @@ -25,3 +25,119 @@ 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