From 203784161195e85344d9f0b4280271f91c89598b Mon Sep 17 00:00:00 2001 From: phixxy Date: Sat, 4 May 2024 12:56:34 -0700 Subject: [PATCH] beta version of llama working --- cogs/llama.py | 170 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 cogs/llama.py diff --git a/cogs/llama.py b/cogs/llama.py new file mode 100644 index 0000000..bf52831 --- /dev/null +++ b/cogs/llama.py @@ -0,0 +1,170 @@ +import os +import time +import json +import logging +import random +import asyncio +import aiofiles +import aiohttp +import discord +from discord.ext import commands, tasks + +class Llama(commands.Cog): + + def __init__(self, bot): + self.bot = bot + self.bot_id = 1097302679836971038 + self.admin_id = 242018983241318410 + self.working_dir = "tmp/chatgpt/" + self.data_dir = "data/chatgpt/" + self.http_session = self.create_aiohttp_session() + self.logger = logging.getLogger("bot") + self.server = "http://192.168.0.34:11434" + + + def create_aiohttp_session(self): + return aiohttp.ClientSession() + + def folder_setup(self): + try: + folders = [ + self.working_dir, + self.data_dir, + self.data_dir + "config", + self.data_dir + "logs", + ] + + for folder in folders: + if not os.path.exists(folder): + os.mkdir(folder) + + except Exception as e: + self.logger.exception(f"LLama failed to make directories: {e}") + + async def question_llama(self, topic, model="mistral"): + data = { + "model": "mistral", + "prompt": topic, + "stream": False + } + + url = f"{self.server}/api/generate" + + try: + async with self.http_session.post(url, json=data) as resp: + response_data = await resp.json() + if resp.status != 200: + self.logger.error(f"Error occurred in answer_question: {response_data}") + return "Error occurred in answer_question" + response = response_data['response'] + return response + + except Exception as error: + self.logger.exception("Error occurred in question_llama") + return "Error occurred in question_llama" + + def create_channel_config(self, filepath): + config_dict = { + "personality":"average", + "channel_topic":"casual", + "chat_enabled":False, + "chat_history_len":5, + "react_to_msgs":False, + "log_images":False, + } + + with open(filepath,"w") as f: + json.dump(config_dict,f) + self.logger.debug("Wrote ChatGPT config variables to file.") + + async def get_channel_config(self, channel_id): + filepath = f"{self.data_dir}config/{channel_id}.json" + if not os.path.exists(filepath): + self.create_channel_config(filepath) + with open(filepath, "r") as f: + config_dict = json.loads(f.readline()) + return config_dict + + def edit_channel_config(self, channel_id, key, value): + config_file = f"{self.data_dir}config/{channel_id}.json" + with open(config_file, 'r') as f: + config_data = json.load(f) + config_data[key] = value + with open(config_file, "w") as f: + json.dump(config_data, f) + + def read_db(self,filepath): + with open(filepath,"r") as fileobj: + db_content = json.load(fileobj) + return db_content + + def save_to_db(self,filepath,db_content): + with open(filepath,"w") as fileobj: + json.dump(db_content,fileobj,indent=4) + + @commands.command( + description="Chat", + help="Enable or disable bot chat in this channel. Usage !chat (enable|disable)", + brief="Enable or disable bot chat" + ) + async def llama_chat(self, ctx, message): + if "enable" in message: + self.edit_channel_config(ctx.channel.id, "llama_enabled", True) + await ctx.send("Chat Enabled") + elif "disable" in message: + self.edit_channel_config(ctx.channel.id, "llama_enabled", False) + await ctx.send("Chat Disabled") + else: + await ctx.send("Usage: !chat (enable|disable)") + + async def log_chat_and_get_history(self, message, logfile, channel_vars): + log_line = '' + if message.attachments and channel_vars.get("log_images", False) and not message.author.bot: + #log_image MUST BE ADDED TO THE JSON FILES + for attachment in message.attachments: + image_description = await self.view_image(attachment.url) + image_description = image_description.replace("\n"," ") + log_line += attachment.url + " " + image_description + " " + log_line += message.content + log_line = message.author.name + ": " + log_line +"\n" + chat_history = "" + self.logger.debug("Logging: " + log_line, end="") + with open(logfile, "a", encoding="utf-8") as f: + f.write(log_line) + with open(logfile, "r", encoding="utf-8") as f: + for line in (f.readlines() [-int(channel_vars["chat_history_len"]):]): + chat_history += line + return chat_history + + async def chat_response(self, message, channel_vars, chat_history_string): + async with message.channel.typing(): + await asyncio.sleep(1) + prompt = f"You are a {channel_vars['personality']} chat bot named Sparkytron 3000 created by @phixxy.com. Your personality should be {channel_vars['personality']}. You are currently in a {channel_vars['channel_topic']} chatroom. The message history is: {chat_history_string}\nSparkytron 3000: " + response = await self.question_llama(prompt) + if "Sparkytron 3000:" in response[0:17]: + response = response.replace("Sparkytron 3000:", "") + max_len = 1999 + if len(response) > max_len: + messages=[response[y-max_len:y] for y in range(max_len, len(response)+max_len,max_len)] + else: + messages=[response] + for response_message in messages: + await message.channel.send(response_message) + + @commands.Cog.listener() + async def on_message(self, message: discord.Message): + # Log Chat + # Todo, make a logging cog to handle this stuff later + logfile = f"{self.data_dir}logs/{message.channel.id}.log" + channel_vars = await self.get_channel_config(message.channel.id) + chat_history_string = await self.log_chat_and_get_history(message, logfile, channel_vars) + # Chat Response + if channel_vars["llama_enabled"] and not message.author.bot or self.bot_id in [x.id for x in message.mentions]: + if message.content and message.content[0] != "!": + await self.chat_response(message, channel_vars, chat_history_string) + elif not message.content: + await self.chat_response(message, channel_vars, chat_history_string) + + +async def setup(bot): + await bot.add_cog(Llama(bot)) \ No newline at end of file