discordbot/cogs/upload_cog.py
Slipstream 8788cea989
feat: Add upload cog, update AI model, and refine Teto prompt
- Introduce `upload_cog.py` for handling file uploads.
- Update AI model from `google/gemini-2.5-flash-preview` to `google/gemini-2.5-flash-preview-05-20` for improved performance and stability.
- Refine Teto's roleplay system prompt:
    - Add explicit 18+ age disclaimer for all roleplay characters.
    - Include context about prepended usernames for better AI understanding.
    - Adjust Teto's perceived age to 18 for consistency.
    - Remove outdated or unnecessary prompt instructions.
2025-05-21 11:14:06 -06:00

271 lines
11 KiB
Python

import discord
from discord.ext import commands
from discord import app_commands
import aiohttp
import io
import os
import json
import asyncio
import base64
from typing import Optional, Dict, Any, Union
class UploadCog(commands.Cog, name="Upload"):
"""Cog for interacting with the upload API"""
def __init__(self, bot: commands.Bot):
self.bot = bot
self.api_base_url = "https://upload.slipstreamm.dev/upload"
self.session = None
self.captcha_cache = {} # Store captcha IDs temporarily
print("UploadCog initialized!")
# Create command group
self.upload_group = app_commands.Group(
name="upload",
description="Commands for interacting with the upload API",
guild_only=False
)
# Register commands
self.register_commands()
# Add command group to the bot's tree
self.bot.tree.add_command(self.upload_group)
async def cog_load(self):
"""Called when the cog is loaded."""
self.session = aiohttp.ClientSession()
print("UploadCog session created")
async def cog_unload(self):
"""Called when the cog is unloaded."""
if self.session:
await self.session.close()
print("UploadCog session closed")
def register_commands(self):
"""Register all commands for this cog"""
# --- Text Captcha Command ---
text_captcha_command = app_commands.Command(
name="textcaptcha",
description="Generate a text captcha for file uploads",
callback=self.upload_text_captcha_callback,
parent=self.upload_group
)
self.upload_group.add_command(text_captcha_command)
# --- Image Captcha Command ---
image_captcha_command = app_commands.Command(
name="imagecaptcha",
description="Generate an image captcha for file uploads",
callback=self.upload_image_captcha_callback,
parent=self.upload_group
)
self.upload_group.add_command(image_captcha_command)
# --- Upload File Command ---
upload_file_command = app_commands.Command(
name="file",
description="Upload a file with captcha verification",
callback=self.upload_file_callback,
parent=self.upload_group
)
app_commands.describe(
file="The file to upload",
captcha_id="The ID of the captcha challenge",
captcha_solution="The solution to the captcha challenge",
expires_after="Time in seconds until the file expires (default: 86400 - 24 hours)"
)(upload_file_command)
self.upload_group.add_command(upload_file_command)
async def _make_api_request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
"""Make a request to the API"""
if not self.session:
self.session = aiohttp.ClientSession()
# Ensure the endpoint starts with a slash
if not endpoint.startswith("/"):
endpoint = f"/{endpoint}"
url = f"{self.api_base_url}{endpoint}"
try:
if method.upper() == "GET":
async with self.session.get(url, **kwargs) as response:
if response.status == 200:
return await response.json()
else:
error_text = await response.text()
raise Exception(f"API request failed: {response.status} - {error_text}")
elif method.upper() == "POST":
async with self.session.post(url, **kwargs) as response:
if response.status in (200, 201):
return await response.json()
else:
error_text = await response.text()
raise Exception(f"API request failed: {response.status} - {error_text}")
else:
raise ValueError(f"Unsupported HTTP method: {method}")
except Exception as e:
print(f"Error making API request to {url}: {e}")
raise
async def upload_text_captcha_callback(self, interaction: discord.Interaction):
"""Generate a text captcha for file uploads"""
await interaction.response.defer(ephemeral=True)
try:
# Make API request to generate text captcha
captcha_data = await self._make_api_request("GET", "/api/captcha/text")
# Store captcha ID in cache
captcha_id = captcha_data.get("captcha_id")
user_id = interaction.user.id
self.captcha_cache[user_id] = captcha_id
# Create embed with captcha information
embed = discord.Embed(
title="Text Captcha Challenge",
description=captcha_data.get("message", "Please solve the captcha challenge"),
color=discord.Color.blue()
)
embed.add_field(name="Captcha ID", value=f"`{captcha_id}`", inline=False)
embed.add_field(name="Challenge", value=f"`{captcha_data.get('challenge', 'N/A')}`", inline=False)
# Add instructions for using the captcha
instructions = (
"**How to use this captcha:**\n"
"1. Copy the Captcha ID above\n"
"2. Type the challenge text as your solution\n"
"3. Use the `/upload file` command with your file, the Captcha ID, and your solution"
)
embed.add_field(name="Instructions", value=instructions, inline=False)
embed.set_footer(text="This captcha is valid for 10 minutes")
await interaction.followup.send(embed=embed, ephemeral=True)
except Exception as e:
await interaction.followup.send(f"Error generating text captcha: {e}", ephemeral=True)
async def upload_image_captcha_callback(self, interaction: discord.Interaction):
"""Generate an image captcha for file uploads"""
await interaction.response.defer(ephemeral=True)
try:
# Make API request to generate image captcha
captcha_data = await self._make_api_request("GET", "/api/captcha/image")
# Store captcha ID in cache
captcha_id = captcha_data.get("captcha_id")
user_id = interaction.user.id
self.captcha_cache[user_id] = captcha_id
# Create embed with captcha information
embed = discord.Embed(
title="Image Captcha Challenge",
description=captcha_data.get("message", "Please solve the captcha challenge"),
color=discord.Color.blue()
)
embed.add_field(name="Captcha ID", value=f"`{captcha_id}`", inline=False)
# Extract the base64 image data
image_data = captcha_data.get("image", "")
if image_data.startswith("data:image/png;base64,"):
image_data = image_data.replace("data:image/png;base64,", "")
# Convert base64 to bytes
image_bytes = io.BytesIO(base64.b64decode(image_data))
image_bytes.seek(0)
# Create file attachment
file = discord.File(fp=image_bytes, filename="captcha.png")
embed.set_image(url="attachment://captcha.png")
# Add instructions for using the captcha
instructions = (
"**How to use this captcha:**\n"
"1. Copy the Captcha ID above\n"
"2. Type the text shown in the image as your solution\n"
"3. Use the `/upload file` command with your file, the Captcha ID, and your solution"
)
embed.add_field(name="Instructions", value=instructions, inline=False)
embed.set_footer(text="This captcha is valid for 10 minutes")
await interaction.followup.send(embed=embed, file=file, ephemeral=True)
else:
embed.add_field(name="Error", value="Invalid image data received", inline=False)
await interaction.followup.send(embed=embed, ephemeral=True)
except Exception as e:
await interaction.followup.send(f"Error generating image captcha: {e}", ephemeral=True)
async def upload_file_callback(self, interaction: discord.Interaction,
file: discord.Attachment,
captcha_id: str,
captcha_solution: str,
expires_after: Optional[int] = 86400):
"""Upload a file with captcha verification"""
await interaction.response.defer(ephemeral=False)
try:
# Validate input
if not file:
await interaction.followup.send("Please provide a file to upload", ephemeral=True)
return
if not captcha_id or not captcha_solution:
await interaction.followup.send("Please provide both captcha ID and solution", ephemeral=True)
return
# Download the file
file_bytes = await file.read()
# Prepare form data
form_data = aiohttp.FormData()
form_data.add_field('file', file_bytes, filename=file.filename, content_type=file.content_type)
form_data.add_field('captcha_id', captcha_id)
form_data.add_field('captcha_solution', captcha_solution)
form_data.add_field('expires_after', str(expires_after))
# Make API request to upload file
upload_data = await self._make_api_request("POST", "/api/upload/third-party", data=form_data)
# Create embed with upload information
embed = discord.Embed(
title="File Uploaded Successfully",
description=f"Your file has been uploaded and will expire in {expires_after} seconds",
color=discord.Color.green()
)
file_id = upload_data.get("id", "unknown")
file_url = f"{self.api_base_url}/uploads/{file_id}"
# Format file size nicely
file_size_bytes = upload_data.get('file_size', 0)
if file_size_bytes < 1024:
file_size_str = f"{file_size_bytes} bytes"
elif file_size_bytes < 1024 * 1024:
file_size_str = f"{file_size_bytes / 1024:.2f} KB"
else:
file_size_str = f"{file_size_bytes / (1024 * 1024):.2f} MB"
embed.add_field(name="File ID", value=file_id, inline=True)
embed.add_field(name="Original Name", value=upload_data.get("original_name", "unknown"), inline=True)
embed.add_field(name="File Size", value=file_size_str, inline=True)
embed.add_field(name="Content Type", value=upload_data.get("content_type", "unknown"), inline=True)
embed.add_field(name="Scan Status", value=upload_data.get("scan_status", "unknown"), inline=True)
embed.add_field(name="File URL", value=file_url, inline=False)
# Add clickable link
embed.description += f"\n\n[Click here to download]({file_url})"
await interaction.followup.send(embed=embed)
except Exception as e:
await interaction.followup.send(f"Error uploading file: {e}", ephemeral=True)
async def setup(bot: commands.Bot):
"""Add the UploadCog to the bot."""
await bot.add_cog(UploadCog(bot))
print("UploadCog setup complete.")