From ccf35b3e8d8bef042ba179873a9b067bdac7ff3e Mon Sep 17 00:00:00 2001 From: cameron Date: Fri, 10 May 2024 08:42:37 -0400 Subject: [PATCH] Bot LLM functionality is now working --- main.py | 1 + plugins/botchat/plugin.py | 96 +++++++++++++++++++++++++++++++---- plugins/botchat/settings.yaml | 1 + 3 files changed, 89 insertions(+), 9 deletions(-) diff --git a/main.py b/main.py index 7617c77..9bf6b8f 100644 --- a/main.py +++ b/main.py @@ -26,6 +26,7 @@ class DankBot(discord.ext.commands.Bot): async def main(): intents = discord.Intents.default() intents.message_content = True + intents.members = True logging.basicConfig(level=logging.INFO) with open(".token") as token_file: diff --git a/plugins/botchat/plugin.py b/plugins/botchat/plugin.py index 6c6bc62..59acd67 100644 --- a/plugins/botchat/plugin.py +++ b/plugins/botchat/plugin.py @@ -1,10 +1,11 @@ # Plugin for bot LLM chat from discord.ext import commands import discord +import io +import aiohttp import yaml import random import os -import llm plugin_folder=os.path.dirname(os.path.realpath(__file__)) prompts_folder=os.path.join(plugin_folder, 'prompts') @@ -12,30 +13,107 @@ default_prompt="default.txt" config_filename=os.path.join(plugin_folder, 'settings.yaml') llm_data = {} +async def prompt_llm(prompt): + print("Prompting LLM") + print(f"PROMPT DATA\n{prompt}") + async with aiohttp.ClientSession(llm_data["api_base"]) as session: + async with session.post("/completion", json={"prompt": prompt, "n_predict": 250}) as resp: + print(f"LLM response status {resp.status}") + response_json=await resp.json() + content=response_json["content"] + return content + +def get_message_contents(msg): + message_text = f"{msg.author.name}: {msg.clean_content}" + print(f"Message contents -- {message_text}") + return message_text + async def get_chat_history(ctx, limit=20): messages = [message async for message in ctx.channel.history(limit=limit)] - plain_messages = list(map(lambda m: f"{m.author.name}: {m.content}", messages)) + plain_messages = list(map(get_message_contents, messages)) plain_messages.reverse() return plain_messages @commands.command(name='llm') async def llm_response(ctx): + await ctx.channel.typing() prompt_file = os.path.join(prompts_folder, default_prompt) with open(prompt_file, 'r') as prompt_file: prompt = prompt_file.read() history_arr = await get_chat_history(ctx) history_str = '\n'.join(history_arr) full_prompt = prompt.replace("", history_str) - response = llm_data["model"].prompt(full_prompt) - print(response) + response = await prompt_llm(full_prompt) + await send_chat_responses(ctx, response) + +async def send_chat_responses(ctx, response_text): + print("Processing chat response") + fullResponseLog = "dank-bot:" + response_text # first response won't include the user + responseLines = fullResponseLog.splitlines() + output_strs = [] + for line in responseLines: + if line.startswith("dank-bot:"): + truncStr = line.replace("dank-bot:","") + output_strs.append(truncStr) + elif line.find(":") > 0 and line.find(":") < 20: + break + else: + output_strs.append(line.strip()) + for outs in output_strs: + final_output_str = await fixup_mentions(ctx, outs) + await ctx.channel.send(final_output_str) + +async def fixup_mentions(ctx, text): + newtext = text + if (isinstance(ctx.channel,discord.DMChannel)): + newtext = newtext.replace(f"@{ctx.author.name}", ctx.author.mention) + elif (isinstance(ctx.channel,discord.GroupChannel)): + for user in ctx.channel.recipients: + newtext = newtext.replace(f"@{user.name}", user.mention) + elif (isinstance(ctx.channel,discord.Thread)): + for user in await ctx.channel.fetch_members(): + member_info = await ctx.channel.guild.fetch_member(user.id) + newtext = newtext.replace(f"@{member_info.name}", member_info.mention) + else: + for user in ctx.channel.members: + newtext = newtext.replace(f"@{user.name}", user.mention) + if ctx.guild != None: + for role in ctx.guild.roles: + newtext = newtext.replace(f"@{role.name}", role.mention) + return newtext + +async def handle_message(ctx): + print("Dank-bot received message") + print(f"Dank-bot ID is {llm_data['bot'].user.id}") + bot_id = llm_data['bot'].user.id + + # First case, bot DMed + if (isinstance(ctx.channel,discord.DMChannel) and ctx.author.id != bot_id): + print("Dank-bot DMed, responding") + await llm_response(ctx) + return + + # Second case, bot mentioned + bot_mentions=list(filter(lambda x: x.id == bot_id, ctx.mentions)) + if (len(bot_mentions) > 0): + print("Dank-bot mentioned, responding") + await llm_response(ctx) + return + + # Other case, random response + random_roll = random.random() + print(f"Dank-bot rolled {random_roll}") + if (random_roll < llm_data['response_probability']): + print(f"{random_roll} < {llm_data['response_probability']}, responding") + await llm_response(ctx) + return async def setup(bot): with open(config_filename, 'r') as conf_file: yaml_config = yaml.safe_load(conf_file) - model = llm.get_model("gpt-3.5-turbo-instruct") - model.key = yaml_config["api_key"] - model.api_base = yaml_config["api_base"] - model.completion = True - llm_data["model"] = model + llm_data["api_base"] = yaml_config["api_base"] + llm_data["response_probability"] = yaml_config["response_probability"] bot.add_command(llm_response) + bot.add_listener(handle_message, "on_message") + llm_data["bot"] = bot print("LLM interface initialized") diff --git a/plugins/botchat/settings.yaml b/plugins/botchat/settings.yaml index bbf37ef..aa07d46 100644 --- a/plugins/botchat/settings.yaml +++ b/plugins/botchat/settings.yaml @@ -1,2 +1,3 @@ api_base: "http://192.168.1.204:5000" api_key: "empty" +response_probability: 0.05