浏览代码

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

Co-authored-by: Yeuoly <admin@srmxy.cn>
Richards Tu 1 年之前
父节点
当前提交
d7f0056e2d

+ 7 - 0
api/core/tools/provider/builtin/tavily/tavily.py

@@ -16,6 +16,13 @@ class TavilyProvider(BuiltinToolProviderController):
                 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:

+ 19 - 62
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,43 @@ 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
+        if 'exclude_domains' in params and isinstance(params['exclude_domains'], str) and params['exclude_domains'] != 'None':
+            params['exclude_domains'] = params['exclude_domains'].split()
+        else:
+            params['exclude_domains'] = []
+        if 'include_domains' in params and isinstance(params['include_domains'], str) and params['include_domains'] != 'None':
+            params['include_domains'] = params['include_domains'].split()
+        else:
+            params['include_domains'] = []
+        
         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 +105,14 @@ class TavilySearchTool(BuiltinTool):
             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)

+ 135 - 0
api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml

@@ -25,3 +25,138 @@ parameters:
       pt_BR: used for searching
     llm_description: key words for searching
     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