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

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

View file

@ -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 = []