diff --git a/server/sources/llamacpp.py b/server/sources/llamacpp.py index 3e17445..758be3f 100644 --- a/server/sources/llamacpp.py +++ b/server/sources/llamacpp.py @@ -28,7 +28,8 @@ class LlamacppLLM(GeneratorLLM): output = self.llm.create_chat_completion( messages = history ) - self.state.current_buffer = output + with self.state.lock: + self.state.current_buffer = output except Exception as e: self.logger.error(f"Error: {e}") finally: diff --git a/sources/utility.py b/sources/utility.py index 27eebb4..9e502ae 100644 --- a/sources/utility.py +++ b/sources/utility.py @@ -6,6 +6,32 @@ import threading import itertools import time +global thinking_toggle +thinking_toggle = False + +def get_color_map(): + if platform.system().lower() != "windows": + color_map = { + "success": Fore.GREEN, + "failure": Fore.RED, + "status": Fore.LIGHTGREEN_EX, + "code": Fore.LIGHTBLUE_EX, + "warning": Fore.YELLOW, + "output": Fore.LIGHTCYAN_EX, + "info": Fore.CYAN + } + else: + color_map = { + "success": "green", + "failure": "red", + "status": "light_green", + "code": "light_blue", + "warning": "yellow", + "output": "cyan", + "info": "black" + } + return color_map + def pretty_print(text, color="info"): """ Print text with color formatting. @@ -22,54 +48,38 @@ def pretty_print(text, color="info"): - "output": Cyan - "default": Black (Windows only) """ - if platform.system().lower() != "windows": - color_map = { - "success": Fore.GREEN, - "failure": Fore.RED, - "status": Fore.LIGHTGREEN_EX, - "code": Fore.LIGHTBLUE_EX, - "warning": Fore.YELLOW, - "output": Fore.LIGHTCYAN_EX, - "info": Fore.CYAN - } - if color not in color_map: - print(text) - pretty_print(f"Invalid color {color} in pretty_print", "warning") - return - print(f"{color_map[color]}{text}{Fore.RESET}") - print(' ' * 10) # prevent cut-off - else: - color_map = { - "success": "green", - "failure": "red", - "status": "light_green", - "code": "light_blue", - "warning": "yellow", - "output": "cyan", - "default": "black" - } - if color not in color_map: - color = "default" - print(colored(text, color_map[color])) - print(' ' * 10) # prevent cut-off + global thinking_toggle + thinking_toggle = False + color_map = get_color_map() + if color not in color_map: + color = "info" + print(colored(text, color_map[color])) + print(' ' * 10) # prevent cut-off -def animate_thinking(text, color="status", duration=2): +def animate_thinking(text, color="status", duration=30): + """ + Animate a thinking spinner while a task is being executed. + It use a daemon thread to run the animation. This will not block the main thread. + Color are the same as pretty_print. + """ + global thinking_toggle + thinking_toggle = False def _animate(): - color_map = { - "success": (Fore.GREEN, "green"), - "failure": (Fore.RED, "red"), - "status": (Fore.LIGHTGREEN_EX, "light_green"), - "code": (Fore.LIGHTBLUE_EX, "light_blue"), - "warning": (Fore.YELLOW, "yellow"), - "output": (Fore.LIGHTCYAN_EX, "cyan"), - "default": (Fore.RESET, "black"), - "info": (Fore.CYAN, "cyan") - } - fore_color, term_color = color_map.get(color, color_map["default"]) - spinner = itertools.cycle(['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']) + global thinking_toggle + color_map = get_color_map() + fore_color, term_color = color_map.get(color, color_map["info"]) + #spinner = itertools.cycle(['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']) + spinner = itertools.cycle([ + '█▒▒▒▒▒', '██▒▒▒▒', '███▒▒▒', '████▒▒', '█████▒', '██████', + '█████▒', '████▒▒', '███▒▒▒', '██▒▒▒▒', '█▒▒▒▒▒', '▒█▒▒▒█', + '▒▒█▒█▒', '▒▒▒██▒', '▒█▒█▒▒', '█▒▒▒▒▒' + ]) end_time = time.time() + duration while time.time() < end_time: + if not thinking_toggle: + # stop if another text is printed + break symbol = next(spinner) if platform.system().lower() != "windows": print(f"{fore_color}{symbol} {text}{Fore.RESET}", flush=True) @@ -77,9 +87,9 @@ def animate_thinking(text, color="status", duration=2): print(colored(f"{symbol} {text}", term_color), flush=True) time.sleep(0.1) print("\033[1A\033[K", end="", flush=True) - animation_thread = threading.Thread(target=_animate) + thinking_toggle = True + animation_thread = threading.Thread(target=_animate, daemon=True) animation_thread.start() - animation_thread.join() def timer_decorator(func): """