Bot LLM functionality is now working
This commit is contained in:
parent
ac9cd831a5
commit
ccf35b3e8d
1
main.py
1
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:
|
||||
|
|
|
|||
|
|
@ -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("<CONVHISTORY>", 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")
|
||||
|
|
|
|||
|
|
@ -1,2 +1,3 @@
|
|||
api_base: "http://192.168.1.204:5000"
|
||||
api_key: "empty"
|
||||
response_probability: 0.05
|
||||
|
|
|
|||
Loading…
Reference in New Issue