[dev] print char logo in shell env

memory_scope/chat/cli_memory_chat.py

@@ -13,6 +13,7 @@ from memory_scope.prompts.memory_chat_prompt import MEMORY_PROMPT, SYSTEM_PROMPT
 from memory_scope.scheme.message import Message
 from memory_scope.scheme.model_response import ModelResponse, ModelResponseGen
 from memory_scope.utils.logger import Logger
+from memory_scope.utils.tool_functions import char_logo
 
 
 class CliMemoryChat(BaseMemoryChat):
@@ -28,8 +29,13 @@ class CliMemoryChat(BaseMemoryChat):
         self.stream: bool = stream
         self.kwargs: dict = kwargs
 
+        self._logo = char_logo("MemoryScope")
         self.logger = Logger.get_logger()
 
+    def print_logo(self):
+        for line in self._logo:
+            print(line)
+
     @property
     def memory_service(self) -> BaseMemoryService:
         if isinstance(self._memory_service, str):
@@ -63,11 +69,12 @@ class CliMemoryChat(BaseMemoryChat):
         self.memory_service.add_messages(new_message)
         related_memories: List[str] = self.memory_service.read_memory()
         system_message: Message = self.get_system_prompt(related_memories, time_created)
+        result = self.generation_model.call(messages=[system_message, new_message], stream=self.stream)
         if self.stream:
-            for result in self.generation_model.call(messages=[system_message, new_message], stream=self.stream):
-                yield result
-
-        self.memory_service.add_messages(result.message)
+            for _ in result:
+                yield _
+        else:
+            return result
 
     def process_commands(self, query: str) -> bool:
         continue_run = True
@@ -81,8 +88,8 @@ class CliMemoryChat(BaseMemoryChat):
         elif query == "help":
             questionary.print("CLI commands", "bold")
             for cmd, desc in self.USER_COMMANDS.items():
-                questionary.print(cmd, "bold")
-                questionary.print(f" {desc}")
+                questionary.print(text=f" /{cmd}:", style="bold")
+                questionary.print(text=f"  {desc}")
 
         elif query == "stream":
             self.stream = bool(args[0])
@@ -109,21 +116,18 @@ class CliMemoryChat(BaseMemoryChat):
         return continue_run
 
     def run(self):
-        self.USER_COMMANDS.update({f"/{k}": v for k, v in self.memory_service.op_description_dict.items()})
+        self.print_logo()
+        self.USER_COMMANDS.update(self.memory_service.op_description_dict)
 
         while True:
             try:
-                query = questionary.text(
-                    message="Please enter your message or command:",
-                    multiline=False,
-                    qmark=">",
-                ).ask()
-                query: str = query.strip()
+                query = questionary.text(message="user:", multiline=False, qmark=">").ask()
 
-                if query == "":
-                    questionary.print("Empty input received. Please try again!")
+                if not query:
                     continue
 
+                query: str = query.strip()
+
                 # handle cli / commands with memory ops
                 if query.startswith("/"):
                     if self.process_commands(query=query):
@@ -131,6 +135,9 @@ class CliMemoryChat(BaseMemoryChat):
                     else:
                         break
 
+                msg = None
+                questionary.print("> ", end="", style="fg:yellow")
+                questionary.print("assistant: ", end="", style="bold")
                 if self.stream:
                     for msg in self.chat_with_memory(query=query):
                         questionary.print(msg.delta, end="")
@@ -138,6 +145,7 @@ class CliMemoryChat(BaseMemoryChat):
                 else:
                     msg = self.chat_with_memory(query=query)
                     questionary.print(msg.message.content)
+                self.memory_service.add_messages(msg.message)
 
             except KeyboardInterrupt:
                 questionary.print("User interrupt occurred.")
@@ -147,6 +155,9 @@ class CliMemoryChat(BaseMemoryChat):
                     break
 
             except Exception as e:
-                questionary.print(f"An exception occurred when running cli memory chat. args={e.args}")
-                self.logger.exception(f"An exception occurred when running cli memory chat. args={e.args}")
+                line = f"An exception occurred when running cli memory chat. args={e.args}"
+                questionary.print(line)
+                self.logger.exception(line)
                 continue
+
+        questionary.print(f"A memory writing thread is still running, please be patient and wait!")

memory_scope/cli.py

@@ -1,6 +1,6 @@
 import sys
 
-sys.path.append(".")
+sys.path.append(".")  # noqa: E402
 
 import json
 from concurrent.futures import ThreadPoolExecutor

memory_scope/memory/operation/summary_memory.py

@@ -39,8 +39,13 @@ class SummaryMemory(BaseWorkflow, BaseOperation):
 
     def _loop_operation(self):
         while self._loop_switch:
-            time.sleep(self.interval_time)
-            self.run_operation()
+            for _ in range(self.interval_time):
+                if self._loop_switch:
+                    time.sleep(1)
+                else:
+                    break
+            if self._loop_switch:
+                self.run_operation()
 
     def run_operation_backend(self):
         if not self._loop_switch:

memory_scope/memory/operation/write_memory.py

@@ -66,8 +66,13 @@ class WriteMemory(BaseWorkflow, BaseOperation):
 
     def _loop_operation(self):
         while self._loop_switch:
-            time.sleep(self.interval_time)
-            self.run_operation()
+            for _ in range(self.interval_time):
+                if self._loop_switch:
+                    time.sleep(1)
+                else:
+                    break
+            if self._loop_switch:
+                self.run_operation()
 
     def run_operation_backend(self):
         if not self._loop_switch:

memory_scope/memory/worker/dummy_worker.py

@@ -1,3 +1,5 @@
+import datetime
+
 from memory_scope.constants.common_constants import RESULT, WORKFLOW_NAME
 from memory_scope.memory.worker.base_worker import BaseWorker
 
@@ -6,4 +8,5 @@ class DummyWorker(BaseWorker):
     def _run(self):
         workflow_name = self.get_context(WORKFLOW_NAME)
         self.logger.info(f"enter workflow={workflow_name}.dummy_worker!")
-        self.set_context(RESULT, f"test {workflow_name}")
+        ts = int(datetime.datetime.now().timestamp())
+        self.set_context(RESULT, f"test {workflow_name} ts={ts}")

memory_scope/utils/tool_functions.py

@@ -1,11 +1,19 @@
+import random
 import re
+import time
 from copy import deepcopy
 from datetime import datetime
 from importlib import import_module
+from typing import get_args
+
+import pyfiglet
+from termcolor import colored, COLORS
+from termcolor._types import Color
 
 from memory_scope.constants.common_constants import WEEKDAYS
 from memory_scope.enumeration.message_role_enum import MessageRoleEnum
 
+ALL_COLORS = get_args(COLORS)
 
 def under_line_to_hump(underline_str):
     sub = re.sub(r"(_\w)", lambda x: x.group(1)[1].upper(), underline_str)
@@ -68,11 +76,10 @@ def get_datetime_info_dict(parse_dt: datetime):
     }
 
 
-def time_to_formatted_str(
-    time: datetime | str | int | float = None,
-    date_format: str = "%Y%m%d",  # e.g. %Y%m%d -> "20240528", add %H:%M:%S
-    string_format: str = "",
-) -> str:
+def time_to_formatted_str(time: datetime | str | int | float = None,
+                          date_format: str = "%Y%m%d",  # e.g. %Y%m%d -> "20240528", add %H:%M:%S
+                          string_format: str = "") -> str:
+
     if isinstance(time, str | int | float):
         if isinstance(time, str):
             time = float(time)
@@ -89,3 +96,22 @@ def time_to_formatted_str(
         return_str = string_format.format(**get_datetime_info_dict(current_dt))
 
     return return_str
+
+
+def char_logo(words: str, seed: int = time.time_ns(), color: Color = None):
+    font = pyfiglet.Figlet()
+    rendered_text = font.renderText(words)
+    colored_lines = []
+    all_colors = list(COLORS.keys())
+    random.seed = seed
+    for line in rendered_text.splitlines():
+        line_color = color
+        if line_color is None:
+            random.shuffle(all_colors)
+            line_color = all_colors[0]
+        colored_line = ""
+        for char in line:
+            colored_char = colored(char, line_color, attrs=['bold'])
+            colored_line += colored_char
+        colored_lines.append(colored_line)
+    return colored_lines