Spaces:
Running
Running
import discord | |
from discord import app_commands | |
import os | |
import requests | |
import asyncio | |
import aiohttp # Use aiohttp for asynchronous HTTP requests | |
import gradio as gr # Import Gradio | |
import google.generativeai as genai | |
from io import BytesIO | |
import PIL.Image | |
import base64 | |
from typing import List | |
# --- Environment Variables & Setup --- | |
DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN") | |
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") | |
# --- Discord Bot Setup --- | |
intents = discord.Intents.default() | |
client = discord.Client(intents=intents) | |
tree = app_commands.CommandTree(client) | |
genai.configure(api_key=GEMINI_API_KEY) | |
model = genai.GenerativeModel("gemini-2.0-flash-exp") | |
if not DISCORD_BOT_TOKEN or not GEMINI_API_KEY: | |
raise ValueError("Both DISCORD_BOT_TOKEN and GEMINI_API_KEY must be set.") | |
async def hello_command(interaction): | |
await interaction.response.send_message("Hello there!") | |
async def generate_command( | |
interaction: discord.Interaction, | |
prompt: str, | |
): | |
try: | |
# Defer the interaction to allow time for processing | |
await interaction.response.defer() | |
genai.configure(api_key=GEMINI_API_KEY) | |
model = genai.GenerativeModel("gemini-2.0-flash-exp") | |
content = [] # List to store the content parts | |
# Handle attachments if the command is used in reply to a message | |
if interaction.message.reference: | |
ref_message = await interaction.channel.fetch_message(interaction.message.reference.message_id) | |
if ref_message.attachments: # Check if the replied-to message has attachments | |
for attachment in ref_message.attachments: | |
if not attachment.content_type or not attachment.content_type.startswith("image/"): | |
await interaction.followup.send(f"The attachment {attachment.filename} is not an image.", ephemeral=True) | |
continue | |
# Download each attachment | |
async with aiohttp.ClientSession() as session: | |
async with session.get(attachment.url) as resp: | |
if resp.status == 200: | |
image_bytes = await resp.read() | |
base64_image = base64.b64encode(image_bytes).decode("utf-8") | |
content.append({"mime_type": "image/jpeg", "data": base64_image}) | |
else: | |
await interaction.followup.send( | |
f"Failed to download image: {attachment.filename}, status code: {resp.status}" | |
) | |
return | |
# Add the text prompt to the content | |
content.append(prompt) | |
# Call the Generative AI model | |
response = model.generate_content(content, stream=True) | |
current_message = None | |
current_message_content = "" | |
# Process streaming responses | |
for part in response: | |
text_chunk = part.text | |
if len(current_message_content) + len(text_chunk) <= 2000: | |
current_message_content += text_chunk | |
if current_message: | |
await current_message.edit(content=current_message_content) | |
else: | |
current_message = await interaction.followup.send(content=current_message_content) | |
else: | |
if current_message: | |
await current_message.edit(content=current_message_content) | |
else: | |
await interaction.followup.send(content=current_message_content) | |
current_message_content = text_chunk | |
current_message = await interaction.followup.send(content=current_message_content) | |
if current_message_content: | |
if current_message: | |
await current_message.edit(content=current_message_content) | |
else: | |
await interaction.followup.send(content=current_message_content) | |
except Exception as e: | |
print(e) | |
await interaction.followup.send(f"An error occurred: {e}") | |
async def process_image(interaction: discord.Interaction): | |
try: | |
# Ensure the command is replying to a message | |
if not interaction.message.reference: | |
await interaction.response.send_message("Please reply to a message with an image to use this command.", ephemeral=True) | |
return | |
# Fetch the original message being replied to | |
ref_message = await interaction.channel.fetch_message(interaction.message.reference.message_id) | |
# Check if the original message has attachments | |
if not ref_message.attachments: | |
await interaction.response.send_message("The replied-to message does not contain any attachments.", ephemeral=True) | |
return | |
# Process the first attachment (assuming it's an image) | |
attachment = ref_message.attachments[0] | |
if not attachment.content_type.startswith("image/"): | |
await interaction.response.send_message("The attachment is not an image.", ephemeral=True) | |
return | |
# Defer the response to allow time for processing | |
await interaction.response.defer() | |
# Download and process the image as needed | |
async with aiohttp.ClientSession() as session: | |
async with session.get(attachment.url) as resp: | |
if resp.status == 200: | |
image_bytes = await resp.read() | |
# Process the image bytes as needed | |
# For example, pass them to your AI model | |
else: | |
await interaction.followup.send(f"Failed to download the image: {attachment.filename}") | |
return | |
# After processing, send a response | |
await interaction.followup.send("Image processed successfully.") | |
except Exception as e: | |
print(f"An error occurred: {e}") | |
await interaction.followup.send(f"An error occurred: {e}") | |
async def on_ready(): | |
await tree.sync() | |
print("Bot is ready!") | |
client.event(on_ready) | |
# --- Gradio Interface --- | |
def echo_text(text): | |
return text | |
def run_gradio(): | |
gr.Interface( | |
fn=echo_text, | |
inputs="text", | |
outputs="text", | |
live=False, | |
title="Minimal Gradio Interface", | |
).launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True) | |
async def main(): | |
bot_task = asyncio.create_task(client.start(DISCORD_BOT_TOKEN)) | |
gradio_task = asyncio.to_thread(run_gradio) | |
await asyncio.gather(bot_task, gradio_task) | |
if __name__ == "__main__": | |
asyncio.run(main()) | |