From cc134578a7c163d72d9ae23685dc9f3721841295 Mon Sep 17 00:00:00 2001 From: Dream Hunter Date: Sun, 19 Feb 2023 15:52:17 +0800 Subject: [PATCH] feat: update chatgpt (#33) --- router/chatgpt.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/router/chatgpt.py b/router/chatgpt.py index 7e3e39a..e9bb83b 100644 --- a/router/chatgpt.py +++ b/router/chatgpt.py @@ -1,6 +1,7 @@ import logging from functools import lru_cache +from pydantic import BaseModel from fastapi import status, APIRouter from fastapi.responses import PlainTextResponse, JSONResponse from revChatGPT.V1 import Chatbot @@ -12,19 +13,28 @@ _logger = logging.getLogger(__name__) +class ChatgptMessage(BaseModel): + token: str + text: str + chat_id: str = "Default" + + @ratelimit(10, 10) @router.post("/chatgpt", response_class=PlainTextResponse, tags=["chatgpt"]) -def get_chatgpt_message(token: str, text: str, chat_id: str = "Default") -> str: - if token != Tools.get_api_token(): +def get_chatgpt_message(chatgpt_message: ChatgptMessage) -> str: + if chatgpt_message.token != Tools.get_api_token(): return JSONResponse( status_code=status.HTTP_400_BAD_REQUEST, content={"message": "token is not correct"} ) - chatbot = get_chatbot(chat_id) - res = "" - for data in chatbot.ask(text): - res = data["message"] - return res + chatbot = get_chatbot(chatgpt_message.chat_id) + try: + res = "" + for data in chatbot.ask(chatgpt_message.text): + res = data["message"] + return res + except Exception as e: + get_chatbot.cache_clear() @lru_cache()