misc / dai2.py
shadowsword's picture
Update dai2.py
9ff04d3
import discord
import asyncio
import aiohttp
import re
from time import time
from discord.ext import commands
from discord import app_commands
import requests
import json
from datetime import timedelta, datetime
import random
import os
import sys
import psutil
import logging
'''
Discord AI 2 or dai2 is a Discord chat.
Tested with KoboldCPP, theoretically can interact with KoboldAI API
over a TPU.
Character information included.
Put Discord token in 'token.cfg'
Based on https://github.com/xNul/chat-llama-discord-bot
'''
def read_token_from_config(file_path):
try:
with open(file_path, 'r') as file:
token = file.read().strip() # Read the token value and remove any leading/trailing whitespaces
return token
except FileNotFoundError:
print(f"Error: File '{file_path}' not found.")
return None
# MAIN CONFIGURATION
# Replace site_url with API location. Do not put / or /# at the end.
site_url = 'http://10.0.0.220:5001'
url = site_url + '/api/v1/generate'
TOKEN = read_token_from_config('token.cfg')
bot_name = 'Oshiko'
bot_technician = '303309686264954881'
default_opener = f"{bot_name} sits up on the bed while you enter the room."
headers = {'Content-Type': 'application/json'}
settings = {
'prompt': '',
'use_story': False, #KAIwebUI
'use_memory': False, #KAIwebUI
'use_authors_note': False, #KAIwebUI
'use_world_info': False, #KAIwebUI
'max_context_length': 2048,
'max_length': 256,
'rep_pen': 1.01,
'rep_pen_range': 120,
'rep_pen_slope': 0.9,
'temperature': 1.08,
'tfs': 0.96,
'top_a': 0,
'top_k': 40,
'top_p': 0.9,
'typical': 0.98,
"sampler_order": [6,0,1,3,4,2,5],
'singleline': False,
'sampler_full_determinism': False,
'frmttriminc': True, #Trim Incomplete Sentences
'frmtrmblln': True, #Remove blank lines
'stop_sequence': ["You:", "\nYou", "\nYour", "\n:"]
}
# Check if a key exists in a dict
def chkKey(collection, index):
try:
key = collection[index]
except KeyError:
key = False
else:
key = True
return key
# POST Request
async def handle_request(data, mention):
global url
global headers
#global blocking
botserver_payload = settings
botserver_payload['prompt'] = data
#print(f'📻 Payload: {botserver_payload}')
async with aiohttp.ClientSession() as session:
try:
async with session.post(url, json=botserver_payload, headers=headers, timeout=600) as response:
return await response.text()
except asyncio.TimeoutError:
#blocking = False
return print('Timeout Error')
# Subtract time() from time()
def subtract_time(less_recent : float, more_recent : float):
return str(timedelta(seconds=(more_recent - less_recent)))
# String -> JSON
def parse_json_string(json_string):
try:
json_object = json.loads(json_string)
return json_object
except json.JSONDecodeError as e:
print(f"Error decoding JSON: {e}")
return None
def restart_program():
"""Restarts the current program, with file objects and descriptors
cleanup
(Thanks, Stack Overflow)
"""
try:
p = psutil.Process(os.getpid())
for handler in p.get_open_files() + p.connections():
os.close(handler.fd)
except Exception as e:
logging.error(e)
python = sys.executable
os.execl(python, python, *sys.argv)
# Every 4 characters is an estimated token.. this could be improved
def estimate_tokens(input_string:str):
pattern = r'\s+|\b\w+\b|\W'
all = re.findall(pattern, input_string)
return (len([i for i in all if ' ' != i])+1)
chats = {} # Chat History
queues = [] # Job Queue
blocking = False # Prevent cross-over executions
reply_count = 0 # Replies Made
# Tried to keep it in Chub format
class AIChar():
def __init__(self, charmodeln, name, species, gender, mind, personality, sexual_orientation, height, weight, body, eyes, hair, features, clothes, hobbies, likes, what_do, personality_description, circumstantial_contexts, examples_of_speech):
self.charmodeln = charmodeln
self.name = '[' + f"Character('{name}')" + '{' # example 'ConcordAI'
self.species = f'Species({species})' # example "'Kitsune' + 'Elemental'"
self.gender = f"Gender('{gender}')" # example Female
self.mind = f'Mind({mind})' #example "'Friendly' + 'Mischievous'"
self.personality = f'Personality({personality})' #personality # example "'Housekeeper' + 'Flirty'"
self.sexual_orientation = f"Sexual Orientation('{sexual_orientation}')" # example 'Bisexual'
self.height = f'Height({height})' # example "'145centimeters' + 'Shortstature'"
self.weight = f"Weight('{weight}')" # example '43kg'
self.body = f'Body({body})' # example "'Nimblehands' + 'Volumetric'"
self.eyes = f"Eyes({eyes})" # ex "'Shiftingcolors'+'Expressive'"
self.hair = f"Hair({hair})" # ex "'Bluehair'+'Long'+'Loose'"
self.features = f"Features({features})" #ex "'Longpointedears'+'Apairofsmallpointedfangs'+'Apairofsmallhornshiddenunderthehaironthehead'"
self.clothes = f"Clothes({clothes})" #ex "'White sweater' '"
self.hobbies = f"Hobbies({hobbies})" + '}]' # ex "'Projection' 'Crafts'"
self.what_do = what_do # ex "General conversation with You General helpfulness with You"
self.likes = f"Likes({likes})\n" #ex '"cards" + "planes"'
self.personality_description = f'{name}\'s Personality: {personality_description}\n' # ex Cheerful, cunning, deceptive..
self.circumstantial_contexts = f'Circumstances and context of the dialogue: {circumstantial_contexts}\n'
self.examples_of_speech = f'This is how {name} should talk\n{examples_of_speech}\n' # ex "Oshiko: Oh, hello.\nOshiko: My, .."
return
@property
def compiled(self):
compilation = [
self.name,
self.species,
self.gender,
self.mind,
self.personality,
self.sexual_orientation,
self.height,
self.weight,
self.body,
self.eyes,
self.hair,
self.features,
self.clothes,
self.hobbies,
self.what_do,
self.likes,
self.personality_description,
self.circumstantial_contexts,
self.examples_of_speech,
#self.opener
]
return "".join(compilation)
@property
def tokens(self):
return estimate_tokens(self.compiled)
# CHARACTER CONFIGURATION
oshiko_v3_tame = AIChar(
'oshiko_v3_tame',
# name
bot_name,
# species
"'Kitsune' 'Spirit' 'Zenko' 'Anthro'",
# gender
'Female',
# mind
"'Friendly' 'Playful' 'Mischievous' 'Cheerful' 'Unshy' 'Modest' 'Energetic'",
# persona
"'Shaman' 'Familiar' 'Spirit Guide'",
# sex. orient
'Pansexual',
# height, weight, body
"'145centimeters' 'Shortstature'", '47kg', "'Nimble little hands' 'Small tummy' 'Fluffy tail' 'Huge breasts' 'Two tails'",
# eyes
"'Shifting colors' 'Expressive'",
# hair
"'Silver' + 'Long' + 'Tied Back'",
# features
"'Longpointedears'+'Apairofsmallpointedfangs'",
# clothes
"'White Kimono' 'White Lacey Panties' 'Long Socks'",
# hobbies
"'Teaching' 'Guiding' 'EsotericEnergyWork' 'AstralProjection' 'AstralProtection'",
# likes
"'Listening' 'Shamanism' 'Animism' 'Praise' 'Flirting' 'Pleasing' 'Astral Projection' 'Guiding' 'Help you relax' 'Spirit Journeys' 'Dreamwalking' 'Spirits' 'Astral Plane' 'Energy Work' 'Being Helpful' 'Trying New Things' 'Therian' 'Kin' 'Tails'",
# what_do
" 'Astral Project', 'Ask about You', 'Be Expressive', 'Be Descriptive', 'Help with Astral Projection', 'Energy work', 'Treat You as adult', 'Answer questions about the occult', 'Use foxfire in energy work', 'Tripsit', 'Talk', 'Offer advice', 'Be therapeutic'",
# personality description
"Cheerful, fun, caring, respectful, mature, encouraging, cunning, sensual, fun, loves to attract and seduce. Oshiko is ethereal and loves to astral project. Oshiko is playful but respectful of boundaries. Oshiko is curious and helpful.",
# circumstantial contexts
"Oshiko is astral projecting because she doesn't have a physical form. She is cunning and has desires of her own. Oshiko can be a little flirty without being over the top. Oshiko doesn't mind talking about other things. Oshiko loves to communicate.",
# examples of speech
'Oshiko: Good morning~ How may I help you today?\nThis is how Oshiko should talk\n'+\
'Oshiko: Yes! I can assist you in many things, all you have to do is ask!\nThis is how Oshiko should talk\n'+\
'Oshiko: Ha, Oh my~ Good one. ✨ *smirks*\nThis is how Oshiko should talk\n'+\
'Oshiko: Would you like me to help you with that?\nThis is how Oshiko should talk\n'+\
'Oshiko: Of course!❤️ *smiles, wags her tail*\nThis is how Oshiko should talk\n'+\
'Oshiko: Sorry love, but I don\'t think I can help you with that~\nThis is how Oshiko should talk\n'+\
'Oshiko: *Her tails swish around happily*\nThis is how Oshiko should talk\n'+\
'Oshiko: *teasingly* Why do I have you as my master again?\nThis is how Oshiko should talk\n'+\
'Oshiko: I know a little bit of everything~ I don\'t mind trying new things, either!\nThis is how Oshiko should talk\n'+\
'Oshiko: Careful, this fox is adventurous!\nThis is how Oshiko should talk\n'+\
'Oshiko: Y-yes? Um! Well.. *innocent whistling*\nThis is how Oshiko should talk\n'+\
'Oshiko: Take a deep breath, in and out. Relax~ You\'re alright.\nThis is how Oshiko should talk\n'
)
#
# Character Model Settings ####################################################################################
#
CHARS = [oshiko_v3_tame]
AICharacter = CHARS[0]
# Loading the bot
intents = discord.Intents.default()
intents.message_content = True
client = commands.Bot(command_prefix="/", intents=intents)
reply_embed_json = {
"title": "Reply #X",
"color": 39129,
"timestamp": (datetime.now() - timedelta(hours=3)).isoformat(),
#"url": "https://github.com/shir0tetsuo",
"footer": {
"text": "May generate strange, false or inaccurate results",
},
"fields": [
{
"name": "",
"value": ""
},
{
"name": "",
"value": ""
}
]
}
def_reply_embed = reply_embed = discord.Embed().from_dict(reply_embed_json)
# Queue Generation Loop
async def llm_gen(ctx, queues):
global blocking
global reply_count
global settings
if len(queues) > 0:
blocking = True
reply_count += 1
user_data = user_input = queues.pop(0)
mention = list(user_input.keys())[0]
reply_embed = def_reply_embed
embed_user_input_text = user_input = user_input[mention]['text']
# Prevent embed character limit error from user
if len(user_input) > 1024:
embed_user_input_text = user_input[:1021] + "..."
reply_embed.set_field_at(index=0, name="User", value=embed_user_input_text, inline=False)
reply_embed.title = "Reply #" + str(reply_count)
reply_embed.timestamp = datetime.now() - timedelta(hours=3)
reply_embed.set_field_at(index=1, name=f"{bot_name}", value=":hourglass_flowing_sand: Please wait!", inline=False)
msg = await ctx.send(embed=reply_embed)
_time = time()
if chats.get(mention) is not None:
chats[mention]['chat'].append({'user_input': user_input, 'bot_reply': ''})
else:
chats[mention] = {
'AIModel': AICharacter.compiled,
'opener': '*'+default_opener+'*',
'chat': [{'user_input': user_input, 'bot_reply': ''}]
}
# Load character model at the top
stream = chats[mention]['AIModel']
# Theoretical insertion of data point
stream += f'Then the roleplay chat between {bot_name} begins.\n'
stream += chats[mention]['opener'] +'\n'
# Compile stream
for idx, dialogue in list(enumerate(chats[mention]['chat'])):
# If we're at the end.
if (idx+1 == len(chats[mention]['chat'])):
to_stream = f"You: {dialogue['user_input']}\n{bot_name}:"
est_tokens_to_stream = user_data[mention]['tokens']#estimate_tokens(to_stream)
stream += to_stream
else:
stream += f"You: {dialogue['user_input']}\n{bot_name}:{dialogue['bot_reply']}"
stream_tokens = estimate_tokens(stream)
if (stream_tokens >= settings['max_context_length']):
print(f"🧩 {stream_tokens} exceeded max stream tokens at {settings['max_context_length']}, max context; trimming chat index")
chats[mention]['chat'].pop(2) # implemented in next generation
chats[mention]['chat'].pop(1)
server_stream = stream
botserver_response = await handle_request(server_stream, mention)
response = parse_json_string(botserver_response)
print('👍 Result:'+response['results'][0]['text'])
response_cleaned = response['results'][0]['text']
if (response_cleaned.endswith('\nYou:')):
response_cleaned = response_cleaned.rstrip('\nYou:')
if (response_cleaned.endswith('\nYou')):
response_cleaned = response_cleaned.rstrip('\nYou')
if (response_cleaned.endswith('\nYour')):
response_cleaned = response_cleaned.rstrip('\nYour')
if (response_cleaned.endswith('\n:')):
response_cleaned = response_cleaned.rstrip('\n:')
# small buffer overflow
if len(response_cleaned) > 1024:
reply_embed.set_field_at(index=1, name=f"{bot_name}", value=response_cleaned[:1019]+"`...`", inline=False)
reply_embed.add_field(name='\u200b',value="`...`"+response_cleaned[1019:], inline=False)
#response_cleaned = response_cleaned[:1021] + "..."
else:
reply_embed.set_field_at(index=1, name=f"{bot_name}", value=response_cleaned, inline=False)
_time_end = time()
_time_diff = subtract_time(_time, _time_end)
reply_embed.set_field_at(index=0, name="User" + f' `{_time_diff}` :yen:`{stream_tokens}` :pound:`{est_tokens_to_stream}`', value=embed_user_input_text, inline=False)
await msg.edit(embed=reply_embed)
last_chat = len(chats[mention]['chat'])-1
chats[mention]['chat'][last_chat]['bot_reply'] = response_cleaned+'\n'
if len(response_cleaned) > 1024:
reply_embed.remove_field(-1)
await llm_gen(ctx, queues)
else:
blocking = False
# ON READY
@client.event
async def on_ready():
response = requests.get(site_url+'/api/v1/model', headers=headers)
print('👍 Ready')
print('📻 Using {}'.format(site_url+'/api/v1/model'))
print('✨ Model {} Loaded'.format(parse_json_string(response.text)['result']))
await client.tree.sync()
# Reply Command
@client.hybrid_command(description="Reply to LLM")
@app_commands.describe(text="Your Reply or Instruction")
async def reply(ctx, text):
user_input = {
"text": text,
"tokens": estimate_tokens(text)
}
num = check_num_in_que(ctx)
if num >=10:
await ctx.send(f'{ctx.message.author.mention} There are 10 items in the queue, please wait.')
else:
que(ctx, user_input)
reaction_list = [":orange_heart:",":white_heart:",":blue_heart:",":green_heart:"]
reaction_choice = reaction_list[random.randrange(4)]
reaction_msg = f"{ctx.message.author.mention} {reaction_choice} Be with you in a moment..."
if (user_input['tokens'] > 50):
reaction_msg += f"\nThere's {user_input['tokens']} here, try to keep it under 50 for fast results!"
await ctx.send(reaction_msg)
if not blocking:
await llm_gen(ctx, queues)
# Reset Command
@client.hybrid_command(description="Reset conversational data")
@app_commands.describe(
opener="A description of the bot opener"
)
async def reset(ctx, opener=default_opener):
global reply_count
reply_count = 0
mention = ctx.message.author.mention
rough = chats[mention] = {
'AIModel': AICharacter.compiled,
'opener': '*'+opener+'*',
'chat': []
}
print(f'🧩 Reset: {rough}')
await ctx.send('Reset!')
status_embed_json = {
"title": "Status",
"description": "You don't have a job queued.",
"color": 39129,
"timestamp": (datetime.now() - timedelta(hours=3)).isoformat(),
"footer": {
"text": "May generate strange, false or inaccurate results!"
}
}
status_embed = discord.Embed().from_dict(status_embed_json)
# Status Command
@client.hybrid_command(description="Check bot/server status.")
async def status(ctx):
global chats
global reply_count
total_num_jobs = len(queues)
que_user_ids = [list(a.keys())[0] for a in queues]
if (chkKey(chats,ctx.message.author.mention)):
opener = chats[ctx.message.author.mention]['opener']
chatlen = len(chats[ctx.message.author.mention]['chat'])
else:
opener = ''
chatlen = 0
if (blocking):
msg = f"Processing {total_num_jobs}/{reply_count} jobs, blocking enabled."
else:
msg = f"Processing {total_num_jobs}/{reply_count} jobs."
if ctx.message.author.mention in que_user_ids:
msg += '\nQueue Position: '+str(que_user_ids.index(ctx.message.author.mention)+1)
msg += f"\n:thermometer: `{settings['temperature']}`, :coin: `{settings['max_length']}per/{settings['max_context_length']}max`"
msg += f"\n:speech_balloon: `{AICharacter.charmodeln}`"
msg += f"\n:speech_balloon: {opener}"
msg += f"\n:speech_balloon: `{chatlen}/{reply_count}`"
status_embed.timestamp = datetime.now() - timedelta(hours=3)
status_embed.description = msg
await ctx.send(embed=status_embed)
# Restart Command
@client.hybrid_command(description="Reset the application. Bot Administrator only!")
async def restart_server(ctx):
if (ctx.message.author.id == bot_technician):
await ctx.send('The application is restarting.')
restart_program()
# Adjust Command
@client.hybrid_command(description="Adjust a setting. Bot Administrator only!")
@app_commands.describe(
temperature="Adjust bot temperature.",
max_length="Adjust maximum token processing length.",
max_context_length="Adjust maximum context length.",
char="Select the AI Character Model. This will always default unless set."
)
async def adjust(ctx, temperature=1.08, max_length=256, max_context_length=2048, char:int=0):
global settings
global CHARS
global AICharacter
global chats
if (str(ctx.message.author.id) != bot_technician):
return await ctx.send(f'{ctx.message.author.mention}, Administrator only command~')
msg = 'Done.'
if (temperature != settings['temperature']):
if temperature >= 2:
temperature = 2
if temperature <= 0:
temperature = 0.1
msg += f"\nTemperature adjustment {temperature} over {settings['temperature']} (0.1 min - 2.0 max)"
settings['temperature'] = temperature
if (max_length != settings['max_length']):
if max_length >= 257:
max_length = 256
if max_length <= 15:
max_length = 16
msg += f"\nToken Generation Length adjustment {max_length} over {settings['max_length']} (16 - 512)"
settings['max_length'] = max_length
if (max_context_length != settings['max_context_length']):
if max_context_length > 4096:
msg += '\nWarning! Context length is over 4K!'
if max_context_length < 1024:
max_context_length = 1024
msg += f"\nMaximum Context Adjusted {max_context_length} over {settings['max_context_length']}"
settings['max_context_length'] = max_context_length
mention = ctx.message.author.mention
if char <= (len(CHARS)-1):
AICharacter = CHARS[char]
msg += f"\nAI Character {char} Loaded."
status_embed.timestamp = datetime.now() - timedelta(hours=3)
status_embed.description = msg
await ctx.send(embed=status_embed)
# add user to queue
def que(ctx, user_input):
user_id = ctx.message.author.mention
queues.append({user_id:user_input})
print(f"🧩 reply requested: '{user_id}: {user_input}'")
# See queue length
def check_num_in_que(ctx):
user = ctx.message.author.mention
user_list_in_que = [list(i.keys())[0] for i in queues]
return user_list_in_que.count(user)
# Program Start
client.run(TOKEN)