From 2750b5fa7864868b4652936275b2500a0dbd2142 Mon Sep 17 00:00:00 2001
From: xory <xory@tutamail.com>
Date: Fri, 15 Aug 2025 13:31:56 +0000
Subject: [PATCH] i guess this counts as img & pdf recog

---
 main.py | 50 ++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 42 insertions(+), 8 deletions(-)

diff --git a/main.py b/main.py
index 1663b4f..1aec5cc 100644
--- a/main.py
+++ b/main.py
@@ -5,6 +5,7 @@ from discord import app_commands
 from discord.ext import commands
 from tools import searxng, open_url, run_command
 from traceback import print_exc
+from typing import Optional
 import asyncio
 import os
 import io
@@ -38,28 +39,61 @@ async def generation(interaction: discord.Interaction) -> None:
 @bot.tree.command(name="ask", description="ai thing yes 👍")
 @app_commands.allowed_installs(guilds=False, users=True)
 @app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
-async def ask(interaction: discord.Interaction, prompt: str) -> None:
+async def ask(interaction: discord.Interaction, prompt: str, attachment: Optional[discord.Attachment]) -> None:
     await interaction.response.defer()
+
+    attachment_text: str | None = None
+    attachment_data: bytes | None = None
+    content_type: str | None = None
+    if attachment:
+        attachment_data = await attachment.read()
+        content_type = attachment.content_type
+        if content_type and ('text/plain' in content_type or 'text/markdown' in content_type):
+            attachment_text = attachment_data.decode("utf-8")
+
     response: types.GenerateContentResponse | None = None
     for _ in range(5):
         try:
-            response = await client.aio.models.generate_content(
-                model="gemini-2.5-flash",
-                contents=[
-                    prompt
-                ],
-                config=config
-            )
+            if not attachment:
+                response = await client.aio.models.generate_content(
+                    model="gemini-2.5-flash",
+                    contents=[
+                        prompt
+                    ],
+                    config=config
+                )
+            elif attachment and attachment_text:
+                response = await client.aio.models.generate_content(
+                    model="gemini-2.5-flash",
+                    contents=[
+                        prompt,
+                        types.Part.from_text(text=attachment_text)
+                    ],
+                    config=config
+                )
+            elif content_type and attachment_data and ("image/" in content_type or "application/pdf" in content_type):
+                response = await client.aio.models.generate_content(
+                    model="gemini-2.5-flash",
+                    contents=[
+                        prompt,
+                        types.Part.from_bytes(
+                            data=attachment_data, mime_type=content_type)
+                    ],
+                    config=config
+                )
             break
         except:
             print_exc()
             await asyncio.sleep(10)
             continue
+
     if not response:
+        print("e")
         await interaction.edit_original_response(content="`[E] API Error. Please try again later.`")
         return
     if not response.text:
         await interaction.edit_original_response(content="`[E] Model returned no response`")
+
     generation: str = response.text or ""
     if len(generation) > 2000:
         buffer = io.BytesIO(generation.encode("utf-8"))