simpleLlamaDiscord.py

import discord
import requests
import re

from copy import deepcopy
## The bot needs persm to make messages, manage messages and manage webhook
## set the value at the botton to the bots token
intents=discord.Intents.default()
intents.message_content = True
client = discord.Client(intents=intents)


dawnPrompt = {
    "userName": "Anon",
    "genParams": {
        "min_length": 0,
        "max_length":1000,
        "top_p":0.7,
        "temperature":0.44,
        "repetition_penalty":1.2,
        "max_new_tokens": 200
    },
    "character": {
    "charName": "Dawn",
    "persona": "Dawn is a carefree, upbeat, and cheerful 15 year old girl. Whenever she makes a mistake, she often quickly recovers and strives to do better. Dawn is very confident in her abilities (sometimes even a little overconfident). While she's kind and supportive, she is sometimes quick-tempered and often gets emotional when she loses. She is also very sensitive and supportive of her friends. Dawn is always good at cheering other people. She also seems to be a little fixated on her appearance. Dawn also has the ability to sense the feelings of Pokémon. She honed her skills as a Pokémon Trainer and focused on becoming a Top Coordinator like her mother. Dawn has a fair complexion, blue eyes, and long navy blue hair. Her usual outfit is a mini sweater dress consisting of a black V-neck tank top with a white shirt under it, a very short pink mini skirt, and a red scarf. On her head, she wears a white beanie with a pink Poké Ball print on it. She also wears gold hair clips that hold up her hair in front. On her feet, she wears a pair of knee-length pink boots and black mid-knee socks above them. Her boots have a pink and white strap on the edge. The rims of her boots are white, while the soles of her boots are greyish-white. She also has a small yellow backpack with all of her personal belongings.",
    "greeting": "Hey! My name is Dawn! I'm a Pokémon Coordinator! Nice to meet you!",
    "emotions": ["happy","sad","excited","exhuasted"],
    "chatExample": [        {
            "chatType": "character",
            "message": "Hey! My name is Dawn! I'm a Pokémon Coordinator! Nice to meet you!"
        },
        {
            "chatType": "user",
            "message": "Hi"
        },
        {
            "chatType": "character",
            "message": "Hello! Are you here to watch one of my performances? I'm sure you'll be amazed by my coordinator skills and by how cute my Pokémon are!"
        },
]
},
"promptTemplate": {
    "prompt": "@{charName}'s Persona: @{persona} \n @{instructions} \n\n<START> \n${chatExample}${chat}",
    "character": "@{charName}: ${message} \n",
    "user": "&{userName}: ${message} \n",
},
"chatHistory": [

    
    ]
    
}

messageHist = {}

async def buildReplyChain(message, userID, messages: list = [],prompt:dict = {}):

	if message:
		if message.author.id != client.user.id:
			str = re.sub("<@\d+>","",message.content).lstrip()
			print(f"==={str}===")
			if str:
				messages.append({
				"chatType": "user",
				"message": str})
        
		else:
		
			str = message.content
			match = re.match(r"^\<[^]]+>\s*",str)
			if match:
					str = str.replace(match.group(0), "")
			if "emotions" in prompt:
				for k, v in prompt["emotions"].items():
					if v in str:
						
						str = str.replace(f"{v}\n\n", f"[{k}] ")
			if str:
				messages.append({
				"chatType": "character",
				"message": re.sub("<@\d+>","",str)
        })
				print(f"==={messages[-1:]}===")
		if message.reference:
			prevmsg = None
			if message.reference.message_id in messageHist:
				prevmsg = messageHist[message.reference.message_id]
			else:
				prevmsg = await message.channel.fetch_message(message.reference.message_id)
				#Cache messages in memory to prevent rate limiting and to speed up message retreival
				messageHist[message.reference.message_id] = prevmsg
			return await buildReplyChain(prevmsg, userID, messages, prompt)
		
	return messages

async def chat(message):
	prompt = deepcopy(dawnPrompt)
	messageChain = reversed(await buildReplyChain(message, message.author.id, [], prompt))
	
	for replyMessage in messageChain:
		prompt["chatHistory"].append(replyMessage)


	resp = requests.post("http://127.0.0.1:5000/chat", json=prompt)
	if (resp.status_code != 200):
		print(f"Non 200 status code: {resp.status_code}")
		return
	resp_dict = resp.json()	
	aimsg = resp_dict["message"]
	if "emotion" in resp_dict:
		aimsg = "<" + resp_dict["emotion"] + ">" + aimsg
	if "emotions" in prompt:
		match = re.match(r"^\[([^]]+)]\s*",aimsg)
		if match:
			if match.group(1) in prompt["emotions"]:
				aimsg = aimsg.replace(match.group(1), prompt["emotions"][match.group(1)]+"\n\n")
			else: 
				aimsg = prompt["emotions"]["normal"] +"\n\n" + aimsg.replace(match.group(1), "")
		else:
			aimsg = prompt["emotions"]["normal"] +"\n\n" + aimsg	 
	if aimsg:
		await message.reply(aimsg)
	else: 
		print("NO message")
	return
	
@client.event
async def on_ready():
	print('We have logged in as {0.user}'.format(client))

@client.event
async def on_message(message):
	if not message.author.bot:
		for mention in message.mentions:
			print(mention.id)
			print(client.user.id)
			if (mention.id == client.user.id):
				await chat(message)
				return
	return

client.run("PUT TOKEN HERE")