Idimatic REST API: Telemetry

This commit is contained in:
Dinesh Yeduguru 2025-01-15 20:54:34 -08:00
parent 8d30ecb91a
commit 749bc74509
3 changed files with 39 additions and 23 deletions

View file

@ -169,39 +169,51 @@ class QueryCondition(BaseModel):
value: Any value: Any
class QueryTracesResponse(BaseModel):
data: List[Trace]
class QuerySpansResponse(BaseModel):
data: List[Span]
class QuerySpanTreeResponse(BaseModel):
data: Dict[str, SpanWithStatus]
@runtime_checkable @runtime_checkable
class Telemetry(Protocol): class Telemetry(Protocol):
@webmethod(route="/telemetry/log-event") @webmethod(route="/telemetry/events", method="POST")
async def log_event( async def log_event(
self, event: Event, ttl_seconds: int = DEFAULT_TTL_DAYS * 86400 self, event: Event, ttl_seconds: int = DEFAULT_TTL_DAYS * 86400
) -> None: ... ) -> None: ...
@webmethod(route="/telemetry/query-traces", method="POST") @webmethod(route="/telemetry/traces", method="GET")
async def query_traces( async def query_traces(
self, self,
attribute_filters: Optional[List[QueryCondition]] = None, attribute_filters: Optional[List[QueryCondition]] = None,
limit: Optional[int] = 100, limit: Optional[int] = 100,
offset: Optional[int] = 0, offset: Optional[int] = 0,
order_by: Optional[List[str]] = None, order_by: Optional[List[str]] = None,
) -> List[Trace]: ... ) -> QueryTracesResponse: ...
@webmethod(route="/telemetry/query-span-tree", method="POST") @webmethod(route="/telemetry/spans/{span_id}/tree", method="GET")
async def query_span_tree( async def query_span_tree(
self, self,
span_id: str, span_id: str,
attributes_to_return: Optional[List[str]] = None, attributes_to_return: Optional[List[str]] = None,
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> Dict[str, SpanWithStatus]: ... ) -> QuerySpanTreeResponse: ...
@webmethod(route="/telemetry/query-spans", method="POST") @webmethod(route="/telemetry/spans", method="GET")
async def query_spans( async def query_spans(
self, self,
attribute_filters: List[QueryCondition], attribute_filters: List[QueryCondition],
attributes_to_return: List[str], attributes_to_return: List[str],
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> List[Span]: ... ) -> QuerySpansResponse: ...
@webmethod(route="/telemetry/save-spans-to-dataset", method="POST") @webmethod(route="/telemetry/spans/export", method="POST")
async def save_spans_to_dataset( async def save_spans_to_dataset(
self, self,
attribute_filters: List[QueryCondition], attribute_filters: List[QueryCondition],

View file

@ -21,13 +21,13 @@ from llama_stack.apis.telemetry import (
Event, Event,
MetricEvent, MetricEvent,
QueryCondition, QueryCondition,
QuerySpanTreeResponse,
QueryTracesResponse,
SpanEndPayload, SpanEndPayload,
SpanStartPayload, SpanStartPayload,
SpanStatus, SpanStatus,
SpanWithStatus,
StructuredLogEvent, StructuredLogEvent,
Telemetry, Telemetry,
Trace,
UnstructuredLogEvent, UnstructuredLogEvent,
) )
from llama_stack.distribution.datatypes import Api from llama_stack.distribution.datatypes import Api
@ -241,12 +241,14 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
limit: Optional[int] = 100, limit: Optional[int] = 100,
offset: Optional[int] = 0, offset: Optional[int] = 0,
order_by: Optional[List[str]] = None, order_by: Optional[List[str]] = None,
) -> List[Trace]: ) -> QueryTracesResponse:
return await self.trace_store.query_traces( return QueryTracesResponse(
attribute_filters=attribute_filters, data=await self.trace_store.query_traces(
limit=limit, attribute_filters=attribute_filters,
offset=offset, limit=limit,
order_by=order_by, offset=offset,
order_by=order_by,
)
) )
async def query_span_tree( async def query_span_tree(
@ -254,9 +256,11 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
span_id: str, span_id: str,
attributes_to_return: Optional[List[str]] = None, attributes_to_return: Optional[List[str]] = None,
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> Dict[str, SpanWithStatus]: ) -> QuerySpanTreeResponse:
return await self.trace_store.get_span_tree( return QuerySpanTreeResponse(
span_id=span_id, data=await self.trace_store.get_span_tree(
attributes_to_return=attributes_to_return, span_id=span_id,
max_depth=max_depth, attributes_to_return=attributes_to_return,
max_depth=max_depth,
)
) )

View file

@ -7,7 +7,7 @@
from typing import List, Optional from typing import List, Optional
from llama_stack.apis.datasetio import DatasetIO from llama_stack.apis.datasetio import DatasetIO
from llama_stack.apis.telemetry import QueryCondition, Span from llama_stack.apis.telemetry import QueryCondition, QuerySpansResponse, Span
class TelemetryDatasetMixin: class TelemetryDatasetMixin:
@ -48,7 +48,7 @@ class TelemetryDatasetMixin:
attribute_filters: List[QueryCondition], attribute_filters: List[QueryCondition],
attributes_to_return: List[str], attributes_to_return: List[str],
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> List[Span]: ) -> QuerySpansResponse:
traces = await self.query_traces(attribute_filters=attribute_filters) traces = await self.query_traces(attribute_filters=attribute_filters)
spans = [] spans = []