open-webui/backend/open_webui/retrieval/web/yandex.py

import base64
import io
import json
import logging
import os
from typing import List, Optional
from xml.etree import ElementTree as ET
from xml.etree.ElementTree import Element

import requests
from fastapi import Request
from open_webui.env import FORWARD_SESSION_INFO_HEADER_CHAT_ID
from open_webui.retrieval.web.main import SearchResult, get_filtered_results
from open_webui.utils.headers import include_user_info_headers

log = logging.getLogger(__name__)


def xml_element_contents_to_string(element: Element) -> str:
    if element is None:
        return ''
    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 '')

    return ''.join(buffer)


def search_yandex(
    request: Request,
    yandex_search_url: str,
    yandex_search_api_key: str,
    yandex_search_config: str,
    query: str,
    count: int,
    filter_list: Optional[List[str]] = None,
    user=None,
) -> List[SearchResult]:
    try:
        headers = {
            # LICENSE covers this Open WebUI user-agent identifier.
            # Do not alter, remove, obscure, or replace it except as LICENSE permits:
            # https://docs.openwebui.com/license.
            '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)
        if chat_id:
            headers[FORWARD_SESSION_INFO_HEADER_CHAT_ID] = str(chat_id)

        payload = {} if yandex_search_config == '' else json.loads(yandex_search_config)

        if type(payload.get('query', None)) != dict:
            payload['query'] = {}

        if 'searchType' not in payload['query']:
            payload['query']['searchType'] = 'SEARCH_TYPE_RU'

        payload['query']['queryText'] = query

        if type(payload.get('groupSpec', None)) != dict:
            payload['groupSpec'] = {}

        if 'groupMode' not in payload['groupSpec']:
            payload['groupSpec']['groupMode'] = 'GROUP_MODE_DEEP'

        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),
            headers=headers,
            json=payload,
        )

        response.raise_for_status()

        response_body = response.json()
        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'))

        doc_root = ET.parse(io.BytesIO(search_result_body_bytes))

        results = []

        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')),
                }
            )

        results = get_filtered_results(results, filter_list)

        results = [
            SearchResult(
                link=result.get('url'),
                title=result.get('title'),
                snippet=result.get('snippet'),
            )
            for result in results[:count]
        ]

        log.info(f'Yandex search results: {results}')

        return results
    except Exception as e:
        log.error(f'Error in search: {e}')

        return []


if __name__ == '__main__':
    from fastapi import FastAPI
    from starlette.datastructures import Headers

    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(),
            },
            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',
        3,
    )

    print(result)