151 lines
5.5 KiB
Python
151 lines
5.5 KiB
Python
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, open_url, run_command
|
|
from traceback import print_exc
|
|
from typing import Optional
|
|
import asyncio
|
|
import os
|
|
import io
|
|
import discord
|
|
|
|
load_dotenv()
|
|
|
|
client = genai.Client(api_key=os.getenv("GEM_API_KEY"))
|
|
config = types.GenerateContentConfig(
|
|
tools=[searxng, open_url, run_command],
|
|
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]) -> 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:
|
|
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)
|