forked from rejnronuz/muzovkant-cogdump
341 lines
12 KiB
Python
341 lines
12 KiB
Python
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)) |