feat : simpler, better terminal color

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:

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):
     """