jarvis-models/src/blackbox/websearch.py

import datetime
from typing import Any, Coroutine

from fastapi import Request, Response, status
from fastapi.responses import JSONResponse
from openai import OpenAI
from .blackbox import Blackbox

import logging
from ..log.logging_time import logging_time
import requests
import json

logger = logging.getLogger
DEFAULT_COLLECTION_ID = "123"

from injector import singleton
@singleton
class WebSearch(Blackbox):

    def __call__(self, *args, **kwargs):
        return self.processing(*args, **kwargs)

    def valid(self, *args, **kwargs) -> bool:
        data = args[0]
        return isinstance(data, list)

    # @logging_time(logger=logger)
    def processing(self, question: str, settings: dict) -> str:

        if settings is None:
            settings = {}

        # from googlesearch import search

        # question = "要搜索的关键词"
        # for url in search(question, num_results=10):
        #     print(url)

        url = "https://google.serper.dev/search"

        payload = json.dumps({
            "q": question,
            "location": "China", # 限制所在位置为中国
            "gl": "cn", # 限制国家为中国
            "hl": "zh-cn", # 限制搜索结果为中文
            "tbs": "qdr:y" # 限制搜索结果为一年内的
        })

        headers = {
        'X-API-KEY': '00c0f5144e44721bd0cfed219e2b3256bb3dd5fc',
        'Content-Type': 'application/json'
        }

        response = requests.request("POST", url, headers=headers, data=payload)

        print("web search results:", response.json())

        return response.json()


    async def fast_api_handler(self, request: Request) -> Response:
        try:
            data = await request.json()
        except:
            return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)

        user_question = data.get("question")
        setting = data.get("settings")

        return JSONResponse(
            content={"response": self.processing(user_question, setting)},
            status_code=status.HTTP_200_OK)