address feedback

llama_stack/providers/utils/telemetry/sqlite.py

@@ -11,8 +11,8 @@ from typing import List, Optional
 import aiosqlite
 
 from llama_stack.apis.telemetry import (
-    MaterializedSpan,
     QueryCondition,
+    SpanWithChildren,
     Trace,
     TraceStore,
 )
@@ -24,56 +24,76 @@ class SQLiteTraceStore(TraceStore):
 
     async def query_traces(
         self,
-        attribute_conditions: Optional[List[QueryCondition]] = None,
-        attribute_keys_to_return: Optional[List[str]] = None,
+        attribute_filters: Optional[List[QueryCondition]] = None,
+        attributes_to_return: Optional[List[str]] = None,
         limit: Optional[int] = 100,
         offset: Optional[int] = 0,
         order_by: Optional[List[str]] = None,
     ) -> List[Trace]:
-        # Build the SQL query with attribute selection
-        select_clause = """
-            SELECT DISTINCT t.trace_id, t.root_span_id, t.start_time, t.end_time
-        """
-        if attribute_keys_to_return:
-            for key in attribute_keys_to_return:
-                select_clause += (
-                    f", json_extract(s.attributes, '$.{key}') as attr_{key}"
-                )
+        print(attribute_filters, attributes_to_return, limit, offset, order_by)
 
-        query = (
-            select_clause
-            + """
-            FROM traces t
-            JOIN spans s ON t.trace_id = s.trace_id
-        """
-        )
-        params = []
+        def build_attribute_select() -> str:
+            if not attributes_to_return:
+                return ""
+            return "".join(
+                f", json_extract(s.attributes, '$.{key}') as attr_{key}"
+                for key in attributes_to_return
+            )
 
-        # Add attribute conditions if present
-        if attribute_conditions:
-            conditions = []
-            for condition in attribute_conditions:
-                conditions.append(
-                    f"json_extract(s.attributes, '$.{condition.key}') {condition.op} ?"
-                )
-                params.append(condition.value)
-            if conditions:
-                query += " WHERE " + " AND ".join(conditions)
+        def build_where_clause() -> tuple[str, list]:
+            if not attribute_filters:
+                return "", []
+
+            conditions = [
+                f"json_extract(s.attributes, '$.{condition.key}') {condition.op} ?"
+                for condition in attribute_filters
+            ]
+            params = [condition.value for condition in attribute_filters]
+            where_clause = " WHERE " + " AND ".join(conditions)
+            return where_clause, params
+
+        def build_order_clause() -> str:
+            if not order_by:
+                return ""
 
-        # Add ordering
-        if order_by:
             order_clauses = []
             for field in order_by:
-                desc = False
-                if field.startswith("-"):
-                    field = field[1:]
-                    desc = True
-                order_clauses.append(f"t.{field} {'DESC' if desc else 'ASC'}")
-            query += " ORDER BY " + ", ".join(order_clauses)
+                desc = field.startswith("-")
+                clean_field = field[1:] if desc else field
+                order_clauses.append(f"t.{clean_field} {'DESC' if desc else 'ASC'}")
+            return " ORDER BY " + ", ".join(order_clauses)
 
-        # Add limit and offset
-        query += f" LIMIT {limit} OFFSET {offset}"
+        # Build the main query
+        base_query = """
+            WITH matching_traces AS (
+                SELECT DISTINCT t.trace_id
+                FROM traces t
+                JOIN spans s ON t.trace_id = s.trace_id
+                {where_clause}
+            ),
+            filtered_traces AS (
+                SELECT t.trace_id, t.root_span_id, t.start_time, t.end_time
+                {attribute_select}
+                FROM matching_traces mt
+                JOIN traces t ON mt.trace_id = t.trace_id
+                LEFT JOIN spans s ON t.trace_id = s.trace_id
+                {order_clause}
+            )
+            SELECT DISTINCT trace_id, root_span_id, start_time, end_time
+            FROM filtered_traces
+            LIMIT {limit} OFFSET {offset}
+        """
 
+        where_clause, params = build_where_clause()
+        query = base_query.format(
+            attribute_select=build_attribute_select(),
+            where_clause=where_clause,
+            order_clause=build_order_clause(),
+            limit=limit,
+            offset=offset,
+        )
+
+        # Execute query and return results
         async with aiosqlite.connect(self.conn_string) as conn:
             conn.row_factory = aiosqlite.Row
             async with conn.execute(query, params) as cursor:
@@ -91,15 +111,15 @@ class SQLiteTraceStore(TraceStore):
     async def get_materialized_span(
         self,
         span_id: str,
-        attribute_keys_to_return: Optional[List[str]] = None,
+        attributes_to_return: Optional[List[str]] = None,
         max_depth: Optional[int] = None,
-    ) -> MaterializedSpan:
+    ) -> SpanWithChildren:
         # Build the attributes selection
         attributes_select = "s.attributes"
-        if attribute_keys_to_return:
+        if attributes_to_return:
             json_object = ", ".join(
                 f"'{key}', json_extract(s.attributes, '$.{key}')"
-                for key in attribute_keys_to_return
+                for key in attributes_to_return
             )
             attributes_select = f"json_object({json_object})"
 
@@ -135,7 +155,7 @@ class SQLiteTraceStore(TraceStore):
                 root_span = None
 
                 for row in rows:
-                    span = MaterializedSpan(
+                    span = SpanWithChildren(
                         span_id=row["span_id"],
                         trace_id=row["trace_id"],
                         parent_span_id=row["parent_span_id"],