From 63e5200e2fd5f0072c19b831de9655159373dc65 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 10 Apr 2025 08:46:12 -0700 Subject: [PATCH] refac --- backend/open_webui/retrieval/web/sougou.py | 32 +++++++++++++++------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/backend/open_webui/retrieval/web/sougou.py b/backend/open_webui/retrieval/web/sougou.py index a0726372c..af7957c4f 100644 --- a/backend/open_webui/retrieval/web/sougou.py +++ b/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] ]