Idiomatic REST API: Telemetry (#786)

# What does this PR do?

Changes Telemetry API to follow more idiomatic REST


- [ ] Addresses issue (#issue)


## Test Plan

TBD, once i get an approval for rest endpoints

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,18 +48,18 @@ 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 = []
 
-        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

@@ -10,7 +10,7 @@ from typing import Dict, List, Optional, Protocol
 
 import aiosqlite
 
-from llama_stack.apis.telemetry import QueryCondition, SpanWithStatus, Trace
+from llama_stack.apis.telemetry import QueryCondition, Span, SpanWithStatus, Trace
 
 
 class TraceStore(Protocol):
@@ -167,3 +167,23 @@ class SQLiteTraceStore(TraceStore):
                     spans_by_id[span.span_id] = span
 
                 return spans_by_id
+
+    async def get_trace(self, trace_id: str) -> Trace:
+        query = "SELECT * FROM traces WHERE trace_id = ?"
+        async with aiosqlite.connect(self.conn_string) as conn:
+            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:
+        query = "SELECT * FROM spans WHERE trace_id = ? AND span_id = ?"
+        async with aiosqlite.connect(self.conn_string) as conn:
+            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)