huixiangdou/service/web_search.py

# Copyright (c) OpenMMLab. All rights reserved.
"""Web search utils."""
import argparse
import json
import os

import pytoml
import requests
from bs4 import BeautifulSoup as BS
from loguru import logger
from readability import Document


class Article:

    def __init__(self, content: str = '', source=''):
        self.content = content
        self.source = source

    def __str__(self):
        return self.content

    def __len__(self):
        return len(self.content)

    def cut(self, start_index, end_index):
        self.source = self.source[start_index:end_index]


class WebSearch:
    """This class provides functionality to perform web search operations.

    Attributes:
        config_path (str): Path to the configuration file.
        retry (int): Number of times to retry a request before giving up.

    Methods:
        load_urls(): Loads URLs from config file.
        load_key(): Retrieves API key from the config file.
        load_save_dir(): Gets the directory path for saving results.
        google(query: str, max_article:int): Performs Google search for the given query and returns top max_article results.  # noqa E501
        save_search_result(query:str, articles: list): Saves the search result into a text file.  # noqa E501
        get(query: str, max_article=1): Searches with cache. If the query already exists in the cache, return the cached result.  # noqa E501
    """

    def __init__(self, config_path: dict, retry: int = 2) -> None:
        """Initializes the WebSearch object with the given config path and
        retry count."""
        self.config_path = config_path
        self.retry = retry

    def load_urls(self):
        """Reads the configuration file and fetches the ordered URLs.

        In case of an error, logs the exception and returns an empty list.
        """
        try:
            with open(self.config_path, encoding='utf8') as f:
                config = pytoml.load(f)
                return config['web_search']['domain_partial_order']
        except Exception as e:
            logger.error(str(e))
        return []

    def load_key(self):
        """Attempts to load the API key from the configuration file.

        Raises an Exception if it fails to find a valid API key.
        """
        with open(self.config_path, encoding='utf8') as f:
            config = pytoml.load(f)
            api_key = config['web_search']['x_api_key']
            # if len(api_key) > 0 and 'YOUR-API-KEY' not in api_key:
            return api_key
        raise Exception(
            'web_search X-API-KEY not found, please input your API key')

    def load_save_dir(self):
        """Attempts to read the directory where the search results are stored
        from the configuration file.

        Returns None if it fails.
        """
        try:
            with open(self.config_path, encoding='utf8') as f:
                config = pytoml.load(f)
                return config['web_search']['save_dir']
        except Exception:
            pass
        return None

    def fetch_url(self, query: str, target_link: str):
        if not target_link.startswith('http'):
            return None

        logger.info(f'extract: {target_link}')
        try:
            response = requests.get(target_link, timeout=30)

            doc = Document(response.text)
            content_html = doc.summary()
            title = doc.short_title()
            soup = BS(content_html, 'html.parser')

            if len(soup.text) < 4 * len(query):
                return None
            content = '{} {}'.format(title, soup.text)
            content = content.replace('\n\n', '\n')
            content = content.replace('\n\n', '\n')
            content = content.replace('  ', ' ')

            return Article(content=content, source=target_link)
        except Exception as e:
            logger.error('fetch_url {}'.format(str(e)))
        return None

    def google(self, query: str, max_article):
        """Executes a google search based on the provided query.

        Parses the response and extracts the relevant URLs based on the
        priority defined in the configuration file. Performs a GET request on
        these URLs and extracts the title and content of the page. The content
        is cleaned and added to the articles list. Returns a list of articles.
        """
        url = 'https://google.serper.dev/search'

        payload = json.dumps({'q': f'{query}', 'hl': 'zh-cn'})
        headers = {
            'X-API-KEY': self.load_key(),
            'Content-Type': 'application/json'
        }
        response = requests.request('POST',
                                    url,
                                    headers=headers,
                                    data=payload,
                                    timeout=15)  # noqa E501
        jsonobj = json.loads(response.text)

        keys = self.load_urls()
        urls = {}
        normal_urls = []

        for organic in jsonobj['organic']:
            link = ''
            logger.debug(organic)

            if 'link' in organic:
                link = organic['link']
            else:
                link = organic['sitelinks'][0]['link']

            for key in keys:
                if key in link:
                    if key not in urls:
                        urls[key] = [link]
                    else:
                        urls[key].append(link)
                    break
                else:
                    normal_urls.append(link)

        logger.debug(f'gather urls: {urls}')

        links = []
        for key in keys:
            if key in urls:
                links += urls[key]

        target_links = links[0:max_article]

        logger.debug(f'target_links:{target_links}')

        articles = []
        for target_link in target_links:
            # network with exponential backoff
            a = self.fetch_url(query=query, target_link=target_link)
            if a is not None:
                articles.append(a)

        return articles

    def save_search_result(self, query: str, articles: list):
        """Writes the search results (articles) for the provided query into a
        text file.

        If the directory does not exist, it creates one. In case of an error,
        logs a warning message.
        """
        try:
            save_dir = self.load_save_dir()
            if save_dir is None:
                return

            if not os.path.exists(save_dir):
                os.makedirs(save_dir)

            filepath = os.path.join(save_dir, query)

            text = ''
            if len(articles) > 0:
                texts = [str(a) for a in articles]
                text = '\n\n'.join(texts)
            with open(filepath, 'w', encoding='utf8') as f:
                f.write(text)
        except Exception as e:
            logger.warning(f'error while saving search result {str(e)}')

    def get(self, query: str, max_article=1):
        """Executes a google search with cache.

        If the query already exists in the cache, returns the cached result. If
        an exception occurs during the process, retries the request based on
        the retry count. Sleeps for a random time interval between retries.
        """
        query = query.strip()
        query = query[0:32]
        try:
            articles = self.google(query=query, max_article=max_article)
            self.save_search_result(query=query, articles=articles)
            return articles, None
        except Exception as e:
            logger.error(('web_search exception', query, str(e)))
            return [], Exception('search fail, please check TOKEN')
        return [], None


def parse_args():
    """Parses command-line arguments for web search."""
    parser = argparse.ArgumentParser(description='Web search.')
    parser.add_argument('--keywords',
                        type=str,
                        help='Keywords for search and parse.')
    parser.add_argument(
        '--config_path',
        default='config.ini',
        help='Feature store configuration path. Default value is config.ini')
    args = parser.parse_args()
    return args


def fetch_web_content(target_link: str):
    """Fetches and parses the content of the target URL.

    Extracts the main content and title from the HTML of the page. Returns the
    title and content as a single string.
    """
    response = requests.get(target_link, timeout=60)

    doc = Document(response.text)
    content_html = doc.summary()
    title = doc.short_title()
    soup = BS(content_html, 'html.parser')
    ret = '{} {}'.format(title, soup.text)
    return ret


if __name__ == '__main__':
    parser = parse_args()
    s = WebSearch(config_path=parser.config_path)

    print(s.get('mmdeploy 安装教程'))
    print(
        fetch_web_content(
            'https://mmdeploy.readthedocs.io/zh-cn/latest/get_started.html'))