This commit is contained in:
Timothy Jaeryang Baek 2025-04-10 08:46:12 -07:00
parent 1865e29bfc
commit 63e5200e2f

View File

@ -2,11 +2,6 @@ import logging
import json import json
from typing import Optional, List 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.retrieval.web.main import SearchResult, get_filtered_results
from open_webui.env import SRC_LOG_LEVELS from open_webui.env import SRC_LOG_LEVELS
@ -22,24 +17,41 @@ def search_sougou(
count: int, count: int,
filter_list: Optional[List[str]] = None, filter_list: Optional[List[str]] = None,
) -> List[SearchResult]: ) -> 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: try:
cred = credential.Credential(sougou_api_sid, sougou_api_sk) cred = credential.Credential(sougou_api_sid, sougou_api_sk)
http_profile = HttpProfile() http_profile = HttpProfile()
http_profile.endpoint = "tms.tencentcloudapi.com" http_profile.endpoint = "tms.tencentcloudapi.com"
client_profile = ClientProfile() client_profile = ClientProfile()
client_profile.http_profile = http_profile client_profile.http_profile = http_profile
params = json.dumps({"Query": query, 'Cnt': 20}) params = json.dumps({"Query": query, "Cnt": 20})
common_client = CommonClient("tms", "2020-12-29", cred, "", profile=client_profile) common_client = CommonClient(
"tms", "2020-12-29", cred, "", profile=client_profile
)
results = [ 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: if filter_list:
sorted_results = get_filtered_results(sorted_results, filter_list) sorted_results = get_filtered_results(sorted_results, filter_list)
return [ return [
SearchResult( 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] for result in sorted_results[:count]
] ]