fix: use logger for console telemetry

currently `print` is being used with custom formatting to achieve telemetry output in the console_span_processor

This causes telemetry not to show up in log files when using `LLAMA_STACK_LOG_FILE`. During testing it looks like telemetry is not being captures when it is

switch to using Rich formatting with the logger and then strip the formatting off when a log file is being uses so the formatting looks normal

Signed-off-by: Charlie Doern <cdoern@redhat.com>

llama_stack/log.py

@@ -6,6 +6,7 @@
 
 import logging
 import os
+import re
 import sys
 from logging.config import dictConfig
 
@@ -30,6 +31,7 @@ CATEGORIES = [
     "eval",
     "tools",
     "client",
+    "telemetry",
 ]
 
 # Initialize category levels with default level
@@ -113,6 +115,11 @@ def parse_environment_config(env_config: str) -> dict[str, int]:
     return category_levels
 
 
+def strip_rich_markup(text):
+    """Remove Rich markup tags like [dim], [bold magenta], etc."""
+    return re.sub(r"\[/?[a-zA-Z0-9 _#=,]+\]", "", text)
+
+
 class CustomRichHandler(RichHandler):
     def __init__(self, *args, **kwargs):
         kwargs["console"] = Console(width=150)
@@ -131,6 +138,19 @@ class CustomRichHandler(RichHandler):
                 self.markup = original_markup
 
 
+class CustomFileHandler(logging.FileHandler):
+    def __init__(self, filename, mode="a", encoding=None, delay=False):
+        super().__init__(filename, mode, encoding, delay)
+        # Default formatter to match console output
+        self.default_formatter = logging.Formatter("%(asctime)s %(name)s:%(lineno)d %(category)s: %(message)s")
+        self.setFormatter(self.default_formatter)
+
+    def emit(self, record):
+        if hasattr(record, "msg"):
+            record.msg = strip_rich_markup(str(record.msg))
+        super().emit(record)
+
+
 def setup_logging(category_levels: dict[str, int], log_file: str | None) -> None:
     """
     Configure logging based on the provided category log levels and an optional log file.
@@ -167,8 +187,7 @@ def setup_logging(category_levels: dict[str, int], log_file: str | None) -> None
     # Add a file handler if log_file is set
     if log_file:
         handlers["file"] = {
-            "class": "logging.FileHandler",
+            "()": CustomFileHandler,
-            "formatter": "rich",
             "filename": log_file,
             "mode": "a",
             "encoding": "utf-8",

llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py

@@ -11,19 +11,9 @@ from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanProcessor
 from opentelemetry.trace.status import StatusCode
 
-# Colors for console output
+from llama_stack.log import get_logger
-COLORS = {
+
-    "reset": "\033[0m",
+logger = get_logger(name="console_span_processor", category="telemetry")
-    "bold": "\033[1m",
-    "dim": "\033[2m",
-    "red": "\033[31m",
-    "green": "\033[32m",
-    "yellow": "\033[33m",
-    "blue": "\033[34m",
-    "magenta": "\033[35m",
-    "cyan": "\033[36m",
-    "white": "\033[37m",
-}
 
 
 class ConsoleSpanProcessor(SpanProcessor):
@@ -35,34 +25,21 @@ class ConsoleSpanProcessor(SpanProcessor):
             return
 
         timestamp = datetime.fromtimestamp(span.start_time / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
+        logger.info(f"[dim]{timestamp}[/dim] [bold magenta][START][/bold magenta] [dim]{span.name}[/dim]")
-        print(
-            f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
-            f"{COLORS['magenta']}[START]{COLORS['reset']} "
-            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
-        )
 
     def on_end(self, span: ReadableSpan) -> None:
         if span.attributes and span.attributes.get("__autotraced__"):
             return
 
         timestamp = datetime.fromtimestamp(span.end_time / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
+        span_context = f"[dim]{timestamp}[/dim] [bold magenta][END][/bold magenta] [dim]{span.name}[/dim]"
-        span_context = (
-            f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
-            f"{COLORS['magenta']}[END]{COLORS['reset']} "
-            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
-        )
-
         if span.status.status_code == StatusCode.ERROR:
-            span_context += f"{COLORS['reset']} {COLORS['red']}[ERROR]{COLORS['reset']}"
+            span_context += " [bold red][ERROR][/bold red]"
         elif span.status.status_code != StatusCode.UNSET:
-            span_context += f"{COLORS['reset']} [{span.status.status_code}]"
+            span_context += f" [{span.status.status_code}]"
-
         duration_ms = (span.end_time - span.start_time) / 1e6
-        span_context += f"{COLORS['reset']} ({duration_ms:.2f}ms)"
+        span_context += f" ({duration_ms:.2f}ms)"
-
+        logger.info(span_context)
-        print(span_context)
 
         if self.print_attributes and span.attributes:
             for key, value in span.attributes.items():
@@ -71,31 +48,26 @@ class ConsoleSpanProcessor(SpanProcessor):
                 str_value = str(value)
                 if len(str_value) > 1000:
                     str_value = str_value[:997] + "..."
-                print(f"    {COLORS['dim']}{key}: {str_value}{COLORS['reset']}")
+                logger.info(f"    [dim]{key}[/dim]: {str_value}")
 
         for event in span.events:
             event_time = datetime.fromtimestamp(event.timestamp / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
             severity = event.attributes.get("severity", "info")
             message = event.attributes.get("message", event.name)
-            if isinstance(message, dict | list):
+            if isinstance(message, dict) or isinstance(message, list):
                 message = json.dumps(message, indent=2)
-
+            severity_color = {
-            severity_colors = {
+                "error": "red",
-                "error": f"{COLORS['bold']}{COLORS['red']}",
+                "warn": "yellow",
-                "warn": f"{COLORS['bold']}{COLORS['yellow']}",
+                "info": "white",
-                "info": COLORS["white"],
+                "debug": "dim",
-                "debug": COLORS["dim"],
+            }.get(severity, "white")
-            }
+            logger.info(f" {event_time} [bold {severity_color}][{severity.upper()}][/bold {severity_color}] {message}")
-            msg_color = severity_colors.get(severity, COLORS["white"])
-
-            print(f" {event_time} {msg_color}[{severity.upper()}] {message}{COLORS['reset']}")
-
             if event.attributes:
                 for key, value in event.attributes.items():
                     if key.startswith("__") or key in ["message", "severity"]:
                         continue
-                    print(f"   {COLORS['dim']}{key}: {value}{COLORS['reset']}")
+                    logger.info(f"/r[dim]{key}[/dim]: {value}")
 
     def shutdown(self) -> None:
         """Shutdown the processor."""