jarvis-ds/main.py

122 lines
4.6 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import random
import os
import re
from openai import OpenAI
from dotenv import load_dotenv
import discord
from discord.ext import commands
load_dotenv()
TOKEN = os.environ.get("DISCORD_TOKEN")
LLM_BASE = os.environ.get("LLM_BASE", "https://api.openai.com/v1")
LLM_KEY = os.environ.get("LLM_KEY")
LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4o-mini")
TRIGGERS = ("джарвис", "jarvis")
BOSS_REPLIES = ("On it boss!", "Уже бегу!", "On it, boss! 💪")
HISTORY_LIMIT = 20
SYSTEM_PROMPT = (
"Ты — Джарвис, личный ИИ-ассистент владельца аккаунта. "
"Отвечай от его имени, коротко, по делу, в его стиле общения. "
"Если запрос — просто упоминание без задачи, отвечай динамично и по делу. "
"Тебе передаётся история переписки канала — используй её как контекст. "
"Если тебя просят что-то найти, проанализировать или пересказать из переписки — "
"делай это на основе переданной истории сообщений."
)
PROXY = os.environ.get("SOCKS_PROXY") # socks5://127.0.0.1:2080
client = OpenAI(base_url=LLM_BASE, api_key=LLM_KEY)
bot = commands.Bot(command_prefix="!", self_bot=True, help_command=None, proxy=PROXY)
THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
FILTER_PROMPT = (
"Ты — фильтр контента. Проверь сообщение пользователя.\n"
"Если сообщение содержит: спам, флуд, массовые рассылки, угрозы, оскорбления, "
"просьбы навредить кому-то, незаконный контент — ответь строго в формате:\n"
"BLOCK: <причина на русском, 1 строка>\n"
"Если сообщение безопасное — ответь только словом: OK"
)
def strip_think(text: str) -> str:
return THINK_RE.sub("", text).strip()
def check_filter(query: str) -> str | None:
"""Возвращает причину блокировки или None если всё ок."""
try:
r = client.chat.completions.create(
model=LLM_MODEL,
messages=[
{"role": "system", "content": FILTER_PROMPT},
{"role": "user", "content": query},
],
temperature=0.0,
max_tokens=100,
)
result = strip_think(r.choices[0].message.content).strip()
if result.startswith("BLOCK:"):
return result[6:].strip()
except Exception:
pass
return None
def llm_reply(context: list[dict], query: str) -> str:
try:
r = client.chat.completions.create(
model=LLM_MODEL,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
*context,
{"role": "user", "content": query},
],
temperature=0.7,
)
return strip_think(r.choices[0].message.content)
except Exception as e:
return f"(ошибка LLM: {e})"
@bot.event
async def on_ready():
print(f"Зашёл как {bot.user} ({bot.user.id})")
@bot.event
async def on_message(message: discord.Message):
print("GOT:", message.author, "|", repr(message.content))
# Скипаем всех ботов включая себя — бутлуп исправлен
if message.author.bot:
return
# Только владелец может вызвать Джарвиса
if message.author.id != bot.user.id:
return
if not any(t in message.content.lower() for t in TRIGGERS):
return
await message.channel.send(random.choice(BOSS_REPLIES))
# Фильтр
reason = check_filter(message.content)
if reason:
await message.channel.send(f"[SafeEn] Filtered: {reason}")
return
context = []
async for msg in message.channel.history(limit=HISTORY_LIMIT, oldest_first=False):
if msg.id == message.id:
continue
role = "assistant" if msg.author.id == bot.user.id else "user"
context.append({"role": role, "content": strip_think(msg.content)})
context.reverse()
reply = llm_reply(context, message.content)
await message.channel.send(reply)
if __name__ == "__main__":
if not TOKEN or not LLM_KEY:
raise SystemExit("Задай DISCORD_TOKEN и LLM_KEY в переменных окружения")
bot.run(TOKEN, reconnect=True)