From be8689af56393941528506f52abb26b4c8826d67 Mon Sep 17 00:00:00 2001 From: Dory Date: Sat, 16 Aug 2025 21:52:21 -0700 Subject: [PATCH] Initial commit of bot; 100% vibe-coded. --- .gitignore | 122 ++++++++++++++++++++++++++++++++++++++++++++ bot.py | 126 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_bot.py | 112 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 360 insertions(+) create mode 100644 .gitignore create mode 100644 bot.py create mode 100644 tests/test_bot.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3825e91 --- /dev/null +++ b/.gitignore @@ -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/ diff --git a/bot.py b/bot.py new file mode 100644 index 0000000..1c693f4 --- /dev/null +++ b/bot.py @@ -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) diff --git a/tests/test_bot.py b/tests/test_bot.py new file mode 100644 index 0000000..772229e --- /dev/null +++ b/tests/test_bot.py @@ -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()