Skip to content

Commit

Permalink
feat(link_ai_bot.py): add support for creating images using OpenAI's …
Browse files Browse the repository at this point in the history
…DALL-E API
  • Loading branch information
lanvent committed Jun 10, 2023
1 parent 419a3e5 commit b25e843
Showing 1 changed file with 16 additions and 3 deletions.
19 changes: 16 additions & 3 deletions bot/linkai/link_ai_bot.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,15 @@

from bot.bot import Bot
from bot.chatgpt.chat_gpt_session import ChatGPTSession
from bot.openai.open_ai_image import OpenAIImage
from bot.session_manager import SessionManager
from bridge.context import Context
from bridge.context import Context, ContextType
from bridge.reply import Reply, ReplyType
from common.log import logger
from config import conf


class LinkAIBot(Bot):
class LinkAIBot(Bot, OpenAIImage):
# authentication failed
AUTH_FAILED_CODE = 401
NO_QUOTA_CODE = 406
Expand All @@ -24,7 +25,19 @@ def __init__(self):
self.sessions = SessionManager(ChatGPTSession, model=conf().get("model") or "gpt-3.5-turbo")

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):
if retry_count >= 2:
Expand Down

0 comments on commit b25e843

Please sign in to comment.