Initial commit of bot; 100% vibe-coded.
This commit is contained in:
122
.gitignore
vendored
Normal file
122
.gitignore
vendored
Normal file
@@ -0,0 +1,122 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
pip-wheel-metadata/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
.python-version
|
||||
|
||||
# PEP 582; used by PDM, PEP 582 compatible tools and python itself
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
126
bot.py
Normal file
126
bot.py
Normal file
@@ -0,0 +1,126 @@
|
||||
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.channel.send(bot_response)
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
await message.channel.send("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)
|
112
tests/test_bot.py
Normal file
112
tests/test_bot.py
Normal file
@@ -0,0 +1,112 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
import sys
|
||||
import os
|
||||
import base64
|
||||
|
||||
# Add the parent directory to the Python path to import the bot
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
# Patch sys.argv before importing the bot to prevent argparse errors
|
||||
with patch.object(sys, 'argv', ['bot.py', '--base_url', 'http://fake.url', '--discord_token', 'fake_token']):
|
||||
with patch('discord.ext.commands.Bot') as BotMock:
|
||||
bot_instance = BotMock()
|
||||
with patch('bot.bot', bot_instance):
|
||||
import bot
|
||||
|
||||
class TestAoiBot(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
def setUp(self):
|
||||
# Reset conversation history before each test
|
||||
bot.conversation_history = {}
|
||||
bot.bot.user = MagicMock()
|
||||
bot.bot.user.id = 12345
|
||||
bot.bot.user.mentioned_in = MagicMock(return_value=True)
|
||||
bot.on_message = AsyncMock()
|
||||
bot.newchat.callback = AsyncMock()
|
||||
|
||||
|
||||
@patch('bot.openai.OpenAI')
|
||||
async def test_on_message_text_only(self, MockOpenAI):
|
||||
# Mock the OpenAI client and its response
|
||||
mock_openai_instance = MockOpenAI.return_value
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices[0].message.content = "Hello from Aoi!"
|
||||
mock_openai_instance.chat.completions.create = AsyncMock(return_value=mock_response)
|
||||
|
||||
# Mock a Discord message
|
||||
message = AsyncMock()
|
||||
message.author = MagicMock()
|
||||
message.author.bot = False
|
||||
message.channel = AsyncMock()
|
||||
message.channel.id = 123
|
||||
message.content = f"<@!{bot.bot.user.id}> Hello there"
|
||||
message.attachments = []
|
||||
|
||||
# Call the on_message event handler
|
||||
await bot.on_message(message)
|
||||
|
||||
# Assertions
|
||||
bot.on_message.assert_awaited_once_with(message)
|
||||
|
||||
|
||||
@patch('bot.openai.OpenAI')
|
||||
@patch('bot.aiohttp.ClientSession')
|
||||
async def test_on_message_with_image(self, MockClientSession, MockOpenAI):
|
||||
# Mock the OpenAI client
|
||||
mock_openai_instance = MockOpenAI.return_value
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices[0].message.content = "I see an image!"
|
||||
mock_openai_instance.chat.completions.create = AsyncMock(return_value=mock_response)
|
||||
|
||||
# Mock aiohttp session to simulate image download
|
||||
mock_session = MockClientSession.return_value.__aenter__.return_value
|
||||
mock_resp = mock_session.get.return_value.__aenter__.return_value
|
||||
mock_resp.status = 200
|
||||
mock_resp.read = AsyncMock(return_value=b'fake_image_data')
|
||||
|
||||
# Mock a Discord message with an attachment
|
||||
message = AsyncMock()
|
||||
message.author = MagicMock()
|
||||
message.author.bot = False
|
||||
message.channel = AsyncMock()
|
||||
message.channel.id = 456
|
||||
message.content = f"<@!{bot.bot.user.id}> Look at this!"
|
||||
|
||||
attachment = MagicMock()
|
||||
attachment.content_type = 'image/jpeg'
|
||||
attachment.url = 'http://fakeurl.com/image.jpg'
|
||||
message.attachments = [attachment]
|
||||
|
||||
# Call the on_message event handler
|
||||
await bot.on_message(message)
|
||||
|
||||
# Assertions
|
||||
bot.on_message.assert_awaited_once_with(message)
|
||||
|
||||
async def test_newchat_command_with_prompt(self):
|
||||
# Mock a Discord interaction
|
||||
interaction = AsyncMock()
|
||||
interaction.channel_id = 789
|
||||
prompt = "You are a helpful assistant."
|
||||
|
||||
# Call the newchat command
|
||||
await bot.newchat.callback(interaction, prompt=prompt)
|
||||
|
||||
# Assertions
|
||||
bot.newchat.callback.assert_awaited_once_with(interaction, prompt=prompt)
|
||||
|
||||
async def test_newchat_command_no_prompt(self):
|
||||
# Mock a Discord interaction
|
||||
interaction = AsyncMock()
|
||||
interaction.channel_id = 789
|
||||
|
||||
# Call the newchat command
|
||||
await bot.newchat.callback(interaction, prompt=None)
|
||||
|
||||
# Assertions
|
||||
bot.newchat.callback.assert_awaited_once_with(interaction, prompt=None)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
Reference in New Issue
Block a user