refactor + preparing bot for public use

Signed-off-by: salmonik <raisersoft@proton.me>

1 files changed, 198 insertions, 0 deletions

diff --git a/cogs/assistant.py b/cogs/assistant.py
new file mode 100644
index 0000000..71f5150
--- /dev/null
+++ b/cogs/assistant.py
@@ -0,0 +1,198 @@
+# This file is part of Senator Sperma
+# Licensed under the GPL-3.0 License
+# Copyright (C) 2026 Salmonik @ salmonik@tuta.io
+# https://codeberg.org/salmonik/senator-sperma
+
+import re
+import discord
+from discord.ext import commands
+import requests
+import json
+import logging
+import shutil
+import os
+from tempfile import NamedTemporaryFile
+from cogs.downloader import Downloader
+
+logger = logging.getLogger(__name__)
+
+class Jarvis(commands.Cog):
+    def __init__(self, bot):
+        self.bot = bot
+        self.keyword = os.getenv("ASSISTANT_KEYWORD", "jarvis").lower()
+
+    @commands.Cog.listener()
+    async def on_message(self, message):
+        if message.author == self.bot.user:
+            return
+
+        logger.info(message.content)
+
+        if message.mentions:
+            if message.mentions[0].id == self.bot.user.id:
+                await self.assistant(message)
+
+        elif self.keyword in message.content.lower():
+            await self.assistant(message)
+
+    @staticmethod
+    async def assistant(message):
+        headers = {
+            "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
+            "Content-Type": "application/json"
+        }
+
+        chain = await Jarvis._build_message_chain(message)
+        messages = Jarvis._build_messages(message, chain)
+        payload = Jarvis._build_payload(messages)
+
+        response_data = await Jarvis._make_api_request(payload, headers)
+        response_message, action = Jarvis._parse_response(response_data)
+
+        channel = message.channel
+        await message.reply(response_message)
+
+        if action == "download":
+            await Jarvis._handle_download_action(message, channel)
+
+    @staticmethod
+    async def _build_message_chain(message):
+        chain = []
+        current = message
+        max_depth = 5
+
+        while current and len(chain) < max_depth:
+            chain.append(current)
+            if not current.reference:
+                break
+            try:
+                if current.reference.resolved:
+                    current = current.reference.resolved
+                else:
+                    current = await current.channel.fetch_message(current.reference.message_id)
+            except discord.NotFound:
+                break
+
+        return list(reversed(chain))
+
+    @staticmethod
+    def _build_messages(message, chain=None):
+        messages = [{
+            "role": "system",
+            "content": "Jesteś Senator Sperma Armostrong, także znany jako Jarvis."
+                       "Twoim zadaniem, po za asystowaniem użytkownika, jest mówienie programowi nadrzędnemu jak sie zachować"
+                       "pole 'message' może być tylko wypełnione w języku Polskim."
+                       "zachowuj sie jak Jarvis, bądź osobistym asystentem użytkownika, nie używaj emotikon"
+                       "nie pytaj sie użytkownika czy chcesz abyś coś zrobił, użytkownik ci sam powie abyś coś wykonał"
+        }]
+
+        if chain is None:
+            chain = [message]
+
+        for msg in chain:
+            messages.append({
+                "role": "user",
+                "content": f"{msg.author.display_name}: {msg.content}"
+            })
+
+        return messages
+
+    @staticmethod
+    def _build_payload(messages):
+        return {
+            "model": "arcee-ai/trinity-large-preview:free",
+            "messages": messages,
+            "response_format": Jarvis._get_response_schema(),
+            "plugins": [{"id": "response-healing"}]
+        }
+
+    @staticmethod
+    def _get_response_schema():
+        return {
+            "type": "json_schema",
+            "json_schema": {
+                "name": "response",
+                "strict": True,
+                "schema": {
+                    "type": "object",
+                    "properties": {
+                        "message": {
+                            "type": "string",
+                            "description": "the message that user will receive",
+                        },
+                        "action": {
+                            "type": "string",
+                            "description": "the action to do, either download, compress or none"
+                        }
+                    },
+                    "required": ["message", "action"],
+                    "additionalProperties": False,
+                }
+            }
+        }
+
+    @staticmethod
+    async def _make_api_request(payload, headers):
+        url = "https://openrouter.ai/api/v1/chat/completions"
+        response = requests.post(url=url, headers=headers, json=payload)
+        response_data = response.json()
+        logger.info(response_data)
+        return response_data
+
+    @staticmethod
+    def _parse_response(response_data):
+        content = response_data['choices'][0]['message']['content']
+        inner_json_str = content.strip('```json\n').strip('```').strip()
+        inner_data = json.loads(inner_json_str)
+
+        response_message = inner_data['message']
+        action = inner_data['action']
+
+        logger.info(response_message)
+        logger.info(action)
+
+        return response_message, action
+
+    @staticmethod
+    async def _handle_download_action(message, channel):
+        url_match = Jarvis._extract_url(message.content)
+        if not url_match:
+            await channel.send("Nie znaleziono poprawnego linku w odpowiedzi.")
+            return
+
+        response = await Downloader.download_standalone(url=url_match)
+        await Jarvis._process_download_response(response, channel)
+
+    @staticmethod
+    def _extract_url(content):
+        match = re.search(r'https://(?:[A-Za-z0-9-]+\.)*[A-Za-z0-9-]{1,63}\.[A-Za-z]{2,63}/\S*', content)
+        return match.group(0) if match else None
+
+    @staticmethod
+    async def _process_download_response(response, channel):
+        if not isinstance(response, dict):
+            return
+
+        if response["status"] == "success":
+            await Jarvis._handle_successful_download(response, channel)
+        else:
+            await Jarvis._handle_failed_download(response)
+
+    @staticmethod
+    async def _handle_successful_download(response, channel):
+        if response["type"] == "url":
+            await channel.send(f"pobioron i wysyłon: {response['content']}")
+        elif response["type"] == "file":
+            file_data = response["content"]
+            await channel.send("wysyłon:", file=discord.File(file_data))
+            shutil.rmtree(response["temp_dir"])
+
+    @staticmethod
+    async def _handle_failed_download(response):
+        logger.error(f"Error: {response['content']}")
+        if 'details' in response:
+            logger.error(f"Details: {response['details']}")
+        shutil.rmtree(response["temp_dir"])
+
+def setup(bot):
+    bot.add_cog(Jarvis(bot))