384 lines
15 KiB
Python
384 lines
15 KiB
Python
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 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://api.greenpt.ai/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,
|
|
max_tokens=None,
|
|
)
|
|
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
|
|
|
|
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 not reply.strip():
|
|
await message.reply("*(The assistant returned an empty response.)*")
|
|
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()
|
|
max_tokens = await self.config.guild(guild).max_tokens()
|
|
|
|
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,
|
|
}
|
|
|
|
if max_tokens:
|
|
payload["max_tokens"] = max_tokens
|
|
|
|
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:
|
|
choice = data["choices"][0]["message"]
|
|
reply = choice.get("content") or ""
|
|
except (KeyError, IndexError):
|
|
return "Unexpected API response."
|
|
|
|
if not reply.strip():
|
|
return "The model returned an empty response. Try rephrasing or switching models."
|
|
|
|
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="models")
|
|
async def assistant_models(self, ctx: commands.Context):
|
|
"""List available models sorted by cost (cheapest first)."""
|
|
# Pricing per 1M tokens (EUR) from GreenPT docs
|
|
MODEL_PRICING = {
|
|
"deepseek-v4-flash-0731": {"input": 0.14, "output": 0.35},
|
|
"minimax-m2.5": {"input": 0.17, "output": 0.99},
|
|
"mistral-small-3.2-24b-instruct-2506": {"input": 0.20, "output": 0.40},
|
|
"gpt-oss-120b": {"input": 0.20, "output": 0.70},
|
|
"voxtral-small-24b-2507": {"input": 0.20, "output": 0.45},
|
|
"green-l": {"input": 0.25, "output": 0.80},
|
|
"green-l-raw": {"input": 0.25, "output": 0.80},
|
|
"qwen3-coder-30b-a3b-instruct": {"input": 0.25, "output": 0.95},
|
|
"gemma-3-27b-it": {"input": 0.30, "output": 0.60},
|
|
"qwen3.6-35b-a3b": {"input": 0.30, "output": 1.80},
|
|
"green-r": {"input": 0.35, "output": 0.95},
|
|
"green-r-raw": {"input": 0.35, "output": 0.95},
|
|
"gemma4": {"input": 0.50, "output": 1.50},
|
|
"devstral-2-123b-instruct-2512": {"input": 0.50, "output": 2.40},
|
|
"kimi-k2.6": {"input": 0.66, "output": 3.75},
|
|
"qwen3.5-397b-a17b": {"input": 0.70, "output": 4.35},
|
|
"kimi-k2.7-code": {"input": 0.79, "output": 3.85},
|
|
"qwen3-235b-a22b-instruct-2507": {"input": 0.90, "output": 2.70},
|
|
"glm-5.2": {"input": 1.10, "output": 4.40},
|
|
"llama-3.3-70b-instruct": {"input": 1.10, "output": 1.10},
|
|
"mistral-medium-3.5-128b": {"input": 1.80, "output": 9.00},
|
|
"kimi-k3": {"input": 3.30, "output": 16.50},
|
|
}
|
|
|
|
api_key = await self.config.api_key()
|
|
if not api_key or not self.session:
|
|
await ctx.send("API key not set. Use `[p]assistant setkey` first.")
|
|
return
|
|
|
|
# Fetch live model list from API
|
|
try:
|
|
async with self.session.get(
|
|
f"{self.BASE_URL}/models",
|
|
headers={"Authorization": f"Bearer {api_key}"},
|
|
timeout=aiohttp.ClientTimeout(total=15),
|
|
) as resp:
|
|
if resp.status != 200:
|
|
await ctx.send(f"Failed to fetch models (HTTP {resp.status}).")
|
|
return
|
|
data = await resp.json()
|
|
except Exception as exc:
|
|
await ctx.send(f"Request failed: {exc}")
|
|
return
|
|
|
|
models = data.get("data", [])
|
|
if not models:
|
|
await ctx.send("No models returned by the API.")
|
|
return
|
|
|
|
# Build list with pricing, sorted by combined cost (input + output)
|
|
current_model = await self.config.guild(ctx.guild).model()
|
|
lines = []
|
|
model_entries = []
|
|
for m in models:
|
|
model_id = m["id"]
|
|
pricing = MODEL_PRICING.get(model_id)
|
|
if pricing:
|
|
total = pricing["input"] + pricing["output"]
|
|
model_entries.append((total, model_id, pricing))
|
|
else:
|
|
# Unknown pricing, put at the end
|
|
model_entries.append((9999, model_id, None))
|
|
|
|
model_entries.sort(key=lambda x: x[0])
|
|
|
|
for total, model_id, pricing in model_entries:
|
|
marker = " ◀" if model_id == current_model else ""
|
|
if pricing:
|
|
lines.append(
|
|
f"`{model_id}` — €{pricing['input']:.2f} in / €{pricing['output']:.2f} out{marker}"
|
|
)
|
|
else:
|
|
lines.append(f"`{model_id}` — pricing unknown{marker}")
|
|
|
|
embed = discord.Embed(
|
|
title="Available Models (cheapest first)",
|
|
description="\n".join(lines),
|
|
color=await ctx.embed_color(),
|
|
)
|
|
embed.set_footer(text="Prices in EUR per 1M tokens. ◀ = current model.")
|
|
await ctx.send(embed=embed)
|
|
|
|
@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="maxlength")
|
|
async def assistant_maxlength(self, ctx: commands.Context, tokens: int = None):
|
|
"""Set max response length in tokens. Omit to clear (unlimited)."""
|
|
if tokens is None:
|
|
await self.config.guild(ctx.guild).max_tokens.set(None)
|
|
await ctx.send("Max response length cleared (unlimited).")
|
|
elif tokens < 1:
|
|
await ctx.send("Must be a positive number.")
|
|
else:
|
|
await self.config.guild(ctx.guild).max_tokens.set(tokens)
|
|
await ctx.send(f"Max response length set to **{tokens}** tokens.")
|
|
|
|
@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()
|
|
max_tokens = await cfg.max_tokens()
|
|
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="Max Tokens", value=str(max_tokens) if max_tokens else "Unlimited", inline=True)
|
|
embed.add_field(name="System Prompt", value=sys_prompt[:200], inline=False)
|
|
|
|
await ctx.send(embed=embed)
|