diff --git a/memory_scope/chat/cli_memory_chat.py b/memory_scope/chat/cli_memory_chat.py index 53cfa650..5e0ca20f 100644 --- a/memory_scope/chat/cli_memory_chat.py +++ b/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!") diff --git a/memory_scope/cli.py b/memory_scope/cli.py index a76c61db..a0cdda2b 100644 --- a/memory_scope/cli.py +++ b/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 diff --git a/memory_scope/memory/operation/summary_memory.py b/memory_scope/memory/operation/summary_memory.py index 78cc84fb..bef933eb 100644 --- a/memory_scope/memory/operation/summary_memory.py +++ b/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: diff --git a/memory_scope/memory/operation/write_memory.py b/memory_scope/memory/operation/write_memory.py index eecc45c5..362ec10f 100644 --- a/memory_scope/memory/operation/write_memory.py +++ b/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: diff --git a/memory_scope/memory/worker/dummy_worker.py b/memory_scope/memory/worker/dummy_worker.py index 5ae68efc..2564b4eb 100644 --- a/memory_scope/memory/worker/dummy_worker.py +++ b/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}") diff --git a/memory_scope/utils/tool_functions.py b/memory_scope/utils/tool_functions.py index aeb8dfbc..a9d0a236 100644 --- a/memory_scope/utils/tool_functions.py +++ b/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