import aiohttp import discord from redbot.core import commands, Config, checks from redbot.core.bot import Red 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 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 isinstance(message.channel, discord.TextChannel): return guild = message.guild if guild is None: return enabled = await self.config.guild(guild).enabled() if not enabled: return api_key = await self.config.api_key() if not api_key: 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 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).""" await ctx.send_modal(APIKeyModal(self.config)) @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)