Spaces:
Running
Running
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 | |
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) | |
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 | |
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 | |
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 | |
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 | |
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 | |
async def restart_server(ctx): | |
if (ctx.message.author.id == bot_technician): | |
await ctx.send('The application is restarting.') | |
restart_program() | |
# Adjust Command | |
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) | |