diff --git a/llama_stack/apis/telemetry/telemetry.py b/llama_stack/apis/telemetry/telemetry.py index d04cb67e3..94d51f9c6 100644 --- a/llama_stack/apis/telemetry/telemetry.py +++ b/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], diff --git a/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py b/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py index e2e318375..a9d94d857 100644 --- a/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py +++ b/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, + ) ) diff --git a/llama_stack/providers/utils/telemetry/dataset_mixin.py b/llama_stack/providers/utils/telemetry/dataset_mixin.py index e488f2475..c52ec7d89 100644 --- a/llama_stack/providers/utils/telemetry/dataset_mixin.py +++ b/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 = []