Skip to content

Commit

Permalink
feat: Groq module + /model command + some bug fix
Browse files Browse the repository at this point in the history
  • Loading branch information
@F33RNI
F33RNI committed Apr 19, 2024
1 parent c778231 commit e1ecac0
Show file tree
Hide file tree
Showing 22 changed files with 671 additions and 21 deletions.
200 changes: 194 additions & 6 deletions bot_handler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@
BOT_COMMAND_CHAT = "chat"
BOT_COMMAND_MODULE = "module"
BOT_COMMAND_STYLE = "style"
BOT_COMMAND_MODEL = "model"
BOT_COMMAND_CLEAR = "clear"
BOT_COMMAND_LANG = "lang"
BOT_COMMAND_CHAT_ID = "chatid"
Expand Down Expand Up @@ -121,6 +122,8 @@ def __init__(
logging_queue: multiprocessing.Queue,
queue_handler_: queue_handler.QueueHandler,
modules: Dict,
web_cooldown_timer: multiprocessing.Value,
web_request_lock: multiprocessing.Lock,
):
self.config = config
self.config_file = config_file
Expand All @@ -130,6 +133,10 @@ def __init__(
self.queue_handler = queue_handler_
self.modules = modules

# LMAO
self.web_cooldown_timer = web_cooldown_timer
self.web_request_lock = web_request_lock

self.prevent_shutdown_flag = multiprocessing.Value(c_bool, False)

self._application = None
Expand Down Expand Up @@ -172,6 +179,7 @@ def start_bot(self):
self._application.add_handler(CaptionCommandHandler(BOT_COMMAND_CHAT, self.bot_module_request))
self._application.add_handler(CaptionCommandHandler(BOT_COMMAND_MODULE, self.bot_command_module))
self._application.add_handler(CaptionCommandHandler(BOT_COMMAND_STYLE, self.bot_command_style))
self._application.add_handler(CaptionCommandHandler(BOT_COMMAND_MODEL, self.bot_command_model))
self._application.add_handler(CaptionCommandHandler(BOT_COMMAND_CLEAR, self.bot_command_clear))
self._application.add_handler(CaptionCommandHandler(BOT_COMMAND_LANG, self.bot_command_lang))
self._application.add_handler(CaptionCommandHandler(BOT_COMMAND_CHAT_ID, self.bot_command_chatid))
Expand Down Expand Up @@ -302,15 +310,13 @@ async def query_callback(self, update: Update, context: ContextTypes.DEFAULT_TYP
return

# Parse data from markup
action, data_, reply_message_id = data_.split("|")
action, data_, argument_ = data_.split("|")
if not action:
raise Exception("No action in callback data")
if not data_:
data_ = None
if not reply_message_id:
reply_message_id = None
else:
reply_message_id = int(reply_message_id.strip())
if not argument_:
argument_ = None

# Get user
banned, user = await self._user_get_check(update, context, prompt_language_selection=False)
Expand All @@ -329,6 +335,12 @@ async def query_callback(self, update: Update, context: ContextTypes.DEFAULT_TYP

# Regenerate request
if action == "regenerate":
# Parse message ID
if not argument_:
reply_message_id = None
else:
reply_message_id = int(argument_.strip())

# Get last message ID
reply_message_id_last = self.users_handler.get_key(0, "reply_message_id_last", user=user)
if reply_message_id_last is None or reply_message_id_last != reply_message_id:
Expand Down Expand Up @@ -364,6 +376,12 @@ async def query_callback(self, update: Update, context: ContextTypes.DEFAULT_TYP

# Continue generating
elif action == "continue":
# Parse message ID
if not argument_:
reply_message_id = None
else:
reply_message_id = int(argument_.strip())

# Get last message ID
reply_message_id_last = self.users_handler.get_key(0, "reply_message_id_last", user=user)
if reply_message_id_last is None or reply_message_id_last != reply_message_id:
Expand All @@ -385,6 +403,12 @@ async def query_callback(self, update: Update, context: ContextTypes.DEFAULT_TYP

# Send suggestion
elif action == "suggestion":
# Parse message ID
if not argument_:
reply_message_id = None
else:
reply_message_id = int(argument_.strip())

# Get last message ID
reply_message_id_last = self.users_handler.get_key(0, "reply_message_id_last", user=user)
if reply_message_id_last is None or reply_message_id_last != reply_message_id:
Expand Down Expand Up @@ -421,6 +445,12 @@ async def query_callback(self, update: Update, context: ContextTypes.DEFAULT_TYP

# Stop generating
elif action == "stop":
# Parse message ID
if not argument_:
reply_message_id = None
else:
reply_message_id = int(argument_.strip())

# Get last message ID
reply_message_id_last = self.users_handler.get_key(0, "reply_message_id_last", user=user)
if reply_message_id_last is None or reply_message_id_last != reply_message_id:
Expand Down Expand Up @@ -466,6 +496,10 @@ async def query_callback(self, update: Update, context: ContextTypes.DEFAULT_TYP
elif action == "style":
await self._bot_command_style_raw(data_, user, context)

# Change model
elif action == "model":
await self._bot_command_model_raw(data_, argument_, user, context)

# Change language
elif action == "lang":
await self._bot_command_lang_raw(data_, user, context)
Expand Down Expand Up @@ -752,8 +786,20 @@ async def bot_command_restart(self, update: Update, context: ContextTypes.DEFAUL
continue
logging.info(f"Trying to load and initialize {module_name} module")
try:
use_web = (
module_name.startswith("lmao_")
and module_name in self.config.get("modules").get("lmao_web_for_modules", [])
and "lmao_web_api_url" in self.config.get("modules")
)
module = module_wrapper_global.ModuleWrapperGlobal(
module_name, self.config, self.messages, self.users_handler, self.logging_queue
module_name,
self.config,
self.messages,
self.users_handler,
self.logging_queue,
use_web=use_web,
web_cooldown_timer=self.web_cooldown_timer,
web_request_lock=self.web_request_lock,
)
self.modules[module_name] = module
reload_logs += f"Intialized and loaded {module_name} module\n"
Expand Down Expand Up @@ -1042,6 +1088,148 @@ async def _bot_command_style_raw(self, style: str or None, user: Dict, context:
context,
)

async def bot_command_model(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
"""/model commands callback
Args:
update (Update): update object from bot's callback
context (ContextTypes.DEFAULT_TYPE): context object from bot's callback
"""
# Get user
banned, user = await self._user_get_check(update, context)
if user is None:
return
user_id = user.get("user_id")
user_name = self.users_handler.get_key(0, "user_name", "", user=user)
lang_id = self.users_handler.get_key(0, "lang_id", user=user)

# Log command
logging.info(f"/model command from {user_name} ({user_id})")

# Exit if banned
if banned:
return

module_id = self.users_handler.get_key(0, "module", self.config.get("modules").get("default"), user=user)

model = None

# User specified model
if context.args and len(context.args) >= 1:
try:
model = context.args[0].strip().lower()

# Get available models
current_module_id = self.users_handler.get_key(
0, "module", self.config.get("modules").get("default"), user=user
)
available_models = self.config.get(current_module_id).get("models", [])

# Get current model
model_current = self.config.get(module_id).get("model_default")
model_current = self.users_handler.get_key(0, f"{module_id}_model", model_current, user=user)

# Check
if not model_current or len(available_models) == 0:
await _send_safe(
user_id,
self.messages.get_message("model_no_models", lang_id=lang_id),
context,
)
return

# Check
if model not in available_models:
raise Exception(f"No model {model} in {' '.join(available_models)}")
except Exception as e:
logging.error("Error retrieving requested model", exc_info=e)
await _send_safe(
user["user_id"],
self.messages.get_message("model_change_error", lang_id=lang_id).format(error_text=str(e)),
context,
)
return

# Change model or ask the user
await self._bot_command_model_raw(module_id, model, user, context)

async def _bot_command_model_raw(
self, module_id: str or None, model: str or None, user: Dict, context: ContextTypes.DEFAULT_TYPE
) -> None:
"""Changes model of module
Args:
module_id (str or None): id of module to change model of
model (str or None): model name or None to ask user
user (Dict): user's data as dictionary
context (ContextTypes.DEFAULT_TYPE): context object from bot's callback
"""
user_id = user.get("user_id")
lang_id = self.users_handler.get_key(0, "lang_id", user=user)

# Extract current user's module and model
module_icon_names = self.messages.get_message("modules", lang_id=lang_id)
if not module_id:
module_id = self.users_handler.get_key(0, "module", self.config.get("modules").get("default"), user=user)
current_module_name = module_icon_names.get(module_id).get("name")
current_module_icon = module_icon_names.get(module_id).get("icon")
current_module_name = f"{current_module_icon} {current_module_name}"

# Get available models
available_models = self.config.get(module_id).get("models", [])

# Get current model
model_current = self.config.get(module_id).get("model_default")
model_current = self.users_handler.get_key(0, f"{module_id}_model", model_current, user=user)

# Check
if not model_current or len(available_models) == 0:
await _send_safe(
user_id,
self.messages.get_message("model_no_models", lang_id=lang_id),
context,
)
return

# Ask user
if not model:
buttons = []
for model_ in available_models:
buttons.append(InlineKeyboardButton(model_, callback_data=f"model|{module_id}|{model_}"))

await _send_safe(
user_id,
self.messages.get_message("model_select", lang_id=lang_id).format(
module_name=current_module_name, current_model=model_current
),
context,
reply_markup=InlineKeyboardMarkup(bot_sender.build_menu(buttons)),
)
return

# Change model
try:
# Change model of user
self.users_handler.set_key(user_id, f"{module_id}_model", model)

# Send confirmation
await _send_safe(
user_id,
self.messages.get_message("model_changed", lang_id=lang_id).format(
module_name=current_module_name, changed_model=model
),
context,
)

# Error changing model
except Exception as e:
logging.error("Error changing model", exc_info=e)
await _send_safe(
user_id,
self.messages.get_message("model_change_error", lang_id=lang_id).format(error_text=str(e)),
context,
)

########################################
# General (non-modules) commands below #
########################################
Expand Down
8 changes: 8 additions & 0 deletions bot_sender.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,14 @@ def build_markup(
)
buttons.append(button_style)

# Add change model button
if request_response.module_name in module_wrapper_global.MODULES_WITH_MODELS:
button_model = InlineKeyboardButton(
messages_.get_message("button_model_change", user_id=user_id),
callback_data=f"model|{request_response.module_name}|",
)
buttons.append(button_model)

# Add change module button for all modules
button_module = InlineKeyboardButton(
messages_.get_message("button_module", user_id=user_id),
Expand Down
9 changes: 7 additions & 2 deletions config.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,15 @@
"config_version": 7,

"__comment02__": "General config of enabled modules",
"__comment03__": "Available modules: lmao_chatgpt, lmao_ms_copilot, ms_copilot, ms_copilot_designer, gemini",
"__comment03__": "Available modules: lmao_chatgpt, lmao_ms_copilot, ms_copilot, ms_copilot_designer, gemini, groq",
"__comment04__": "NOTE: ms_copilot and ms_copilot_designer modules are DEPRECATED and will be removed soon",
"modules": {
"__comment01__": "Enabled modules",
"enabled": [
"lmao_chatgpt",
"lmao_ms_copilot",
"gemini"
"gemini",
"groq"
],

"__comment02__": "Default (initial) module for handling user messages",
Expand Down Expand Up @@ -99,6 +100,10 @@
"command": "module",
"description": "↕ Change module to chat with"
},
{
"command": "model",
"description": "⚙️ Change model of module (ex. For 🔴 Groq module)"
},
{
"command": "clear",
"description": "🧹 Clear chat history"
Expand Down
12 changes: 5 additions & 7 deletions google_ai_module.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,10 +155,12 @@ def process_request(self, request_response: RequestResponseContainer) -> None:
module_config = self.config.get(_NAME)

# Cool down
if time.time() - self._last_request_time.value <= module_config.get("cooldown_seconds"):
time_to_wait = module_config.get("cooldown_seconds") - (time.time() - self._last_request_time.value)
if time.time() - self._last_request_time.value <= module_config.get("user_cooldown_seconds"):
time_to_wait = module_config.get("user_cooldown_seconds") - (
time.time() - self._last_request_time.value
)
logging.warning(f"Too frequent requests. Waiting {time_to_wait} seconds...")
time.sleep(self._last_request_time.value + module_config.get("cooldown_seconds") - time.time())
time.sleep(self._last_request_time.value + module_config.get("user_cooldown_seconds") - time.time())
self._last_request_time.value = time.time()

response = None
Expand Down Expand Up @@ -224,10 +226,6 @@ def process_request(self, request_response: RequestResponseContainer) -> None:
# Save conversation ID
self.users_handler.set_key(request_response.user_id, f"{_NAME}_conversation_id", conversation_id)

# Error
except Exception as e:
raise e

finally:
self.processing_flag.value = False

Expand Down

0 comments on commit e1ecac0

Please sign in to comment.