49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
import random
|
|
|
|
import discord
|
|
from discord.ext import commands
|
|
|
|
from app import cfg
|
|
from app.system import SystemPrompt
|
|
from app.var import CustomError
|
|
from pkg.llm import LanguageModel
|
|
from pkg.llm.data import Payload
|
|
|
|
|
|
class AurochsIntellect(commands.Cog):
|
|
def __init__(self, bot: commands.Bot, llm_api: str):
|
|
self.bot = bot
|
|
self.llm = LanguageModel(llm_api, CustomError.LLM_EXCEPTION)
|
|
|
|
self.prompts = [p.value.variable for p in SystemPrompt]
|
|
self.weights = [p.value.weight for p in SystemPrompt]
|
|
|
|
@commands.Cog.listener()
|
|
async def on_message(self, message: discord.Message):
|
|
assert self.bot.user is not None
|
|
|
|
if message.author == self.bot.user or not self.bot.user.mentioned_in(message):
|
|
return
|
|
|
|
async with message.channel.typing():
|
|
clean_content = message.content.replace(f'<@{self.bot.user.id}>', '').strip()
|
|
|
|
system_prompt = random.choices(self.prompts, self.weights, k=1)[0]
|
|
|
|
payload = Payload(
|
|
system_prompt=system_prompt.format(message.author.display_name),
|
|
user_message=clean_content,
|
|
temperature=0.7,
|
|
max_tokens=200,
|
|
stream=False
|
|
)
|
|
|
|
response = await self.llm.make_request(payload)
|
|
await message.reply(response)
|
|
|
|
@commands.command()
|
|
async def status(self, context: commands.Context):
|
|
await context.send("The herd is accounted for.")
|
|
|
|
async def setup(bot: commands.Bot):
|
|
await bot.add_cog(AurochsIntellect(bot, cfg.LLM_API)) |