Skip to content

Commit

Permalink
Merge pull request #19 from F33RNI/next
Browse files Browse the repository at this point in the history 
Next
  • Loading branch information
@F33RNI
F33RNI authored Mar 10, 2023
2 parents 4fef5c9 + 27488e7 commit 2549e95
Show file tree
Hide file tree
Showing 8 changed files with 555 additions and 263 deletions.
2 changes: 1 addition & 1 deletion .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,4 @@
/*.zip
/test.py
/logs/
/chats.json
/chats/
240 changes: 199 additions & 41 deletions AIHandler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.
"""
import asyncio
import logging
import os
import queue
import threading
import time
Expand All @@ -30,11 +32,13 @@
ERROR_CHATGPT_DISABLED = 'ChatGPT module is disabled in settings.json'
ERROR_DALLE_DISABLED = 'DALL-E module is disabled in settings.json'

CONVERSATION_DIR_OR_FILE = 'conversations'


class AIHandler:
def __init__(self, settings, chats_file, authenticator):
def __init__(self, settings, chats_dir, authenticator):
self.settings = settings
self.chats_file = chats_file
self.chats_dir = chats_dir
self.authenticator = authenticator

# Loop running flag
Expand All @@ -59,6 +63,10 @@ def thread_start(self):
Starts background thread
:return:
"""
# Create chats directory if not exists
if not os.path.exists(self.chats_dir):
os.makedirs(self.chats_dir)

# Set flag
self.loop_running = True

Expand All @@ -74,7 +82,7 @@ def get_chat(self, chat_id: int):
:return: (conversation_id, parent_id)
"""
logging.info('Loading conversation_id for chat_id ' + str(chat_id))
chats = load_json(self.chats_file)
chats = load_json(os.path.join(self.chats_dir, 'chats.json'))
if chats is not None and str(chat_id) in chats:
chat = chats[str(chat_id)]
conversation_id = None
Expand All @@ -98,7 +106,7 @@ def set_chat(self, chat_id: int, conversation_id=None, parent_id=None):
"""
logging.info('Saving conversation_id ' + str(conversation_id) + ' and parent_id '
+ str(parent_id) + ' for chat_id ' + str(chat_id))
chats = load_json(self.chats_file)
chats = load_json(os.path.join(self.chats_dir, 'chats.json'))
if chats is not None:
if str(chat_id) in chats:
# Save or delete conversation_id
Expand All @@ -122,7 +130,104 @@ def set_chat(self, chat_id: int, conversation_id=None, parent_id=None):
chats[str(chat_id)]['parent_id'] = parent_id
else:
chats = {}
save_json(self.chats_file, chats)
save_json(os.path.join(self.chats_dir, 'chats.json'), chats)

def save_conversation(self, chatbot_, conversation_id) -> None:
"""
Saves conversation
:param chatbot_:
:param conversation_id:
:return:
"""
logging.info('Saving conversation ' + conversation_id)
# API type 0
if int(self.settings['modules']['chatgpt_api_type']) == 0:
conversations_file = CONVERSATION_DIR_OR_FILE + '.json'
chatbot_.conversations.save(os.path.join(self.chats_dir, conversations_file))

# API type 3
elif int(self.settings['modules']['chatgpt_api_type']) == 3:
if not os.path.exists(os.path.join(self.chats_dir, CONVERSATION_DIR_OR_FILE)):
os.makedirs(os.path.join(self.chats_dir, CONVERSATION_DIR_OR_FILE))
conversation_file = os.path.join(self.chats_dir, CONVERSATION_DIR_OR_FILE, conversation_id + '.json')
chatbot_.save(conversation_file, conversation_id)

def load_conversation(self, chatbot_, conversation_id) -> None:
"""
Loads conversation
:param chatbot_:
:param conversation_id:
:return:
"""
logging.info('Loading conversation ' + conversation_id)
try:
# API type 0
if int(self.settings['modules']['chatgpt_api_type']) == 0:
conversations_file = CONVERSATION_DIR_OR_FILE + '.json'
chatbot_.conversations.load(os.path.join(self.chats_dir, conversations_file))

# API type 3
elif int(self.settings['modules']['chatgpt_api_type']) == 3:
conversation_file = os.path.join(self.chats_dir, CONVERSATION_DIR_OR_FILE, conversation_id + '.json')
if os.path.exists(conversation_file):
chatbot_.load(conversation_file, conversation_id)
else:
logging.warning('File ' + conversation_file + ' not exists!')
except Exception as e:
logging.warning('Error loading conversation ' + conversation_id + ' ' + str(e))

def delete_conversation(self, conversation_id) -> None:
"""
Deletes conversation
:param conversation_id:
:return:
"""
logging.info('Deleting conversation ' + conversation_id)
try:
# API type 0
if int(self.settings['modules']['chatgpt_api_type']) == 0:
self.authenticator.chatbot.reset()
try:
conversations_file = os.path.join(self.chats_dir, CONVERSATION_DIR_OR_FILE + '.json')
self.authenticator.chatbot.conversations.remove_conversation(conversations_file)
self.authenticator.chatbot.conversations.save(conversations_file)
except Exception as e:
logging.error('Error deleting conversation ' + conversation_id, e)

# API type 1
elif int(self.settings['modules']['chatgpt_api_type']) == 1:
self.authenticator.chatbot.reset_chat()
try:
self.authenticator.chatbot.delete_conversation(conversation_id)
except Exception as e:
logging.error('Error deleting conversation ' + conversation_id, e)

# API type 2
elif int(self.settings['modules']['chatgpt_api_type']) == 2:
self.authenticator.chatbot.reset_chat()
try:
self.authenticator.chatbot.conversations.remove(conversation_id)
except Exception as e:
logging.error('Error deleting conversation ' + conversation_id, e)

# API type 3
elif int(self.settings['modules']['chatgpt_api_type']) == 3:
self.authenticator.chatbot.reset(conversation_id)

# Wrong API type
else:
raise Exception('Wrong chatgpt_api_type')

# Delete conversation file if exists
try:
conversation_file = os.path.join(self.chats_dir, CONVERSATION_DIR_OR_FILE, conversation_id + '.json')
if os.path.exists(conversation_file):
logging.info('Removing ' + conversation_file + ' file')
os.remove(conversation_file)
except Exception as e:
logging.error('Error removing conversation file for conversation ' + conversation_id, e)
except Exception as e:
logging.warning('Error loading conversation ' + conversation_id + ' ' + str(e))

def gpt_loop(self):
"""
Expand Down Expand Up @@ -153,17 +258,46 @@ def gpt_loop(self):
api_response = None
raise Exception(ERROR_CHATGPT_DISABLED)

# API type 0
if self.authenticator.api_type == 0:
# Get conversation_id
conversation_id, parent_id = self.get_chat(container.chat_id)
# Too many requests in 1 hour
if self.authenticator.chatbot_too_many_requests:
raise Exception(Authenticator.TOO_MANY_REQUESTS_MESSAGE)

# Get chatbot from Authenticator class
# Wait for chatbot
chatbot = self.authenticator.chatbot
while not self.authenticator.chatbot_working or chatbot is None:
time.sleep(1)
chatbot = self.authenticator.chatbot

# Reset chat
# Too many requests in 1 hour
if self.authenticator.chatbot_too_many_requests:
raise Exception(Authenticator.TOO_MANY_REQUESTS_MESSAGE)

# Lock chatbot
self.authenticator.chatbot_locked = True

# Get conversation_id and parent_id
conversation_id, parent_id = self.get_chat(container.chat_id)

# Log request
logging.info('Asking: ' + str(container.request)
+ ', conversation_id: ' + str(conversation_id) + ', parent_id: ' + str(parent_id))

# API type 0
if int(self.settings['modules']['chatgpt_api_type']) == 0:
# Reset current chat
chatbot.reset()

# Try to load conversation
if conversation_id is not None:
self.load_conversation(chatbot, conversation_id)

# Try to load conversation
if conversation_id is not None:
try:
chatbot.load_conversation(conversation_id)
except Exception as e:
logging.warning('Error loading conversation with id: ' + conversation_id + ' ' + str(e))

# Ask
for data in chatbot.ask_stream(str(container.request), conversation_id=conversation_id):
# Initialize response
Expand All @@ -178,6 +312,7 @@ def gpt_loop(self):
if conversation_id is None:
conversation_id = str(uuid.uuid4())
chatbot.save_conversation(conversation_id)
self.save_conversation(chatbot, conversation_id)
self.set_chat(container.chat_id, conversation_id, parent_id)
except Exception as e:
logging.warning('Error saving conversation! ' + str(e))
Expand All @@ -186,32 +321,8 @@ def gpt_loop(self):
api_response = api_response.replace('<|im_end|>', '').replace('<|im_start|>', '')

# API type 1
elif self.authenticator.api_type == 1:
# Too many requests in 1 hour
if self.authenticator.chatbot_too_many_requests:
raise Exception(Authenticator.TOO_MANY_REQUESTS_MESSAGE)

# Wait for chatbot
chatbot = self.authenticator.chatbot
while not self.authenticator.chatbot_working or chatbot is None:
time.sleep(1)
chatbot = self.authenticator.chatbot

# Too many requests in 1 hour
if self.authenticator.chatbot_too_many_requests:
raise Exception(Authenticator.TOO_MANY_REQUESTS_MESSAGE)

# Lock chatbot
self.authenticator.chatbot_locked = True

# Get conversation_id and parent_id
conversation_id, parent_id = self.get_chat(container.chat_id)

# Log request
logging.info('Asking: ' + str(container.request)
+ ', conversation_id: ' + str(conversation_id) + ', parent_id: ' + str(parent_id))

# Reset chat
elif int(self.settings['modules']['chatgpt_api_type']) == 1:
# Reset current chat
chatbot.reset_chat()

# Ask
Expand All @@ -230,18 +341,65 @@ def gpt_loop(self):
parent_id = data['parent_id']

# Log conversation id and parent id
logging.info('Current conversation_id: ' + conversation_id + ', parent_id: ' + parent_id)
logging.info('Current conversation_id: ' + str(conversation_id)
+ ', parent_id: ' + str(parent_id))

# Save conversation id
self.set_chat(container.chat_id, conversation_id, parent_id)

# API type 2
elif int(self.settings['modules']['chatgpt_api_type']) == 2:
# Make async function sync
api_responses_ = []
conversation_ids_ = []

async def ask_async(request_, conversation_id_):
async for data_ in chatbot.ask(request_, conversation_id=conversation_id_):
# Get last response
api_responses_.append(data_['message'])

# Store conversation_id
if 'conversation_id' in data_ and data_['conversation_id'] is not None:
conversation_ids_.append(data_['conversation_id'])

# Ask
loop = asyncio.new_event_loop()
loop.run_until_complete(ask_async(str(container.request), conversation_id))

# Log conversation id and parent id
logging.info('Current conversation_id: ' + str(conversation_id)
+ ', parent_id: ' + str(parent_id))

# Save conversation id
self.set_chat(container.chat_id, conversation_id, parent_id)

# API type 3
elif int(self.settings['modules']['chatgpt_api_type']) == 3:
# Generate conversation ID
if conversation_id is None:
conversation_id = str(uuid.uuid4())

# Try to load conversation
else:
self.load_conversation(chatbot, conversation_id)

# Ask
for data in chatbot.ask(str(container.request), convo_id=conversation_id):
# Initialize response
if api_response is None:
api_response = ''

# Append response
api_response += str(data)

# Save conversation id
self.save_conversation(chatbot, conversation_id)
self.set_chat(container.chat_id, conversation_id)

# Wrong api type
else:
raise Exception('Wrong chatgpt_api_type')

# Log response
logging.info(str(api_response))

# DALL-E
else:
# Check if ChatGPT is enabled
Expand Down
Loading

0 comments on commit 2549e95

Please sign in to comment.