Merge branch 'main' into implement-search-for-PGVector

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -28,7 +28,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import interleaved_con
 EMBEDDING_MODELS = {}
 
 
-log = get_logger(name=__name__, category="inference")
+log = get_logger(name=__name__, category="providers::utils")
 
 
 class SentenceTransformerEmbeddingMixin:

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -54,7 +54,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
 )
 
-logger = get_logger(name=__name__, category="inference")
+logger = get_logger(name=__name__, category="providers::utils")
 
 
 class LiteLLMOpenAIMixin(
@@ -429,28 +429,6 @@ class LiteLLMOpenAIMixin(
         )
         return await litellm.acompletion(**params)
 
-    async def batch_completion(
-        self,
-        model_id: str,
-        content_batch: list[InterleavedContent],
-        sampling_params: SamplingParams | None = None,
-        response_format: ResponseFormat | None = None,
-        logprobs: LogProbConfig | None = None,
-    ):
-        raise NotImplementedError("Batch completion is not supported for OpenAI Compat")
-
-    async def batch_chat_completion(
-        self,
-        model_id: str,
-        messages_batch: list[list[Message]],
-        sampling_params: SamplingParams | None = None,
-        tools: list[ToolDefinition] | None = None,
-        tool_config: ToolConfig | None = None,
-        response_format: ResponseFormat | None = None,
-        logprobs: LogProbConfig | None = None,
-    ):
-        raise NotImplementedError("Batch chat completion is not supported for OpenAI Compat")
-
     async def check_model_availability(self, model: str) -> bool:
         """
         Check if a specific model is available via LiteLLM for the current

llama_stack/providers/utils/inference/model_registry.py

@@ -17,7 +17,7 @@ from llama_stack.providers.utils.inference import (
     ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR,
 )
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="providers::utils")
 
 
 class RemoteInferenceProviderConfig(BaseModel):

llama_stack/providers/utils/inference/openai_compat.py

@@ -134,7 +134,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
     decode_assistant_message,
 )
 
-logger = get_logger(name=__name__, category="inference")
+logger = get_logger(name=__name__, category="providers::utils")
 
 
 class OpenAICompatCompletionChoiceDelta(BaseModel):

llama_stack/providers/utils/inference/openai_mixin.py

@@ -25,7 +25,7 @@ from llama_stack.apis.inference import (
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.openai_compat import prepare_openai_completion_params
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="providers::utils")
 
 
 class OpenAIMixin(ABC):

llama_stack/providers/utils/inference/prompt_adapter.py

@@ -58,7 +58,7 @@ from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.models.llama.sku_types import ModelFamily, is_multimodal
 from llama_stack.providers.utils.inference import supported_inference_models
 
-log = get_logger(name=__name__, category="inference")
+log = get_logger(name=__name__, category="providers::utils")
 
 
 class ChatCompletionRequestWithRawContent(ChatCompletionRequest):

llama_stack/providers/utils/kvstore/mongodb/mongodb.py

@@ -13,7 +13,7 @@ from llama_stack.providers.utils.kvstore import KVStore
 
 from ..config import MongoDBKVStoreConfig
 
-log = get_logger(name=__name__, category="kvstore")
+log = get_logger(name=__name__, category="providers::utils")
 
 
 class MongoDBKVStoreImpl(KVStore):

llama_stack/providers/utils/kvstore/postgres/postgres.py

@@ -14,7 +14,7 @@ from llama_stack.log import get_logger
 from ..api import KVStore
 from ..config import PostgresKVStoreConfig
 
-log = get_logger(name=__name__, category="kvstore")
+log = get_logger(name=__name__, category="providers::utils")
 
 
 class PostgresKVStoreImpl(KVStore):

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -44,7 +44,7 @@ from llama_stack.providers.utils.memory.vector_store import (
     make_overlapped_chunks,
 )
 
-logger = get_logger(name=__name__, category="memory")
+logger = get_logger(name=__name__, category="providers::utils")
 
 # Constants for OpenAI vector stores
 CHUNK_MULTIPLIER = 5

llama_stack/providers/utils/memory/vector_store.py

@@ -33,7 +33,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 )
 from llama_stack.providers.utils.vector_io.vector_utils import generate_chunk_id
 
-log = get_logger(name=__name__, category="memory")
+log = get_logger(name=__name__, category="providers::utils")
 
 
 class ChunkForDeletion(BaseModel):

llama_stack/providers/utils/scheduler.py

@@ -17,7 +17,7 @@ from pydantic import BaseModel
 
 from llama_stack.log import get_logger
 
-logger = get_logger(name=__name__, category="scheduler")
+logger = get_logger(name=__name__, category="providers::utils")
 
 
 # TODO: revisit the list of possible statuses when defining a more coherent

llama_stack/providers/utils/sqlstore/authorized_sqlstore.py

@@ -17,7 +17,7 @@ from llama_stack.log import get_logger
 from .api import ColumnDefinition, ColumnType, PaginatedResponse, SqlStore
 from .sqlstore import SqlStoreType
 
-logger = get_logger(name=__name__, category="authorized_sqlstore")
+logger = get_logger(name=__name__, category="providers::utils")
 
 # Hardcoded copy of the default policy that our SQL filtering implements
 # WARNING: If default_policy() changes, this constant must be updated accordingly

llama_stack/providers/utils/sqlstore/sqlalchemy_sqlstore.py

@@ -30,7 +30,7 @@ from llama_stack.log import get_logger
 from .api import ColumnDefinition, ColumnType, SqlStore
 from .sqlstore import SqlAlchemySqlStoreConfig
 
-logger = get_logger(name=__name__, category="sqlstore")
+logger = get_logger(name=__name__, category="providers::utils")
 
 TYPE_MAPPING: dict[ColumnType, Any] = {
     ColumnType.INTEGER: Integer,

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,192 @@ class TraceStore(Protocol):
         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 = None,
+        query_type: MetricQueryType = MetricQueryType.RANGE,
+        label_matchers: list[MetricLabelMatcher] | None = None,
+    ) -> QueryMetricsResponse:
+        if end_time is None:
+            end_time = datetime.now(UTC)
+
+        # Build base query
+        if query_type == MetricQueryType.INSTANT:
+            query = """
+                SELECT
+                    se.name,
+                    SUM(CAST(json_extract(se.attributes, '$.value') AS REAL)) as value,
+                    json_extract(se.attributes, '$.unit') as unit,
+                    se.attributes
+                FROM span_events se
+                WHERE se.name = ?
+                AND se.timestamp BETWEEN ? AND ?
+            """
+        else:
+            if granularity:
+                time_format = self._get_time_format_for_granularity(granularity)
+                query = f"""
+                    SELECT
+                        se.name,
+                        SUM(CAST(json_extract(se.attributes, '$.value') AS REAL)) as value,
+                        json_extract(se.attributes, '$.unit') as unit,
+                        se.attributes,
+                        strftime('{time_format}', se.timestamp) as bucket_start
+                    FROM span_events se
+                    WHERE se.name = ?
+                    AND se.timestamp BETWEEN ? AND ?
+                """
+            else:
+                query = """
+                    SELECT
+                        se.name,
+                        json_extract(se.attributes, '$.value') as value,
+                        json_extract(se.attributes, '$.unit') as unit,
+                        se.attributes,
+                        se.timestamp
+                    FROM span_events se
+                    WHERE se.name = ?
+                    AND se.timestamp BETWEEN ? AND ?
+                """
+
+        params = [f"metric.{metric_name}", start_time.isoformat(), end_time.isoformat()]
+
+        # Labels that will be attached to the MetricSeries (preserve matcher labels)
+        all_labels: list[MetricLabel] = []
+        matcher_label_names = set()
+        if label_matchers:
+            for matcher in label_matchers:
+                json_path = f"$.{matcher.name}"
+                if matcher.operator == "=":
+                    query += f" AND json_extract(se.attributes, '{json_path}') = ?"
+                    params.append(matcher.value)
+                elif matcher.operator == "!=":
+                    query += f" AND json_extract(se.attributes, '{json_path}') != ?"
+                    params.append(matcher.value)
+                elif matcher.operator == "=~":
+                    query += f" AND json_extract(se.attributes, '{json_path}') LIKE ?"
+                    params.append(f"%{matcher.value}%")
+                elif matcher.operator == "!~":
+                    query += f" AND json_extract(se.attributes, '{json_path}') NOT LIKE ?"
+                    params.append(f"%{matcher.value}%")
+                # Preserve filter context in output
+                all_labels.append(MetricLabel(name=matcher.name, value=str(matcher.value)))
+                matcher_label_names.add(matcher.name)
+
+        # GROUP BY / ORDER BY logic
+        if query_type == MetricQueryType.RANGE and granularity:
+            group_time_format = self._get_time_format_for_granularity(granularity)
+            query += f" GROUP BY strftime('{group_time_format}', se.timestamp), json_extract(se.attributes, '$.unit')"
+            query += " ORDER BY bucket_start"
+        elif query_type == MetricQueryType.INSTANT:
+            query += " GROUP BY json_extract(se.attributes, '$.unit')"
+        else:
+            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=[])
+
+                data_points = []
+                # We want to add attribute labels, but only those not already present as matcher labels.
+                attr_label_names = set()
+                for row in rows:
+                    # Parse JSON attributes safely, if there are no attributes (weird), just don't add the labels to the result.
+                    try:
+                        attributes = json.loads(row["attributes"] or "{}")
+                    except (TypeError, json.JSONDecodeError):
+                        attributes = {}
+
+                    value = row["value"]
+                    unit = row["unit"] or ""
+
+                    # Add labels from attributes without duplicating matcher labels, if we don't do this, there will be a lot of duplicate label in the result.
+                    for k, v in attributes.items():
+                        if k not in ["value", "unit"] and k not in matcher_label_names and k not in attr_label_names:
+                            all_labels.append(MetricLabel(name=k, value=str(v)))
+                            attr_label_names.add(k)
+
+                    # Determine timestamp
+                    if query_type == MetricQueryType.RANGE and granularity:
+                        try:
+                            bucket_start_raw = row["bucket_start"]
+                        except KeyError as e:
+                            raise ValueError(
+                                "DB did not have a bucket_start time in row when using granularity, this indicates improper formatting"
+                            ) from e
+                        # this value could also be there, but be NULL, I think.
+                        if bucket_start_raw is None:
+                            raise ValueError("bucket_start is None check time format and data")
+                        bucket_start = datetime.fromisoformat(bucket_start_raw)
+                        timestamp = int(bucket_start.timestamp())
+                    elif query_type == MetricQueryType.INSTANT:
+                        timestamp = int(datetime.now(UTC).timestamp())
+                    else:
+                        try:
+                            timestamp_raw = row["timestamp"]
+                        except KeyError as e:
+                            raise ValueError(
+                                "DB did not have a timestamp in row, this indicates improper formatting"
+                            ) from e
+                        # this value could also be there, but be NULL, I think.
+                        if timestamp_raw is None:
+                            raise ValueError("timestamp is None check time format and data")
+                        timestamp_iso = datetime.fromisoformat(timestamp_raw)
+                        timestamp = int(timestamp_iso.timestamp())
+
+                    data_points.append(
+                        MetricDataPoint(
+                            timestamp=timestamp,
+                            value=value,
+                            unit=unit,
+                        )
+                    )
+
+                metric_series = [MetricSeries(metric=metric_name, labels=all_labels, values=data_points)]
+                return QueryMetricsResponse(data=metric_series)
+
+    def _get_time_format_for_granularity(self, granularity: str | None) -> str:
+        """Get the SQLite strftime format string for a given granularity.
+        Args:
+            granularity: Granularity string (e.g., "1m", "5m", "1h", "1d")
+        Returns:
+            SQLite strftime format string for the granularity
+        """
+        if granularity is None:
+            raise ValueError("granularity cannot be None for this method - use separate logic for no aggregation")
+
+        if granularity.endswith("d"):
+            return "%Y-%m-%d 00:00:00"
+        elif granularity.endswith("h"):
+            return "%Y-%m-%d %H:00:00"
+        elif granularity.endswith("m"):
+            return "%Y-%m-%d %H:%M:00"
+        else:
+            return "%Y-%m-%d %H:%M:00"  # Default to most granular which will give us the most timestamps.
+
     async def query_traces(
         self,
         attribute_filters: list[QueryCondition] | None = None,