Improve stdout logging

Author: cyclotruc

src/process_query.py

@@ -4,10 +4,33 @@
 
 from config import MAX_DISPLAY_SIZE, EXAMPLE_REPOS
 from gitingest import ingest_from_query, clone_repo, parse_query
-from server_utils import logSliderToSize
+from server_utils import logSliderToSize, Colors
 
 templates = Jinja2Templates(directory="templates")
 
+def print_query(query, request, max_file_size, pattern_type, pattern):
+    print(f"{Colors.WHITE}{query['url']:<20}{Colors.END}", end="")
+    if int(max_file_size/1024) != 50:
+        print(f" | {Colors.PURPLE}Size: {int(max_file_size/1024)}kb{Colors.END}", end="")
+    if pattern_type == "include" and pattern != "":
+        print(f" | {Colors.BLUE}Include {pattern}{Colors.END}", end="")
+    elif pattern_type == "exclude" and pattern != "":
+        print(f" | {Colors.YELLOW}Exclude {pattern}{Colors.END}", end="")
+
+
+def print_error(query, request, e, max_file_size, pattern_type, pattern):
+    print(f"{Colors.BROWN}WARNING{Colors.END}: {Colors.RED}<- {Colors.END}", end="")
+    print_query(query, request, max_file_size, pattern_type, pattern)
+    print(f" | {Colors.RED}{e}{Colors.END}")
+
+def print_success(query, request, max_file_size, pattern_type, pattern, summary):
+    estimated_tokens = summary[summary.index("Estimated tokens:") + len("Estimated ") :]
+    print(f"{Colors.GREEN}INFO{Colors.END}: {Colors.GREEN}<- {Colors.END}", end="")
+    print_query(query, request, max_file_size, pattern_type, pattern)
+    print(f" | {Colors.PURPLE}{estimated_tokens}{Colors.END}")
+
+
+
 async def process_query(request: Request, input_text: str, slider_position: int, pattern_type: str = "exclude", pattern: str = "", is_index: bool = False) -> str:
     template = "index.jinja" if is_index else "github.jinja"
     max_file_size = logSliderToSize(slider_position)
@@ -23,11 +46,14 @@ async def process_query(request: Request, input_text: str, slider_position: int,
         summary, tree, content = ingest_from_query(query)
         with open(f"{query['local_path']}.txt", "w") as f:
             f.write(tree + "\n" + content)
-        print(f"{query['slug']:<20}", end="")
-        if pattern and pattern != "":
-            print(f"{pattern_type}[{pattern}]", end="")
-        print(f"\n{query['url']}")
+
+        
+        
     except Exception as e:
+        if 'query' in locals() and query is not None and isinstance(query, dict):
+                print_error(query, request, e, max_file_size, pattern_type, pattern)
+        else:
+            print(f"{Colors.RED}Error: {e}{Colors.END}")
         return templates.TemplateResponse(
             template, 
             {
@@ -43,7 +69,7 @@ async def process_query(request: Request, input_text: str, slider_position: int,
 
     if len(content) > MAX_DISPLAY_SIZE:
         content = f"(Files content cropped to {int(MAX_DISPLAY_SIZE/1000)}k characters, download full ingest to see more)\n" + content[:MAX_DISPLAY_SIZE]
-        
+    print_success(query, request, max_file_size, pattern_type, pattern, summary)
     return templates.TemplateResponse(
         template, 
         {

src/server_utils.py

@@ -1,4 +1,3 @@
-
 ## Rate Limiter
 from slowapi import Limiter
 from slowapi.util import get_remote_address
@@ -12,4 +11,32 @@ def logSliderToSize(position):
     minv = math.log(1)
     maxv = math.log(102400)
 
-    return round(math.exp(minv + (maxv - minv) * pow(position / maxp, 1.5))) * 1024
+    return round(math.exp(minv + (maxv - minv) * pow(position / maxp, 1.5))) * 1024
+
+## Color printing utility
+class Colors:
+    """ANSI color codes"""
+    BLACK = "\033[0;30m"
+    RED = "\033[0;31m"
+    GREEN = "\033[0;32m"
+    BROWN = "\033[0;33m"
+    BLUE = "\033[0;34m"
+    PURPLE = "\033[0;35m"
+    CYAN = "\033[0;36m"
+    LIGHT_GRAY = "\033[0;37m"
+    DARK_GRAY = "\033[1;30m"
+    LIGHT_RED = "\033[1;31m"
+    LIGHT_GREEN = "\033[1;32m"
+    YELLOW = "\033[1;33m"
+    LIGHT_BLUE = "\033[1;34m"
+    LIGHT_PURPLE = "\033[1;35m"
+    LIGHT_CYAN = "\033[1;36m"
+    WHITE = "\033[1;37m"
+    BOLD = "\033[1m"
+    FAINT = "\033[2m"
+    ITALIC = "\033[3m"
+    UNDERLINE = "\033[4m"
+    BLINK = "\033[5m"
+    NEGATIVE = "\033[7m"
+    CROSSED = "\033[9m"
+    END = "\033[0m"