openrouter client
This commit is contained in:
29
cogs/openrouter/config.json
Normal file
29
cogs/openrouter/config.json
Normal file
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"description": "OpenRouter AI client",
|
||||
"usages": [
|
||||
"ai <message>",
|
||||
"ai_reset",
|
||||
"ai_history",
|
||||
"ai_stats",
|
||||
"ai_info",
|
||||
"ai_setmodel <model> (admin)",
|
||||
"ai_setprompt <prompt> (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": []
|
||||
}
|
||||
341
cogs/openrouter/openrouter.py
Normal file
341
cogs/openrouter/openrouter.py
Normal file
@@ -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))
|
||||
Reference in New Issue
Block a user