Console span processor improvements (#577)

Makes the console span processor output spans in less prominent way and
highlight the logs based on severity.


![Screenshot 2024-12-06 at 11 26
46 AM](https://github.com/user-attachments/assets/c3a1b051-85db-4b71-b7a5-7bab5a26f072)

llama_stack/providers/inline/memory/faiss/faiss.py

@@ -27,7 +27,6 @@ from llama_stack.providers.utils.memory.vector_store import (
     BankWithIndex,
     EmbeddingIndex,
 )
-from llama_stack.providers.utils.telemetry import tracing
 
 from .config import FaissImplConfig
 
@@ -95,7 +94,6 @@ class FaissIndex(EmbeddingIndex):
 
         await self.kvstore.delete(f"faiss_index:v1::{self.bank_id}")
 
-    @tracing.span(name="add_chunks")
     async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
         indexlen = len(self.id_by_index)
         for i, chunk in enumerate(chunks):

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

@@ -4,10 +4,12 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import json
 from datetime import datetime
 
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanProcessor
+from opentelemetry.trace.status import StatusCode
 
 # Colors for console output
 COLORS = {
@@ -25,10 +27,11 @@ COLORS = {
 
 
 class ConsoleSpanProcessor(SpanProcessor):
-    """A SpanProcessor that prints spans to the console with color formatting."""
+
+    def __init__(self, print_attributes: bool = False):
+        self.print_attributes = print_attributes
 
     def on_start(self, span: ReadableSpan, parent_context=None) -> None:
-        """Called when a span starts."""
         if span.attributes and span.attributes.get("__autotraced__"):
             return
 
@@ -39,11 +42,10 @@ class ConsoleSpanProcessor(SpanProcessor):
         print(
             f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
             f"{COLORS['magenta']}[START]{COLORS['reset']} "
-            f"{COLORS['cyan']}{span.name}{COLORS['reset']}"
+            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
         )
 
     def on_end(self, span: ReadableSpan) -> None:
-        """Called when a span ends."""
         if span.attributes and span.attributes.get("__autotraced__"):
             return
 
@@ -51,50 +53,60 @@ class ConsoleSpanProcessor(SpanProcessor):
             "%H:%M:%S.%f"
         )[:-3]
 
-        # Build the span context string
         span_context = (
             f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
             f"{COLORS['magenta']}[END]{COLORS['reset']} "
-            f"{COLORS['cyan']}{span.name}{COLORS['reset']} "
+            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
         )
 
-        # Add status if not OK
-        if span.status.status_code != 0:  # UNSET or ERROR
-            status_color = (
-                COLORS["red"] if span.status.status_code == 2 else COLORS["yellow"]
-            )
-            span_context += (
-                f" {status_color}[{span.status.status_code}]{COLORS['reset']}"
-            )
+        if span.status.status_code == StatusCode.ERROR:
+            span_context += f"{COLORS['reset']} {COLORS['red']}[ERROR]{COLORS['reset']}"
+        elif span.status.status_code != StatusCode.UNSET:
+            span_context += f"{COLORS['reset']} [{span.status.status_code}]"
 
-        # Add duration
         duration_ms = (span.end_time - span.start_time) / 1e6
-        span_context += f" {COLORS['dim']}({duration_ms:.2f}ms){COLORS['reset']}"
+        span_context += f"{COLORS['reset']} ({duration_ms:.2f}ms)"
 
-        # Print the main span line
         print(span_context)
 
-        # Print attributes indented
-        if span.attributes:
+        if self.print_attributes and span.attributes:
             for key, value in span.attributes.items():
                 if key.startswith("__"):
                     continue
-                print(f"  {COLORS['dim']}{key}: {value}{COLORS['reset']}")
+                str_value = str(value)
+                if len(str_value) > 1000:
+                    str_value = str_value[:997] + "..."
+                print(f"    {COLORS['dim']}{key}: {str_value}{COLORS['reset']}")
 
-        # Print events indented
         for event in span.events:
             event_time = datetime.utcfromtimestamp(event.timestamp / 1e9).strftime(
                 "%H:%M:%S.%f"
             )[:-3]
+
+            severity = event.attributes.get("severity", "info")
+            message = event.attributes.get("message", event.name)
+            if isinstance(message, (dict, list)):
+                message = json.dumps(message, indent=2)
+
+            severity_colors = {
+                "error": f"{COLORS['bold']}{COLORS['red']}",
+                "warn": f"{COLORS['bold']}{COLORS['yellow']}",
+                "info": COLORS["white"],
+                "debug": COLORS["dim"],
+            }
+            msg_color = severity_colors.get(severity, COLORS["white"])
+
             print(
-                f"  {COLORS['dim']}{event_time}{COLORS['reset']} "
-                f"{COLORS['cyan']}[EVENT]{COLORS['reset']} {event.name}"
+                f" {event_time} "
+                f"{msg_color}[{severity.upper()}] "
+                f"{message}{COLORS['reset']}"
             )
+
             if event.attributes:
                 for key, value in event.attributes.items():
-                    if key.startswith("__"):
+                    if key.startswith("__") or key in ["message", "severity"]:
                         continue
-                    print(f"    {COLORS['dim']}{key}: {value}{COLORS['reset']}")
+                    print(f"   {COLORS['dim']}{key}: {value}{COLORS['reset']}")
 
     def shutdown(self) -> None:
         """Shutdown the processor."""

llama_stack/providers/utils/telemetry/trace_protocol.py

@@ -0,0 +1,129 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+import inspect
+import json
+from functools import wraps
+from typing import Any, AsyncGenerator, Callable, Type, TypeVar
+
+from pydantic import BaseModel
+
+T = TypeVar("T")
+
+
+def serialize_value(value: Any) -> str:
+    """Helper function to serialize values to string representation."""
+    try:
+        if isinstance(value, BaseModel):
+            return value.model_dump_json()
+        elif isinstance(value, list) and value and isinstance(value[0], BaseModel):
+            return json.dumps([item.model_dump_json() for item in value])
+        elif hasattr(value, "to_dict"):
+            return json.dumps(value.to_dict())
+        elif isinstance(value, (dict, list, int, float, str, bool)):
+            return json.dumps(value)
+        else:
+            return str(value)
+    except Exception:
+        return str(value)
+
+
+def trace_protocol(cls: Type[T]) -> Type[T]:
+    """
+    A class decorator that automatically traces all methods in a protocol/base class
+    and its inheriting classes.
+    """
+
+    def trace_method(method: Callable) -> Callable:
+        from llama_stack.providers.utils.telemetry import tracing
+
+        is_async = asyncio.iscoroutinefunction(method)
+        is_async_gen = inspect.isasyncgenfunction(method)
+
+        def create_span_context(self: Any, *args: Any, **kwargs: Any) -> tuple:
+            class_name = self.__class__.__name__
+            method_name = method.__name__
+
+            span_type = (
+                "async_generator" if is_async_gen else "async" if is_async else "sync"
+            )
+            span_attributes = {
+                "__autotraced__": True,
+                "__class__": class_name,
+                "__method__": method_name,
+                "__type__": span_type,
+                "__args__": serialize_value(args),
+            }
+
+            return class_name, method_name, span_attributes
+
+        @wraps(method)
+        async def async_gen_wrapper(
+            self: Any, *args: Any, **kwargs: Any
+        ) -> AsyncGenerator:
+            class_name, method_name, span_attributes = create_span_context(
+                self, *args, **kwargs
+            )
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    count = 0
+                    async for item in method(self, *args, **kwargs):
+                        yield item
+                        count += 1
+                finally:
+                    span.set_attribute("chunk_count", count)
+
+        @wraps(method)
+        async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            class_name, method_name, span_attributes = create_span_context(
+                self, *args, **kwargs
+            )
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    result = await method(self, *args, **kwargs)
+                    span.set_attribute("output", serialize_value(result))
+                    return result
+                except Exception as e:
+                    span.set_attribute("error", str(e))
+                    raise
+
+        @wraps(method)
+        def sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            class_name, method_name, span_attributes = create_span_context(
+                self, *args, **kwargs
+            )
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    result = method(self, *args, **kwargs)
+                    span.set_attribute("output", serialize_value(result))
+                    return result
+                except Exception as _e:
+                    raise
+
+        if is_async_gen:
+            return async_gen_wrapper
+        elif is_async:
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    original_init_subclass = getattr(cls, "__init_subclass__", None)
+
+    def __init_subclass__(cls_child, **kwargs):  # noqa: N807
+        if original_init_subclass:
+            original_init_subclass(**kwargs)
+
+        for name, method in vars(cls_child).items():
+            if inspect.isfunction(method) and not name.startswith("_"):
+                setattr(cls_child, name, trace_method(method))  # noqa: B010
+
+    cls.__init_subclass__ = classmethod(__init_subclass__)
+
+    return cls