Initial commit

ext/aurochs.py

@@ -0,0 +1,49 @@
+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))