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/apis/agents/agents.py

@@ -23,7 +23,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Annotated
 
-from llama_stack.distribution.tracing import trace_protocol
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.common.deployment_types import *  # noqa: F403
 from llama_stack.apis.inference import *  # noqa: F403

llama_stack/apis/inference/inference.py

@@ -21,7 +21,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
-from llama_stack.distribution.tracing import trace_protocol
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.models import *  # noqa: F403

llama_stack/apis/memory/memory.py

@@ -16,7 +16,7 @@ from pydantic import BaseModel, Field
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.memory_banks import *  # noqa: F403
-from llama_stack.distribution.tracing import trace_protocol
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 
 @json_schema_type

llama_stack/apis/memory_banks/memory_banks.py

@@ -20,7 +20,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
 from llama_stack.apis.resource import Resource, ResourceType
-from llama_stack.distribution.tracing import trace_protocol
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 
 @json_schema_type

llama_stack/apis/models/models.py

@@ -10,7 +10,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, ConfigDict, Field
 
 from llama_stack.apis.resource import Resource, ResourceType
-from llama_stack.distribution.tracing import trace_protocol
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 
 class CommonModelFields(BaseModel):

llama_stack/apis/safety/safety.py

@@ -10,7 +10,7 @@ from typing import Any, Dict, List, Protocol, runtime_checkable
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
-from llama_stack.distribution.tracing import trace_protocol
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.shields import *  # noqa: F403

llama_stack/apis/shields/shields.py

@@ -10,7 +10,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_stack.apis.resource import Resource, ResourceType
-from llama_stack.distribution.tracing import trace_protocol
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 
 class CommonShieldFields(BaseModel):

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 == StatusCode.ERROR:
-        if span.status.status_code != 0:  # UNSET or ERROR
+            span_context += f"{COLORS['reset']} {COLORS['red']}[ERROR]{COLORS['reset']}"
-            status_color = (
+        elif span.status.status_code != StatusCode.UNSET:
-                COLORS["red"] if span.status.status_code == 2 else COLORS["yellow"]
+            span_context += f"{COLORS['reset']} [{span.status.status_code}]"
-            )
-            span_context += (
-                f" {status_color}[{span.status.status_code}]{COLORS['reset']}"
-            )
 
-        # 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 self.print_attributes and span.attributes:
-        if 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" {event_time} "
-                f"{COLORS['cyan']}[EVENT]{COLORS['reset']} {event.name}"
+                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."""