feat: implement query_metrics

cdoern · cdoern · commit 786905704812 · 2025-08-07T20:03:58.000-04:00
query_metrics currently has no implementation, meaning once a metric is emitted there is no way in llama stack to query it from the store.

implement query_metrics for the meta_reference provider which follows a similar style to `query_traces`, using the trace_store to format an SQL query and execute it

in this case the parameters for the query are `metric.METRIC_NAME, start_time, and end_time`.

this required client side changes since the client had no `query_metrics` or any associated resources, so any tests here will fail but I will provider manual execution logs for the new tests I am adding

order the metrics by timestamp.

Additionally add `unit` to the `MetricDataPoint` class since this adds much more context to the metric being queried.

Signed-off-by: Charlie Doern &lt;cdoern@redhat.com&gt;
diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html
@@ -15846,12 +15846,16 @@
                     "value": {
                         "type": "number",
                         "description": "The numeric value of the metric at this timestamp"
+                    },
+                    "unit": {
+                        "type": "string"
                     }
                 },
                 "additionalProperties": false,
                 "required": [
                     "timestamp",
-                    "value"
+                    "value",
+                    "unit"
                 ],
                 "title": "MetricDataPoint",
                 "description": "A single data point in a metric time series."
diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml
@@ -11774,10 +11774,13 @@ components:
           type: number
           description: >-
             The numeric value of the metric at this timestamp
+        unit:
+          type: string
       additionalProperties: false
       required:
         - timestamp
         - value
+        - unit
       title: MetricDataPoint
       description: >-
         A single data point in a metric time series.
diff --git a/llama_stack/apis/telemetry/telemetry.py b/llama_stack/apis/telemetry/telemetry.py
@@ -386,6 +386,7 @@ class MetricDataPoint(BaseModel):
 
     timestamp: int
     value: float
+    unit: str
 
 
 @json_schema_type
diff --git a/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py b/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py
@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import datetime
 import logging
 import threading
 from typing import Any
@@ -149,7 +150,36 @@ async def query_metrics(
         query_type: MetricQueryType = MetricQueryType.RANGE,
         label_matchers: list[MetricLabelMatcher] | None = None,
     ) -> QueryMetricsResponse:
-        raise NotImplementedError("Querying metrics is not implemented")
+        """Query metrics from the telemetry store.
+
+        Args:
+            metric_name: The name of the metric to query (e.g., "prompt_tokens")
+            start_time: Start time as Unix timestamp
+            end_time: End time as Unix timestamp (defaults to now if None)
+            granularity: Time granularity for aggregation (not implemented yet)
+            query_type: Type of query (RANGE or INSTANT)
+            label_matchers: Label filters to apply
+
+        Returns:
+            QueryMetricsResponse with metric time series data
+        """
+        # Convert timestamps to datetime objects
+        start_dt = datetime.datetime.fromtimestamp(start_time, datetime.UTC)
+        end_dt = datetime.datetime.fromtimestamp(end_time, datetime.UTC) if end_time else None
+
+        # Use SQLite trace store if available
+        if hasattr(self, "trace_store") and self.trace_store:
+            return await self.trace_store.query_metrics(
+                metric_name=metric_name,
+                start_time=start_dt,
+                end_time=end_dt,
+                granularity=granularity,
+                query_type=query_type,
+                label_matchers=label_matchers,
+            )
+
+        # Fallback to empty response if no trace store
+        return QueryMetricsResponse(data=[])
 
     def _log_unstructured(self, event: UnstructuredLogEvent, ttl_seconds: int) -> None:
         with self._lock:
diff --git a/llama_stack/providers/utils/telemetry/sqlite_trace_store.py b/llama_stack/providers/utils/telemetry/sqlite_trace_store.py
@@ -5,12 +5,23 @@
 # the root directory of this source tree.
 
 import json
-from datetime import datetime
+from datetime import UTC, datetime
 from typing import Protocol
 
 import aiosqlite
 
-from llama_stack.apis.telemetry import QueryCondition, Span, SpanWithStatus, Trace
+from llama_stack.apis.telemetry import (
+    MetricDataPoint,
+    MetricLabel,
+    MetricLabelMatcher,
+    MetricQueryType,
+    MetricSeries,
+    QueryCondition,
+    QueryMetricsResponse,
+    Span,
+    SpanWithStatus,
+    Trace,
+)
 
 
 class TraceStore(Protocol):
@@ -29,11 +40,119 @@ async def get_span_tree(
         max_depth: int | None = None,
     ) -> dict[str, SpanWithStatus]: ...
 
+    async def query_metrics(
+        self,
+        metric_name: str,
+        start_time: datetime,
+        end_time: datetime | None = None,
+        granularity: str | None = "1d",
+        query_type: MetricQueryType = MetricQueryType.RANGE,
+        label_matchers: list[MetricLabelMatcher] | None = None,
+    ) -> QueryMetricsResponse: ...
+
 
 class SQLiteTraceStore(TraceStore):
     def __init__(self, conn_string: str):
         self.conn_string = conn_string
 
+    async def query_metrics(
+        self,
+        metric_name: str,
+        start_time: datetime,
+        end_time: datetime | None = None,
+        granularity: str | None = "1d",
+        query_type: MetricQueryType = MetricQueryType.RANGE,
+        label_matchers: list[MetricLabelMatcher] | None = None,
+    ) -> QueryMetricsResponse:
+        """Query metrics from span events stored in SQLite.
+
+        Args:
+            metric_name: The name of the metric to query (e.g., "prompt_tokens")
+            start_time: Start time for the query range
+            end_time: End time for the query range (defaults to now if None)
+            granularity: Time granularity for aggregation (not implemented yet)
+            query_type: Type of query (RANGE or INSTANT)
+            label_matchers: Label filters to apply
+
+        Returns:
+            QueryMetricsResponse with metric time series data
+        """
+        if end_time is None:
+            end_time = datetime.now(UTC)
+
+        # Build the base query
+        query = """
+            SELECT
+                se.name,
+                se.timestamp,
+                se.attributes
+            FROM span_events se
+            WHERE se.name = ?
+              AND se.timestamp BETWEEN ? AND ?
+        """
+
+        params = [f"metric.{metric_name}", start_time.isoformat(), end_time.isoformat()]
+
+        # Add label matchers if provided
+        if label_matchers:
+            for matcher in label_matchers:
+                if matcher.operator == "=":
+                    query += f" AND json_extract(se.attributes, '$.{matcher.name}') = ?"
+                    params.append(matcher.value)
+                elif matcher.operator == "!=":
+                    query += f" AND json_extract(se.attributes, '$.{matcher.name}') != ?"
+                    params.append(matcher.value)
+                elif matcher.operator == "=~":
+                    query += f" AND json_extract(se.attributes, '$.{matcher.name}') LIKE ?"
+                    params.append(f"%{matcher.value}%")
+                elif matcher.operator == "!~":
+                    query += f" AND json_extract(se.attributes, '$.{matcher.name}') NOT LIKE ?"
+                    params.append(f"%{matcher.value}%")
+
+        query += " ORDER BY se.timestamp"
+
+        # Execute query
+        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()
+
+                if not rows:
+                    return QueryMetricsResponse(data=[])
+
+                # Parse metric data
+                data_points = []
+                labels: list[MetricLabel] = []
+
+                for row in rows:
+                    # Parse JSON attributes
+                    attributes = json.loads(row["attributes"])
+
+                    # Extract metric value and unit
+                    value = attributes.get("value")
+                    unit = attributes.get("unit", "")
+
+                    # Extract labels from attributes
+                    metric_labels = []
+                    for key, val in attributes.items():
+                        if key not in ["value", "unit"]:
+                            metric_labels.append(MetricLabel(name=key, value=str(val)))
+
+                    # Create data point
+                    timestamp = datetime.fromisoformat(row["timestamp"])
+                    data_points.append(
+                        MetricDataPoint(
+                            timestamp=int(timestamp.timestamp()),
+                            value=value,
+                            unit=unit,
+                        )
+                    )
+
+                # Create metric series
+                metric_series = [MetricSeries(metric=metric_name, labels=labels, values=data_points)]
+
+                return QueryMetricsResponse(data=metric_series)
+
     async def query_traces(
         self,
         attribute_filters: list[QueryCondition] | None = None,
