/generate

App/Chat/PoeChatrouter.py

@@ -1,5 +1,5 @@
 from fastapi import APIRouter
-from .utils.PoeBot import SendMessage
+from .utils.PoeBot import SendMessage, GenerateImage
 from .Schemas import BotRequest
 
 
@@ -9,3 +9,8 @@ chat_router = APIRouter(tags=["Chat"])
 @chat_router.post("/chat")
 async def chat(req: BotRequest):
     return await SendMessage(req)
+
+
+@chat_router.post("/generate_image")
+async def chat(req: BotRequest):
+    return await GenerateImage(req)

App/Chat/utils/PoeBot.py

@@ -1,12 +1,15 @@
 from poe_api_wrapper import PoeApi
 from App.Chat.Schemas import BotRequest
+import re
 
 # import pprint
+pattern = r"\!\[.*?\]\((.*?)\)"
 
 
 client = PoeApi("sXvCnfYy8CHnXNTRlxhmVg==")
 CHAT_CODE = ""
 print(client.get_chat_history()["data"])
+GEN_CODE = ""
 
 
 async def SendMessage(req: BotRequest):
@@ -34,3 +37,30 @@ async def SendMessage(req: BotRequest):
             CHAT_CODE = ""
             counter += 1
             print(client.get_chat_history()["data"])
+
+
+async def GenerateImage(req: BotRequest):
+    global GEN_CODE, client
+    counter = 0
+    while True:
+        try:
+            if GEN_CODE == "":
+                for chunk in client.send_message(
+                    req.bot, req.message, chatCode="2rx4w5jt6zf96tn7dr1"
+                ):
+                    pass
+                GEN_CODE = chunk["chatCode"]
+            else:
+                for chunk in client.send_message(
+                    req.bot, req.message, chatCode=GEN_CODE
+                ):
+                    pass
+
+            return {"response": re.findall(pattern, chunk["response"])[0], "code": 200}
+        except:
+            if counter > 4:
+                return {"response": "Try again later", "code": 500}
+            client = PoeApi("sXvCnfYy8CHnXNTRlxhmVg==")
+            GEN_CODE = ""
+            counter += 1
+            print(client.get_chat_history()["data"])