Skip to content

Commit

Permalink
HUUUGE refactoring
Browse files Browse the repository at this point in the history
  • Loading branch information
@F33RNI
F33RNI committed Jun 26, 2023
1 parent 6589c64 commit fffc285
Show file tree
Hide file tree
Showing 14 changed files with 2,270 additions and 963 deletions.
181 changes: 80 additions & 101 deletions BardModule.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,50 +14,58 @@
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.
"""

import asyncio
import ctypes
import logging
import multiprocessing
import os
import uuid
from typing import List, Dict

import Bard

import BotHandler
import UsersHandler
from RequestResponseContainer import RequestResponseContainer


class BardModule:
def __init__(self, config: dict, messages: dict, users_handler: UsersHandler.UsersHandler) -> None:
def __init__(self, config: dict, messages: List[Dict], users_handler: UsersHandler.UsersHandler) -> None:
self.config = config
self.messages = messages
self.users_handler = users_handler

self._enabled = False
self._chatbot = None
self._restart_attempts = 0
self._proxy = None
# All variables here must be multiprocessing
self.processing_flag = multiprocessing.Value(ctypes.c_bool, False)

def initialize(self) -> None:
def initialize(self, proxy=None) -> None:
"""
Initializes Bard bot using this API: https://github.com/acheong08/Bard
:return:
"""
self._enabled = False
self._chatbot = None
self.processing_flag.value = False

try:
# Use manual proxy
if not proxy and self.config["bard"]["proxy"] and self.config["bard"]["proxy"] != "auto":
proxy = self.config["bard"]["proxy"]

# Log
logging.info("Initializing Bard module with proxy {}".format(proxy))

# Set enabled status
self._enabled = self.config["modules"]["bard"]
if not self._enabled:
logging.warning("Bard module disabled in config file!")
return

# Initialize chatbot
self._chatbot = Bard.Chatbot(self.config["bard"]["token"])

# Set proxy
proxy = self.config["bard"]["proxy"]
if proxy and len(proxy) > 1 and proxy.strip().lower() != "auto":
self._proxy = proxy
self._chatbot.session.proxies.update({"http": proxy,
"https": proxy})
if proxy:
self._chatbot = Bard.Chatbot(self.config["bard"]["token"], proxy=proxy)
else:
self._proxy = None

self._chatbot = Bard.Chatbot(self.config["bard"]["token"])
# Done?
if self._chatbot is not None:
logging.info("Bard module initialized")
Expand All @@ -66,24 +74,8 @@ def initialize(self) -> None:

# Error
except Exception as e:
logging.error("Error initializing Bard module!", exc_info=e)
self._enabled = False

def set_proxy(self, proxy: str) -> None:
"""
Sets new proxy from ProxyAutomation
self.config["bard"]["proxy"] must be "auto"
:param proxy: https proxy but in format http://IP:PORT
:return:
"""
if self.config["bard"]["proxy"].strip().lower() != "auto":
return

logging.info("Setting proxy {0} for Bard module".format(proxy))
self._proxy = proxy
if self._enabled and self._chatbot is not None:
self._chatbot.session.proxies.update({"http": proxy,
"https": proxy})
raise e

def process_request(self, request_response: RequestResponseContainer) -> None:
"""
Expand All @@ -94,31 +86,33 @@ def process_request(self, request_response: RequestResponseContainer) -> None:
# Check if we are initialized
if not self._enabled or self._chatbot is None:
logging.error("Bard module not initialized!")
request_response.response = self.messages["response_error"].replace("\\n", "\n") \
lang = UsersHandler.get_key_or_none(request_response.user, "lang", 0)
request_response.response = self.messages[lang]["response_error"].replace("\\n", "\n") \
.format("Bard module not initialized!")
request_response.error = True
self._restart_attempts = 0
self.processing_flag.value = False
return

try:
# Set processing flag
self.processing_flag.value = True

# Get user data
bard_conversation_id = UsersHandler.get_key_or_none(request_response.user, "bard_conversation_id")
bard_response_id = UsersHandler.get_key_or_none(request_response.user, "bard_response_id")
bard_choice_id = UsersHandler.get_key_or_none(request_response.user, "bard_choice_id")

# Increment requests_total for statistics
request_response.user["requests_total"] += 1
self.users_handler.save_user(request_response.user)

# Set conversation id, response id and choice id
if bard_conversation_id and bard_response_id and bard_choice_id:
self._chatbot.conversation_id = bard_conversation_id
self._chatbot.response_id = bard_response_id
self._chatbot.choice_id = bard_choice_id
else:
self._chatbot.conversation_id = ""
self._chatbot.response_id = ""
self._chatbot.choice_id = ""
# Try to load conversation
if bard_conversation_id:
conversation_file = os.path.join(self.config["files"]["conversations_dir"],
bard_conversation_id + ".json")
if os.path.exists(conversation_file):
logging.info("Loading conversation from {}".format(conversation_file))
self._chatbot.load_conversation(conversation_file, bard_conversation_id)
else:
bard_conversation_id = None

# Ask Bard
logging.info("Asking Bard...")
Expand All @@ -133,56 +127,64 @@ def process_request(self, request_response: RequestResponseContainer) -> None:
.format(request_response.user["user_name"], request_response.user["user_id"]))
request_response.response = bard_response["content"]

# Save user data
if self._chatbot.conversation_id and self._chatbot.response_id and self._chatbot.choice_id:
request_response.user["bard_conversation_id"] = self._chatbot.conversation_id
request_response.user["bard_response_id"] = self._chatbot.response_id
request_response.user["bard_choice_id"] = self._chatbot.choice_id
self.users_handler.save_user(request_response.user)
self._chatbot.conversation_id = ""
self._chatbot.response_id = ""
self._chatbot.choice_id = ""
# Generate new conversation id
if not bard_conversation_id:
bard_conversation_id = str(uuid.uuid4()) + "_bard"

# Save conversation
logging.info("Saving conversation to {}".format(bard_conversation_id))
self._chatbot.save_conversation(os.path.join(self.config["files"]["conversations_dir"],
bard_conversation_id + ".json"), bard_conversation_id)

# Reset attempts counter
self._restart_attempts = 0
# Save to user data
request_response.user["bard_conversation_id"] = bard_conversation_id
self.users_handler.save_user(request_response.user)

# Exit requested
except KeyboardInterrupt:
logging.warning("KeyboardInterrupt @ process_request")
self._restart_attempts = 0
return

# Bard or other error
except Exception as e:
# Try to restart
self.restart()
self._restart_attempts += 1
logging.error("Error processing request!", exc_info=e)
error_text = str(e)
if len(error_text) > 100:
error_text = error_text[:100] + "..."

# Try again 1 time
if self._restart_attempts < 2:
self.process_request(request_response)
lang = UsersHandler.get_key_or_none(request_response.user, "lang", 0)
request_response.response = self.messages[lang]["response_error"].replace("\\n", "\n").format(error_text)
request_response.error = True

# Stop restarting and respond with error
else:
request_response.response = self.messages["response_error"].replace("\\n", "\n").format(str(e))
request_response.error = True
self._restart_attempts = 0
# Finish message
BotHandler.async_helper(BotHandler.send_message_async(self.config, self.messages, request_response, end=True))

# Clear processing flag
self.processing_flag.value = False

def clear_conversation_for_user(self, user: dict) -> None:
"""
Clears conversation (chat history) for selected user
:param user:
:return: True if cleared successfully
:return:
"""
if not self._enabled or self._chatbot is None:
return
# Get conversation id
bard_conversation_id = UsersHandler.get_key_or_none(user, "bard_conversation_id")

# Check if we need to clear it
if bard_conversation_id:
# Delete file
try:
conversation_file = os.path.join(self.config["files"]["conversations_dir"],
bard_conversation_id + ".json")
if os.path.exists(conversation_file):
logging.info("Removing {}".format(conversation_file))
os.remove(conversation_file)
except Exception as e:
logging.error("Error removing conversation file!", exc_info=e)

# Reset user data
user["bard_conversation_id"] = ""
user["bard_response_id"] = ""
user["bard_choice_id"] = ""

# Save user
user["bard_conversation_id"] = None
self.users_handler.save_user(user)

def exit(self):
Expand All @@ -192,26 +194,3 @@ def exit(self):
"""
if not self._enabled or self._chatbot is None:
return
self._chatbot.session.close()

def restart(self):
"""
Restarts module and saves proxy
:return:
"""
if not self.config["modules"]["bard"]:
return
logging.info("Restarting Bard module")

# Restart
self.exit()
self.initialize()

# Set proxy
try:
if self._proxy is not None:
self._chatbot.session.proxies.update({"http": self._proxy,
"https": self._proxy})
except Exception as e:
logging.error("Error setting back proxy to Bard module!", exc_info=e)

Loading

0 comments on commit fffc285

Please sign in to comment.