From e1455e655fecdd069c298c358bf6493de8d95a8e Mon Sep 17 00:00:00 2001 From: perforat Date: Wed, 20 May 2026 12:55:38 +0000 Subject: [PATCH] openrouter client --- cogs/openrouter/config.json | 29 +++ cogs/openrouter/openrouter.py | 341 ++++++++++++++++++++++++++++++++++ 2 files changed, 370 insertions(+) create mode 100644 cogs/openrouter/config.json create mode 100644 cogs/openrouter/openrouter.py diff --git a/cogs/openrouter/config.json b/cogs/openrouter/config.json new file mode 100644 index 0000000..f776129 --- /dev/null +++ b/cogs/openrouter/config.json @@ -0,0 +1,29 @@ +{ + "description": "OpenRouter AI client", + "usages": [ + "ai ", + "ai_reset", + "ai_history", + "ai_stats", + "ai_info", + "ai_setmodel (admin)", + "ai_setprompt (admin)", + "ai_setchannel [#channel] (admin)", + "ai_removechannel <#channel> (admin)", + "ai_clearlimit <@user> (admin)", + "ai_globalstats (admin)" + ], + "hidden": false, + + "OPENROUTER_API_KEY": "YOUR_KEY_HERE", + "MODEL": "openai/gpt-4o-mini", + "SYSTEM_PROMPT": "You are a helpful assistant in a Discord server. Be concise and friendly.", + + "MAX_REQUESTS_PER_DAY": 50, + "MAX_TOKENS_PER_USER_PER_DAY": 10000, + "MAX_RESPONSE_TOKENS": 1000, + "MAX_HISTORY_LENGTH": 10, + "COOLDOWN_SECONDS": 10, + + "ALLOWED_CHANNEL_IDS": [] +} \ No newline at end of file diff --git a/cogs/openrouter/openrouter.py b/cogs/openrouter/openrouter.py new file mode 100644 index 0000000..0501209 --- /dev/null +++ b/cogs/openrouter/openrouter.py @@ -0,0 +1,341 @@ +import discord +from discord.ext import commands +import logging +import asyncio +import json +import os +import aiohttp +from datetime import datetime, date + +logger = logging.getLogger(__name__) + + +class AIClient(commands.Cog): + def __init__(self, bot): + self.bot = bot + self.config_path = os.path.join(os.path.dirname(__file__), "config.json") + self.usage_path = os.path.join(os.path.dirname(__file__), "usage.json") + self.config = self._load_config() + self.usage = self._load_usage() + + # user_id (int) -> list of {"role": ..., "content": ...} + self.conversations: dict[int, list[dict]] = {} + # bot message_id -> user_id that owns that conversation thread + self.message_map: dict[int, int] = {} + # user_id -> datetime of last request (for cooldown) + self.last_used: dict[int, datetime] = {} + + # confeeaggeeee + + def _load_config(self) -> dict: + with open(self.config_path, "r", encoding="utf-8") as f: + return json.load(f) + + def _save_config(self) -> None: + with open(self.config_path, "w", encoding="utf-8") as f: + json.dump(self.config, f, indent=4, ensure_ascii=False) + + def _load_usage(self) -> dict: + if os.path.exists(self.usage_path): + with open(self.usage_path, "r", encoding="utf-8") as f: + return json.load(f) + return {} + + def _save_usage(self) -> None: + with open(self.usage_path, "w", encoding="utf-8") as f: + json.dump(self.usage, f, indent=4) + + # helpers + + def _get_user_usage(self, user_id: int) -> dict: + """Return today's usage record for user, resetting if the date changed.""" + uid = str(user_id) + today = str(date.today()) + if uid not in self.usage or self.usage[uid].get("date") != today: + self.usage[uid] = {"date": today, "requests": 0, "tokens": 0} + return self.usage[uid] + + def _check_channel(self, channel_id: int) -> bool: + allowed = self.config.get("ALLOWED_CHANNEL_IDS", []) + return (not allowed) or (channel_id in allowed) + + def _cooldown_remaining(self, user_id: int) -> float: + cooldown = self.config.get("COOLDOWN_SECONDS", 10) + last = self.last_used.get(user_id) + if last is None: + return 0.0 + elapsed = (datetime.now() - last).total_seconds() + remaining = cooldown - elapsed + return max(remaining, 0.0) + + def _trim_history(self, user_id: int) -> None: + max_pairs = self.config.get("MAX_HISTORY_LENGTH", 10) + history = self.conversations.get(user_id, []) + cap = max_pairs * 2 + if len(history) > cap: + self.conversations[user_id] = history[-cap:] + + def _register_bot_message(self, message_id: int, user_id: int) -> None: + self.message_map[message_id] = user_id + if len(self.message_map) > 1000: + for old_id in list(self.message_map.keys())[:200]: + del self.message_map[old_id] + + + async def _call_openrouter(self, messages: list[dict]) -> tuple[str, int]: + # call api + api_key = self.config.get("OPENROUTER_API_KEY", "") + model = self.config.get("MODEL", "openai/gpt-4o-mini") + max_tokens = self.config.get("MAX_RESPONSE_TOKENS", 1000) + + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + "HTTP-Referer": "https://discord.bot", + "X-Title": "Discord Bot", + } + payload = { + "model": model, + "messages": messages, + "max_tokens": max_tokens, + } + + async with aiohttp.ClientSession() as session: + async with session.post( + "https://openrouter.ai/api/v1/chat/completions", + headers=headers, + json=payload, + timeout=aiohttp.ClientTimeout(total=60), + ) as resp: + if resp.status != 200: + body = await resp.text() + raise RuntimeError(f"HTTP {resp.status}: {body[:300]}") + data = await resp.json() + + text = data["choices"][0]["message"]["content"] + tokens = data.get("usage", {}).get("total_tokens", 0) + return text, tokens + + async def _process_request( + self, + message: discord.Message, + user_input: str, + ) -> None: + # logic for !ai + author = message.author + channel = message.channel + + if not self._check_channel(channel.id): + allowed_ids = self.config.get("ALLOWED_CHANNEL_IDS", []) + mention_list = ", ".join(f"<#{cid}>" for cid in allowed_ids) + await channel.send( + f"Ai is only available in: {mention_list}", + delete_after=10, + ) + return + + remaining = self._cooldown_remaining(author.id) + if remaining > 0: + await channel.send( + f"Hit cooldown. Wait for **{remaining:.1f}s**.", + delete_after=10, + ) + return + + usage = self._get_user_usage(author.id) + max_req = self.config.get("MAX_REQUESTS_PER_DAY", 50) + max_tok = self.config.get("MAX_TOKENS_PER_USER_PER_DAY", 10000) + + if usage["requests"] >= max_req: + await channel.send( + f"{author.mention}, daily limit hit " + f"({max_req}). Try tommorow.", + delete_after=15, + ) + return + + if usage["tokens"] >= max_tok: + await channel.send( + f"{author.mention}, daily token limit hit " + f"({max_tok}). Try tommorow.", + delete_after=15, + ) + return + + if author.id not in self.conversations: + self.conversations[author.id] = [] + + self.conversations[author.id].append({"role": "user", "content": user_input}) + self._trim_history(author.id) + + system_prompt = self.config.get( + "SYSTEM_PROMPT", "You are a helpful assistant in a Discord server." + ) + api_messages = [{"role": "system", "content": system_prompt}] + self.conversations[author.id] + + self.last_used[author.id] = datetime.now() + + async with channel.typing(): + try: + response_text, tokens_used = await self._call_openrouter(api_messages) + except Exception as exc: + logger.error("OpenRouter error: %s", exc) + self.conversations[author.id].pop() + await channel.send(f"err: `{exc}`") + return + + self.conversations[author.id].append( + {"role": "assistant", "content": response_text} + ) + + usage["requests"] += 1 + usage["tokens"] += tokens_used + self._save_usage() + + chunks = [response_text[i: i + 1990] for i in range(0, len(response_text), 1990)] + last_sent: discord.Message | None = None + for idx, chunk in enumerate(chunks): + if idx == 0: + last_sent = await message.reply(chunk) + else: + last_sent = await channel.send(chunk) + + if last_sent: + self._register_bot_message(last_sent.id, author.id) + + # cmds + + @commands.command(name="ai") + async def ai(self, ctx: commands.Context, *, prompt: str): + # i really wonder what this does + await self._process_request(ctx.message, prompt) + + @commands.command(name="ai_reset") + async def ai_reset(self, ctx: commands.Context): + # clear your personal conversation history with the ai + self.conversations.pop(ctx.author.id, None) + await ctx.reply("Reset personal conversation history.") + + #!! admin cmds + + @commands.command(name="ai_setchannel") + @commands.has_permissions(administrator=True) + async def ai_setchannel( + self, ctx: commands.Context, channel: discord.TextChannel = None + ): + # whitelist a channel for ai commands, or clear the whitelist if no channel is given + if channel is None: + self.config["ALLOWED_CHANNEL_IDS"] = [] + self._save_config() + await ctx.reply("Whitelist is not active") + return + + ids: list[int] = self.config.setdefault("ALLOWED_CHANNEL_IDS", []) + if channel.id in ids: + await ctx.reply(f"{channel.mention} is already in whitelist") + else: + ids.append(channel.id) + self._save_config() + await ctx.reply(f"Added {channel.mention} into whitelist.") + + @commands.command(name="ai_removechannel") + @commands.has_permissions(administrator=True) + async def ai_removechannel(self, ctx: commands.Context, channel: discord.TextChannel): + # remove channel from whitelist + ids: list[int] = self.config.get("ALLOWED_CHANNEL_IDS", []) + if channel.id not in ids: + await ctx.reply(f"{channel.mention} is not in whitelist") + else: + ids.remove(channel.id) + self._save_config() + await ctx.reply(f"Removed {channel.mention} from whitelist") + + @commands.command(name="ai_clearlimit") + @commands.has_permissions(administrator=True) + async def ai_clearlimit(self, ctx: commands.Context, member: discord.Member): + # reset limit for a specific member + uid = str(member.id) + if uid in self.usage: + del self.usage[uid] + self._save_usage() + await ctx.reply(f"Reset limits for {member.mention}") + + @commands.command(name="ai_globalstats") + @commands.has_permissions(administrator=True) + async def ai_globalstats(self, ctx: commands.Context): + # todays usage + today = str(date.today()) + lines: list[str] = [] + total_req = 0 + total_tok = 0 + + for uid, data in self.usage.items(): + if data.get("date") != today: + continue + member = ctx.guild.get_member(int(uid)) + name = member.display_name if member else f"User {uid}" + req = data["requests"] + tok = data["tokens"] + total_req += req + total_tok += tok + lines.append(f"**{name}** - {req} requests, {tok} tokens") + + if not lines: + await ctx.reply("No water wasted yet.") + return + + lines.sort() + embed = discord.Embed( + title=f"How much water wasted for the date of {today}", + description="\n".join(lines), + color=discord.Color.orange(), + ) + await ctx.reply(embed=embed) + + # reply + + @commands.Cog.listener() + async def on_message(self, message: discord.Message): + # ignore bots + if message.author.bot: + return + + # only handle messages that are a reply to another message + if not message.reference: + return + + ctx = await self.bot.get_context(message) + if ctx.valid: + return + + + ref = message.reference.resolved + if not isinstance(ref, discord.Message): + # Not cached — fetch it + try: + ref = await message.channel.fetch_message(message.reference.message_id) + except (discord.NotFound, discord.HTTPException): + return + + # must be a reply to the bot + if ref.author.id != self.bot.user.id: + return + + # must be a message we registered as part of a conversation + if ref.id not in self.message_map: + return + + await self._process_request(message, message.content) + + # utility + + @staticmethod + def _progress_bar(value: int, maximum: int, length: int = 10) -> str: + if maximum == 0: + return "█" * length + filled = round(length * min(value, maximum) / maximum) + return "█" * filled + "░" * (length - filled) + + +async def setup(bot): + await bot.add_cog(AIClient(bot)) \ No newline at end of file