فهرست منبع

feat(link_ai_bot.py): add support for creating images using OpenAI's DALL-E API

lanvent 2 سال پیش
والد
کامیت
b25e843351
1فایلهای تغییر یافته به همراه16 افزوده شده و 3 حذف شده
  1. 16 3
      bot/linkai/link_ai_bot.py

+ 16 - 3
bot/linkai/link_ai_bot.py

@@ -7,14 +7,15 @@ import requests
 
 
 from bot.bot import Bot
 from bot.bot import Bot
 from bot.chatgpt.chat_gpt_session import ChatGPTSession
 from bot.chatgpt.chat_gpt_session import ChatGPTSession
+from bot.openai.open_ai_image import OpenAIImage
 from bot.session_manager import SessionManager
 from bot.session_manager import SessionManager
-from bridge.context import Context
+from bridge.context import Context, ContextType
 from bridge.reply import Reply, ReplyType
 from bridge.reply import Reply, ReplyType
 from common.log import logger
 from common.log import logger
 from config import conf
 from config import conf
 
 
 
 
-class LinkAIBot(Bot):
+class LinkAIBot(Bot, OpenAIImage):
     # authentication failed
     # authentication failed
     AUTH_FAILED_CODE = 401
     AUTH_FAILED_CODE = 401
     NO_QUOTA_CODE = 406
     NO_QUOTA_CODE = 406
@@ -24,7 +25,19 @@ class LinkAIBot(Bot):
         self.sessions = SessionManager(ChatGPTSession, model=conf().get("model") or "gpt-3.5-turbo")
         self.sessions = SessionManager(ChatGPTSession, model=conf().get("model") or "gpt-3.5-turbo")
 
 
     def reply(self, query, context: Context = None) -> Reply:
     def reply(self, query, context: Context = None) -> Reply:
-        return self._chat(query, context)
+        if context.type == ContextType.TEXT:
+            return self._chat(query, context)
+        elif context.type == ContextType.IMAGE_CREATE:
+            ok, retstring = self.create_img(query, 0)
+            reply = None
+            if ok:
+                reply = Reply(ReplyType.IMAGE_URL, retstring)
+            else:
+                reply = Reply(ReplyType.ERROR, retstring)
+            return reply
+        else:
+            reply = Reply(ReplyType.ERROR, "Bot不支持处理{}类型的消息".format(context.type))
+            return reply
 
 
     def _chat(self, query, context, retry_count=0):
     def _chat(self, query, context, retry_count=0):
         if retry_count >= 2:
         if retry_count >= 2: