Idimatic REST API: Telemetry

llama_stack/apis/telemetry/telemetry.py

@@ -169,39 +169,51 @@ class QueryCondition(BaseModel):
     value: Any
 
 
+class QueryTracesResponse(BaseModel):
+    data: List[Trace]
+
+
+class QuerySpansResponse(BaseModel):
+    data: List[Span]
+
+
+class QuerySpanTreeResponse(BaseModel):
+    data: Dict[str, SpanWithStatus]
+
+
 @runtime_checkable
 class Telemetry(Protocol):
-    @webmethod(route="/telemetry/log-event")
+    @webmethod(route="/telemetry/events", method="POST")
     async def log_event(
         self, event: Event, ttl_seconds: int = DEFAULT_TTL_DAYS * 86400
     ) -> None: ...
 
-    @webmethod(route="/telemetry/query-traces", method="POST")
+    @webmethod(route="/telemetry/traces", method="GET")
     async def query_traces(
         self,
         attribute_filters: Optional[List[QueryCondition]] = None,
         limit: Optional[int] = 100,
         offset: Optional[int] = 0,
         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(
         self,
         span_id: str,
         attributes_to_return: Optional[List[str]] = 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(
         self,
         attribute_filters: List[QueryCondition],
         attributes_to_return: List[str],
         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(
         self,
         attribute_filters: List[QueryCondition],

llama_stack/providers/inline/telemetry/meta_reference/telemetry.py

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

llama_stack/providers/utils/telemetry/dataset_mixin.py

@@ -7,7 +7,7 @@
 from typing import List, Optional
 
 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:
@@ -48,7 +48,7 @@ class TelemetryDatasetMixin:
         attribute_filters: List[QueryCondition],
         attributes_to_return: List[str],
         max_depth: Optional[int] = None,
-    ) -> List[Span]:
+    ) -> QuerySpansResponse:
         traces = await self.query_traces(attribute_filters=attribute_filters)
         spans = []