Merge branch 'main' into add_max_tool_calls_v2

pyproject.toml

@@ -356,6 +356,10 @@ exclude = [
 module = [
     "yaml",
     "fire",
+    "redis.asyncio",
+    "psycopg2",
+    "psycopg2.extras",
+    "psycopg2.extensions",
     "torchtune.*",
     "fairscale.*",
     "torchvision.*",

src/llama_stack/core/conversations/conversations.py

@@ -11,10 +11,9 @@ from typing import Any, Literal
 from pydantic import BaseModel, TypeAdapter
 
 from llama_stack.core.datatypes import AccessRule, StackRunConfig
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
-from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
 from llama_stack_api import (
     Conversation,
     ConversationDeletedResource,
@@ -25,6 +24,7 @@ from llama_stack_api import (
     Conversations,
     Metadata,
 )
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
 
 logger = get_logger(name=__name__, category="openai_conversations")
 

src/llama_stack/core/prompts/prompts.py

@@ -10,7 +10,7 @@ from typing import Any
 from pydantic import BaseModel
 
 from llama_stack.core.datatypes import StackRunConfig
-from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
+from llama_stack.core.storage.kvstore import KVStore, kvstore_impl
 from llama_stack_api import ListPromptsResponse, Prompt, Prompts
 
 

src/llama_stack/core/server/quota.py

@@ -11,9 +11,9 @@ from datetime import UTC, datetime, timedelta
 from starlette.types import ASGIApp, Receive, Scope, Send
 
 from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendType
+from llama_stack.core.storage.kvstore.kvstore import _KVSTORE_BACKENDS, kvstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore.api import KVStore
+from llama_stack_api.internal.kvstore import KVStore
-from llama_stack.providers.utils.kvstore.kvstore import _KVSTORE_BACKENDS, kvstore_impl
 
 logger = get_logger(name=__name__, category="core::server")
 

src/llama_stack/core/stack.py

@@ -385,8 +385,8 @@ def _initialize_storage(run_config: StackRunConfig):
         else:
             raise ValueError(f"Unknown storage backend type: {type}")
 
-    from llama_stack.providers.utils.kvstore.kvstore import register_kvstore_backends
+    from llama_stack.core.storage.kvstore.kvstore import register_kvstore_backends
-    from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+    from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 
     register_kvstore_backends(kv_backends)
     register_sqlstore_backends(sql_backends)

src/llama_stack/core/storage/datatypes.py

@@ -12,6 +12,8 @@ from typing import Annotated, Literal
 
 from pydantic import BaseModel, Field, field_validator
 
+from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR
+
 
 class StorageBackendType(StrEnum):
     KV_REDIS = "kv_redis"
@@ -256,15 +258,24 @@ class ResponsesStoreReference(InferenceStoreReference):
 
 class ServerStoresConfig(BaseModel):
     metadata: KVStoreReference | None = Field(
-        default=None,
+        default=KVStoreReference(
+            backend="kv_default",
+            namespace="registry",
+        ),
         description="Metadata store configuration (uses KV backend)",
     )
     inference: InferenceStoreReference | None = Field(
-        default=None,
+        default=InferenceStoreReference(
+            backend="sql_default",
+            table_name="inference_store",
+        ),
         description="Inference store configuration (uses SQL backend)",
     )
     conversations: SqlStoreReference | None = Field(
-        default=None,
+        default=SqlStoreReference(
+            backend="sql_default",
+            table_name="openai_conversations",
+        ),
         description="Conversations store configuration (uses SQL backend)",
     )
     responses: ResponsesStoreReference | None = Field(
@@ -272,13 +283,21 @@ class ServerStoresConfig(BaseModel):
         description="Responses store configuration (uses SQL backend)",
     )
     prompts: KVStoreReference | None = Field(
-        default=None,
+        default=KVStoreReference(backend="kv_default", namespace="prompts"),
         description="Prompts store configuration (uses KV backend)",
     )
 
 
 class StorageConfig(BaseModel):
     backends: dict[str, StorageBackendConfig] = Field(
+        default={
+            "kv_default": SqliteKVStoreConfig(
+                db_path=f"${{env.SQLITE_STORE_DIR:={DISTRIBS_BASE_DIR}}}/kvstore.db",
+            ),
+            "sql_default": SqliteSqlStoreConfig(
+                db_path=f"${{env.SQLITE_STORE_DIR:={DISTRIBS_BASE_DIR}}}/sql_store.db",
+            ),
+        },
         description="Named backend configurations (e.g., 'default', 'cache')",
     )
     stores: ServerStoresConfig = Field(

src/llama_stack/core/storage/kvstore/__init__.py

@@ -4,4 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from llama_stack_api.internal.kvstore import KVStore as KVStore
+
 from .kvstore import *  # noqa: F401, F403

src/llama_stack/core/storage/kvstore/kvstore.py

@@ -13,11 +13,19 @@ from __future__ import annotations
 
 import asyncio
 from collections import defaultdict
+from datetime import datetime
+from typing import cast
 
-from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig, StorageBackendType
+from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig
+from llama_stack_api.internal.kvstore import KVStore
 
-from .api import KVStore
+from .config import (
-from .config import KVStoreConfig
+    KVStoreConfig,
+    MongoDBKVStoreConfig,
+    PostgresKVStoreConfig,
+    RedisKVStoreConfig,
+    SqliteKVStoreConfig,
+)
 
 
 def kvstore_dependencies():
@@ -33,7 +41,7 @@ def kvstore_dependencies():
 
 class InmemoryKVStoreImpl(KVStore):
     def __init__(self):
-        self._store = {}
+        self._store: dict[str, str] = {}
 
     async def initialize(self) -> None:
         pass
@@ -41,7 +49,7 @@ class InmemoryKVStoreImpl(KVStore):
     async def get(self, key: str) -> str | None:
         return self._store.get(key)
 
-    async def set(self, key: str, value: str) -> None:
+    async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         self._store[key] = value
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
@@ -70,7 +78,8 @@ def register_kvstore_backends(backends: dict[str, StorageBackendConfig]) -> None
     _KVSTORE_INSTANCES.clear()
     _KVSTORE_LOCKS.clear()
     for name, cfg in backends.items():
-        _KVSTORE_BACKENDS[name] = cfg
+        typed_cfg = cast(KVStoreConfig, cfg)
+        _KVSTORE_BACKENDS[name] = typed_cfg
 
 
 async def kvstore_impl(reference: KVStoreReference) -> KVStore:
@@ -94,19 +103,20 @@ async def kvstore_impl(reference: KVStoreReference) -> KVStore:
         config = backend_config.model_copy()
         config.namespace = reference.namespace
 
-        if config.type == StorageBackendType.KV_REDIS.value:
+        impl: KVStore
+        if isinstance(config, RedisKVStoreConfig):
             from .redis import RedisKVStoreImpl
 
             impl = RedisKVStoreImpl(config)
-        elif config.type == StorageBackendType.KV_SQLITE.value:
+        elif isinstance(config, SqliteKVStoreConfig):
             from .sqlite import SqliteKVStoreImpl
 
             impl = SqliteKVStoreImpl(config)
-        elif config.type == StorageBackendType.KV_POSTGRES.value:
+        elif isinstance(config, PostgresKVStoreConfig):
             from .postgres import PostgresKVStoreImpl
 
             impl = PostgresKVStoreImpl(config)
-        elif config.type == StorageBackendType.KV_MONGODB.value:
+        elif isinstance(config, MongoDBKVStoreConfig):
             from .mongodb import MongoDBKVStoreImpl
 
             impl = MongoDBKVStoreImpl(config)

src/llama_stack/core/storage/kvstore/mongodb/mongodb.py

@@ -9,8 +9,8 @@ from datetime import datetime
 from pymongo import AsyncMongoClient
 from pymongo.asynchronous.collection import AsyncCollection
 
+from llama_stack.core.storage.kvstore import KVStore
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import KVStore
 
 from ..config import MongoDBKVStoreConfig
 

src/llama_stack/core/storage/kvstore/postgres/postgres.py

@@ -6,12 +6,13 @@
 
 from datetime import datetime
 
-import psycopg2
+import psycopg2  # type: ignore[import-not-found]
-from psycopg2.extras import DictCursor
+from psycopg2.extensions import connection as PGConnection  # type: ignore[import-not-found]
+from psycopg2.extras import DictCursor  # type: ignore[import-not-found]
 
 from llama_stack.log import get_logger
+from llama_stack_api.internal.kvstore import KVStore
 
-from ..api import KVStore
 from ..config import PostgresKVStoreConfig
 
 log = get_logger(name=__name__, category="providers::utils")
@@ -20,12 +21,12 @@ log = get_logger(name=__name__, category="providers::utils")
 class PostgresKVStoreImpl(KVStore):
     def __init__(self, config: PostgresKVStoreConfig):
         self.config = config
-        self.conn = None
+        self._conn: PGConnection | None = None
-        self.cursor = None
+        self._cursor: DictCursor | None = None
 
     async def initialize(self) -> None:
         try:
-            self.conn = psycopg2.connect(
+            self._conn = psycopg2.connect(
                 host=self.config.host,
                 port=self.config.port,
                 database=self.config.db,
@@ -34,11 +35,11 @@ class PostgresKVStoreImpl(KVStore):
                 sslmode=self.config.ssl_mode,
                 sslrootcert=self.config.ca_cert_path,
             )
-            self.conn.autocommit = True
+            self._conn.autocommit = True
-            self.cursor = self.conn.cursor(cursor_factory=DictCursor)
+            self._cursor = self._conn.cursor(cursor_factory=DictCursor)
 
             # Create table if it doesn't exist
-            self.cursor.execute(
+            self._cursor.execute(
                 f"""
                 CREATE TABLE IF NOT EXISTS {self.config.table_name} (
                     key TEXT PRIMARY KEY,
@@ -51,6 +52,11 @@ class PostgresKVStoreImpl(KVStore):
             log.exception("Could not connect to PostgreSQL database server")
             raise RuntimeError("Could not connect to PostgreSQL database server") from e
 
+    def _cursor_or_raise(self) -> DictCursor:
+        if self._cursor is None:
+            raise RuntimeError("Postgres client not initialized")
+        return self._cursor
+
     def _namespaced_key(self, key: str) -> str:
         if not self.config.namespace:
             return key
@@ -58,7 +64,8 @@ class PostgresKVStoreImpl(KVStore):
 
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         key = self._namespaced_key(key)
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"""
             INSERT INTO {self.config.table_name} (key, value, expiration)
             VALUES (%s, %s, %s)
@@ -70,7 +77,8 @@ class PostgresKVStoreImpl(KVStore):
 
     async def get(self, key: str) -> str | None:
         key = self._namespaced_key(key)
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"""
             SELECT value FROM {self.config.table_name}
             WHERE key = %s
@@ -78,12 +86,13 @@ class PostgresKVStoreImpl(KVStore):
             """,
             (key,),
         )
-        result = self.cursor.fetchone()
+        result = cursor.fetchone()
         return result[0] if result else None
 
     async def delete(self, key: str) -> None:
         key = self._namespaced_key(key)
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"DELETE FROM {self.config.table_name} WHERE key = %s",
             (key,),
         )
@@ -92,7 +101,8 @@ class PostgresKVStoreImpl(KVStore):
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
 
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"""
             SELECT value FROM {self.config.table_name}
             WHERE key >= %s AND key < %s
@@ -101,14 +111,15 @@ class PostgresKVStoreImpl(KVStore):
             """,
             (start_key, end_key),
         )
-        return [row[0] for row in self.cursor.fetchall()]
+        return [row[0] for row in cursor.fetchall()]
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
 
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"SELECT key FROM {self.config.table_name} WHERE key >= %s AND key < %s",
             (start_key, end_key),
         )
-        return [row[0] for row in self.cursor.fetchall()]
+        return [row[0] for row in cursor.fetchall()]

src/llama_stack/core/storage/kvstore/redis/redis.py

@@ -6,18 +6,25 @@
 
 from datetime import datetime
 
-from redis.asyncio import Redis
+from redis.asyncio import Redis  # type: ignore[import-not-found]
+
+from llama_stack_api.internal.kvstore import KVStore
 
-from ..api import KVStore
 from ..config import RedisKVStoreConfig
 
 
 class RedisKVStoreImpl(KVStore):
     def __init__(self, config: RedisKVStoreConfig):
         self.config = config
+        self._redis: Redis | None = None
 
     async def initialize(self) -> None:
-        self.redis = Redis.from_url(self.config.url)
+        self._redis = Redis.from_url(self.config.url)
+
+    def _client(self) -> Redis:
+        if self._redis is None:
+            raise RuntimeError("Redis client not initialized")
+        return self._redis
 
     def _namespaced_key(self, key: str) -> str:
         if not self.config.namespace:
@@ -26,30 +33,37 @@ class RedisKVStoreImpl(KVStore):
 
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         key = self._namespaced_key(key)
-        await self.redis.set(key, value)
+        client = self._client()
+        await client.set(key, value)
         if expiration:
-            await self.redis.expireat(key, expiration)
+            await client.expireat(key, expiration)
 
     async def get(self, key: str) -> str | None:
         key = self._namespaced_key(key)
-        value = await self.redis.get(key)
+        client = self._client()
+        value = await client.get(key)
         if value is None:
             return None
-        await self.redis.ttl(key)
+        await client.ttl(key)
-        return value
+        if isinstance(value, bytes):
+            return value.decode("utf-8")
+        if isinstance(value, str):
+            return value
+        return str(value)
 
     async def delete(self, key: str) -> None:
         key = self._namespaced_key(key)
-        await self.redis.delete(key)
+        await self._client().delete(key)
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
+        client = self._client()
         cursor = 0
         pattern = start_key + "*"  # Match all keys starting with start_key prefix
-        matching_keys = []
+        matching_keys: list[str | bytes] = []
         while True:
-            cursor, keys = await self.redis.scan(cursor, match=pattern, count=1000)
+            cursor, keys = await client.scan(cursor, match=pattern, count=1000)
 
             for key in keys:
                 key_str = key.decode("utf-8") if isinstance(key, bytes) else key
@@ -61,7 +75,7 @@ class RedisKVStoreImpl(KVStore):
 
         # Then fetch all values in a single MGET call
         if matching_keys:
-            values = await self.redis.mget(matching_keys)
+            values = await client.mget(matching_keys)
             return [
                 value.decode("utf-8") if isinstance(value, bytes) else value for value in values if value is not None
             ]
@@ -70,7 +84,18 @@ class RedisKVStoreImpl(KVStore):
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
         """Get all keys in the given range."""
-        matching_keys = await self.redis.zrangebylex(self.namespace, f"[{start_key}", f"[{end_key}")
+        start_key = self._namespaced_key(start_key)
-        if not matching_keys:
+        end_key = self._namespaced_key(end_key)
-            return []
+        client = self._client()
-        return [k.decode("utf-8") for k in matching_keys]
+        cursor = 0
+        pattern = start_key + "*"
+        result: list[str] = []
+        while True:
+            cursor, keys = await client.scan(cursor, match=pattern, count=1000)
+            for key in keys:
+                key_str = key.decode("utf-8") if isinstance(key, bytes) else str(key)
+                if start_key <= key_str <= end_key:
+                    result.append(key_str)
+            if cursor == 0:
+                break
+        return result

src/llama_stack/core/storage/kvstore/sqlite/sqlite.py

@@ -10,8 +10,8 @@ from datetime import datetime
 import aiosqlite
 
 from llama_stack.log import get_logger
+from llama_stack_api.internal.kvstore import KVStore
 
-from ..api import KVStore
 from ..config import SqliteKVStoreConfig
 
 logger = get_logger(name=__name__, category="providers::utils")

src/llama_stack/core/storage/sqlstore/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+
+from llama_stack_api.internal.sqlstore import (
+    ColumnDefinition as ColumnDefinition,
+)
+from llama_stack_api.internal.sqlstore import (
+    ColumnType as ColumnType,
+)
+from llama_stack_api.internal.sqlstore import (
+    SqlStore as SqlStore,
+)
+
+from .sqlstore import *  # noqa: F401,F403

src/llama_stack/core/storage/sqlstore/authorized_sqlstore.py

@@ -14,8 +14,8 @@ from llama_stack.core.datatypes import User
 from llama_stack.core.request_headers import get_authenticated_user
 from llama_stack.core.storage.datatypes import StorageBackendType
 from llama_stack.log import get_logger
-
+from llama_stack_api import PaginatedResponse
-from .api import ColumnDefinition, ColumnType, PaginatedResponse, SqlStore
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType, SqlStore
 
 logger = get_logger(name=__name__, category="providers::utils")
 

src/llama_stack/core/storage/sqlstore/sqlalchemy_sqlstore.py

@@ -29,8 +29,7 @@ from sqlalchemy.sql.elements import ColumnElement
 from llama_stack.core.storage.datatypes import SqlAlchemySqlStoreConfig
 from llama_stack.log import get_logger
 from llama_stack_api import PaginatedResponse
-
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType, SqlStore
-from .api import ColumnDefinition, ColumnType, SqlStore
 
 logger = get_logger(name=__name__, category="providers::utils")
 

src/llama_stack/core/storage/sqlstore/sqlstore.py

@@ -16,8 +16,7 @@ from llama_stack.core.storage.datatypes import (
     StorageBackendConfig,
     StorageBackendType,
 )
-
+from llama_stack_api.internal.sqlstore import SqlStore
-from .api import SqlStore
 
 sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"]
 

src/llama_stack/core/store/registry.py

@@ -12,8 +12,8 @@ import pydantic
 
 from llama_stack.core.datatypes import RoutableObjectWithProvider
 from llama_stack.core.storage.datatypes import KVStoreReference
+from llama_stack.core.storage.kvstore import KVStore, kvstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
 
 logger = get_logger(__name__, category="core::registry")
 

src/llama_stack/distributions/starter/starter.py

@@ -17,6 +17,8 @@ from llama_stack.core.datatypes import (
     ToolGroupInput,
     VectorStoresConfig,
 )
+from llama_stack.core.storage.kvstore.config import PostgresKVStoreConfig
+from llama_stack.core.storage.sqlstore.sqlstore import PostgresSqlStoreConfig
 from llama_stack.core.utils.dynamic import instantiate_class_type
 from llama_stack.distributions.template import DistributionTemplate, RunConfigSettings
 from llama_stack.providers.inline.files.localfs.config import LocalfsFilesImplConfig
@@ -35,8 +37,6 @@ from llama_stack.providers.remote.vector_io.pgvector.config import (
 )
 from llama_stack.providers.remote.vector_io.qdrant.config import QdrantVectorIOConfig
 from llama_stack.providers.remote.vector_io.weaviate.config import WeaviateVectorIOConfig
-from llama_stack.providers.utils.kvstore.config import PostgresKVStoreConfig
-from llama_stack.providers.utils.sqlstore.sqlstore import PostgresSqlStoreConfig
 from llama_stack_api import RemoteProviderSpec
 
 

src/llama_stack/distributions/template.py

@@ -35,13 +35,13 @@ from llama_stack.core.storage.datatypes import (
     SqlStoreReference,
     StorageBackendType,
 )
+from llama_stack.core.storage.kvstore.config import SqliteKVStoreConfig
+from llama_stack.core.storage.kvstore.config import get_pip_packages as get_kv_pip_packages
+from llama_stack.core.storage.sqlstore.sqlstore import SqliteSqlStoreConfig
+from llama_stack.core.storage.sqlstore.sqlstore import get_pip_packages as get_sql_pip_packages
 from llama_stack.core.utils.dynamic import instantiate_class_type
 from llama_stack.core.utils.image_types import LlamaStackImageType
 from llama_stack.providers.utils.inference.model_registry import ProviderModelEntry
-from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
-from llama_stack.providers.utils.kvstore.config import get_pip_packages as get_kv_pip_packages
-from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
-from llama_stack.providers.utils.sqlstore.sqlstore import get_pip_packages as get_sql_pip_packages
 from llama_stack_api import DatasetPurpose, ModelType
 
 

src/llama_stack/providers/inline/agents/meta_reference/agents.py

@@ -6,8 +6,8 @@
 
 
 from llama_stack.core.datatypes import AccessRule
+from llama_stack.core.storage.kvstore import InmemoryKVStoreImpl, kvstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import InmemoryKVStoreImpl, kvstore_impl
 from llama_stack.providers.utils.responses.responses_store import ResponsesStore
 from llama_stack_api import (
     Agents,

src/llama_stack/providers/inline/batches/reference/__init__.py

@@ -7,7 +7,7 @@
 from typing import Any
 
 from llama_stack.core.datatypes import AccessRule, Api
-from llama_stack.providers.utils.kvstore import kvstore_impl
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack_api import Files, Inference, Models
 
 from .batches import ReferenceBatchesImpl

src/llama_stack/providers/inline/batches/reference/batches.py

@@ -16,8 +16,8 @@ from typing import Any, Literal
 from openai.types.batch import BatchError, Errors
 from pydantic import BaseModel
 
+from llama_stack.core.storage.kvstore import KVStore
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import KVStore
 from llama_stack_api import (
     Batches,
     BatchObject,

src/llama_stack/providers/inline/datasetio/localfs/datasetio.py

@@ -5,8 +5,8 @@
 # the root directory of this source tree.
 from typing import Any
 
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_uri
-from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.utils.pagination import paginate_records
 from llama_stack_api import Dataset, DatasetIO, DatasetsProtocolPrivate, PaginatedResponse
 

src/llama_stack/providers/inline/eval/meta_reference/eval.py

@@ -8,8 +8,8 @@ from typing import Any
 
 from tqdm import tqdm
 
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.providers.utils.common.data_schema_validator import ColumnName
-from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack_api import (
     Agents,
     Benchmark,

src/llama_stack/providers/inline/files/localfs/files.py

@@ -13,11 +13,10 @@ from fastapi import Depends, File, Form, Response, UploadFile
 
 from llama_stack.core.datatypes import AccessRule
 from llama_stack.core.id_generation import generate_object_id
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.files.form_data import parse_expires_after
-from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
-from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
 from llama_stack_api import (
     ExpiresAfter,
     Files,
@@ -28,6 +27,7 @@ from llama_stack_api import (
     Order,
     ResourceNotFoundError,
 )
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
 
 from .config import LocalfsFilesImplConfig
 

src/llama_stack/providers/inline/vector_io/faiss/faiss.py

@@ -14,9 +14,8 @@ import faiss  # type: ignore[import-untyped]
 import numpy as np
 from numpy.typing import NDArray
 
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import kvstore_impl
-from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
 from llama_stack_api import (
@@ -32,6 +31,7 @@ from llama_stack_api import (
     VectorStoreNotFoundError,
     VectorStoresProtocolPrivate,
 )
+from llama_stack_api.internal.kvstore import KVStore
 
 from .config import FaissVectorIOConfig
 

src/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -14,9 +14,8 @@ import numpy as np
 import sqlite_vec  # type: ignore[import-untyped]
 from numpy.typing import NDArray
 
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import kvstore_impl
-from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import (
     RERANKER_TYPE_RRF,
@@ -35,6 +34,7 @@ from llama_stack_api import (
     VectorStoreNotFoundError,
     VectorStoresProtocolPrivate,
 )
+from llama_stack_api.internal.kvstore import KVStore
 
 logger = get_logger(name=__name__, category="vector_io")
 

src/llama_stack/providers/registry/agents.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 
-from llama_stack.providers.utils.kvstore import kvstore_dependencies
+from llama_stack.core.storage.kvstore import kvstore_dependencies
 from llama_stack_api import (
     Api,
     InlineProviderSpec,

src/llama_stack/providers/registry/files.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack.providers.utils.sqlstore.sqlstore import sql_store_pip_packages
+from llama_stack.core.storage.sqlstore.sqlstore import sql_store_pip_packages
 from llama_stack_api import Api, InlineProviderSpec, ProviderSpec, RemoteProviderSpec
 
 

src/llama_stack/providers/remote/datasetio/huggingface/huggingface.py

@@ -6,7 +6,7 @@
 from typing import Any
 from urllib.parse import parse_qs, urlparse
 
-from llama_stack.providers.utils.kvstore import kvstore_impl
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.providers.utils.pagination import paginate_records
 from llama_stack_api import Dataset, DatasetIO, DatasetsProtocolPrivate, PaginatedResponse
 

src/llama_stack/providers/remote/files/openai/files.py

@@ -10,10 +10,9 @@ from typing import Annotated, Any
 from fastapi import Depends, File, Form, Response, UploadFile
 
 from llama_stack.core.datatypes import AccessRule
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
 from llama_stack.providers.utils.files.form_data import parse_expires_after
-from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
-from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
 from llama_stack_api import (
     ExpiresAfter,
     Files,
@@ -24,6 +23,7 @@ from llama_stack_api import (
     Order,
     ResourceNotFoundError,
 )
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
 from openai import OpenAI
 
 from .config import OpenAIFilesImplConfig

src/llama_stack/providers/remote/files/s3/files.py

@@ -19,10 +19,9 @@ if TYPE_CHECKING:
 
 from llama_stack.core.datatypes import AccessRule
 from llama_stack.core.id_generation import generate_object_id
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
 from llama_stack.providers.utils.files.form_data import parse_expires_after
-from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
-from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
 from llama_stack_api import (
     ExpiresAfter,
     Files,
@@ -33,6 +32,7 @@ from llama_stack_api import (
     Order,
     ResourceNotFoundError,
 )
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
 
 from .config import S3FilesImplConfig
 

src/llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -11,10 +11,9 @@ from urllib.parse import urlparse
 import chromadb
 from numpy.typing import NDArray
 
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
-from llama_stack.providers.utils.kvstore import kvstore_impl
-from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
 from llama_stack_api import (
@@ -27,6 +26,7 @@ from llama_stack_api import (
     VectorStore,
     VectorStoresProtocolPrivate,
 )
+from llama_stack_api.internal.kvstore import KVStore
 
 from .config import ChromaVectorIOConfig as RemoteChromaVectorIOConfig
 

src/llama_stack/providers/remote/vector_io/milvus/milvus.py

@@ -11,10 +11,9 @@ from typing import Any
 from numpy.typing import NDArray
 from pymilvus import AnnSearchRequest, DataType, Function, FunctionType, MilvusClient, RRFRanker, WeightedRanker
 
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
-from llama_stack.providers.utils.kvstore import kvstore_impl
-from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import (
     RERANKER_TYPE_WEIGHTED,
@@ -34,6 +33,7 @@ from llama_stack_api import (
     VectorStoreNotFoundError,
     VectorStoresProtocolPrivate,
 )
+from llama_stack_api.internal.kvstore import KVStore
 
 from .config import MilvusVectorIOConfig as RemoteMilvusVectorIOConfig
 

src/llama_stack/providers/remote/vector_io/pgvector/pgvector.py

@@ -13,10 +13,9 @@ from psycopg2 import sql
 from psycopg2.extras import Json, execute_values
 from pydantic import BaseModel, TypeAdapter
 
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
-from llama_stack.providers.utils.kvstore import kvstore_impl
-from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
 from llama_stack.providers.utils.vector_io.vector_utils import WeightedInMemoryAggregator, sanitize_collection_name
@@ -31,6 +30,7 @@ from llama_stack_api import (
     VectorStoreNotFoundError,
     VectorStoresProtocolPrivate,
 )
+from llama_stack_api.internal.kvstore import KVStore
 
 from .config import PGVectorVectorIOConfig
 

src/llama_stack/providers/remote/vector_io/qdrant/qdrant.py

@@ -13,9 +13,9 @@ from numpy.typing import NDArray
 from qdrant_client import AsyncQdrantClient, models
 from qdrant_client.models import PointStruct
 
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
-from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
 from llama_stack_api import (

src/llama_stack/providers/remote/vector_io/weaviate/weaviate.py

@@ -13,9 +13,8 @@ from weaviate.classes.init import Auth
 from weaviate.classes.query import Filter, HybridFusion
 
 from llama_stack.core.request_headers import NeedsRequestProviderData
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import kvstore_impl
-from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import (
     RERANKER_TYPE_RRF,
@@ -35,6 +34,7 @@ from llama_stack_api import (
     VectorStoreNotFoundError,
     VectorStoresProtocolPrivate,
 )
+from llama_stack_api.internal.kvstore import KVStore
 
 from .config import WeaviateVectorIOConfig
 

src/llama_stack/providers/utils/inference/inference_store.py

@@ -10,6 +10,8 @@ from sqlalchemy.exc import IntegrityError
 
 from llama_stack.core.datatypes import AccessRule
 from llama_stack.core.storage.datatypes import InferenceStoreReference, StorageBackendType
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
 from llama_stack.log import get_logger
 from llama_stack_api import (
     ListOpenAIChatCompletionResponse,
@@ -18,10 +20,7 @@ from llama_stack_api import (
     OpenAIMessageParam,
     Order,
 )
-
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
-from ..sqlstore.api import ColumnDefinition, ColumnType
-from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from ..sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
 
 logger = get_logger(name=__name__, category="inference")
 

src/llama_stack/providers/utils/kvstore/sqlite/config.py

@@ -1,20 +0,0 @@
-# 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.
-
-from pydantic import BaseModel, Field
-
-from llama_stack_api import json_schema_type
-
-
-@json_schema_type
-class SqliteControlPlaneConfig(BaseModel):
-    db_path: str = Field(
-        description="File path for the sqlite database",
-    )
-    table_name: str = Field(
-        default="llamastack_control_plane",
-        description="Table into which all the keys will be placed",
-    )

src/llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -17,7 +17,6 @@ from pydantic import TypeAdapter
 
 from llama_stack.core.id_generation import generate_object_id
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.vector_store import (
     ChunkForDeletion,
     content_from_data_and_mime_type,
@@ -53,6 +52,7 @@ from llama_stack_api import (
     VectorStoreSearchResponse,
     VectorStoreSearchResponsePage,
 )
+from llama_stack_api.internal.kvstore import KVStore
 
 EMBEDDING_DIMENSION = 768
 

src/llama_stack/providers/utils/responses/responses_store.py

@@ -6,6 +6,8 @@
 
 from llama_stack.core.datatypes import AccessRule
 from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqlStoreReference
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
 from llama_stack.log import get_logger
 from llama_stack_api import (
     ListOpenAIResponseInputItem,
@@ -17,10 +19,7 @@ from llama_stack_api import (
     OpenAIResponseObjectWithInput,
     Order,
 )
-
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
-from ..sqlstore.api import ColumnDefinition, ColumnType
-from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from ..sqlstore.sqlstore import sqlstore_impl
 
 logger = get_logger(name=__name__, category="openai_responses")
 

src/llama_stack/providers/utils/sqlstore/api.py

@@ -1,140 +0,0 @@
-# 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.
-
-from collections.abc import Mapping, Sequence
-from enum import Enum
-from typing import Any, Literal, Protocol
-
-from pydantic import BaseModel
-
-from llama_stack_api import PaginatedResponse
-
-
-class ColumnType(Enum):
-    INTEGER = "INTEGER"
-    STRING = "STRING"
-    TEXT = "TEXT"
-    FLOAT = "FLOAT"
-    BOOLEAN = "BOOLEAN"
-    JSON = "JSON"
-    DATETIME = "DATETIME"
-
-
-class ColumnDefinition(BaseModel):
-    type: ColumnType
-    primary_key: bool = False
-    nullable: bool = True
-    default: Any = None
-
-
-class SqlStore(Protocol):
-    """
-    A protocol for a SQL store.
-    """
-
-    async def create_table(self, table: str, schema: Mapping[str, ColumnType | ColumnDefinition]) -> None:
-        """
-        Create a table.
-        """
-        pass
-
-    async def insert(self, table: str, data: Mapping[str, Any] | Sequence[Mapping[str, Any]]) -> None:
-        """
-        Insert a row or batch of rows into a table.
-        """
-        pass
-
-    async def upsert(
-        self,
-        table: str,
-        data: Mapping[str, Any],
-        conflict_columns: list[str],
-        update_columns: list[str] | None = None,
-    ) -> None:
-        """
-        Insert a row and update specified columns when conflicts occur.
-        """
-        pass
-
-    async def fetch_all(
-        self,
-        table: str,
-        where: Mapping[str, Any] | None = None,
-        where_sql: str | None = None,
-        limit: int | None = None,
-        order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
-        cursor: tuple[str, str] | None = None,
-    ) -> PaginatedResponse:
-        """
-        Fetch all rows from a table with optional cursor-based pagination.
-
-        :param table: The table name
-        :param where: Simple key-value WHERE conditions
-        :param where_sql: Raw SQL WHERE clause for complex queries
-        :param limit: Maximum number of records to return
-        :param order_by: List of (column, order) tuples for sorting
-        :param cursor: Tuple of (key_column, cursor_id) for pagination (None for first page)
-                      Requires order_by with exactly one column when used
-        :return: PaginatedResult with data and has_more flag
-
-        Note: Cursor pagination only supports single-column ordering for simplicity.
-        Multi-column ordering is allowed without cursor but will raise an error with cursor.
-        """
-        pass
-
-    async def fetch_one(
-        self,
-        table: str,
-        where: Mapping[str, Any] | None = None,
-        where_sql: str | None = None,
-        order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
-    ) -> dict[str, Any] | None:
-        """
-        Fetch one row from a table.
-        """
-        pass
-
-    async def update(
-        self,
-        table: str,
-        data: Mapping[str, Any],
-        where: Mapping[str, Any],
-    ) -> None:
-        """
-        Update a row in a table.
-        """
-        pass
-
-    async def delete(
-        self,
-        table: str,
-        where: Mapping[str, Any],
-    ) -> None:
-        """
-        Delete a row from a table.
-        """
-        pass
-
-    async def add_column_if_not_exists(
-        self,
-        table: str,
-        column_name: str,
-        column_type: ColumnType,
-        nullable: bool = True,
-    ) -> None:
-        """
-        Add a column to an existing table if the column doesn't already exist.
-
-        This is useful for table migrations when adding new functionality.
-        If the table doesn't exist, this method should do nothing.
-        If the column already exists, this method should do nothing.
-
-        :param table: Table name
-        :param column_name: Name of the column to add
-        :param column_type: Type of the column to add
-        :param nullable: Whether the column should be nullable (default: True)
-        """
-        pass

src/llama_stack_api/internal/__init__.py

@@ -3,3 +3,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+
+# Internal subpackage for shared interfaces that are not part of the public API.
+
+__all__: list[str] = []

src/llama_stack_api/internal/kvstore.py

@@ -9,6 +9,8 @@ from typing import Protocol
 
 
 class KVStore(Protocol):
+    """Protocol for simple key/value storage backends."""
+
     # TODO: make the value type bytes instead of str
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None: ...
 
@@ -19,3 +21,6 @@ class KVStore(Protocol):
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]: ...
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: ...
+
+
+__all__ = ["KVStore"]

src/llama_stack_api/internal/sqlstore.py

@@ -0,0 +1,79 @@
+# 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.
+
+from collections.abc import Mapping, Sequence
+from enum import Enum
+from typing import Any, Literal, Protocol
+
+from pydantic import BaseModel
+
+from llama_stack_api import PaginatedResponse
+
+
+class ColumnType(Enum):
+    INTEGER = "INTEGER"
+    STRING = "STRING"
+    TEXT = "TEXT"
+    FLOAT = "FLOAT"
+    BOOLEAN = "BOOLEAN"
+    JSON = "JSON"
+    DATETIME = "DATETIME"
+
+
+class ColumnDefinition(BaseModel):
+    type: ColumnType
+    primary_key: bool = False
+    nullable: bool = True
+    default: Any = None
+
+
+class SqlStore(Protocol):
+    """Protocol for common SQL-store functionality."""
+
+    async def create_table(self, table: str, schema: Mapping[str, ColumnType | ColumnDefinition]) -> None: ...
+
+    async def insert(self, table: str, data: Mapping[str, Any] | Sequence[Mapping[str, Any]]) -> None: ...
+
+    async def upsert(
+        self,
+        table: str,
+        data: Mapping[str, Any],
+        conflict_columns: list[str],
+        update_columns: list[str] | None = None,
+    ) -> None: ...
+
+    async def fetch_all(
+        self,
+        table: str,
+        where: Mapping[str, Any] | None = None,
+        where_sql: str | None = None,
+        limit: int | None = None,
+        order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
+        cursor: tuple[str, str] | None = None,
+    ) -> PaginatedResponse: ...
+
+    async def fetch_one(
+        self,
+        table: str,
+        where: Mapping[str, Any] | None = None,
+        where_sql: str | None = None,
+        order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
+    ) -> dict[str, Any] | None: ...
+
+    async def update(self, table: str, data: Mapping[str, Any], where: Mapping[str, Any]) -> None: ...
+
+    async def delete(self, table: str, where: Mapping[str, Any]) -> None: ...
+
+    async def add_column_if_not_exists(
+        self,
+        table: str,
+        column_name: str,
+        column_type: ColumnType,
+        nullable: bool = True,
+    ) -> None: ...
+
+
+__all__ = ["ColumnDefinition", "ColumnType", "SqlStore"]

tests/integration/files/test_files.py

@@ -175,7 +175,7 @@ def test_expires_after_requests(openai_client):
 
 
 @pytest.mark.xfail(message="User isolation broken for current providers, must be fixed.")
-@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
+@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
 def test_files_authentication_isolation(mock_get_authenticated_user, llama_stack_client):
     """Test that users can only access their own files."""
     from llama_stack_client import NotFoundError
@@ -275,7 +275,7 @@ def test_files_authentication_isolation(mock_get_authenticated_user, llama_stack
         raise e
 
 
-@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
+@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
 def test_files_authentication_shared_attributes(
     mock_get_authenticated_user, llama_stack_client, provider_type_is_openai
 ):
@@ -335,7 +335,7 @@ def test_files_authentication_shared_attributes(
         raise e
 
 
-@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
+@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
 def test_files_authentication_anonymous_access(
     mock_get_authenticated_user, llama_stack_client, provider_type_is_openai
 ):

tests/integration/providers/utils/sqlstore/test_authorized_sqlstore.py

@@ -13,14 +13,14 @@ import pytest
 from llama_stack.core.access_control.access_control import default_policy
 from llama_stack.core.datatypes import User
 from llama_stack.core.storage.datatypes import SqlStoreReference
-from llama_stack.providers.utils.sqlstore.api import ColumnType
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import (
-from llama_stack.providers.utils.sqlstore.sqlstore import (
     PostgresSqlStoreConfig,
     SqliteSqlStoreConfig,
     register_sqlstore_backends,
     sqlstore_impl,
 )
+from llama_stack_api.internal.sqlstore import ColumnType
 
 
 def get_postgres_config():
@@ -96,7 +96,7 @@ async def cleanup_records(sql_store, table_name, record_ids):
 
 
 @pytest.mark.parametrize("backend_config", BACKEND_CONFIGS)
-@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
+@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
 async def test_authorized_store_attributes(mock_get_authenticated_user, authorized_store, request):
     """Test that JSON column comparisons work correctly for both PostgreSQL and SQLite"""
     backend_name = request.node.callspec.id
@@ -190,7 +190,7 @@ async def test_authorized_store_attributes(mock_get_authenticated_user, authoriz
 
 
 @pytest.mark.parametrize("backend_config", BACKEND_CONFIGS)
-@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
+@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
 async def test_user_ownership_policy(mock_get_authenticated_user, authorized_store, request):
     """Test that 'user is owner' policies work correctly with record ownership"""
     from llama_stack.core.access_control.datatypes import AccessRule, Action, Scope

tests/unit/conversations/test_conversations.py

@@ -23,7 +23,7 @@ from llama_stack.core.storage.datatypes import (
     SqlStoreReference,
     StorageConfig,
 )
-from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack_api import OpenAIResponseInputMessageContentText, OpenAIResponseMessage
 
 
@@ -38,6 +38,9 @@ async def service():
             },
             stores=ServerStoresConfig(
                 conversations=SqlStoreReference(backend="sql_test", table_name="openai_conversations"),
+                metadata=None,
+                inference=None,
+                prompts=None,
             ),
         )
         register_sqlstore_backends({"sql_test": storage.backends["sql_test"]})
@@ -142,6 +145,9 @@ async def test_policy_configuration():
             },
             stores=ServerStoresConfig(
                 conversations=SqlStoreReference(backend="sql_test", table_name="openai_conversations"),
+                metadata=None,
+                inference=None,
+                prompts=None,
             ),
         )
         register_sqlstore_backends({"sql_test": storage.backends["sql_test"]})

tests/unit/core/test_stack_validation.py

@@ -10,8 +10,9 @@ from unittest.mock import AsyncMock
 
 import pytest
 
-from llama_stack.core.datatypes import QualifiedModel, SafetyConfig, StackRunConfig, StorageConfig, VectorStoresConfig
+from llama_stack.core.datatypes import QualifiedModel, SafetyConfig, StackRunConfig, VectorStoresConfig
 from llama_stack.core.stack import validate_safety_config, validate_vector_stores_config
+from llama_stack.core.storage.datatypes import ServerStoresConfig, StorageConfig
 from llama_stack_api import Api, ListModelsResponse, ListShieldsResponse, Model, ModelType, Shield
 
 
@@ -21,7 +22,15 @@ class TestVectorStoresValidation:
         run_config = StackRunConfig(
             image_name="test",
             providers={},
-            storage=StorageConfig(backends={}, stores={}),
+            storage=StorageConfig(
+                backends={},
+                stores=ServerStoresConfig(
+                    metadata=None,
+                    inference=None,
+                    conversations=None,
+                    prompts=None,
+                ),
+            ),
             vector_stores=VectorStoresConfig(
                 default_provider_id="faiss",
                 default_embedding_model=QualifiedModel(
@@ -41,7 +50,15 @@ class TestVectorStoresValidation:
         run_config = StackRunConfig(
             image_name="test",
             providers={},
-            storage=StorageConfig(backends={}, stores={}),
+            storage=StorageConfig(
+                backends={},
+                stores=ServerStoresConfig(
+                    metadata=None,
+                    inference=None,
+                    conversations=None,
+                    prompts=None,
+                ),
+            ),
             vector_stores=VectorStoresConfig(
                 default_provider_id="faiss",
                 default_embedding_model=QualifiedModel(

tests/unit/files/test_files.py

@@ -9,11 +9,11 @@ import pytest
 
 from llama_stack.core.access_control.access_control import default_policy
 from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
+from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack.providers.inline.files.localfs import (
     LocalfsFilesImpl,
     LocalfsFilesImplConfig,
 )
-from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack_api import OpenAIFilePurpose, Order, ResourceNotFoundError
 
 

tests/unit/fixtures.py

@@ -6,9 +6,9 @@
 
 import pytest
 
+from llama_stack.core.storage.kvstore.config import SqliteKVStoreConfig
+from llama_stack.core.storage.kvstore.sqlite import SqliteKVStoreImpl
 from llama_stack.core.store.registry import CachedDiskDistributionRegistry, DiskDistributionRegistry
-from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
-from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
 
 
 @pytest.fixture(scope="function")

tests/unit/prompts/prompts/conftest.py

@@ -18,7 +18,7 @@ from llama_stack.core.storage.datatypes import (
     SqlStoreReference,
     StorageConfig,
 )
-from llama_stack.providers.utils.kvstore import register_kvstore_backends
+from llama_stack.core.storage.kvstore import register_kvstore_backends
 
 
 @pytest.fixture

tests/unit/providers/agents/meta_reference/test_openai_responses.py

@@ -17,6 +17,7 @@ from openai.types.chat.chat_completion_chunk import (
 
 from llama_stack.core.access_control.access_control import default_policy
 from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
+from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack.providers.inline.agents.meta_reference.responses.openai_responses import (
     OpenAIResponsesImpl,
 )
@@ -24,7 +25,6 @@ from llama_stack.providers.utils.responses.responses_store import (
     ResponsesStore,
     _OpenAIResponseObjectWithInputAndMessages,
 )
-from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack_api.agents import Order
 from llama_stack_api.inference import (
     OpenAIAssistantMessageParam,

tests/unit/providers/batches/conftest.py

@@ -13,9 +13,9 @@ from unittest.mock import AsyncMock
 import pytest
 
 from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
+from llama_stack.core.storage.kvstore import kvstore_impl, register_kvstore_backends
 from llama_stack.providers.inline.batches.reference.batches import ReferenceBatchesImpl
 from llama_stack.providers.inline.batches.reference.config import ReferenceBatchesImplConfig
-from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends
 
 
 @pytest.fixture

tests/unit/providers/files/conftest.py

@@ -9,8 +9,8 @@ import pytest
 from moto import mock_aws
 
 from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
+from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack.providers.remote.files.s3 import S3FilesImplConfig, get_adapter_impl
-from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
 
 
 class MockUploadFile:

tests/unit/providers/files/test_s3_files_auth.py

@@ -18,11 +18,11 @@ async def test_listing_hides_other_users_file(s3_provider, sample_text_file):
     user_a = User("user-a", {"roles": ["team-a"]})
     user_b = User("user-b", {"roles": ["team-b"]})
 
-    with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
+    with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
         mock_get_user.return_value = user_a
         uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
 
-    with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
+    with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
         mock_get_user.return_value = user_b
         listed = await s3_provider.openai_list_files()
         assert all(f.id != uploaded.id for f in listed.data)
@@ -41,11 +41,11 @@ async def test_cannot_access_other_user_file(s3_provider, sample_text_file, op):
     user_a = User("user-a", {"roles": ["team-a"]})
     user_b = User("user-b", {"roles": ["team-b"]})
 
-    with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
+    with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
         mock_get_user.return_value = user_a
         uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
 
-    with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
+    with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
         mock_get_user.return_value = user_b
         with pytest.raises(ResourceNotFoundError):
             await op(s3_provider, uploaded.id)
@@ -56,11 +56,11 @@ async def test_shared_role_allows_listing(s3_provider, sample_text_file):
     user_a = User("user-a", {"roles": ["shared-role"]})
     user_b = User("user-b", {"roles": ["shared-role"]})
 
-    with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
+    with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
         mock_get_user.return_value = user_a
         uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
 
-    with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
+    with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
         mock_get_user.return_value = user_b
         listed = await s3_provider.openai_list_files()
         assert any(f.id == uploaded.id for f in listed.data)
@@ -79,10 +79,10 @@ async def test_shared_role_allows_access(s3_provider, sample_text_file, op):
     user_x = User("user-x", {"roles": ["shared-role"]})
     user_y = User("user-y", {"roles": ["shared-role"]})
 
-    with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
+    with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
         mock_get_user.return_value = user_x
         uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
 
-    with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
+    with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
         mock_get_user.return_value = user_y
         await op(s3_provider, uploaded.id)

tests/unit/providers/vector_io/conftest.py

@@ -11,13 +11,13 @@ import numpy as np
 import pytest
 
 from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
+from llama_stack.core.storage.kvstore import register_kvstore_backends
 from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
 from llama_stack.providers.inline.vector_io.faiss.faiss import FaissIndex, FaissVectorIOAdapter
 from llama_stack.providers.inline.vector_io.sqlite_vec import SQLiteVectorIOConfig
 from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import SQLiteVecIndex, SQLiteVecVectorIOAdapter
 from llama_stack.providers.remote.vector_io.pgvector.config import PGVectorVectorIOConfig
 from llama_stack.providers.remote.vector_io.pgvector.pgvector import PGVectorIndex, PGVectorVectorIOAdapter
-from llama_stack.providers.utils.kvstore import register_kvstore_backends
 from llama_stack_api import Chunk, ChunkMetadata, QueryChunksResponse, VectorStore
 
 EMBEDDING_DIMENSION = 768
@@ -279,7 +279,7 @@ async def pgvector_vec_adapter(unique_kvstore_config, mock_inference_api, embedd
         ) as mock_check_version:
             mock_check_version.return_value = "0.5.1"
 
-            with patch("llama_stack.providers.utils.kvstore.kvstore_impl") as mock_kvstore_impl:
+            with patch("llama_stack.core.storage.kvstore.kvstore_impl") as mock_kvstore_impl:
                 mock_kvstore = AsyncMock()
                 mock_kvstore_impl.return_value = mock_kvstore
 

tests/unit/registry/test_registry.py

@@ -9,12 +9,12 @@ import pytest
 
 from llama_stack.core.datatypes import VectorStoreWithOwner
 from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
+from llama_stack.core.storage.kvstore import kvstore_impl, register_kvstore_backends
 from llama_stack.core.store.registry import (
     KEY_FORMAT,
     CachedDiskDistributionRegistry,
     DiskDistributionRegistry,
 )
-from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends
 from llama_stack_api import Model, VectorStore
 
 

tests/unit/server/test_quota.py

@@ -15,7 +15,7 @@ from starlette.middleware.base import BaseHTTPMiddleware
 from llama_stack.core.datatypes import QuotaConfig, QuotaPeriod
 from llama_stack.core.server.quota import QuotaMiddleware
 from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
-from llama_stack.providers.utils.kvstore import register_kvstore_backends
+from llama_stack.core.storage.kvstore import register_kvstore_backends
 
 
 @pytest.fixture

tests/unit/server/test_resolver.py

@@ -24,8 +24,8 @@ from llama_stack.core.storage.datatypes import (
     SqlStoreReference,
     StorageConfig,
 )
-from llama_stack.providers.utils.kvstore import register_kvstore_backends
+from llama_stack.core.storage.kvstore import register_kvstore_backends
-from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack_api import Inference, InlineProviderSpec, ProviderSpec
 
 

tests/unit/utils/inference/test_inference_store.py

@@ -9,8 +9,8 @@ import time
 import pytest
 
 from llama_stack.core.storage.datatypes import InferenceStoreReference, SqliteSqlStoreConfig
+from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack.providers.utils.inference.inference_store import InferenceStore
-from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack_api import (
     OpenAIAssistantMessageParam,
     OpenAIChatCompletion,

tests/unit/utils/kvstore/test_sqlite_memory.py

@@ -5,8 +5,8 @@
 # the root directory of this source tree.
 
 
-from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+from llama_stack.core.storage.kvstore.config import SqliteKVStoreConfig
-from llama_stack.providers.utils.kvstore.sqlite.sqlite import SqliteKVStoreImpl
+from llama_stack.core.storage.kvstore.sqlite.sqlite import SqliteKVStoreImpl
 
 
 async def test_memory_kvstore_persistence_behavior():

tests/unit/utils/responses/test_responses_store.py

@@ -11,8 +11,8 @@ from uuid import uuid4
 import pytest
 
 from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
+from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack.providers.utils.responses.responses_store import ResponsesStore
-from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack_api import OpenAIMessageParam, OpenAIResponseInput, OpenAIResponseObject, OpenAIUserMessageParam, Order
 
 

tests/unit/utils/sqlstore/test_sqlstore.py

@@ -9,9 +9,9 @@ from tempfile import TemporaryDirectory
 
 import pytest
 
-from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
+from llama_stack.core.storage.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
-from llama_stack.providers.utils.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
+from llama_stack.core.storage.sqlstore.sqlstore import SqliteSqlStoreConfig
-from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
 
 
 async def test_sqlite_sqlstore():

tests/unit/utils/test_authorized_sqlstore.py

@@ -10,13 +10,13 @@ from unittest.mock import patch
 from llama_stack.core.access_control.access_control import default_policy, is_action_allowed
 from llama_stack.core.access_control.datatypes import Action
 from llama_stack.core.datatypes import User
-from llama_stack.providers.utils.sqlstore.api import ColumnType
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore, SqlRecord
-from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore, SqlRecord
+from llama_stack.core.storage.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
-from llama_stack.providers.utils.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
+from llama_stack.core.storage.sqlstore.sqlstore import SqliteSqlStoreConfig
-from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
+from llama_stack_api.internal.sqlstore import ColumnType
 
 
-@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
+@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
 async def test_authorized_fetch_with_where_sql_access_control(mock_get_authenticated_user):
     """Test that fetch_all works correctly with where_sql for access control"""
     with TemporaryDirectory() as tmp_dir:
@@ -78,7 +78,7 @@ async def test_authorized_fetch_with_where_sql_access_control(mock_get_authentic
         assert row["title"] == "User Document"
 
 
-@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
+@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
 async def test_sql_policy_consistency(mock_get_authenticated_user):
     """Test that SQL WHERE clause logic exactly matches is_action_allowed policy logic"""
     with TemporaryDirectory() as tmp_dir:
@@ -164,7 +164,7 @@ async def test_sql_policy_consistency(mock_get_authenticated_user):
             )
 
 
-@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
+@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
 async def test_authorized_store_user_attribute_capture(mock_get_authenticated_user):
     """Test that user attributes are properly captured during insert"""
     with TemporaryDirectory() as tmp_dir: