Timothy Jaeryang Baek 2 months ago
parent
commit
63e5200e2f
1 changed files with 22 additions and 10 deletions
  1. 22 10
      backend/open_webui/retrieval/web/sougou.py

+ 22 - 10
backend/open_webui/retrieval/web/sougou.py

@@ -2,11 +2,6 @@ import logging
 import json
 from typing import Optional, List
 
-from tencentcloud.common.common_client import CommonClient
-from tencentcloud.common import credential
-from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
-from tencentcloud.common.profile.client_profile import ClientProfile
-from tencentcloud.common.profile.http_profile import HttpProfile
 
 from open_webui.retrieval.web.main import SearchResult, get_filtered_results
 from open_webui.env import SRC_LOG_LEVELS
@@ -22,24 +17,41 @@ def search_sougou(
     count: int,
     filter_list: Optional[List[str]] = None,
 ) -> List[SearchResult]:
+    from tencentcloud.common.common_client import CommonClient
+    from tencentcloud.common import credential
+    from tencentcloud.common.exception.tencent_cloud_sdk_exception import (
+        TencentCloudSDKException,
+    )
+    from tencentcloud.common.profile.client_profile import ClientProfile
+    from tencentcloud.common.profile.http_profile import HttpProfile
+
     try:
         cred = credential.Credential(sougou_api_sid, sougou_api_sk)
         http_profile = HttpProfile()
         http_profile.endpoint = "tms.tencentcloudapi.com"
         client_profile = ClientProfile()
         client_profile.http_profile = http_profile
-        params = json.dumps({"Query": query, 'Cnt': 20})
-        common_client = CommonClient("tms", "2020-12-29", cred, "", profile=client_profile)
+        params = json.dumps({"Query": query, "Cnt": 20})
+        common_client = CommonClient(
+            "tms", "2020-12-29", cred, "", profile=client_profile
+        )
         results = [
-            json.loads(page) for page in common_client.call_json("SearchPro", json.loads(params))["Response"]["Pages"]
+            json.loads(page)
+            for page in common_client.call_json("SearchPro", json.loads(params))[
+                "Response"
+            ]["Pages"]
         ]
-        sorted_results = sorted(results, key=lambda x: x.get("scour", 0.0), reverse=True)
+        sorted_results = sorted(
+            results, key=lambda x: x.get("scour", 0.0), reverse=True
+        )
         if filter_list:
             sorted_results = get_filtered_results(sorted_results, filter_list)
 
         return [
             SearchResult(
-                link=result.get("url"), title=result.get("title"), snippet=result.get("passage")
+                link=result.get("url"),
+                title=result.get("title"),
+                snippet=result.get("passage"),
             )
             for result in sorted_results[:count]
         ]