consolidate telemetry to meta reference inline

llama_stack/providers/utils/telemetry/sqlite.py

@@ -0,0 +1,157 @@
+# 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.
+
+import json
+from datetime import datetime
+from typing import List, Optional
+
+import aiosqlite
+
+from llama_stack.apis.telemetry import (
+    MaterializedSpan,
+    QueryCondition,
+    Trace,
+    TraceStore,
+)
+
+
+class SQLiteTraceStore(TraceStore):
+    def __init__(self, conn_string: str):
+        self.conn_string = conn_string
+
+    async def query_traces(
+        self,
+        attribute_conditions: Optional[List[QueryCondition]] = None,
+        attribute_keys_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}"
+                )
+
+        query = (
+            select_clause
+            + """
+            FROM traces t
+            JOIN spans s ON t.trace_id = s.trace_id
+        """
+        )
+        params = []
+
+        # 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)
+
+        # 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)
+
+        # Add limit and offset
+        query += f" LIMIT {limit} OFFSET {offset}"
+
+        async with aiosqlite.connect(self.conn_string) as conn:
+            conn.row_factory = aiosqlite.Row
+            async with conn.execute(query, params) as cursor:
+                rows = await cursor.fetchall()
+                return [
+                    Trace(
+                        trace_id=row["trace_id"],
+                        root_span_id=row["root_span_id"],
+                        start_time=datetime.fromisoformat(row["start_time"]),
+                        end_time=datetime.fromisoformat(row["end_time"]),
+                    )
+                    for row in rows
+                ]
+
+    async def get_materialized_span(
+        self,
+        span_id: str,
+        attribute_keys_to_return: Optional[List[str]] = None,
+        max_depth: Optional[int] = None,
+    ) -> MaterializedSpan:
+        # Build the attributes selection
+        attributes_select = "s.attributes"
+        if attribute_keys_to_return:
+            json_object = ", ".join(
+                f"'{key}', json_extract(s.attributes, '$.{key}')"
+                for key in attribute_keys_to_return
+            )
+            attributes_select = f"json_object({json_object})"
+
+        # SQLite CTE query with filtered attributes
+        query = f"""
+        WITH RECURSIVE span_tree AS (
+            SELECT s.*, 1 as depth, {attributes_select} as filtered_attributes
+            FROM spans s
+            WHERE s.span_id = ?
+
+            UNION ALL
+
+            SELECT s.*, st.depth + 1, {attributes_select} as filtered_attributes
+            FROM spans s
+            JOIN span_tree st ON s.parent_span_id = st.span_id
+            WHERE (? IS NULL OR st.depth < ?)
+        )
+        SELECT *
+        FROM span_tree
+        ORDER BY depth, start_time
+        """
+
+        async with aiosqlite.connect(self.conn_string) as conn:
+            conn.row_factory = aiosqlite.Row
+            async with conn.execute(query, (span_id, max_depth, max_depth)) as cursor:
+                rows = await cursor.fetchall()
+
+                if not rows:
+                    raise ValueError(f"Span {span_id} not found")
+
+                # Build span tree
+                spans_by_id = {}
+                root_span = None
+
+                for row in rows:
+                    span = MaterializedSpan(
+                        span_id=row["span_id"],
+                        trace_id=row["trace_id"],
+                        parent_span_id=row["parent_span_id"],
+                        name=row["name"],
+                        start_time=datetime.fromisoformat(row["start_time"]),
+                        end_time=datetime.fromisoformat(row["end_time"]),
+                        attributes=json.loads(row["filtered_attributes"]),
+                        status=row["status"].lower(),
+                        children=[],
+                    )
+
+                    spans_by_id[span.span_id] = span
+
+                    if span.span_id == span_id:
+                        root_span = span
+                    elif span.parent_span_id in spans_by_id:
+                        spans_by_id[span.parent_span_id].children.append(span)
+
+                return root_span