Compare commits
3 Commits
dd47cc05a6
...
fd08420f26
Author | SHA1 | Date | |
---|---|---|---|
fd08420f26 | |||
ad9c069993 | |||
8c750b2fbb |
47
bot.py
47
bot.py
@@ -6,6 +6,7 @@ import argparse
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from llm_client import Conversation
|
||||
from database import Database
|
||||
|
||||
# --- Configuration ---
|
||||
DEFAULT_AVATAR = "https://cdn.discordapp.com/avatars/1406466525858369716/f1dfeaf2a1c361dbf981e2e899c7f981?size=256"
|
||||
@@ -23,10 +24,15 @@ parser.add_argument(
|
||||
args = parser.parse_args()
|
||||
|
||||
# --- Bot Setup ---
|
||||
class AoiBot(commands.Bot):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
async def setup_hook(self):
|
||||
self.db = Database.get()
|
||||
intents = discord.Intents.default()
|
||||
intents.messages = True
|
||||
intents.message_content = True
|
||||
bot = commands.Bot(command_prefix="/", intents=intents)
|
||||
bot = AoiBot(command_prefix="/", intents=intents)
|
||||
|
||||
|
||||
# --- Helpers ---
|
||||
@@ -75,12 +81,12 @@ async def on_message(message):
|
||||
|
||||
bot_tag = f'<@{bot.user.id}>'
|
||||
channel = message.channel
|
||||
conversation = await Conversation.get(channel.id, args.base_url)
|
||||
conversation = await Conversation.get(channel.id, args.base_url, bot.db)
|
||||
user_message = message.content
|
||||
if user_message.startswith(bot_tag):
|
||||
user_message = user_message[len(bot_tag):]
|
||||
user_message = user_message.replace(bot_tag, conversation.bot_name).strip()
|
||||
print(f'> {message.author.name}: {user_message}')
|
||||
print(f'{channel.id}> {message.author.name}: {user_message}')
|
||||
|
||||
media = []
|
||||
if message.attachments:
|
||||
@@ -91,13 +97,16 @@ async def on_message(message):
|
||||
async with channel.typing():
|
||||
response = await conversation.generate(user_message, media)
|
||||
for old_message_id in conversation.last_messages:
|
||||
old_message = await channel.fetch_message(old_message_id)
|
||||
await old_message.clear_reaction("🔁")
|
||||
await old_message.clear_reaction("❌")
|
||||
try:
|
||||
old_message = await channel.fetch_message(old_message_id)
|
||||
await old_message.clear_reaction("🔁")
|
||||
await old_message.clear_reaction("❌")
|
||||
except (discord.NotFound, discord.Forbidden):
|
||||
pass # Ignore if message is not found or we don't have perms
|
||||
conversation.last_messages = await discord_send(
|
||||
channel, response, conversation.bot_name,
|
||||
)
|
||||
conversation.save()
|
||||
await conversation.save()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
await message.reply("Sorry, I had a little hiccup. Baka!")
|
||||
@@ -108,7 +117,7 @@ async def on_reaction_add(reaction, user):
|
||||
return
|
||||
message = reaction.message
|
||||
channel = message.channel
|
||||
conversation = await Conversation.get(channel.id, args.base_url)
|
||||
conversation = await Conversation.get(channel.id, args.base_url, bot.db)
|
||||
if message.id not in conversation.last_messages:
|
||||
await reaction.clear()
|
||||
return
|
||||
@@ -124,17 +133,18 @@ async def on_reaction_add(reaction, user):
|
||||
except (discord.NotFound, discord.Forbidden) as e:
|
||||
# don't do anything if any message in the list is not found
|
||||
await reaction.clear()
|
||||
return
|
||||
for message in messages:
|
||||
await message.delete()
|
||||
|
||||
if reaction.emoji == "❌":
|
||||
conversation.pop()
|
||||
await conversation.pop()
|
||||
elif reaction.emoji == "🔁":
|
||||
response = await conversation.regenerate()
|
||||
conversation.last_messages = await discord_send(
|
||||
channel, response, conversation.bot_name,
|
||||
)
|
||||
conversation.save()
|
||||
await conversation.save()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
await message.reply("Sorry, I had a little hiccup. Baka!")
|
||||
@@ -148,11 +158,26 @@ async def on_reaction_add(reaction, user):
|
||||
async def newchat(interaction: discord.Interaction, prompt: str = None):
|
||||
await interaction.response.defer()
|
||||
channel_id = interaction.channel_id
|
||||
conversation = await Conversation.create(channel_id, args.base_url, prompt)
|
||||
conversation = await Conversation.create(
|
||||
channel_id, args.base_url, bot.db, prompt
|
||||
)
|
||||
await interaction.followup.send(
|
||||
f'Starting a new chat with {conversation.bot_name}: "{prompt}"'
|
||||
)
|
||||
|
||||
@bot.tree.command(
|
||||
name="changeprompt",
|
||||
description="Change the current chat's system prompt."
|
||||
)
|
||||
async def changeprompt(interaction: discord.Interaction, prompt: str):
|
||||
await interaction.response.defer()
|
||||
channel_id = interaction.channel_id
|
||||
conversation = await Conversation.get(channel_id, args.base_url, bot.db)
|
||||
await conversation.update_prompt(prompt)
|
||||
await interaction.followup.send(
|
||||
f'Now chatting with {conversation.bot_name}: "{prompt}"'
|
||||
)
|
||||
|
||||
|
||||
# --- Running the Bot ---
|
||||
if __name__ == "__main__":
|
||||
|
18
database.py
18
database.py
@@ -2,12 +2,17 @@ import sqlite3
|
||||
import json
|
||||
|
||||
class Database:
|
||||
def __init__(self, db_path='conversations.db'):
|
||||
self.db_path = db_path
|
||||
self.conn = sqlite3.connect(self.db_path)
|
||||
self.create_table()
|
||||
def __init__(self, db_conn):
|
||||
self.conn = db_conn
|
||||
self._create_table()
|
||||
|
||||
def create_table(self):
|
||||
@classmethod
|
||||
def get(cls, db_path='conversations.db'):
|
||||
"""Creates and returns a connected Database instance."""
|
||||
print(f"Initializing DB connection to: {db_path}")
|
||||
return Database(sqlite3.connect(db_path))
|
||||
|
||||
def _create_table(self):
|
||||
with self.conn:
|
||||
self.conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
@@ -18,7 +23,7 @@ class Database:
|
||||
)
|
||||
""")
|
||||
|
||||
def get(self, conversation_id):
|
||||
def get_conversation(self, conversation_id):
|
||||
with self.conn:
|
||||
cursor = self.conn.cursor()
|
||||
cursor.execute("SELECT history, bot_name, last_messages FROM conversations WHERE id = ?", (conversation_id,))
|
||||
@@ -41,4 +46,3 @@ class Database:
|
||||
with self.conn:
|
||||
self.conn.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
|
||||
|
||||
db = Database()
|
||||
|
@@ -1,7 +1,7 @@
|
||||
import aiohttp
|
||||
import base64
|
||||
from openai import AsyncOpenAI
|
||||
from database import db
|
||||
from database import Database
|
||||
|
||||
API_KEY = "eh"
|
||||
MODEL = "p620"
|
||||
@@ -13,12 +13,13 @@ NAME_PROMPT = "reply with your name, nothing else, no punctuation"
|
||||
|
||||
|
||||
class Conversation:
|
||||
def __init__(self, client, name, prompt, convo_id):
|
||||
def __init__(self, client, name, prompt, convo_id, db):
|
||||
self.history = [{"role": "system", "content": prompt}]
|
||||
self.bot_name = name
|
||||
self.last_messages = []
|
||||
self.client = client
|
||||
self.id = convo_id
|
||||
self.db = db
|
||||
|
||||
def __str__(self):
|
||||
return (
|
||||
@@ -26,29 +27,32 @@ class Conversation:
|
||||
f"{self.history}"
|
||||
)
|
||||
|
||||
def save(self):
|
||||
db.save(self.id, self.history, self.bot_name, self.last_messages)
|
||||
async def save(self):
|
||||
self.db.save(
|
||||
self.id, self.history, self.bot_name, self.last_messages
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def get(cls, key, base_url):
|
||||
convo_data = db.get(key)
|
||||
async def get(cls, key, base_url, db):
|
||||
convo_data = db.get_conversation(key)
|
||||
if convo_data:
|
||||
history, bot_name, last_messages = convo_data
|
||||
client = AsyncOpenAI(base_url=base_url, api_key=API_KEY)
|
||||
convo = cls(client, bot_name, history[0]['content'], key)
|
||||
convo = cls(client, bot_name, history[0]['content'], key, db)
|
||||
convo.history = history
|
||||
convo.last_messages = last_messages
|
||||
return convo
|
||||
return await Conversation.create(key, base_url)
|
||||
return await Conversation.create(key, base_url, db)
|
||||
|
||||
@classmethod
|
||||
async def create(cls, key, base_url, prompt=None):
|
||||
async def create(cls, key, base_url, db, prompt=None):
|
||||
client = AsyncOpenAI(base_url=base_url, api_key=API_KEY)
|
||||
if not prompt:
|
||||
convo = cls(client, DEFAULT_NAME, DEFAULT_SYSTEM_PROMPT, key)
|
||||
convo = cls(client, DEFAULT_NAME, DEFAULT_SYSTEM_PROMPT, key, db)
|
||||
else:
|
||||
convo = cls(client, await cls.get_name(client, prompt), prompt, key)
|
||||
convo.save()
|
||||
name = await cls.get_name(client, prompt)
|
||||
convo = cls(client, name, prompt, key, db)
|
||||
await convo.save()
|
||||
return convo
|
||||
|
||||
@classmethod
|
||||
@@ -68,10 +72,15 @@ class Conversation:
|
||||
{"role": "assistant", "content": assistant},
|
||||
])
|
||||
|
||||
def pop(self):
|
||||
async def pop(self):
|
||||
if len(self.history) >= 3:
|
||||
self.history = self.history[:-2]
|
||||
self.save()
|
||||
await self.save()
|
||||
|
||||
async def update_prompt(self, prompt):
|
||||
self.history[0] = {"role": "system", "content": prompt}
|
||||
self.bot_name = await self.get_name(self.client, prompt)
|
||||
await self.save()
|
||||
|
||||
async def generate(self, text, media=tuple()):
|
||||
# prepare text part
|
||||
|
Reference in New Issue
Block a user