passing tests

llama_stack/apis/telemetry/telemetry.py

@@ -204,7 +204,7 @@ class Telemetry(Protocol):
     async def get_span(self, trace_id: str, span_id: str) -> Span: ...
 
     @webmethod(route="/telemetry/spans/{span_id}/tree", method="GET")
-    async def query_span_tree(
+    async def get_span_tree(
         self,
         span_id: str,
         attributes_to_return: Optional[List[str]] = None,

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

@@ -259,7 +259,7 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
     async def get_span(self, trace_id: str, span_id: str) -> Span:
         return await self.trace_store.get_span(trace_id, span_id)
 
-    async def query_span_tree(
+    async def get_span_tree(
         self,
         span_id: str,
         attributes_to_return: Optional[List[str]] = None,

llama_stack/providers/utils/telemetry/dataset_mixin.py

@@ -49,17 +49,17 @@ class TelemetryDatasetMixin:
         attributes_to_return: List[str],
         max_depth: Optional[int] = None,
     ) -> QuerySpansResponse:
-        traces = await self.query_traces(attribute_filters=attribute_filters).data
+        traces = await self.query_traces(attribute_filters=attribute_filters)
         spans = []
 
-        for trace in traces:
-            spans_by_id = await self.query_span_tree(
+        for trace in traces.data:
+            spans_by_id_resp = 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():
+            for span in spans_by_id_resp.data.values():
                 if span.attributes and all(
                     attr in span.attributes and span.attributes[attr] is not None
                     for attr in attributes_to_return
@@ -76,4 +76,4 @@ class TelemetryDatasetMixin:
                         )
                     )
 
-        return spans
+        return QuerySpansResponse(data=spans)

llama_stack/providers/utils/telemetry/sqlite_trace_store.py

@@ -174,6 +174,8 @@ class SQLiteTraceStore(TraceStore):
             conn.row_factory = aiosqlite.Row
             async with conn.execute(query, (trace_id,)) as cursor:
                 row = await cursor.fetchone()
+                if row is None:
+                    raise ValueError(f"Trace {trace_id} not found")
                 return Trace(**row)
 
     async def get_span(self, trace_id: str, span_id: str) -> Span:
@@ -182,4 +184,6 @@ class SQLiteTraceStore(TraceStore):
             conn.row_factory = aiosqlite.Row
             async with conn.execute(query, (trace_id, span_id)) as cursor:
                 row = await cursor.fetchone()
+                if row is None:
+                    raise ValueError(f"Span {span_id} not found")
                 return Span(**row)