phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-02 08:44:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

console span processor improvements

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-12-06 11:19:21 -08:00
parent 084ec337af
commit 208cded50e
10 changed files with 43 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

2
llama_stack/apis/agents/agents.py
View file

@ -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

2
llama_stack/apis/inference/inference.py
View file

@ -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

2
llama_stack/apis/memory/memory.py
View file

@ -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

2
llama_stack/apis/memory_banks/memory_banks.py
View file

@ -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

2
llama_stack/apis/models/models.py
View file

@ -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):

2
llama_stack/apis/safety/safety.py
View file

@ -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

2
llama_stack/apis/shields/shields.py
View file

@ -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):

2
llama_stack/providers/inline/memory/faiss/faiss.py
View file

@ -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):

65
llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py
View file

@ -4,6 +4,7 @@
# 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
@ -25,10 +26,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
@ -37,13 +39,10 @@ class ConsoleSpanProcessor(SpanProcessor):
)[:-3]
print(
f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
f"{COLORS['magenta']}[START]{COLORS['reset']} "
f"{COLORS['cyan']}{span.name}{COLORS['reset']}"
f"{COLORS['dim']}{timestamp} " f"[START] " f"{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,48 +50,54 @@ 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']} "
)
span_context = f"{COLORS['dim']}{timestamp} [END] {span.name}"
# 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 == 2: # ERROR
span_context += f" {COLORS['red']}[ERROR]{COLORS['reset']}"
elif span.status.status_code != 0: # UNSET
span_context += f" {COLORS['dim']}[{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" ({duration_ms:.2f}ms){COLORS['reset']}"
# 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']}")

0
llama_stack/distribution/tracing.py → llama_stack/providers/utils/telemetry/trace_protocol.py
View file