Update Telemetry API so OpenAPI generation can work (#640)

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.)
This commit is contained in:
Ashwin Bharambe 2024-12-16 13:00:14 -08:00 committed by GitHub
parent 78e2bfbe7a
commit 2e5bfcd42a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 349 additions and 473 deletions

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, SpanWithChildren
from llama_stack.apis.telemetry import QueryCondition, Span
class TelemetryDatasetMixin:
@ -53,19 +53,18 @@ class TelemetryDatasetMixin:
spans = []
for trace in traces:
span_tree = await self.get_span_tree(
spans_by_id = await self.get_span_tree(
span_id=trace.root_span_id,
attributes_to_return=attributes_to_return,
max_depth=max_depth,
)
def extract_spans(span: SpanWithChildren) -> List[Span]:
result = []
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
):
result.append(
spans.append(
Span(
trace_id=trace.root_span_id,
span_id=span.span_id,
@ -77,11 +76,4 @@ class TelemetryDatasetMixin:
)
)
for child in span.children:
result.extend(extract_spans(child))
return result
spans.extend(extract_spans(span_tree))
return spans