282 lines
11 KiB
Python
282 lines
11 KiB
Python
import logging
|
|
import aiohttp
|
|
import discord
|
|
from redbot.core import commands, Config, checks
|
|
from redbot.core.bot import Red
|
|
|
|
log = logging.getLogger("red.assistant")
|
|
|
|
|
|
class APIKeyModal(discord.ui.Modal):
|
|
def __init__(self, config):
|
|
super().__init__(title="API Key")
|
|
self.config = config
|
|
self.key_input = discord.ui.TextInput(
|
|
label="API Key",
|
|
style=discord.TextStyle.short,
|
|
placeholder="Paste your API key here",
|
|
required=True,
|
|
)
|
|
self.add_item(self.key_input)
|
|
|
|
async def on_submit(self, interaction: discord.Interaction):
|
|
await self.config.api_key.set(self.key_input.value)
|
|
await interaction.response.send_message("API key saved.", ephemeral=True)
|
|
|
|
|
|
class APIKeyButton(discord.ui.View):
|
|
def __init__(self, config):
|
|
super().__init__(timeout=60)
|
|
self.config = config
|
|
|
|
@discord.ui.button(label="Set API Key", style=discord.ButtonStyle.primary)
|
|
async def set_key(self, interaction: discord.Interaction, button: discord.ui.Button):
|
|
await interaction.response.send_modal(APIKeyModal(self.config))
|
|
self.stop()
|
|
|
|
|
|
class Assistant(commands.Cog):
|
|
"""AI chatbot assistant. Set a chat channel or get responses when the bot is pinged."""
|
|
|
|
BASE_URL = "https://greenpt.com/v1"
|
|
|
|
def __init__(self, bot: Red):
|
|
self.bot = bot
|
|
self.session: aiohttp.ClientSession | None = None
|
|
self.config = Config.get_conf(self, identifier=987654321)
|
|
self.config.register_guild(
|
|
enabled=True,
|
|
chat_channel=None,
|
|
model="gemma4",
|
|
temperature=0.7,
|
|
system_prompt="You are a helpful assistant.",
|
|
question_mode=False,
|
|
)
|
|
self.config.register_global(api_key=None)
|
|
self.history: dict[int, list[dict]] = {}
|
|
|
|
async def cog_load(self):
|
|
self.session = aiohttp.ClientSession()
|
|
|
|
async def cog_unload(self):
|
|
if self.session:
|
|
await self.session.close()
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Listeners
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@commands.Cog.listener()
|
|
async def on_message(self, message: discord.Message):
|
|
if message.author.bot:
|
|
return
|
|
if not message.guild:
|
|
return
|
|
|
|
guild = message.guild
|
|
|
|
if self.bot.user in message.mentions:
|
|
log.info("[ASSISTANT DEBUG] Mention detected from %s in #%s", message.author, message.channel)
|
|
|
|
enabled = await self.config.guild(guild).enabled()
|
|
if not enabled:
|
|
if self.bot.user in message.mentions:
|
|
log.info("[ASSISTANT DEBUG] Guild %s is disabled", guild.id)
|
|
return
|
|
|
|
api_key = await self.config.api_key()
|
|
if not api_key:
|
|
if self.bot.user in message.mentions:
|
|
log.info("[ASSISTANT DEBUG] No API key set")
|
|
return
|
|
|
|
chat_channel_id = await self.config.guild(guild).chat_channel()
|
|
question_mode = await self.config.guild(guild).question_mode()
|
|
|
|
is_chat_channel = chat_channel_id and message.channel.id == chat_channel_id
|
|
is_ping = self.bot.user in message.mentions
|
|
|
|
if not is_chat_channel and not is_ping:
|
|
return
|
|
|
|
log.info("[ASSISTANT DEBUG] Processing message, is_ping=%s, is_chat_channel=%s", is_ping, is_chat_channel)
|
|
|
|
content = message.content
|
|
if is_ping:
|
|
for mention in (f"<@{self.bot.user.id}>", f"<@!{self.bot.user.id}>"):
|
|
content = content.replace(mention, "").strip()
|
|
|
|
if not content:
|
|
return
|
|
|
|
if question_mode and not content.rstrip().endswith("?"):
|
|
if is_ping:
|
|
await message.reply("Question mode is on — messages must end with **?**.")
|
|
return
|
|
|
|
async with message.channel.typing():
|
|
reply = await self._chat(guild, message.channel.id, content)
|
|
|
|
if reply is None:
|
|
await message.reply("No response from the assistant.")
|
|
return
|
|
|
|
if len(reply) <= 2000:
|
|
await message.reply(reply)
|
|
else:
|
|
chunks = [reply[i : i + 2000] for i in range(0, len(reply), 2000)]
|
|
for chunk in chunks:
|
|
await message.reply(chunk)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API
|
|
# ---------------------------------------------------------------------------
|
|
|
|
async def _chat(self, guild: discord.Guild, channel_id: int, user_msg: str) -> str | None:
|
|
api_key = await self.config.api_key()
|
|
if not api_key or not self.session:
|
|
return None
|
|
|
|
model = await self.config.guild(guild).model()
|
|
temperature = await self.config.guild(guild).temperature()
|
|
system_prompt = await self.config.guild(guild).system_prompt()
|
|
|
|
key = (guild.id, channel_id)
|
|
if key not in self.history:
|
|
self.history[key] = []
|
|
|
|
self.history[key].append({"role": "user", "content": user_msg})
|
|
|
|
messages = [{"role": "system", "content": system_prompt}] + self.history[key][-20:]
|
|
|
|
payload = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
}
|
|
|
|
try:
|
|
async with self.session.post(
|
|
f"{self.BASE_URL}/chat/completions",
|
|
json=payload,
|
|
headers={
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
timeout=aiohttp.ClientTimeout(total=60),
|
|
) as resp:
|
|
if resp.status != 200:
|
|
body = await resp.text()
|
|
return f"API error {resp.status}: {body[:300]}"
|
|
data = await resp.json()
|
|
except Exception as exc:
|
|
return f"Request failed: {exc}"
|
|
|
|
try:
|
|
reply = data["choices"][0]["message"]["content"]
|
|
except (KeyError, IndexError):
|
|
return "Unexpected API response."
|
|
|
|
self.history[key].append({"role": "assistant", "content": reply})
|
|
|
|
if len(self.history[key]) > 40:
|
|
self.history[key] = self.history[key][-20:]
|
|
|
|
return reply
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Commands
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@commands.group(name="assistant", invoke_without_command=True)
|
|
@commands.guild_only()
|
|
async def assistant(self, ctx: commands.Context):
|
|
"""AI assistant settings."""
|
|
await ctx.send_help(ctx.command)
|
|
|
|
@assistant.command(name="setkey")
|
|
@checks.is_owner()
|
|
async def assistant_setkey(self, ctx: commands.Context):
|
|
"""Set the API key (opens a private modal)."""
|
|
view = APIKeyButton(self.config)
|
|
await ctx.send("Click below to set the API key:", view=view)
|
|
|
|
@assistant.command(name="toggle")
|
|
async def assistant_toggle(self, ctx: commands.Context):
|
|
"""Toggle the assistant on or off for this server."""
|
|
current = await self.config.guild(ctx.guild).enabled()
|
|
await self.config.guild(ctx.guild).enabled.set(not current)
|
|
state = "enabled" if not current else "disabled"
|
|
await ctx.send(f"Assistant **{state}** for this server.")
|
|
|
|
@assistant.command(name="channel")
|
|
async def assistant_channel(self, ctx: commands.Context, channel: discord.TextChannel | None = None):
|
|
"""Set the dedicated chat channel. Omit to clear."""
|
|
if channel is None:
|
|
await self.config.guild(ctx.guild).chat_channel.set(None)
|
|
await ctx.send("Chat channel cleared.")
|
|
else:
|
|
await self.config.guild(ctx.guild).chat_channel.set(channel.id)
|
|
await ctx.send(f"Chat channel set to {channel.mention}.")
|
|
|
|
@assistant.command(name="model")
|
|
async def assistant_model(self, ctx: commands.Context, model: str):
|
|
"""Set the model (e.g. gemma4, green-l, green-r)."""
|
|
await self.config.guild(ctx.guild).model.set(model)
|
|
await ctx.send(f"Model set to **{model}**.")
|
|
|
|
@assistant.command(name="temp")
|
|
async def assistant_temp(self, ctx: commands.Context, temperature: float):
|
|
"""Set the temperature (0.0 - 2.0)."""
|
|
if not 0.0 <= temperature <= 2.0:
|
|
await ctx.send("Must be between 0.0 and 2.0.")
|
|
return
|
|
await self.config.guild(ctx.guild).temperature.set(temperature)
|
|
await ctx.send(f"Temperature set to **{temperature}**.")
|
|
|
|
@assistant.command(name="system")
|
|
async def assistant_system(self, ctx: commands.Context, *, prompt: str):
|
|
"""Set the system prompt."""
|
|
await self.config.guild(ctx.guild).system_prompt.set(prompt)
|
|
await ctx.send("System prompt updated.")
|
|
|
|
@assistant.command(name="question")
|
|
async def assistant_question(self, ctx: commands.Context):
|
|
"""Toggle question mode (messages must end with ?)."""
|
|
current = await self.config.guild(ctx.guild).question_mode()
|
|
await self.config.guild(ctx.guild).question_mode.set(not current)
|
|
state = "ON" if not current else "OFF"
|
|
await ctx.send(f"Question mode **{state}**.")
|
|
|
|
@assistant.command(name="clear")
|
|
async def assistant_clear(self, ctx: commands.Context):
|
|
"""Clear conversation history for this channel."""
|
|
key = (ctx.guild.id, ctx.channel.id)
|
|
self.history.pop(key, None)
|
|
await ctx.send("History cleared for this channel.")
|
|
|
|
@assistant.command(name="settings")
|
|
async def assistant_settings(self, ctx: commands.Context):
|
|
"""Show current settings."""
|
|
cfg = self.config.guild(ctx.guild)
|
|
enabled = await cfg.enabled()
|
|
chat_channel = await cfg.chat_channel()
|
|
model = await cfg.model()
|
|
temp = await cfg.temperature()
|
|
sys_prompt = await cfg.system_prompt()
|
|
q_mode = await cfg.question_mode()
|
|
has_key = await self.config.api_key() is not None
|
|
|
|
channel_str = f"<#{chat_channel}>" if chat_channel else "None"
|
|
|
|
embed = discord.Embed(title="Assistant Settings", color=await ctx.embed_color())
|
|
embed.add_field(name="Enabled", value=str(enabled), inline=True)
|
|
embed.add_field(name="Chat Channel", value=channel_str, inline=True)
|
|
embed.add_field(name="API Key", value="Set" if has_key else "Not set", inline=True)
|
|
embed.add_field(name="Model", value=model, inline=True)
|
|
embed.add_field(name="Temperature", value=str(temp), inline=True)
|
|
embed.add_field(name="Question Mode", value=str(q_mode), inline=True)
|
|
embed.add_field(name="System Prompt", value=sys_prompt[:200], inline=False)
|
|
|
|
await ctx.send(embed=embed)
|