From a2b5b10625126b7e98643af9fe40569c9e004047 Mon Sep 17 00:00:00 2001 From: salmonik Date: Sat, 18 Apr 2026 00:24:50 +0200 Subject: refactor + preparing bot for public use Signed-off-by: salmonik --- cogs/assistant.py | 198 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 198 insertions(+) create mode 100644 cogs/assistant.py (limited to 'cogs/assistant.py') 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)) -- cgit v1.2.3