1 files changed, 128 insertions, 90 deletions

diff --git a/cogs/jarvis.py b/cogs/jarvis.py
index 0096b38..b199c8a 100644
--- a/cogs/jarvis.py
+++ b/cogs/jarvis.py
@@ -29,126 +29,164 @@ class Jarvis(commands.Cog):
 
         if message.mentions:
             if message.mentions[0].id == self.bot.user.id:
-                content = await self.get_original(message)
-                await self.jarvis(message, content)
+                await self.jarvis(message)
 
         elif "jarvis" in message.content.lower():
-            content = await self.get_original(message)
-            await self.jarvis(message, content)
+            await self.jarvis(message)
 
     @staticmethod
-    async def jarvis(message, content):
-        url = "https://openrouter.ai/api/v1/chat/completions"
+    async def jarvis(message):
         headers = {
-
-            "Authorization": f"Bearer {os.getenv("OPENROUTER_API_KEY")}",
+            "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
             "Content-Type": "application/json"
         }
         
-        # Build messages array
-        messages = [
-            {
-                "role": "system",
-                "content": "your task is to tell the program what to do and how to respond after receiving a message.\n"
-                           "message field has to be filled out in polish.\n"
-                           "act like jarvis, be a personal assistant to user, do not use emojis\n"
-                           "do not ask the user if they want to do an action. the user will tell you that"
-            }
-        ]
+        messages = Jarvis._build_messages(message)
+        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 channel.send(response_message)
+        
+        if action == "download":
+            await Jarvis._handle_download_action(message, channel)
+
+    @staticmethod
+    def _build_messages(message):
+        messages = [{
+            "role": "system",
+            "content": "your task is to tell the program what to do and how to respond after receiving a message.\n"
+                       "message field has to be filled out in polish.\n"
+                       "act like jarvis, be a personal assistant to user, do not use emojis\n"
+                       "do not ask the user if they want to do an action. the user will tell you that"
+        }]
         
-        # Add original message with display name
         messages.append({
             "role": "user",
             "content": f"[from {message.author.display_name}]: {message.content}"
         })
         
-        # If it's a reply, add the original message and reply as separate messages
-        if message.reference and (message.type == discord.MessageType.reply or message.reference.resolved):
-            if message.reference.resolved:
-                original_msg = message.reference.resolved
-                messages.append({
-                    "role": "user",
-                    "content": f"[original message from {original_msg.author.display_name}]: {original_msg.content}"
-                })
-                messages.append({
-                    "role": "user", 
-                    "content": f"[reply from {message.author.display_name}]: {message.content}"
-                })
-        
-        payload = {
+        if Jarvis._is_reply(message):
+            messages.extend(Jarvis._build_reply_messages(message))
+            
+        return messages
+
+    @staticmethod
+    def _is_reply(message):
+        return message.reference and (message.type == discord.MessageType.reply or message.reference.resolved)
+
+    @staticmethod
+    def _build_reply_messages(message):
+        original_msg = message.reference.resolved
+        return [
+            {
+                "role": "user",
+                "content": f"[original message from {original_msg.author.display_name}]: {original_msg.content}"
+            },
+            {
+                "role": "user", 
+                "content": f"[reply from {message.author.display_name}]: {message.content}"
+            }
+        ]
+
+    @staticmethod
+    def _build_payload(messages):
+        return {
             "model": "arcee-ai/trinity-large-preview:free",
             "messages": messages,
-            "response_format": {
-                "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"
-                            }
+            "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",
                         },
-                        "required": ["message", "action"],
-                        "additionalProperties": False,
+                        "action": {
+                            "type": "string",
+                            "description": "the action to do, either download, compress or none"
                         }
-                    }
-                },
-            "plugins": [
-                {"id": "response-healing"}
-            ]
+                    },
+                    "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)
-        # Parse the JSON response
         response_data = response.json()
         logger.info(response_data)
-        # Extract the content string
-        content = response_data['choices'][0]['message']['content']
+        return response_data
 
-        # Remove the code block markers and parse the inner JSON
+    @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)
-
-        # Now you can access the message and action
-        response_message = inner_data['message']  # "Proszę podać plik do kompresji lub link do pobrania."
-        action = inner_data['action']  # "none OR download OR compress"
-        channel = message.channel
+        
+        response_message = inner_data['message']
+        action = inner_data['action']
+        
         logger.info(response_message)
         logger.info(action)
+        
+        return response_message, action
 
-        await channel.send(response_message)
+    @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
 
-        if action == "download":
-            # Extract the URL from the message using regex
-            url_match = re.search(r'https://(?:[A-Za-z0-9-]+\.)*[A-Za-z0-9-]{1,63}\.[A-Za-z]{2,63}/\S*', message.content).group(0)
-            if not url_match:
-                await channel.send("Nie znaleziono poprawnego linku w odpowiedzi.")
-                return
-
-            response = await Downloader.download_standalone(url=url_match)
-            if isinstance(response, dict):
-                if response["status"] == "success":
-                    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"])
-                else:
-                    logger.error(f"Error: {response['content']}")
-                    if 'details' in response:
-                        logger.error(f"Details: {response['details']}")
-                    shutil.rmtree(response["temp_dir"])
-
-            elif action == "compress":
-                pass
+        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"])
 
     @staticmethod
     async def get_original(message):