From 485b85ee76329c1d45473102486abd3e5d1388f5 Mon Sep 17 00:00:00 2001
From: TheInvisibleMage <64934262+TheInvisibleMage@users.noreply.github.com>
Date: Sun, 31 Dec 2023 16:03:23 +1100
Subject: [PATCH] Superboogav2 Quick Fixes (#5089)

---
 extensions/superboogav2/chat_handler.py     | 28 ++++++---------------
 extensions/superboogav2/notebook_handler.py |  4 +--
 extensions/superboogav2/script.py           |  4 +--
 3 files changed, 11 insertions(+), 25 deletions(-)

diff --git a/extensions/superboogav2/chat_handler.py b/extensions/superboogav2/chat_handler.py
index 215f7bdb..419b9264 100644
--- a/extensions/superboogav2/chat_handler.py
+++ b/extensions/superboogav2/chat_handler.py
@@ -1,14 +1,14 @@
 """
 This module is responsible for modifying the chat prompt and history.
 """
-import json
 import re
 
 import extensions.superboogav2.parameters as parameters
 
-from modules import chat
+from modules import chat, shared
 from modules.text_generation import get_encoded_length
 from modules.logging_colors import logger
+from modules.chat import load_character_memoized
 from extensions.superboogav2.utils import create_context_text, create_metadata_source
 
 from .data_processor import process_and_add_to_collector
@@ -17,14 +17,6 @@ from .chromadb import ChromaCollector
 
 CHAT_METADATA = create_metadata_source('automatic-chat-insert')
 
-INSTRUCT_MODE = 'instruct'
-CHAT_INSTRUCT_MODE = 'chat-instruct'
-
-
-def _is_instruct_mode(state: dict):
-    mode = state.get('mode')
-    return mode == INSTRUCT_MODE or mode == CHAT_INSTRUCT_MODE
-
 
 def _remove_tag_if_necessary(user_input: str):
     if not parameters.get_is_manual():
@@ -51,17 +43,11 @@ def _format_single_exchange(name, text):
 
 
 def _get_names(state: dict):
-    if _is_instruct_mode(state):
-        user_name = state['name1_instruct']
-        bot_name = state['name2_instruct']
-    else:
-        user_name = state['name1']
-        bot_name = state['name2']
-
-    if not user_name:
-        user_name = 'User'
-    if not bot_name:
-        bot_name = 'Assistant'
+    default_char = shared.settings.get('character', "Assistant")
+    default_user = shared.settings.get('name1', "You")
+    character = state.get('character', default_char)
+    user_name = state.get('name1', default_user)
+    user_name, bot_name, _, _, _ = load_character_memoized(character, user_name, '')
 
     return user_name, bot_name
 
diff --git a/extensions/superboogav2/notebook_handler.py b/extensions/superboogav2/notebook_handler.py
index 9faadfed..7b864349 100644
--- a/extensions/superboogav2/notebook_handler.py
+++ b/extensions/superboogav2/notebook_handler.py
@@ -16,9 +16,9 @@ def _remove_special_tokens(string):
     return re.sub(pattern, '', string)
 
 
-def input_modifier_internal(string, collector):
+def input_modifier_internal(string, collector, is_chat):
     # Sanity check.
-    if shared.is_chat():
+    if is_chat:
         return string
 
     # Find the user input
diff --git a/extensions/superboogav2/script.py b/extensions/superboogav2/script.py
index 0870ab4c..66f56e29 100644
--- a/extensions/superboogav2/script.py
+++ b/extensions/superboogav2/script.py
@@ -167,8 +167,8 @@ def custom_generate_chat_prompt(user_input, state, **kwargs):
     return custom_generate_chat_prompt_internal(user_input, state, collector, **kwargs)
 
 
-def input_modifier(string):
-    return input_modifier_internal(string, collector)
+def input_modifier(string, state, is_chat=False):
+    return input_modifier_internal(string, collector, is_chat)
 
 
 def ui():