From 589d690699e20dff8db8145482fd1073b5d0a188 Mon Sep 17 00:00:00 2001 From: crazywoola <100913391+crazywoola@users.noreply.github.com> Date: Sun, 14 Apr 2024 10:13:56 +0800 Subject: [PATCH] Revert "Update yaml and py file in Tavily Tool" (#3464) --- .../tools/provider/builtin/tavily/tavily.py | 7 -- .../builtin/tavily/tools/tavily_search.py | 72 +++++++++-- .../builtin/tavily/tools/tavily_search.yaml | 116 ------------------ 3 files changed, 62 insertions(+), 133 deletions(-) diff --git a/api/core/tools/provider/builtin/tavily/tavily.py b/api/core/tools/provider/builtin/tavily/tavily.py index 3a4a0854562f6..a013d41fcf869 100644 --- a/api/core/tools/provider/builtin/tavily/tavily.py +++ b/api/core/tools/provider/builtin/tavily/tavily.py @@ -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: 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 a4a9003f17c4c..9a4d27376b6d7 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 +from typing import Any, Optional import requests @@ -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]: @@ -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) \ No newline at end of file + return self.create_text_message(text=results) 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 aed1c93eabb26..ccdb9408fcbc8 100644 --- a/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml +++ b/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml @@ -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