Implement possibility to answer to a different channel

2023-04-12 17:56:31 +02:00
parent defe598651
commit b1ece64874
6 changed files with 74 additions and 58 deletions
@@ -32,8 +32,11 @@ def parse_response(content: str) -> Dict:
 def parse_maybe_json(json_string):
    if json_string is None:
        return None
-    json_string = json_string.strip()
-
+    if isinstance(json_string, list):
+        return ' '.join([str(x) for x in json_string])
+    if isinstance(json_string, dict):
+        return ' '.join([str(x) for x in json_string.values()])
+    json_string = str(json_string).strip()
    try:
        parsed_json = parse_response(json_string)
    except Exception:
@@ -70,9 +73,17 @@ class AIMessage(AIMessageBase):


 class AIResponse(AIMessageBase):
-    def __init__(self, answer: Optional[str], answer_needed: bool, staff: Optional[str], picture: Optional[str], hack: bool) -> None:
+    def __init__(self,
+                 answer: Optional[str],
+                 answer_needed: bool,
+                 channel: Optional[str],
+                 staff: Optional[str],
+                 picture: Optional[str],
+                 hack: bool
+                 ) -> None:
        self.answer = answer
        self.answer_needed = answer_needed
+        self.channel = channel
        self.staff = staff
        self.picture = picture
        self.hack = hack
@@ -117,7 +128,7 @@ class AIResponder(object):
        raise RuntimeError(f"Failed to generate image {repr(description)} after multiple retries")

    async def post_process(self, message: AIMessage, response: Dict[str, Any]) -> AIResponse:
-        for fld in ('answer', 'staff', 'picture', 'hack'):
+        for fld in ('answer', 'channel', 'staff', 'picture', 'hack'):
            if str(response.get(fld)).strip().lower() in \
               ('none', '', 'null', '"none"', '"null"', "'none'", "'null'"):
                response[fld] = None
@@ -134,11 +145,15 @@ class AIResponder(object):
            response['answer'] = re.sub(r'\[[^\]]*\]\(([^\)]*)\)', r'\1', response['answer'])
            if message.direct or message.user in message.message:
                response['answer_needed'] = True
-        return AIResponse(response['answer'],
-                          response['answer_needed'],
-                          parse_maybe_json(response['staff']),
-                          parse_maybe_json(response['picture']),
-                          response['hack'])
+        response_message = AIResponse(response['answer'],
+                                      response['answer_needed'],
+                                      response['channel'],
+                                      parse_maybe_json(response['staff']),
+                                      parse_maybe_json(response['picture']),
+                                      response['hack'])
+        if response_message.staff is not None and response_message.answer is not None:
+            response_message.answer_needed = True
+        return response_message

    def short_path(self, message: AIMessage, limit: int) -> bool:
        if message.direct or 'short-path' not in self.config:
@@ -224,7 +239,7 @@ class AIResponder(object):
    async def send(self, message: AIMessage) -> AIResponse:
        limit = self.config["history-limit"]
        if self.short_path(message, limit):
-            return AIResponse(None, False, None, None, False)
+            return AIResponse(None, False, None, None, None, False)
        retries = 3
        while retries > 0:
            messages = self._message(message, limit)
@@ -8,6 +8,7 @@ from discord.ext import commands
 from watchdog.observers import Observer
 from watchdog.events import FileSystemEventHandler
 from .ai_responder import AIResponder, AIMessage
+from typing import Optional


 class ConfigFileHandler(FileSystemEventHandler):
@@ -55,13 +56,8 @@ class FjerkroaBot(commands.Bot):

    async def on_ready(self):
        print(f"We have logged in as {self.user}")
-        self.staff_channel = None
-        self.welcome_channel = None
-        for guild in self.guilds:
-            if self.staff_channel is None and self.config['staff-channel'] is not None:
-                self.staff_channel = discord.utils.get(guild.channels, name=self.config['staff-channel'])
-            if self.welcome_channel is None and self.config['welcome-channel'] is not None:
-                self.welcome_channel = discord.utils.get(guild.channels, name=self.config['welcome-channel'])
+        self.staff_channel = self.channel_by_name(self.config['staff-channel'])
+        self.welcome_channel = self.channel_by_name(self.config['welcome-channel'])

    async def on_member_join(self, member):
        logging.info(f"User {member.name} joined")
@@ -72,6 +68,8 @@ class FjerkroaBot(commands.Bot):
    async def on_message(self, message: Message) -> None:
        if self.user is not None and message.author.id == self.user.id:
            return
+        if not isinstance(message.channel, TextChannel):
+            return
        message_content = str(message.content).strip()
        if len(message_content) < 1:
            return
@@ -82,6 +80,18 @@ class FjerkroaBot(commands.Bot):
        msg = AIMessage(message.author.name, message_content, channel_name, self.user in message.mentions)
        await self.respond(msg, message.channel)

+    def channel_by_name(self,
+                        channel_name: Optional[str],
+                        fallback_channel: Optional[TextChannel] = None
+                        ) -> Optional[TextChannel]:
+        if channel_name is None:
+            return fallback_channel
+        for guild in self.guilds:
+            channel = discord.utils.get(guild.channels, name=channel_name)
+            if channel is not None and isinstance(channel, TextChannel):
+                return channel
+        return fallback_channel
+
    async def respond(self, message: AIMessage, channel: TextChannel) -> None:
        try:
            channel_name = str(channel.name)
@@ -97,20 +107,22 @@ class FjerkroaBot(commands.Bot):
            airesponder = self.airesponder
        async with channel.typing():
            response = await airesponder.send(message)
-            if response.hack:
-                logging.warning(f"User {message.user} tried to hack the system.")
-                if response.staff is None:
-                    response.staff = f"User {message.user} try to hack the AI."
-            if response.staff is not None and self.staff_channel is not None:
-                async with self.staff_channel.typing():
-                    await self.staff_channel.send(response.staff, suppress_embeds=True)
-            if not response.answer_needed:
-                return
+        if response.hack:
+            logging.warning(f"User {message.user} tried to hack the system.")
+            if response.staff is None:
+                response.staff = f"User {message.user} try to hack the AI."
+        answer_channel = self.channel_by_name(response.channel, channel)
+        if response.staff is not None and self.staff_channel is not None:
+            async with self.staff_channel.typing():
+                await self.staff_channel.send(response.staff, suppress_embeds=True)
+        if not response.answer_needed or answer_channel is None:
+            return
+        async with answer_channel.typing():
            if response.picture is not None:
                images = [discord.File(fp=await airesponder.draw(response.picture), filename="image.png")]
-                await channel.send(response.answer, files=images, suppress_embeds=True)
+                await answer_channel.send(response.answer, files=images, suppress_embeds=True)
            else:
-                await channel.send(response.answer, suppress_embeds=True)
+                await answer_channel.send(response.answer, suppress_embeds=True)

    async def close(self):
        self.observer.stop()