refac
This commit is contained in:
@@ -20,14 +20,14 @@ log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def xml_element_contents_to_string(element: Element) -> str:
|
||||
buffer = [element.text if element.text else ""]
|
||||
buffer = [element.text if element.text else '']
|
||||
|
||||
for child in element:
|
||||
buffer.append(xml_element_contents_to_string(child))
|
||||
|
||||
buffer.append(element.tail if element.tail else "")
|
||||
buffer.append(element.tail if element.tail else '')
|
||||
|
||||
return "".join(buffer)
|
||||
return ''.join(buffer)
|
||||
|
||||
|
||||
def search_yandex(
|
||||
@@ -42,42 +42,38 @@ def search_yandex(
|
||||
) -> List[SearchResult]:
|
||||
try:
|
||||
headers = {
|
||||
"User-Agent": "Open WebUI (https://github.com/open-webui/open-webui) RAG Bot",
|
||||
"Authorization": f"Api-Key {yandex_search_api_key}",
|
||||
'User-Agent': 'Open WebUI (https://github.com/open-webui/open-webui) RAG Bot',
|
||||
'Authorization': f'Api-Key {yandex_search_api_key}',
|
||||
}
|
||||
|
||||
if user is not None:
|
||||
headers = include_user_info_headers(headers, user)
|
||||
|
||||
chat_id = getattr(request.state, "chat_id", None)
|
||||
chat_id = getattr(request.state, 'chat_id', None)
|
||||
if chat_id:
|
||||
headers[FORWARD_SESSION_INFO_HEADER_CHAT_ID] = str(chat_id)
|
||||
|
||||
payload = {} if yandex_search_config == "" else json.loads(yandex_search_config)
|
||||
payload = {} if yandex_search_config == '' else json.loads(yandex_search_config)
|
||||
|
||||
if type(payload.get("query", None)) != dict:
|
||||
payload["query"] = {}
|
||||
if type(payload.get('query', None)) != dict:
|
||||
payload['query'] = {}
|
||||
|
||||
if "searchType" not in payload["query"]:
|
||||
payload["query"]["searchType"] = "SEARCH_TYPE_RU"
|
||||
if 'searchType' not in payload['query']:
|
||||
payload['query']['searchType'] = 'SEARCH_TYPE_RU'
|
||||
|
||||
payload["query"]["queryText"] = query
|
||||
payload['query']['queryText'] = query
|
||||
|
||||
if type(payload.get("groupSpec", None)) != dict:
|
||||
payload["groupSpec"] = {}
|
||||
if type(payload.get('groupSpec', None)) != dict:
|
||||
payload['groupSpec'] = {}
|
||||
|
||||
if "groupMode" not in payload["groupSpec"]:
|
||||
payload["groupSpec"]["groupMode"] = "GROUP_MODE_DEEP"
|
||||
if 'groupMode' not in payload['groupSpec']:
|
||||
payload['groupSpec']['groupMode'] = 'GROUP_MODE_DEEP'
|
||||
|
||||
payload["groupSpec"]["groupsOnPage"] = count
|
||||
payload["groupSpec"]["docsInGroup"] = 1
|
||||
payload['groupSpec']['groupsOnPage'] = count
|
||||
payload['groupSpec']['docsInGroup'] = 1
|
||||
|
||||
response = requests.post(
|
||||
(
|
||||
"https://searchapi.api.cloud.yandex.net/v2/web/search"
|
||||
if yandex_search_url == ""
|
||||
else yandex_search_url
|
||||
),
|
||||
('https://searchapi.api.cloud.yandex.net/v2/web/search' if yandex_search_url == '' else yandex_search_url),
|
||||
headers=headers,
|
||||
json=payload,
|
||||
)
|
||||
@@ -85,29 +81,21 @@ def search_yandex(
|
||||
response.raise_for_status()
|
||||
|
||||
response_body = response.json()
|
||||
if "rawData" not in response_body:
|
||||
raise Exception(f"No `rawData` in response body: {response_body}")
|
||||
if 'rawData' not in response_body:
|
||||
raise Exception(f'No `rawData` in response body: {response_body}')
|
||||
|
||||
search_result_body_bytes = base64.decodebytes(
|
||||
bytes(response_body["rawData"], "utf-8")
|
||||
)
|
||||
search_result_body_bytes = base64.decodebytes(bytes(response_body['rawData'], 'utf-8'))
|
||||
|
||||
doc_root = ET.parse(io.BytesIO(search_result_body_bytes))
|
||||
|
||||
results = []
|
||||
|
||||
for group in doc_root.findall("response/results/grouping/group"):
|
||||
for group in doc_root.findall('response/results/grouping/group'):
|
||||
results.append(
|
||||
{
|
||||
"url": xml_element_contents_to_string(group.find("doc/url")).strip(
|
||||
"\n"
|
||||
),
|
||||
"title": xml_element_contents_to_string(
|
||||
group.find("doc/title")
|
||||
).strip("\n"),
|
||||
"snippet": xml_element_contents_to_string(
|
||||
group.find("doc/passages/passage")
|
||||
),
|
||||
'url': xml_element_contents_to_string(group.find('doc/url')).strip('\n'),
|
||||
'title': xml_element_contents_to_string(group.find('doc/title')).strip('\n'),
|
||||
'snippet': xml_element_contents_to_string(group.find('doc/passages/passage')),
|
||||
}
|
||||
)
|
||||
|
||||
@@ -115,49 +103,47 @@ def search_yandex(
|
||||
|
||||
results = [
|
||||
SearchResult(
|
||||
link=result.get("url"),
|
||||
title=result.get("title"),
|
||||
snippet=result.get("snippet"),
|
||||
link=result.get('url'),
|
||||
title=result.get('title'),
|
||||
snippet=result.get('snippet'),
|
||||
)
|
||||
for result in results[:count]
|
||||
]
|
||||
|
||||
log.info(f"Yandex search results: {results}")
|
||||
log.info(f'Yandex search results: {results}')
|
||||
|
||||
return results
|
||||
except Exception as e:
|
||||
log.error(f"Error in search: {e}")
|
||||
log.error(f'Error in search: {e}')
|
||||
|
||||
return []
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if __name__ == '__main__':
|
||||
from starlette.datastructures import Headers
|
||||
from fastapi import FastAPI
|
||||
|
||||
result = search_yandex(
|
||||
Request(
|
||||
{
|
||||
"type": "http",
|
||||
"asgi.version": "3.0",
|
||||
"asgi.spec_version": "2.0",
|
||||
"method": "GET",
|
||||
"path": "/internal",
|
||||
"query_string": b"",
|
||||
"headers": Headers({}).raw,
|
||||
"client": ("127.0.0.1", 12345),
|
||||
"server": ("127.0.0.1", 80),
|
||||
"scheme": "http",
|
||||
"app": FastAPI(),
|
||||
'type': 'http',
|
||||
'asgi.version': '3.0',
|
||||
'asgi.spec_version': '2.0',
|
||||
'method': 'GET',
|
||||
'path': '/internal',
|
||||
'query_string': b'',
|
||||
'headers': Headers({}).raw,
|
||||
'client': ('127.0.0.1', 12345),
|
||||
'server': ('127.0.0.1', 80),
|
||||
'scheme': 'http',
|
||||
'app': FastAPI(),
|
||||
},
|
||||
None,
|
||||
),
|
||||
os.environ.get("YANDEX_WEB_SEARCH_URL", ""),
|
||||
os.environ.get("YANDEX_WEB_SEARCH_API_KEY", ""),
|
||||
os.environ.get(
|
||||
"YANDEX_WEB_SEARCH_CONFIG", '{"query": {"searchType": "SEARCH_TYPE_COM"}}'
|
||||
),
|
||||
"TOP movies of the past year",
|
||||
os.environ.get('YANDEX_WEB_SEARCH_URL', ''),
|
||||
os.environ.get('YANDEX_WEB_SEARCH_API_KEY', ''),
|
||||
os.environ.get('YANDEX_WEB_SEARCH_CONFIG', '{"query": {"searchType": "SEARCH_TYPE_COM"}}'),
|
||||
'TOP movies of the past year',
|
||||
3,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user