diff --git a/llama_stack/apis/agents/agents.py b/llama_stack/apis/agents/agents.py index d2243c96f..6e41df4f6 100644 --- a/llama_stack/apis/agents/agents.py +++ b/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 diff --git a/llama_stack/apis/inference/inference.py b/llama_stack/apis/inference/inference.py index 85b29a147..233cd1b50 100644 --- a/llama_stack/apis/inference/inference.py +++ b/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 diff --git a/llama_stack/apis/memory/memory.py b/llama_stack/apis/memory/memory.py index b75df8a1a..2f3a94956 100644 --- a/llama_stack/apis/memory/memory.py +++ b/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 diff --git a/llama_stack/apis/memory_banks/memory_banks.py b/llama_stack/apis/memory_banks/memory_banks.py index 0b8b2563f..a17e8e48d 100644 --- a/llama_stack/apis/memory_banks/memory_banks.py +++ b/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 diff --git a/llama_stack/apis/models/models.py b/llama_stack/apis/models/models.py index 2c0f1ee21..cb9cb1117 100644 --- a/llama_stack/apis/models/models.py +++ b/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): diff --git a/llama_stack/apis/safety/safety.py b/llama_stack/apis/safety/safety.py index 41058f107..26ae45ae7 100644 --- a/llama_stack/apis/safety/safety.py +++ b/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 diff --git a/llama_stack/apis/shields/shields.py b/llama_stack/apis/shields/shields.py index b28605727..8d4d5f9fd 100644 --- a/llama_stack/apis/shields/shields.py +++ b/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): diff --git a/llama_stack/providers/inline/memory/faiss/faiss.py b/llama_stack/providers/inline/memory/faiss/faiss.py index dfefefeb8..78de13120 100644 --- a/llama_stack/providers/inline/memory/faiss/faiss.py +++ b/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): diff --git a/llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py b/llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py index 6c4d7e8d4..2f00b21b8 100644 --- a/llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py +++ b/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.""" diff --git a/llama_stack/distribution/tracing.py b/llama_stack/providers/utils/telemetry/trace_protocol.py similarity index 100% rename from llama_stack/distribution/tracing.py rename to llama_stack/providers/utils/telemetry/trace_protocol.py