Pārlūkot izejas kodu

Update search_toolkit: add addit_search

StrayWarrior 1 nedēļu atpakaļ
vecāks
revīzija
041b40641b
1 mainītis faili ar 28 papildinājumiem un 0 dzēšanām
  1. 28 0
      pqai_agent/toolkit/search_toolkit.py

+ 28 - 0
pqai_agent/toolkit/search_toolkit.py

@@ -188,6 +188,33 @@ class SearchToolkit(BaseToolkit):
         except Exception as e:
             return {"error": f"Bing scraping error: {e!s}"}
 
+    def aiddit_search(self, keyword: str) -> Dict[str, Any]:
+        r"""Search using Aiddit API.
+
+        Args:
+            keyword (str): The search keyword.
+        Returns:
+            Dict[str, Any]: A dictionary containing search results.
+        """
+        url = "http://smcp-api.aiddit.com/mcp/custom/search"
+        headers = {
+            "Content-Type": "application/json",
+        }
+        data = {
+            "keyword": keyword
+        }
+        try:
+            response = requests.post(url, headers=headers, json=data)
+            response.raise_for_status()
+            resp_json = response.json()
+            if resp_json.get('code') != 0:
+                return {"error": f"Aiddit search error: {resp_json.get('message', 'Unknown error')}"}
+            resp_data = resp_json['data']
+            results = resp_data.get('results', [])[:5]
+            return {'results': results}  # Limit to 5 results
+        except requests.RequestException as e:
+            return {"error": f"Aiddit search error: {e!s}"}
+
     def get_tools(self) -> List[FunctionTool]:
         r"""Returns a list of FunctionTool objects representing the
         functions in the toolkit.
@@ -199,4 +226,5 @@ class SearchToolkit(BaseToolkit):
         return [
             FunctionTool(self.search_baidu),
             FunctionTool(self.search_bing),
+            FunctionTool(self.aiddit_search),
         ]