Support drawing.

fjerkroa_bot/ai_responder.py

@@ -1,5 +1,8 @@
 import json
 import openai
+import aiohttp
+import logging
+from io import BytesIO
 from typing import Optional, List, Dict, Any, Tuple
 
 
@@ -43,6 +46,16 @@ class AIResponder(object):
             messages.append(msg)
         return messages, history
 
+    async def draw(self, description: str) -> BytesIO:
+        while True:
+            try:
+                response = await openai.Image.acreate(prompt=description, n=1, size="512x512")
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(response['data'][0]['url']) as image:
+                        return BytesIO(await image.read())
+            except Exception as err:
+                logging.warning(f"Failed to generate image {repr(description)}: {repr(err)}")
+
     async def send(self, message: AIMessage) -> AIResponse:
         limit = self.config["history-limit"]
         while True:
@@ -54,13 +67,20 @@ class AIResponder(object):
                                                              top_p=self.config["top-p"],
                                                              presence_penalty=self.config["presence-penalty"],
                                                              frequency_penalty=self.config["frequency-penalty"])
+                answer = result['choices'][0]['message']
+                response = json.loads(answer['content'])
             except openai.error.InvalidRequestError as err:
-                if 'maximum context length is' in str(err) and limit > 2:
+                if 'maximum context length is' in str(err) and limit > 4:
                     limit -= 1
                     continue
                 raise err
-            answer = result['choices'][0]['message']
-            response = json.loads(answer['content'])
+            except Exception as err:
+                logging.warning(f"failed to generate response: {repr(err)}")
+                continue
             history.append(answer)
             self.history = history
-            return response
+            return AIResponse(response['answer'],
+                              response['answer_needed'],
+                              response['staff'],
+                              response['picture'],
+                              response['hack'])

fjerkroa_bot/discord_bot.py

@@ -59,11 +59,15 @@ class FjerkroaBot(commands.Bot):
         response = await self.airesponder.send(msg)
         if response.staff is not None:
             async with self.staff_channel.typing():
-                self.staff_channel.send(response.staff)
+                await self.staff_channel.send(response.staff)
         if not response.answer_needed:
             return
         async with message.channel.typing():
-            message.channel.send(response.answer)
+            if response.picture is not None:
+                images = [discord.File(fp=await self.airesponder.draw(response.picture), filename="image.png")]
+                await message.channel.send(response.answer, files=images)
+            else:
+                await message.channel.send(response.answer)
 
     async def close(self):
         self.observer.stop()