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, Literal 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: Optional[bool], model: Optional[Literal["pro", "flash", "flash-lite"]]) -> 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 model_name: str = "" match model: case "lite": model_name = "gemini-flash-lite-latest" case "pro": model_name = "gemini-2.5-pro" case _: model_name = "gemini-flash-latest" print(f"[d] using {model_name}") # TODO: delete this for _ in range(5): try: if not attachment: response = await client.aio.models.generate_content( model=model_name, contents=[ prompt ], config=config ) elif attachment and attachment_text: response = await client.aio.models.generate_content( model=model_name, 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=model_name, 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="genimage", description="ai image gen woooo") @app_commands.allowed_installs(guilds=False, users=True) @app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True) async def gen_image(interaction: discord.Interaction, prompt: str) -> None: await interaction.response.defer() response: types.GenerateImagesResponse | None for _ in range(5): try: response = await client.models.generate_images( model="imagen-4.0-generate-001", prompt=prompt, config=types.GenerateImagesConfig( number_of_images=1, ) ) except: print_exc() await asyncio.sleep(10) continue if not response: await interaction.edit_original_response(content="`[E] API error. Please try again later.`") await interaction.edit_original_response(content="Done", attachments=response.generated_images) @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)