feat: Add new compact MetricInResponse type

This commit is contained in:
Dinesh Yeduguru 2025-03-12 11:54:23 -07:00
parent 0fdb15bcc7
commit 8b63aba7a4
5 changed files with 21 additions and 26 deletions

View file

@ -4630,8 +4630,6 @@
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"trace_id",
"span_id",
"timestamp", "timestamp",
"type", "type",
"metric", "metric",
@ -8450,8 +8448,6 @@
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"trace_id",
"span_id",
"timestamp", "timestamp",
"type", "type",
"payload" "payload"
@ -8524,8 +8520,6 @@
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"trace_id",
"span_id",
"timestamp", "timestamp",
"type", "type",
"message", "message",

View file

@ -3149,8 +3149,6 @@ components:
type: string type: string
additionalProperties: false additionalProperties: false
required: required:
- trace_id
- span_id
- timestamp - timestamp
- type - type
- metric - metric
@ -5760,8 +5758,6 @@ components:
$ref: '#/components/schemas/StructuredLogPayload' $ref: '#/components/schemas/StructuredLogPayload'
additionalProperties: false additionalProperties: false
required: required:
- trace_id
- span_id
- timestamp - timestamp
- type - type
- payload - payload
@ -5804,8 +5800,6 @@ components:
$ref: '#/components/schemas/LogSeverity' $ref: '#/components/schemas/LogSeverity'
additionalProperties: false additionalProperties: false
required: required:
- trace_id
- span_id
- timestamp - timestamp
- type - type
- message - message

View file

@ -75,9 +75,9 @@ class LogSeverity(Enum):
class EventCommon(BaseModel): class EventCommon(BaseModel):
trace_id: str trace_id: Optional[str] = None
span_id: str span_id: Optional[str] = None
timestamp: datetime timestamp: Optional[datetime] = None
attributes: Optional[Dict[str, Primitive]] = Field(default_factory=dict) attributes: Optional[Dict[str, Primitive]] = Field(default_factory=dict)
@ -93,7 +93,14 @@ class MetricEvent(EventCommon):
type: Literal[EventType.METRIC.value] = EventType.METRIC.value type: Literal[EventType.METRIC.value] = EventType.METRIC.value
metric: str # this would be an enum metric: str # this would be an enum
value: Union[int, float] value: Union[int, float]
unit: str unit: Optional[str] = None
@json_schema_type
class MetricInResponse(BaseModel):
metric: str
value: Union[int, float]
unit: Optional[str] = None
# This is a short term solution to allow inference API to return metrics # This is a short term solution to allow inference API to return metrics
@ -117,7 +124,7 @@ class MetricEvent(EventCommon):
class MetricResponseMixin(BaseModel): class MetricResponseMixin(BaseModel):
metrics: Optional[List[MetricEvent]] = None metrics: Optional[List[MetricInResponse]] = None
@json_schema_type @json_schema_type

View file

@ -48,7 +48,7 @@ from llama_stack.apis.scoring import (
ScoringFnParams, ScoringFnParams,
) )
from llama_stack.apis.shields import Shield from llama_stack.apis.shields import Shield
from llama_stack.apis.telemetry import MetricEvent, Telemetry from llama_stack.apis.telemetry import MetricEvent, MetricInResponse, Telemetry
from llama_stack.apis.tools import ( from llama_stack.apis.tools import (
RAGDocument, RAGDocument,
RAGQueryConfig, RAGQueryConfig,
@ -206,12 +206,12 @@ class InferenceRouter(Inference):
completion_tokens: int, completion_tokens: int,
total_tokens: int, total_tokens: int,
model: Model, model: Model,
) -> List[MetricEvent]: ) -> List[MetricInResponse]:
metrics = self._construct_metrics(prompt_tokens, completion_tokens, total_tokens, model) metrics = self._construct_metrics(prompt_tokens, completion_tokens, total_tokens, model)
if self.telemetry: if self.telemetry:
for metric in metrics: for metric in metrics:
await self.telemetry.log_event(metric) await self.telemetry.log_event(metric)
return metrics return [MetricInResponse(metric=metric.metric, value=metric.value) for metric in metrics]
async def _count_tokens( async def _count_tokens(
self, self,

View file

@ -153,20 +153,20 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
else: else:
print(f"Warning: No active span found for span_id {span_id}. Dropping event: {event}") print(f"Warning: No active span found for span_id {span_id}. Dropping event: {event}")
def _get_or_create_counter(self, name: str, unit: str) -> metrics.Counter: def _get_or_create_counter(self, name: str, unit: Optional[str] = None) -> metrics.Counter:
if name not in _GLOBAL_STORAGE["counters"]: if name not in _GLOBAL_STORAGE["counters"]:
_GLOBAL_STORAGE["counters"][name] = self.meter.create_counter( _GLOBAL_STORAGE["counters"][name] = self.meter.create_counter(
name=name, name=name,
unit=unit, unit=unit or "",
description=f"Counter for {name}", description=f"Counter for {name}",
) )
return _GLOBAL_STORAGE["counters"][name] return _GLOBAL_STORAGE["counters"][name]
def _get_or_create_gauge(self, name: str, unit: str) -> metrics.ObservableGauge: def _get_or_create_gauge(self, name: str, unit: Optional[str] = None) -> metrics.ObservableGauge:
if name not in _GLOBAL_STORAGE["gauges"]: if name not in _GLOBAL_STORAGE["gauges"]:
_GLOBAL_STORAGE["gauges"][name] = self.meter.create_gauge( _GLOBAL_STORAGE["gauges"][name] = self.meter.create_gauge(
name=name, name=name,
unit=unit, unit=unit or "",
description=f"Gauge for {name}", description=f"Gauge for {name}",
) )
return _GLOBAL_STORAGE["gauges"][name] return _GLOBAL_STORAGE["gauges"][name]
@ -181,11 +181,11 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
up_down_counter = self._get_or_create_up_down_counter(event.metric, event.unit) up_down_counter = self._get_or_create_up_down_counter(event.metric, event.unit)
up_down_counter.add(event.value, attributes=event.attributes) up_down_counter.add(event.value, attributes=event.attributes)
def _get_or_create_up_down_counter(self, name: str, unit: str) -> metrics.UpDownCounter: def _get_or_create_up_down_counter(self, name: str, unit: Optional[str] = None) -> metrics.UpDownCounter:
if name not in _GLOBAL_STORAGE["up_down_counters"]: if name not in _GLOBAL_STORAGE["up_down_counters"]:
_GLOBAL_STORAGE["up_down_counters"][name] = self.meter.create_up_down_counter( _GLOBAL_STORAGE["up_down_counters"][name] = self.meter.create_up_down_counter(
name=name, name=name,
unit=unit, unit=unit or "",
description=f"UpDownCounter for {name}", description=f"UpDownCounter for {name}",
) )
return _GLOBAL_STORAGE["up_down_counters"][name] return _GLOBAL_STORAGE["up_down_counters"][name]