huginn/main.py

from google import genai
from google.genai import types
from dotenv import load_dotenv
from discord import app_commands
from discord.ext import commands
from tools import searxng, run_command, open_url
from traceback import print_exc
from typing import Optional
import asyncio
import os
import io
import discord

load_dotenv()

sysprompt = ""
with open("sysprompt.md", "r") as sysprompt_file:
    sysprompt = sysprompt_file.read()
client = genai.Client(api_key=os.getenv("GEM_API_KEY"))
config = types.GenerateContentConfig(
    tools=[searxng, open_url, run_command],
    system_instruction=sysprompt,
    safety_settings=[
        types.SafetySetting(
            category=types.HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY,
            threshold=types.HarmBlockThreshold.BLOCK_NONE
        ),
        types.SafetySetting(
            category=types.HarmCategory.HARM_CATEGORY_HARASSMENT,
            threshold=types.HarmBlockThreshold.BLOCK_NONE
        ),
        types.SafetySetting(
            category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
            threshold=types.HarmBlockThreshold.BLOCK_NONE
        ),
        types.SafetySetting(
            category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
            threshold=types.HarmBlockThreshold.BLOCK_NONE
        ),
        types.SafetySetting(
            category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
            threshold=types.HarmBlockThreshold.BLOCK_NONE
        ),

    ]
)
intents = discord.Intents.default()
intents.message_content = True
bot = commands.Bot(intents=intents, command_prefix="-")


@bot.tree.command(name="test", description="Test command")
@app_commands.allowed_installs(guilds=False, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def test(interaction: discord.Interaction) -> None:
    await interaction.response.send_message("hai :3")


@bot.tree.command(name="generation", description="generation")
@app_commands.allowed_installs(guilds=False, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def generation(interaction: discord.Interaction) -> None:
    await interaction.response.send_message("generation")


@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, attachment: Optional[discord.Attachment], ephemeral: bool) -> None:
    if not ephemeral:
        await interaction.response.defer()
    else:
        await interaction.response.defer(ephemeral=True)

    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:
            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"))
        buffer.seek(0)
        file = discord.File(buffer, filename="output.md")
        await interaction.edit_original_response(content="Sent as a file", attachments=[file])
    else:
        await interaction.edit_original_response(content=generation)


@bot.tree.command(name="sync", description="bot.tree.sync()")
@app_commands.allowed_installs(guilds=False, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def sync(interaction: discord.Interaction) -> None:
    if interaction.user.id == 1139850599085645844:
        await interaction.response.send_message("`[I] Syncing...`")
        await bot.tree.sync()
        await interaction.edit_original_response(content="`[I] Syncing... OK`")
    else:
        await interaction.response.send_message("`[E] 403 Forbidden`")


@bot.event
async def on_ready():
    await bot.tree.sync()
    print("Logged in!")


api_key: str | None = os.getenv("DSC_API_KEY")
if not api_key:
    raise RuntimeError
bot.run(api_key)