mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-28 02:53:30 +00:00
2e5bfcd42a
We cannot use recursive types because not only does our OpenAPI generator not like them, even if it did, it is not easy for all client languages to automatically construct proper APIs (especially considering garbage collection) around them. For now, we can return a `Dict[str, SpanWithStatus]` instead of `SpanWithChildren` and rely on the client to reconstruct the tree. Also fixed a super subtle issue with the OpenAPI generation process (monkey-patching of json_schema_type wasn't working because of import reordering.)
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
from typing import List, Optional
|
|
|
|
from llama_stack.apis.datasetio import DatasetIO
|
|
from llama_stack.apis.telemetry import QueryCondition, Span
|
|
|
|
|
|
class TelemetryDatasetMixin:
|
|
"""Mixin class that provides dataset-related functionality for telemetry providers."""
|
|
|
|
datasetio_api: DatasetIO
|
|
|
|
async def save_spans_to_dataset(
|
|
self,
|
|
attribute_filters: List[QueryCondition],
|
|
attributes_to_save: List[str],
|
|
dataset_id: str,
|
|
max_depth: Optional[int] = None,
|
|
) -> None:
|
|
spans = await self.query_spans(
|
|
attribute_filters=attribute_filters,
|
|
attributes_to_return=attributes_to_save,
|
|
max_depth=max_depth,
|
|
)
|
|
|
|
rows = [
|
|
{
|
|
"trace_id": span.trace_id,
|
|
"span_id": span.span_id,
|
|
"parent_span_id": span.parent_span_id,
|
|
"name": span.name,
|
|
"start_time": span.start_time,
|
|
"end_time": span.end_time,
|
|
**{attr: span.attributes.get(attr) for attr in attributes_to_save},
|
|
}
|
|
for span in spans
|
|
]
|
|
|
|
await self.datasetio_api.append_rows(dataset_id=dataset_id, rows=rows)
|
|
|
|
async def query_spans(
|
|
self,
|
|
attribute_filters: List[QueryCondition],
|
|
attributes_to_return: List[str],
|
|
max_depth: Optional[int] = None,
|
|
) -> List[Span]:
|
|
traces = await self.query_traces(attribute_filters=attribute_filters)
|
|
spans = []
|
|
|
|
for trace in traces:
|
|
spans_by_id = await self.get_span_tree(
|
|
span_id=trace.root_span_id,
|
|
attributes_to_return=attributes_to_return,
|
|
max_depth=max_depth,
|
|
)
|
|
|
|
for span in spans_by_id.values():
|
|
if span.attributes and all(
|
|
attr in span.attributes and span.attributes[attr] is not None
|
|
for attr in attributes_to_return
|
|
):
|
|
spans.append(
|
|
Span(
|
|
trace_id=trace.root_span_id,
|
|
span_id=span.span_id,
|
|
parent_span_id=span.parent_span_id,
|
|
name=span.name,
|
|
start_time=span.start_time,
|
|
end_time=span.end_time,
|
|
attributes=span.attributes,
|
|
)
|
|
)
|
|
|
|
return spans
|