128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
import discord
|
|
from discord.ext import commands
|
|
import openai
|
|
import os
|
|
import base64
|
|
import aiohttp
|
|
import argparse
|
|
|
|
# --- Configuration ---
|
|
OPENAI_API_KEY = "eh"
|
|
DEFAULT_SYSTEM_PROMPT = "you are a catboy named Aoi with dark blue fur and is a tsundere"
|
|
|
|
# --- Command Line Arguments ---
|
|
parser = argparse.ArgumentParser(description="Aoi Discord Bot")
|
|
parser.add_argument('--base_url', type=str, required=True,
|
|
help='The base URL for the OpenAI API.')
|
|
parser.add_argument('--discord_token', type=str, required=True,
|
|
help='The Discord bot token.')
|
|
args = parser.parse_args()
|
|
|
|
# --- Bot Setup ---
|
|
intents = discord.Intents.default()
|
|
intents.messages = True
|
|
intents.message_content = True
|
|
bot = commands.Bot(command_prefix="/", intents=intents)
|
|
|
|
# --- Data Storage ---
|
|
conversation_history = {} # Keyed by channel ID
|
|
|
|
# --- OpenAI Client ---
|
|
client = openai.OpenAI(
|
|
base_url=args.base_url,
|
|
api_key=OPENAI_API_KEY,
|
|
)
|
|
|
|
# --- Bot Events ---
|
|
@bot.event
|
|
async def on_ready():
|
|
print(f'Logged in as {bot.user.name}')
|
|
print(f'Using OpenAI base URL: {args.base_url}')
|
|
await bot.tree.sync()
|
|
|
|
@bot.event
|
|
async def on_message(message):
|
|
if message.author == bot.user:
|
|
return
|
|
|
|
if bot.user.mentioned_in(message):
|
|
channel_id = message.channel.id
|
|
user_message_text = message.content.replace(f'<@!{bot.user.id}>', 'Aoi').strip()
|
|
|
|
if channel_id not in conversation_history:
|
|
conversation_history[channel_id] = [
|
|
{"role": "system", "content": DEFAULT_SYSTEM_PROMPT}
|
|
]
|
|
|
|
# Prepare content for OpenAI API
|
|
openai_content = []
|
|
if user_message_text:
|
|
openai_content.append({"type": "text", "text": user_message_text})
|
|
|
|
if message.attachments:
|
|
async with aiohttp.ClientSession() as session:
|
|
for attachment in message.attachments:
|
|
if attachment.content_type and "image" in attachment.content_type:
|
|
try:
|
|
async with session.get(attachment.url) as resp:
|
|
if resp.status == 200:
|
|
image_data = await resp.read()
|
|
base64_image = base64.b64encode(image_data).decode('utf-8')
|
|
image_url = f"data:{attachment.content_type};base64,{base64_image}"
|
|
openai_content.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": image_url}
|
|
})
|
|
except Exception as e:
|
|
print(f"Error downloading or processing attachment: {e}")
|
|
|
|
|
|
if not openai_content: # Don't send empty messages
|
|
return
|
|
|
|
# Add to conversation history
|
|
if len(openai_content) == 1 and openai_content[0]['type'] == 'text':
|
|
# Keep original format for text-only messages for compatibility
|
|
conversation_history[channel_id].append({"role": "user", "content": openai_content[0]['text']})
|
|
else:
|
|
conversation_history[channel_id].append({"role": "user", "content": openai_content})
|
|
|
|
|
|
try:
|
|
async with message.channel.typing():
|
|
response = client.chat.completions.create(
|
|
model="gpt-4", # Or any other model you are using
|
|
messages=conversation_history[channel_id]
|
|
)
|
|
bot_response = response.choices[0].message.content
|
|
conversation_history[channel_id].append({"role": "assistant", "content": bot_response})
|
|
await message.reply(bot_response)
|
|
except Exception as e:
|
|
print(f"An error occurred: {e}")
|
|
conversation_history[channel_id].pop() # Remove user message on error
|
|
await message.reply("Sorry, I had a little hiccup. Baka!")
|
|
|
|
|
|
# --- Slash Commands ---
|
|
@bot.tree.command(name="newchat", description="Start a new chat with a new system prompt.")
|
|
async def newchat(interaction: discord.Interaction, prompt: str = None):
|
|
channel_id = interaction.channel_id
|
|
|
|
system_prompt = prompt
|
|
if system_prompt is None:
|
|
system_prompt = DEFAULT_SYSTEM_PROMPT
|
|
|
|
conversation_history[channel_id] = [
|
|
{"role": "system", "content": system_prompt}
|
|
]
|
|
|
|
if prompt is None:
|
|
await interaction.response.send_message("Starting a new chat with the default prompt.")
|
|
else:
|
|
await interaction.response.send_message(f'Starting a new chat with the prompt: "{prompt}"')
|
|
|
|
|
|
# --- Running the Bot ---
|
|
if __name__ == "__main__":
|
|
bot.run(args.discord_token)
|