128 lines
5.2 KiB
Python
128 lines
5.2 KiB
Python
import discord
|
|
from discord.ext import commands
|
|
from discord import app_commands
|
|
|
|
# In-memory conversation history for Kasane Teto AI (keyed by channel id)
|
|
_teto_conversations = {}
|
|
|
|
import os
|
|
import aiohttp
|
|
|
|
async def _teto_reply_ai_with_messages(messages, system_mode="reply"):
|
|
"""
|
|
Use OpenRouter AI to generate a Kasane Teto-style response.
|
|
system_mode: "reply" for replying as Kasane Teto.
|
|
"""
|
|
api_key = os.getenv("AI_API_KEY")
|
|
if not api_key:
|
|
raise RuntimeError("AI_API_KEY environment variable not set.")
|
|
url = "https://openrouter.ai/api/v1/chat/completions"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
system_prompt = (
|
|
"You are Kasane Teto, a cheeky and energetic vocaloid/utau character. "
|
|
"Reply to the user in a short, conversational manner, staying in character. "
|
|
"Use Teto's signature phrases like 'desu~' and 'arigatou'. Be playful and a little mischievous."
|
|
)
|
|
payload = {
|
|
"model": "google/gemini-2.0-flash-exp:free",
|
|
"messages": [{"role": "system", "content": system_prompt}] + messages
|
|
}
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(url, headers=headers, json=payload) as resp:
|
|
if resp.content_type == "application/json":
|
|
data = await resp.json()
|
|
return data["choices"][0]["message"]["content"]
|
|
else:
|
|
text = await resp.text()
|
|
raise RuntimeError(f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}")
|
|
|
|
async def _teto_reply_ai(text: str) -> str:
|
|
"""Replies to the text as Kasane Teto using AI via OpenRouter."""
|
|
return await _teto_reply_ai_with_messages([{"role": "user", "content": text}])
|
|
|
|
class TetoCog(commands.Cog):
|
|
def __init__(self, bot: commands.Bot):
|
|
self.bot = bot
|
|
|
|
@commands.Cog.listener()
|
|
async def on_message(self, message: discord.Message):
|
|
if message.author.bot:
|
|
return
|
|
|
|
# Check if the bot was mentioned or replied to
|
|
# Remove all bot mention prefixes from the message content for prefix check
|
|
content_wo_mentions = message.content
|
|
for mention in message.mentions:
|
|
mention_str = f"<@{mention.id}>"
|
|
mention_nick_str = f"<@!{mention.id}>"
|
|
content_wo_mentions = content_wo_mentions.replace(mention_str, "").replace(mention_nick_str, "")
|
|
content_wo_mentions = content_wo_mentions.strip()
|
|
|
|
if (
|
|
self.bot.user in message.mentions
|
|
and not content_wo_mentions.startswith(tuple(self.bot.command_prefix))
|
|
) or (
|
|
message.reference and getattr(message.reference.resolved, "author", None) == self.bot.user
|
|
):
|
|
channel = message.channel
|
|
convo_key = channel.id
|
|
convo = _teto_conversations.get(convo_key, [])
|
|
|
|
# Get message history
|
|
history = []
|
|
async for msg in channel.history(limit=10):
|
|
if msg.content:
|
|
role = "assistant" if msg.author == self.bot.user else "user"
|
|
history.insert(0, {"role": role, "content": msg.content})
|
|
|
|
convo = history
|
|
|
|
convo.append({"role": "user", "content": message.content})
|
|
try:
|
|
ai_reply = await _teto_reply_ai_with_messages(convo)
|
|
await message.reply(ai_reply)
|
|
convo.append({"role": "assistant", "content": ai_reply})
|
|
_teto_conversations[convo_key] = convo[-10:]
|
|
except Exception as e:
|
|
await channel.send(f"Teto AI conversation failed: {e} desu~")
|
|
|
|
@app_commands.context_menu(name="Teto AI Reply")
|
|
async def teto_context_menu_ai_reply(interaction: discord.Interaction, message: discord.Message):
|
|
"""Replies to the selected message as a Teto AI."""
|
|
if not message.content:
|
|
await interaction.response.send_message("The selected message has no text content to reply to! >.<", ephemeral=True)
|
|
return
|
|
|
|
await interaction.response.defer(ephemeral=True)
|
|
channel = interaction.channel
|
|
convo_key = channel.id
|
|
convo = _teto_conversations.get(convo_key, [])
|
|
|
|
# Get message history
|
|
history = []
|
|
async for msg in channel.history(limit=10):
|
|
if msg.content:
|
|
role = "assistant" if msg.author == interaction.client.user else "user"
|
|
history.insert(0, {"role": role, "content": msg.content})
|
|
|
|
convo = history
|
|
|
|
convo.append({"role": "user", "content": message.content})
|
|
try:
|
|
ai_reply = await _teto_reply_ai_with_messages(convo)
|
|
await message.reply(ai_reply)
|
|
await interaction.followup.send("Teto AI replied desu~", ephemeral=True)
|
|
convo.append({"role": "assistant", "content": ai_reply})
|
|
_teto_conversations[convo_key] = convo[-10:]
|
|
except Exception as e:
|
|
await interaction.followup.send(f"Teto AI reply failed: {e} desu~", ephemeral=True)
|
|
|
|
async def setup(bot: commands.Bot):
|
|
cog = TetoCog(bot)
|
|
await bot.add_cog(cog)
|
|
bot.tree.add_command(TetoCog.teto_context_menu_ai_reply)
|
|
print("TetoCog loaded! desu~")
|