jarvis-models/src/blackbox/tesou.py

from typing import Any, Coroutine

from fastapi import Request, Response, status
from fastapi.responses import JSONResponse
from injector import inject
from ..configuration import TesouConf
from .blackbox import Blackbox

import requests
from injector import singleton

@singleton
class Tesou(Blackbox):
    url: str

    @inject
    def __init__(self, tesou_config: TesouConf):
        self.url = tesou_config.url
    
    def __call__(self, *args, **kwargs):
        return self.processing(*args, **kwargs)

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

    # 用户输入的数据格式为：[{"id": "123", "prompt": "叉烧饭，帮我查询叉烧饭的介绍"}]
    def processing(self, id, prompt)  -> str:
        message = {
            "user_id": id,
            "prompt": prompt,
        }
        response = requests.post(self.url, json=message)
        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_id = data.get("user_id")
        user_prompt = data.get("prompt")
        if user_prompt is None:
            return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)
        return JSONResponse(content={"Response": self.processing(user_id, user_prompt)}, status_code=status.HTTP_200_OK)