feat: chatml-noforce-roles wrapper + cli fix

memgpt/interface.py

@@ -159,8 +159,8 @@ def printd_function_message(icon, msg, color=Fore.RED):
                 if match:
                     function_name = match.group(1)
                     function_args = match.group(2)
-                    if function_name in ["archival_memory_insert", "archival_memory_search"]:
-                        if function_name == "archival_memory_insert":
+                    if function_name in ["archival_memory_insert", "archival_memory_search", "core_memory_replace", "core_memory_append"]:
+                        if function_name in ["archival_memory_insert", "core_memory_append", "core_memory_replace"]:
                             print_function_message("🧠", f"updating memory with {function_name}")
                         elif function_name == "archival_memory_search":
                             print_function_message("🧠", f"searching memory with {function_name}")

memgpt/local_llm/llm_chat_completion_wrappers/chatml.py

@@ -40,6 +40,7 @@ def __init__(
         assistant_prefix_extra='\n{\n  "function":',
         assistant_prefix_extra_first_message='\n{\n  "function": "send_message",',
         allow_custom_roles=True,  # allow roles outside user/assistant
+        use_system_role_in_user=False,  # use the system role on user messages that don't use "type: user_message"
         # allow_function_role=True,  # use function role for function replies?
         allow_function_role=False,  # use function role for function replies?
         no_function_role_role="assistant",  # if no function role, which role to use?
@@ -56,6 +57,7 @@ def __init__(
 
         # role-based
         self.allow_custom_roles = allow_custom_roles
+        self.use_system_role_in_user = use_system_role_in_user
         self.allow_function_role = allow_function_role
         # extras for when the function role is disallowed
         self.no_function_role_role = no_function_role_role
@@ -198,6 +200,13 @@ def chat_completion_to_prompt(self, messages, functions, first_message=False):
                 role_str = message["name"].strip().lower() if (self.allow_custom_roles and "name" in message) else message["role"]
                 msg_str = self._compile_user_message(message)
 
+                if self.use_system_role_in_user:
+                    try:
+                        msg_json = json.loads(message["content"])
+                        if msg_json["type"] != "user_message":
+                            role_str = "system"
+                    except:
+                        pass
                 prompt += f"\n<|im_start|>{role_str}\n{msg_str.strip()}<|im_end|>"
 
             elif message["role"] == "assistant":

memgpt/local_llm/utils.py

@@ -50,6 +50,8 @@ def get_available_wrappers() -> dict:
         # New chatml-based wrappers
         "chatml": chatml.ChatMLInnerMonologueWrapper(),
         "chatml-noforce": chatml.ChatMLOuterInnerMonologueWrapper(),
+        # "chatml-noforce-sysm": chatml.ChatMLOuterInnerMonologueWrapper(use_system_role_in_user=True),
+        "chatml-noforce-roles": chatml.ChatMLOuterInnerMonologueWrapper(use_system_role_in_user=True, allow_function_role=True),
         # With extra hints
         "chatml-hints": chatml.ChatMLInnerMonologueWrapper(assistant_prefix_hint=True),
         "chatml-noforce-hints": chatml.ChatMLOuterInnerMonologueWrapper(assistant_prefix_hint=True),