From ea9664874d7475f9e83431129ccf22d984ec1851 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Thu, 16 Oct 2025 16:24:31 -0700 Subject: [PATCH] feat(stores)!: use backend storage references instead of configs --- llama_stack/cli/stack/_build.py | 7 +- .../core/conversations/conversations.py | 26 +- llama_stack/core/datatypes.py | 40 +-- llama_stack/core/prompts/prompts.py | 12 +- llama_stack/core/routers/__init__.py | 13 +- llama_stack/core/stack.py | 23 +- llama_stack/core/storage/__init__.py | 5 + llama_stack/core/storage/datatypes.py | 283 ++++++++++++++++++ llama_stack/core/store/registry.py | 18 +- llama_stack/distributions/ci-tests/run.yaml | 47 ++- .../distributions/dell/run-with-safety.yaml | 29 +- llama_stack/distributions/dell/run.yaml | 29 +- .../meta-reference-gpu/run-with-safety.yaml | 29 +- .../distributions/meta-reference-gpu/run.yaml | 29 +- .../distributions/nvidia/run-with-safety.yaml | 26 +- llama_stack/distributions/nvidia/run.yaml | 23 +- .../distributions/open-benchmark/run.yaml | 35 +-- .../distributions/postgres-demo/run.yaml | 13 +- .../distributions/starter-gpu/run.yaml | 47 ++- llama_stack/distributions/starter/run.yaml | 47 ++- llama_stack/distributions/watsonx/run.yaml | 32 +- .../inline/agents/meta_reference/agents.py | 4 +- .../inline/agents/meta_reference/config.py | 30 +- .../inline/batches/reference/config.py | 11 +- .../inline/datasetio/localfs/config.py | 14 +- .../inline/eval/meta_reference/config.py | 14 +- .../providers/inline/files/localfs/config.py | 11 +- .../inline/vector_io/chroma/config.py | 11 +- .../inline/vector_io/faiss/config.py | 14 +- .../inline/vector_io/milvus/config.py | 14 +- .../inline/vector_io/qdrant/config.py | 13 +- .../inline/vector_io/sqlite_vec/config.py | 14 +- .../remote/datasetio/huggingface/config.py | 14 +- .../providers/remote/files/s3/config.py | 11 +- .../providers/remote/files/s3/files.py | 2 +- .../remote/vector_io/chroma/config.py | 11 +- .../remote/vector_io/milvus/config.py | 11 +- .../remote/vector_io/pgvector/config.py | 16 +- .../remote/vector_io/qdrant/config.py | 14 +- .../remote/vector_io/weaviate/config.py | 16 +- .../utils/inference/inference_store.py | 35 ++- llama_stack/providers/utils/kvstore/config.py | 151 +--------- .../providers/utils/kvstore/kvstore.py | 54 +++- .../utils/sqlstore/authorized_sqlstore.py | 18 +- .../providers/utils/sqlstore/sqlstore.py | 128 +++----- .../test_persistence_integration.py | 63 ++++ tests/unit/core/test_persistence_config.py | 82 +++++ 47 files changed, 893 insertions(+), 696 deletions(-) create mode 100644 llama_stack/core/storage/__init__.py create mode 100644 llama_stack/core/storage/datatypes.py create mode 100644 tests/integration/test_persistence_integration.py create mode 100644 tests/unit/core/test_persistence_config.py diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py index 471d5cb66..1b60e9bf5 100644 --- a/llama_stack/cli/stack/_build.py +++ b/llama_stack/cli/stack/_build.py @@ -295,12 +295,7 @@ def _generate_run_config( if build_config.external_providers_dir else EXTERNAL_PROVIDERS_DIR, ) - if not run_config.inference_store: - run_config.inference_store = SqliteSqlStoreConfig( - **SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=(DISTRIBS_BASE_DIR / image_name).as_posix(), db_name="inference_store.db" - ) - ) + # Persistence config defaults are handled by PersistenceConfig model validators # build providers dict provider_registry = get_provider_registry(build_config) for api in apis: diff --git a/llama_stack/core/conversations/conversations.py b/llama_stack/core/conversations/conversations.py index d2537c7ee..eb9a3c399 100644 --- a/llama_stack/core/conversations/conversations.py +++ b/llama_stack/core/conversations/conversations.py @@ -4,7 +4,6 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import os import secrets import time from typing import Any @@ -21,16 +20,11 @@ from llama_stack.apis.conversations.conversations import ( Conversations, Metadata, ) -from llama_stack.core.datatypes import AccessRule -from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR +from llama_stack.core.datatypes import AccessRule, StackRunConfig 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 ( - SqliteSqlStoreConfig, - SqlStoreConfig, - sqlstore_impl, -) +from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl logger = get_logger(name=__name__, category="openai_conversations") @@ -38,13 +32,11 @@ logger = get_logger(name=__name__, category="openai_conversations") class ConversationServiceConfig(BaseModel): """Configuration for the built-in conversation service. - :param conversations_store: SQL store configuration for conversations (defaults to SQLite) + :param run_config: Stack run configuration for resolving persistence :param policy: Access control rules """ - conversations_store: SqlStoreConfig = SqliteSqlStoreConfig( - db_path=(DISTRIBS_BASE_DIR / "conversations.db").as_posix() - ) + run_config: StackRunConfig policy: list[AccessRule] = [] @@ -63,14 +55,16 @@ class ConversationServiceImpl(Conversations): self.deps = deps self.policy = config.policy - base_sql_store = sqlstore_impl(config.conversations_store) + # Use conversations store reference from storage config + conversations_ref = config.run_config.storage.conversations + if not conversations_ref: + raise ValueError("storage.conversations must be configured in run config") + + base_sql_store = sqlstore_impl(conversations_ref) self.sql_store = AuthorizedSqlStore(base_sql_store, self.policy) async def initialize(self) -> None: """Initialize the store and create tables.""" - if isinstance(self.config.conversations_store, SqliteSqlStoreConfig): - os.makedirs(os.path.dirname(self.config.conversations_store.db_path), exist_ok=True) - await self.sql_store.create_table( "openai_conversations", { diff --git a/llama_stack/core/datatypes.py b/llama_stack/core/datatypes.py index 94222d49e..766e5b550 100644 --- a/llama_stack/core/datatypes.py +++ b/llama_stack/core/datatypes.py @@ -26,9 +26,8 @@ from llama_stack.apis.tools import ToolGroup, ToolGroupInput, ToolRuntime from llama_stack.apis.vector_dbs import VectorDB, VectorDBInput from llama_stack.apis.vector_io import VectorIO from llama_stack.core.access_control.datatypes import AccessRule +from llama_stack.core.storage.datatypes import KVStoreReference, StorageConfig from llama_stack.providers.datatypes import Api, ProviderSpec -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig -from llama_stack.providers.utils.sqlstore.sqlstore import SqlStoreConfig LLAMA_STACK_BUILD_CONFIG_VERSION = 2 LLAMA_STACK_RUN_CONFIG_VERSION = 2 @@ -356,7 +355,7 @@ class QuotaPeriod(StrEnum): class QuotaConfig(BaseModel): - kvstore: SqliteKVStoreConfig = Field(description="Config for KV store backend (SQLite only for now)") + kvstore: KVStoreReference = Field(description="Config for KV store backend (SQLite only for now)") anonymous_max_requests: int = Field(default=100, description="Max requests for unauthenticated clients per period") authenticated_max_requests: int = Field( default=1000, description="Max requests for authenticated clients per period" @@ -438,18 +437,6 @@ class ServerConfig(BaseModel): ) -class InferenceStoreConfig(BaseModel): - sql_store_config: SqlStoreConfig - max_write_queue_size: int = Field(default=10000, description="Max queued writes for inference store") - num_writers: int = Field(default=4, description="Number of concurrent background writers") - - -class ResponsesStoreConfig(BaseModel): - sql_store_config: SqlStoreConfig - max_write_queue_size: int = Field(default=10000, description="Max queued writes for responses store") - num_writers: int = Field(default=4, description="Number of concurrent background writers") - - class StackRunConfig(BaseModel): version: int = LLAMA_STACK_RUN_CONFIG_VERSION @@ -476,26 +463,11 @@ One or more providers to use for each API. The same provider_type (e.g., meta-re can be instantiated multiple times (with different configs) if necessary. """, ) - metadata_store: KVStoreConfig | None = Field( - default=None, + storage: StorageConfig = Field( description=""" -Configuration for the persistence store used by the distribution registry. If not specified, -a default SQLite store will be used.""", - ) - - inference_store: InferenceStoreConfig | SqlStoreConfig | None = Field( - default=None, - description=""" -Configuration for the persistence store used by the inference API. Can be either a -InferenceStoreConfig (with queue tuning parameters) or a SqlStoreConfig (deprecated). -If not specified, a default SQLite store will be used.""", - ) - - conversations_store: SqlStoreConfig | None = Field( - default=None, - description=""" -Configuration for the persistence store used by the conversations API. -If not specified, a default SQLite store will be used.""", +Storage backend configurations. Each backend is named, and can be referenced by various components +throughout the Stack (both by its core as well as providers). +""", ) # registry of "resources" in the distribution diff --git a/llama_stack/core/prompts/prompts.py b/llama_stack/core/prompts/prompts.py index 26e8f5cef..a819a2c57 100644 --- a/llama_stack/core/prompts/prompts.py +++ b/llama_stack/core/prompts/prompts.py @@ -11,9 +11,8 @@ from pydantic import BaseModel from llama_stack.apis.prompts import ListPromptsResponse, Prompt, Prompts from llama_stack.core.datatypes import StackRunConfig -from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl -from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig class PromptServiceConfig(BaseModel): @@ -41,10 +40,13 @@ class PromptServiceImpl(Prompts): self.kvstore: KVStore async def initialize(self) -> None: - kvstore_config = SqliteKVStoreConfig( - db_path=(DISTRIBS_BASE_DIR / self.config.run_config.image_name / "prompts.db").as_posix() + # Use metadata store backend with prompts-specific namespace + metadata_ref = self.config.run_config.storage.metadata + prompts_ref = KVStoreReference( + namespace="prompts", + backend=metadata_ref.backend if metadata_ref else None, ) - self.kvstore = await kvstore_impl(kvstore_config) + self.kvstore = await kvstore_impl(prompts_ref) def _get_default_key(self, prompt_id: str) -> str: """Get the KVStore key that stores the default version number.""" diff --git a/llama_stack/core/routers/__init__.py b/llama_stack/core/routers/__init__.py index 4463d2460..a0ab8fc3b 100644 --- a/llama_stack/core/routers/__init__.py +++ b/llama_stack/core/routers/__init__.py @@ -6,7 +6,10 @@ from typing import Any -from llama_stack.core.datatypes import AccessRule, RoutedProtocol +from llama_stack.core.datatypes import ( + AccessRule, + RoutedProtocol, +) from llama_stack.core.stack import StackRunConfig from llama_stack.core.store import DistributionRegistry from llama_stack.providers.datatypes import Api, RoutingTable @@ -76,9 +79,13 @@ async def get_auto_router_impl( api_to_dep_impl[dep_name] = deps[dep_api] # TODO: move pass configs to routers instead - if api == Api.inference and run_config.inference_store: + if api == Api.inference: + inference_ref = run_config.storage.inference + if not inference_ref: + raise ValueError("storage.inference must be configured in run config") + inference_store = InferenceStore( - config=run_config.inference_store, + reference=inference_ref, policy=policy, ) await inference_store.initialize() diff --git a/llama_stack/core/stack.py b/llama_stack/core/stack.py index 733b55262..6a68dec7a 100644 --- a/llama_stack/core/stack.py +++ b/llama_stack/core/stack.py @@ -42,6 +42,7 @@ from llama_stack.core.prompts.prompts import PromptServiceConfig, PromptServiceI from llama_stack.core.providers import ProviderImpl, ProviderImplConfig from llama_stack.core.resolver import ProviderRegistry, resolve_impls from llama_stack.core.routing_tables.common import CommonRoutingTableImpl +from llama_stack.core.storage.datatypes import StorageBackendConfig from llama_stack.core.store.registry import create_dist_registry from llama_stack.core.utils.dynamic import instantiate_class_type from llama_stack.log import get_logger @@ -329,6 +330,25 @@ def add_internal_implementations(impls: dict[Api, Any], run_config: StackRunConf impls[Api.conversations] = conversations_impl +def _initialize_storage(run_config: StackRunConfig): + kv_backends: dict[str, StorageBackendConfig] = {} + sql_backends: dict[str, StorageBackendConfig] = {} + for backend_name, backend_config in run_config.storage.backends.items(): + type = backend_config.type.value + if type.startswith("kv_"): + kv_backends[backend_name] = backend_config + elif type.startswith("sql_"): + sql_backends[backend_name] = backend_config + else: + raise ValueError(f"Unknown storage backend type: {type}") + + from llama_stack.providers.utils.kvstore.kvstore import register_kvstore_backends + from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends + + register_kvstore_backends(kv_backends) + register_sqlstore_backends(sql_backends) + + class Stack: def __init__(self, run_config: StackRunConfig, provider_registry: ProviderRegistry | None = None): self.run_config = run_config @@ -347,7 +367,8 @@ class Stack: TEST_RECORDING_CONTEXT.__enter__() logger.info(f"API recording enabled: mode={os.environ.get('LLAMA_STACK_TEST_INFERENCE_MODE')}") - dist_registry, _ = await create_dist_registry(self.run_config.metadata_store, self.run_config.image_name) + _initialize_storage(self.run_config) + dist_registry, _ = await create_dist_registry(self.run_config.storage, self.run_config.image_name) policy = self.run_config.server.auth.access_policy if self.run_config.server.auth else [] internal_impls = {} diff --git a/llama_stack/core/storage/__init__.py b/llama_stack/core/storage/__init__.py new file mode 100644 index 000000000..756f351d8 --- /dev/null +++ b/llama_stack/core/storage/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/llama_stack/core/storage/datatypes.py b/llama_stack/core/storage/datatypes.py new file mode 100644 index 000000000..10046c248 --- /dev/null +++ b/llama_stack/core/storage/datatypes.py @@ -0,0 +1,283 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import re +from abc import abstractmethod +from enum import StrEnum +from pathlib import Path +from typing import Annotated, Literal + +from pydantic import BaseModel, Field, field_validator + + +class StorageBackendType(StrEnum): + KV_REDIS = "kv_redis" + KV_SQLITE = "kv_sqlite" + KV_POSTGRES = "kv_postgres" + KV_MONGODB = "kv_mongodb" + SQL_SQLITE = "sql_sqlite" + SQL_POSTGRES = "sql_postgres" + + +class CommonConfig(BaseModel): + namespace: str | None = Field( + default=None, + description="All keys will be prefixed with this namespace", + ) + default: bool = Field( + default=False, + description="Mark this KV store as the default choice when a reference omits the backend name", + ) + + +class RedisKVStoreConfig(CommonConfig): + type: Literal[StorageBackendType.KV_REDIS] = StorageBackendType.KV_REDIS + host: str = "localhost" + port: int = 6379 + + @property + def url(self) -> str: + return f"redis://{self.host}:{self.port}" + + @classmethod + def pip_packages(cls) -> list[str]: + return ["redis"] + + @classmethod + def sample_run_config(cls): + return { + "type": StorageBackendType.KV_REDIS.value, + "host": "${env.REDIS_HOST:=localhost}", + "port": "${env.REDIS_PORT:=6379}", + } + + +class SqliteKVStoreConfig(CommonConfig): + type: Literal[StorageBackendType.KV_SQLITE] = StorageBackendType.KV_SQLITE + db_path: str = Field( + description="File path for the sqlite database", + ) + + @classmethod + def pip_packages(cls) -> list[str]: + return ["aiosqlite"] + + @classmethod + def sample_run_config(cls, __distro_dir__: str, db_name: str = "kvstore.db"): + return { + "type": StorageBackendType.KV_SQLITE.value, + "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + db_name, + } + + +class PostgresKVStoreConfig(CommonConfig): + type: Literal[StorageBackendType.KV_POSTGRES] = StorageBackendType.KV_POSTGRES + host: str = "localhost" + port: int = 5432 + db: str = "llamastack" + user: str + password: str | None = None + ssl_mode: str | None = None + ca_cert_path: str | None = None + table_name: str = "llamastack_kvstore" + + @classmethod + def sample_run_config(cls, table_name: str = "llamastack_kvstore", **kwargs): + return { + "type": StorageBackendType.KV_POSTGRES.value, + "host": "${env.POSTGRES_HOST:=localhost}", + "port": "${env.POSTGRES_PORT:=5432}", + "db": "${env.POSTGRES_DB:=llamastack}", + "user": "${env.POSTGRES_USER:=llamastack}", + "password": "${env.POSTGRES_PASSWORD:=llamastack}", + "table_name": "${env.POSTGRES_TABLE_NAME:=" + table_name + "}", + } + + @classmethod + @field_validator("table_name") + def validate_table_name(cls, v: str) -> str: + # PostgreSQL identifiers rules: + # - Must start with a letter or underscore + # - Can contain letters, numbers, and underscores + # - Maximum length is 63 bytes + pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*$" + if not re.match(pattern, v): + raise ValueError( + "Invalid table name. Must start with letter or underscore and contain only letters, numbers, and underscores" + ) + if len(v) > 63: + raise ValueError("Table name must be less than 63 characters") + return v + + @classmethod + def pip_packages(cls) -> list[str]: + return ["psycopg2-binary"] + + +class MongoDBKVStoreConfig(CommonConfig): + type: Literal[StorageBackendType.KV_MONGODB] = StorageBackendType.KV_MONGODB + host: str = "localhost" + port: int = 27017 + db: str = "llamastack" + user: str | None = None + password: str | None = None + collection_name: str = "llamastack_kvstore" + + @classmethod + def pip_packages(cls) -> list[str]: + return ["pymongo"] + + @classmethod + def sample_run_config(cls, collection_name: str = "llamastack_kvstore"): + return { + "type": StorageBackendType.KV_MONGODB.value, + "host": "${env.MONGODB_HOST:=localhost}", + "port": "${env.MONGODB_PORT:=5432}", + "db": "${env.MONGODB_DB}", + "user": "${env.MONGODB_USER}", + "password": "${env.MONGODB_PASSWORD}", + "collection_name": "${env.MONGODB_COLLECTION_NAME:=" + collection_name + "}", + } + + +class CommonSqlStoreConfig(BaseModel): + default: bool = Field( + default=False, + description="Mark this SQL store as the default choice when a reference omits the backend name", + ) + + +class SqlAlchemySqlStoreConfig(BaseModel): + @property + @abstractmethod + def engine_str(self) -> str: ... + + # TODO: move this when we have a better way to specify dependencies with internal APIs + @classmethod + def pip_packages(cls) -> list[str]: + return ["sqlalchemy[asyncio]"] + + +class SqliteSqlStoreConfig(SqlAlchemySqlStoreConfig, CommonSqlStoreConfig): + type: Literal[StorageBackendType.SQL_SQLITE] = StorageBackendType.SQL_SQLITE + db_path: str = Field( + description="Database path, e.g. ~/.llama/distributions/ollama/sqlstore.db", + ) + + @property + def engine_str(self) -> str: + return "sqlite+aiosqlite:///" + Path(self.db_path).expanduser().as_posix() + + @classmethod + def sample_run_config(cls, __distro_dir__: str, db_name: str = "sqlstore.db"): + return { + "type": StorageBackendType.SQL_SQLITE.value, + "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + db_name, + } + + @classmethod + def pip_packages(cls) -> list[str]: + return super().pip_packages() + ["aiosqlite"] + + +class PostgresSqlStoreConfig(SqlAlchemySqlStoreConfig): + type: Literal[StorageBackendType.SQL_POSTGRES] = StorageBackendType.SQL_POSTGRES + host: str = "localhost" + port: int = 5432 + db: str = "llamastack" + user: str + password: str | None = None + + @property + def engine_str(self) -> str: + return f"postgresql+asyncpg://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}" + + @classmethod + def pip_packages(cls) -> list[str]: + return super().pip_packages() + ["asyncpg"] + + @classmethod + def sample_run_config(cls, **kwargs): + return { + "type": StorageBackendType.SQL_POSTGRES.value, + "host": "${env.POSTGRES_HOST:=localhost}", + "port": "${env.POSTGRES_PORT:=5432}", + "db": "${env.POSTGRES_DB:=llamastack}", + "user": "${env.POSTGRES_USER:=llamastack}", + "password": "${env.POSTGRES_PASSWORD:=llamastack}", + } + + +# reference = (backend_name, table_name) +class SqlStoreReference(BaseModel): + """A reference to a 'SQL-like' persistent store. A table name must be provided.""" + + table_name: str = Field( + description="Name of the table to use for the SqlStore", + ) + + backend: str | None = Field( + description="Name of backend from persistence.backends, a default will be used if not specified", + default=None, + ) + + +# reference = (backend_name, namespace) +class KVStoreReference(BaseModel): + """A reference to a 'key-value' persistent store. A namespace must be provided.""" + + namespace: str = Field( + description="Key prefix for KVStore backends", + ) + + backend: str | None = Field( + description="Name of backend from persistence.backends, a default will be used if not specified", + default=None, + ) + + +StorageBackendConfig = Annotated[ + RedisKVStoreConfig + | SqliteKVStoreConfig + | PostgresKVStoreConfig + | MongoDBKVStoreConfig + | SqliteSqlStoreConfig + | PostgresSqlStoreConfig, + Field(discriminator="type"), +] + + +class InferenceStoreReference(SqlStoreReference): + """Inference store configuration with queue tuning.""" + + max_write_queue_size: int = Field( + default=10000, + description="Max queued writes for inference store", + ) + num_writers: int = Field( + default=4, + description="Number of concurrent background writers", + ) + + +class StorageConfig(BaseModel): + backends: dict[str, StorageBackendConfig] = Field( + description="Named backend configurations (e.g., 'default', 'cache')", + ) + + # these are stores used natively by the Stack + metadata: KVStoreReference | None = Field( + default=None, + description="Metadata store configuration (uses KVStore backend)", + ) + inference: InferenceStoreReference | None = Field( + default=None, + description="Inference store configuration (uses SqlStore backend)", + ) + conversations: SqlStoreReference | None = Field( + default=None, + description="Conversations store configuration (uses SqlStore backend)", + ) diff --git a/llama_stack/core/store/registry.py b/llama_stack/core/store/registry.py index 04581bab5..ac5d1f1d7 100644 --- a/llama_stack/core/store/registry.py +++ b/llama_stack/core/store/registry.py @@ -11,10 +11,9 @@ from typing import Protocol import pydantic from llama_stack.core.datatypes import RoutableObjectWithProvider -from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR +from llama_stack.core.storage.datatypes import KVStoreReference, StorageConfig from llama_stack.log import get_logger from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig logger = get_logger(__name__, category="core::registry") @@ -191,16 +190,17 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry): async def create_dist_registry( - metadata_store: KVStoreConfig | None, + storage: StorageConfig, image_name: str, ) -> tuple[CachedDiskDistributionRegistry, KVStore]: # instantiate kvstore for storing and retrieving distribution metadata - if metadata_store: - dist_kvstore = await kvstore_impl(metadata_store) - else: - dist_kvstore = await kvstore_impl( - SqliteKVStoreConfig(db_path=(DISTRIBS_BASE_DIR / image_name / "kvstore.db").as_posix()) - ) + # Use metadata store backend with registry-specific namespace + metadata_ref = storage.metadata + registry_ref = KVStoreReference( + namespace="registry", + backend=metadata_ref.backend if metadata_ref else None, + ) + dist_kvstore = await kvstore_impl(registry_ref) dist_registry = CachedDiskDistributionRegistry(dist_kvstore) await dist_registry.initialize() return dist_registry, dist_kvstore diff --git a/llama_stack/distributions/ci-tests/run.yaml b/llama_stack/distributions/ci-tests/run.yaml index a6a6b7c0d..3e592e8d5 100644 --- a/llama_stack/distributions/ci-tests/run.yaml +++ b/llama_stack/distributions/ci-tests/run.yaml @@ -94,29 +94,25 @@ providers: provider_type: inline::faiss config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/faiss_store.db + namespace: vector_io::faiss - provider_id: sqlite-vec provider_type: inline::sqlite-vec config: db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/sqlite_vec.db kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/sqlite_vec_registry.db + namespace: vector_io::sqlite_vec - provider_id: ${env.MILVUS_URL:+milvus} provider_type: inline::milvus config: db_path: ${env.MILVUS_DB_PATH:=~/.llama/distributions/ci-tests}/milvus.db kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/milvus_registry.db + namespace: vector_io::milvus - provider_id: ${env.CHROMADB_URL:+chromadb} provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests/}/chroma_remote_registry.db + namespace: vector_io::chroma_remote - provider_id: ${env.PGVECTOR_DB:+pgvector} provider_type: remote::pgvector config: @@ -126,16 +122,14 @@ providers: user: ${env.PGVECTOR_USER:=} password: ${env.PGVECTOR_PASSWORD:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/pgvector_registry.db + namespace: vector_io::pgvector files: - provider_id: meta-reference-files provider_type: inline::localfs config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/ci-tests/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/files_metadata.db + table_name: files_metadata safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -147,12 +141,11 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses post_training: - provider_id: torchtune-cpu provider_type: inline::torchtune-cpu @@ -163,21 +156,18 @@ providers: provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -207,16 +197,15 @@ providers: provider_type: inline::reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/batches.db + namespace: batches metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/conversations.db models: [] shields: diff --git a/llama_stack/distributions/dell/run-with-safety.yaml b/llama_stack/distributions/dell/run-with-safety.yaml index 5da3cf511..03e42c674 100644 --- a/llama_stack/distributions/dell/run-with-safety.yaml +++ b/llama_stack/distributions/dell/run-with-safety.yaml @@ -27,8 +27,7 @@ providers: config: url: ${env.CHROMADB_URL:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell/}/chroma_remote_registry.db + namespace: vector_io::chroma_remote safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -38,32 +37,28 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses eval: - provider_id: meta-reference provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -87,13 +82,13 @@ providers: - provider_id: rag-runtime provider_type: inline::rag-runtime metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/conversations.db models: - metadata: {} diff --git a/llama_stack/distributions/dell/run.yaml b/llama_stack/distributions/dell/run.yaml index ac0fdc0fa..9e6c03454 100644 --- a/llama_stack/distributions/dell/run.yaml +++ b/llama_stack/distributions/dell/run.yaml @@ -23,8 +23,7 @@ providers: config: url: ${env.CHROMADB_URL:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell/}/chroma_remote_registry.db + namespace: vector_io::chroma_remote safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -34,32 +33,28 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses eval: - provider_id: meta-reference provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -83,13 +78,13 @@ providers: - provider_id: rag-runtime provider_type: inline::rag-runtime metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/conversations.db models: - metadata: {} diff --git a/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml b/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml index 874c5050f..201c0bf7e 100644 --- a/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml +++ b/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml @@ -38,8 +38,7 @@ providers: provider_type: inline::faiss config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/faiss_store.db + namespace: vector_io::faiss safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -49,32 +48,28 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses eval: - provider_id: meta-reference provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -100,13 +95,13 @@ providers: - provider_id: model-context-protocol provider_type: remote::model-context-protocol metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/conversations.db models: - metadata: {} diff --git a/llama_stack/distributions/meta-reference-gpu/run.yaml b/llama_stack/distributions/meta-reference-gpu/run.yaml index 50553d2c7..088abd02d 100644 --- a/llama_stack/distributions/meta-reference-gpu/run.yaml +++ b/llama_stack/distributions/meta-reference-gpu/run.yaml @@ -28,8 +28,7 @@ providers: provider_type: inline::faiss config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/faiss_store.db + namespace: vector_io::faiss safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -39,32 +38,28 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses eval: - provider_id: meta-reference provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -90,13 +85,13 @@ providers: - provider_id: model-context-protocol provider_type: remote::model-context-protocol metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/conversations.db models: - metadata: {} diff --git a/llama_stack/distributions/nvidia/run-with-safety.yaml b/llama_stack/distributions/nvidia/run-with-safety.yaml index e0482f67d..038d03829 100644 --- a/llama_stack/distributions/nvidia/run-with-safety.yaml +++ b/llama_stack/distributions/nvidia/run-with-safety.yaml @@ -29,8 +29,7 @@ providers: provider_type: inline::faiss config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/faiss_store.db + namespace: vector_io::faiss safety: - provider_id: nvidia provider_type: remote::nvidia @@ -41,12 +40,11 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses eval: - provider_id: nvidia provider_type: remote::nvidia @@ -65,8 +63,7 @@ providers: provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/localfs_datasetio.db + namespace: datasetio::localfs - provider_id: nvidia provider_type: remote::nvidia config: @@ -86,16 +83,15 @@ providers: config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/nvidia/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/files_metadata.db + table_name: files_metadata metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/conversations.db models: - metadata: {} diff --git a/llama_stack/distributions/nvidia/run.yaml b/llama_stack/distributions/nvidia/run.yaml index 950782eed..2c08f021a 100644 --- a/llama_stack/distributions/nvidia/run.yaml +++ b/llama_stack/distributions/nvidia/run.yaml @@ -24,8 +24,7 @@ providers: provider_type: inline::faiss config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/faiss_store.db + namespace: vector_io::faiss safety: - provider_id: nvidia provider_type: remote::nvidia @@ -36,12 +35,11 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses eval: - provider_id: nvidia provider_type: remote::nvidia @@ -75,16 +73,15 @@ providers: config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/nvidia/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/files_metadata.db + table_name: files_metadata metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/conversations.db models: [] shields: [] diff --git a/llama_stack/distributions/open-benchmark/run.yaml b/llama_stack/distributions/open-benchmark/run.yaml index a738887b4..84337d935 100644 --- a/llama_stack/distributions/open-benchmark/run.yaml +++ b/llama_stack/distributions/open-benchmark/run.yaml @@ -40,15 +40,13 @@ providers: config: db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/sqlite_vec.db kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/sqlite_vec_registry.db + namespace: vector_io::sqlite_vec - provider_id: ${env.ENABLE_CHROMADB:+chromadb} provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/chroma_remote_registry.db + namespace: vector_io::chroma_remote - provider_id: ${env.ENABLE_PGVECTOR:+pgvector} provider_type: remote::pgvector config: @@ -58,8 +56,7 @@ providers: user: ${env.PGVECTOR_USER:=} password: ${env.PGVECTOR_PASSWORD:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/pgvector_registry.db + namespace: vector_io::pgvector safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -69,32 +66,28 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses eval: - provider_id: meta-reference provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -120,13 +113,13 @@ providers: - provider_id: model-context-protocol provider_type: remote::model-context-protocol metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/conversations.db models: - metadata: {} diff --git a/llama_stack/distributions/postgres-demo/run.yaml b/llama_stack/distributions/postgres-demo/run.yaml index 62faf3f62..d592db5ff 100644 --- a/llama_stack/distributions/postgres-demo/run.yaml +++ b/llama_stack/distributions/postgres-demo/run.yaml @@ -23,8 +23,7 @@ providers: config: url: ${env.CHROMADB_URL:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/postgres-demo}/chroma_remote_registry.db + namespace: vector_io::chroma_remote safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -35,14 +34,14 @@ providers: provider_type: inline::meta-reference config: persistence_store: - type: postgres + type: sql_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} db: ${env.POSTGRES_DB:=llamastack} user: ${env.POSTGRES_USER:=llamastack} password: ${env.POSTGRES_PASSWORD:=llamastack} responses_store: - type: postgres + type: sql_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} db: ${env.POSTGRES_DB:=llamastack} @@ -64,7 +63,7 @@ providers: - provider_id: model-context-protocol provider_type: remote::model-context-protocol metadata_store: - type: postgres + type: kv_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} db: ${env.POSTGRES_DB:=llamastack} @@ -72,14 +71,14 @@ metadata_store: password: ${env.POSTGRES_PASSWORD:=llamastack} table_name: ${env.POSTGRES_TABLE_NAME:=llamastack_kvstore} inference_store: - type: postgres + type: sql_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} db: ${env.POSTGRES_DB:=llamastack} user: ${env.POSTGRES_USER:=llamastack} password: ${env.POSTGRES_PASSWORD:=llamastack} conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/postgres-demo}/conversations.db models: - metadata: {} diff --git a/llama_stack/distributions/starter-gpu/run.yaml b/llama_stack/distributions/starter-gpu/run.yaml index 370d4b516..1e3b3b966 100644 --- a/llama_stack/distributions/starter-gpu/run.yaml +++ b/llama_stack/distributions/starter-gpu/run.yaml @@ -94,29 +94,25 @@ providers: provider_type: inline::faiss config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/faiss_store.db + namespace: vector_io::faiss - provider_id: sqlite-vec provider_type: inline::sqlite-vec config: db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/sqlite_vec.db kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/sqlite_vec_registry.db + namespace: vector_io::sqlite_vec - provider_id: ${env.MILVUS_URL:+milvus} provider_type: inline::milvus config: db_path: ${env.MILVUS_DB_PATH:=~/.llama/distributions/starter-gpu}/milvus.db kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/milvus_registry.db + namespace: vector_io::milvus - provider_id: ${env.CHROMADB_URL:+chromadb} provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu/}/chroma_remote_registry.db + namespace: vector_io::chroma_remote - provider_id: ${env.PGVECTOR_DB:+pgvector} provider_type: remote::pgvector config: @@ -126,16 +122,14 @@ providers: user: ${env.PGVECTOR_USER:=} password: ${env.PGVECTOR_PASSWORD:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/pgvector_registry.db + namespace: vector_io::pgvector files: - provider_id: meta-reference-files provider_type: inline::localfs config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter-gpu/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/files_metadata.db + table_name: files_metadata safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -147,12 +141,11 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses post_training: - provider_id: huggingface-gpu provider_type: inline::huggingface-gpu @@ -166,21 +159,18 @@ providers: provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -210,16 +200,15 @@ providers: provider_type: inline::reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/batches.db + namespace: batches metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/conversations.db models: [] shields: diff --git a/llama_stack/distributions/starter/run.yaml b/llama_stack/distributions/starter/run.yaml index 2f4e7f350..9c618e95d 100644 --- a/llama_stack/distributions/starter/run.yaml +++ b/llama_stack/distributions/starter/run.yaml @@ -94,29 +94,25 @@ providers: provider_type: inline::faiss config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/faiss_store.db + namespace: vector_io::faiss - provider_id: sqlite-vec provider_type: inline::sqlite-vec config: db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/sqlite_vec.db kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/sqlite_vec_registry.db + namespace: vector_io::sqlite_vec - provider_id: ${env.MILVUS_URL:+milvus} provider_type: inline::milvus config: db_path: ${env.MILVUS_DB_PATH:=~/.llama/distributions/starter}/milvus.db kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/milvus_registry.db + namespace: vector_io::milvus - provider_id: ${env.CHROMADB_URL:+chromadb} provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter/}/chroma_remote_registry.db + namespace: vector_io::chroma_remote - provider_id: ${env.PGVECTOR_DB:+pgvector} provider_type: remote::pgvector config: @@ -126,16 +122,14 @@ providers: user: ${env.PGVECTOR_USER:=} password: ${env.PGVECTOR_PASSWORD:=} kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/pgvector_registry.db + namespace: vector_io::pgvector files: - provider_id: meta-reference-files provider_type: inline::localfs config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/files_metadata.db + table_name: files_metadata safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -147,12 +141,11 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses post_training: - provider_id: torchtune-cpu provider_type: inline::torchtune-cpu @@ -163,21 +156,18 @@ providers: provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -207,16 +197,15 @@ providers: provider_type: inline::reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/batches.db + namespace: batches metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/conversations.db models: [] shields: diff --git a/llama_stack/distributions/watsonx/run.yaml b/llama_stack/distributions/watsonx/run.yaml index c3db4eeb8..8d1bfd58f 100644 --- a/llama_stack/distributions/watsonx/run.yaml +++ b/llama_stack/distributions/watsonx/run.yaml @@ -23,8 +23,7 @@ providers: provider_type: inline::faiss config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/faiss_store.db + namespace: vector_io::faiss safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -34,32 +33,28 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/responses_store.db + persistence: + agent_state: + namespace: agents + responses: + table_name: responses eval: - provider_id: meta-reference provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/meta_reference_eval.db + namespace: eval datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/huggingface_datasetio.db + namespace: datasetio::huggingface - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/localfs_datasetio.db + namespace: datasetio::localfs scoring: - provider_id: basic provider_type: inline::basic @@ -90,16 +85,15 @@ providers: config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/watsonx/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/files_metadata.db + table_name: files_metadata metadata_store: - type: sqlite + type: kv_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/registry.db inference_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/inference_store.db conversations_store: - type: sqlite + type: sql_sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/conversations.db models: [] shields: [] diff --git a/llama_stack/providers/inline/agents/meta_reference/agents.py b/llama_stack/providers/inline/agents/meta_reference/agents.py index 810c063e6..c2f6ea640 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agents.py +++ b/llama_stack/providers/inline/agents/meta_reference/agents.py @@ -83,8 +83,8 @@ class MetaReferenceAgentsImpl(Agents): self.policy = policy async def initialize(self) -> None: - self.persistence_store = await kvstore_impl(self.config.persistence_store) - self.responses_store = ResponsesStore(self.config.responses_store, self.policy) + self.persistence_store = await kvstore_impl(self.config.persistence.agent_state) + self.responses_store = ResponsesStore(self.config.persistence.responses, self.policy) await self.responses_store.initialize() self.openai_responses_impl = OpenAIResponsesImpl( inference_api=self.inference_api, diff --git a/llama_stack/providers/inline/agents/meta_reference/config.py b/llama_stack/providers/inline/agents/meta_reference/config.py index 1c392f29c..acda5f463 100644 --- a/llama_stack/providers/inline/agents/meta_reference/config.py +++ b/llama_stack/providers/inline/agents/meta_reference/config.py @@ -8,24 +8,28 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore import KVStoreConfig -from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference, SqlStoreReference + + +class AgentPersistenceConfig(BaseModel): + """Nested persistence configuration for agents.""" + + agent_state: KVStoreReference + responses: SqlStoreReference class MetaReferenceAgentsImplConfig(BaseModel): - persistence_store: KVStoreConfig - responses_store: SqlStoreConfig + persistence: AgentPersistenceConfig @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { - "persistence_store": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="agents_store.db", - ), - "responses_store": SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="responses_store.db", - ), + "persistence": { + "agent_state": KVStoreReference( + namespace="agents", + ).model_dump(exclude_none=True), + "responses": SqlStoreReference( + table_name="responses", + ).model_dump(exclude_none=True), + } } diff --git a/llama_stack/providers/inline/batches/reference/config.py b/llama_stack/providers/inline/batches/reference/config.py index d8d06868b..19078b4f5 100644 --- a/llama_stack/providers/inline/batches/reference/config.py +++ b/llama_stack/providers/inline/batches/reference/config.py @@ -6,13 +6,13 @@ from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference class ReferenceBatchesImplConfig(BaseModel): """Configuration for the Reference Batches implementation.""" - kvstore: KVStoreConfig = Field( + kvstore: KVStoreReference = Field( description="Configuration for the key-value store backend.", ) @@ -33,8 +33,7 @@ class ReferenceBatchesImplConfig(BaseModel): @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="batches.db", - ), + "kvstore": KVStoreReference( + namespace="batches", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/inline/datasetio/localfs/config.py b/llama_stack/providers/inline/datasetio/localfs/config.py index b450e8777..f70fb79ec 100644 --- a/llama_stack/providers/inline/datasetio/localfs/config.py +++ b/llama_stack/providers/inline/datasetio/localfs/config.py @@ -7,20 +7,16 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference class LocalFSDatasetIOConfig(BaseModel): - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="localfs_datasetio.db", - ) + "kvstore": KVStoreReference( + namespace="datasetio::localfs", + ).model_dump(exclude_none=True) } diff --git a/llama_stack/providers/inline/eval/meta_reference/config.py b/llama_stack/providers/inline/eval/meta_reference/config.py index 2a4a29998..210e23201 100644 --- a/llama_stack/providers/inline/eval/meta_reference/config.py +++ b/llama_stack/providers/inline/eval/meta_reference/config.py @@ -7,20 +7,16 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference class MetaReferenceEvalConfig(BaseModel): - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="meta_reference_eval.db", - ) + "kvstore": KVStoreReference( + namespace="eval", + ).model_dump(exclude_none=True) } diff --git a/llama_stack/providers/inline/files/localfs/config.py b/llama_stack/providers/inline/files/localfs/config.py index 6c767af8f..5b437f41c 100644 --- a/llama_stack/providers/inline/files/localfs/config.py +++ b/llama_stack/providers/inline/files/localfs/config.py @@ -8,14 +8,14 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig +from llama_stack.core.storage.datatypes import SqlStoreReference class LocalfsFilesImplConfig(BaseModel): storage_dir: str = Field( description="Directory to store uploaded files", ) - metadata_store: SqlStoreConfig = Field( + metadata_store: SqlStoreReference = Field( description="SQL store configuration for file metadata", ) ttl_secs: int = 365 * 24 * 60 * 60 # 1 year @@ -24,8 +24,7 @@ class LocalfsFilesImplConfig(BaseModel): def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { "storage_dir": "${env.FILES_STORAGE_DIR:=" + __distro_dir__ + "/files}", - "metadata_store": SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="files_metadata.db", - ), + "metadata_store": SqlStoreReference( + table_name="files_metadata", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/inline/vector_io/chroma/config.py b/llama_stack/providers/inline/vector_io/chroma/config.py index a9566f7ff..c373c2241 100644 --- a/llama_stack/providers/inline/vector_io/chroma/config.py +++ b/llama_stack/providers/inline/vector_io/chroma/config.py @@ -8,14 +8,14 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class ChromaVectorIOConfig(BaseModel): db_path: str - kvstore: KVStoreConfig = Field(description="Config for KV store backend") + kvstore: KVStoreReference = Field(description="Config for KV store backend") @classmethod def sample_run_config( @@ -23,8 +23,7 @@ class ChromaVectorIOConfig(BaseModel): ) -> dict[str, Any]: return { "db_path": db_path, - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="chroma_inline_registry.db", - ), + "kvstore": KVStoreReference( + namespace="vector_io::chroma", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/inline/vector_io/faiss/config.py b/llama_stack/providers/inline/vector_io/faiss/config.py index cbcbb1762..c156eee37 100644 --- a/llama_stack/providers/inline/vector_io/faiss/config.py +++ b/llama_stack/providers/inline/vector_io/faiss/config.py @@ -8,22 +8,18 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class FaissVectorIOConfig(BaseModel): - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="faiss_store.db", - ) + "kvstore": KVStoreReference( + namespace="vector_io::faiss", + ).model_dump(exclude_none=True) } diff --git a/llama_stack/providers/inline/vector_io/milvus/config.py b/llama_stack/providers/inline/vector_io/milvus/config.py index 8cbd056be..ba37febaa 100644 --- a/llama_stack/providers/inline/vector_io/milvus/config.py +++ b/llama_stack/providers/inline/vector_io/milvus/config.py @@ -8,25 +8,21 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class MilvusVectorIOConfig(BaseModel): db_path: str - kvstore: KVStoreConfig = Field(description="Config for KV store backend (SQLite only for now)") + kvstore: KVStoreReference = Field(description="Config for KV store backend (SQLite only for now)") consistency_level: str = Field(description="The consistency level of the Milvus server", default="Strong") @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { "db_path": "${env.MILVUS_DB_PATH:=" + __distro_dir__ + "}/" + "milvus.db", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="milvus_registry.db", - ), + "kvstore": KVStoreReference( + namespace="vector_io::milvus", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/inline/vector_io/qdrant/config.py b/llama_stack/providers/inline/vector_io/qdrant/config.py index e15c27ea1..d462c7544 100644 --- a/llama_stack/providers/inline/vector_io/qdrant/config.py +++ b/llama_stack/providers/inline/vector_io/qdrant/config.py @@ -9,23 +9,20 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class QdrantVectorIOConfig(BaseModel): path: str - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { "path": "${env.QDRANT_PATH:=~/.llama/" + __distro_dir__ + "}/" + "qdrant.db", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, db_name="qdrant_registry.db" - ), + "kvstore": KVStoreReference( + namespace="vector_io::qdrant", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/inline/vector_io/sqlite_vec/config.py b/llama_stack/providers/inline/vector_io/sqlite_vec/config.py index 525ed4b1f..e478efb37 100644 --- a/llama_stack/providers/inline/vector_io/sqlite_vec/config.py +++ b/llama_stack/providers/inline/vector_io/sqlite_vec/config.py @@ -8,22 +8,18 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference class SQLiteVectorIOConfig(BaseModel): db_path: str = Field(description="Path to the SQLite database file") - kvstore: KVStoreConfig = Field(description="Config for KV store backend (SQLite only for now)") + kvstore: KVStoreReference = Field(description="Config for KV store backend (SQLite only for now)") @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + "sqlite_vec.db", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="sqlite_vec_registry.db", - ), + "kvstore": KVStoreReference( + namespace="vector_io::sqlite_vec", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/datasetio/huggingface/config.py b/llama_stack/providers/remote/datasetio/huggingface/config.py index 38f933728..1123a103f 100644 --- a/llama_stack/providers/remote/datasetio/huggingface/config.py +++ b/llama_stack/providers/remote/datasetio/huggingface/config.py @@ -7,20 +7,16 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference class HuggingfaceDatasetIOConfig(BaseModel): - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="huggingface_datasetio.db", - ) + "kvstore": KVStoreReference( + namespace="datasetio::huggingface", + ).model_dump(exclude_none=True) } diff --git a/llama_stack/providers/remote/files/s3/config.py b/llama_stack/providers/remote/files/s3/config.py index da20d8668..66e4f6044 100644 --- a/llama_stack/providers/remote/files/s3/config.py +++ b/llama_stack/providers/remote/files/s3/config.py @@ -8,7 +8,7 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig +from llama_stack.core.storage.datatypes import SqlStoreReference class S3FilesImplConfig(BaseModel): @@ -24,7 +24,7 @@ class S3FilesImplConfig(BaseModel): auto_create_bucket: bool = Field( default=False, description="Automatically create the S3 bucket if it doesn't exist" ) - metadata_store: SqlStoreConfig = Field(description="SQL store configuration for file metadata") + metadata_store: SqlStoreReference = Field(description="SQL store configuration for file metadata") @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: @@ -35,8 +35,7 @@ class S3FilesImplConfig(BaseModel): "aws_secret_access_key": "${env.AWS_SECRET_ACCESS_KEY:=}", "endpoint_url": "${env.S3_ENDPOINT_URL:=}", "auto_create_bucket": "${env.S3_AUTO_CREATE_BUCKET:=false}", - "metadata_store": SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="s3_files_metadata.db", - ), + "metadata_store": SqlStoreReference( + table_name="s3_files_metadata", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/files/s3/files.py b/llama_stack/providers/remote/files/s3/files.py index c0e9f81d6..7d76f1721 100644 --- a/llama_stack/providers/remote/files/s3/files.py +++ b/llama_stack/providers/remote/files/s3/files.py @@ -166,7 +166,7 @@ class S3FilesImpl(Files): self._client = _create_s3_client(self._config) await _create_bucket_if_not_exists(self._client, self._config) - self._sql_store = AuthorizedSqlStore(sqlstore_impl(self._config.metadata_store), self.policy) + self._sql_store = AuthorizedSqlStore(sqlstore_impl(self._config.persistence), self.policy) await self._sql_store.create_table( "openai_files", { diff --git a/llama_stack/providers/remote/vector_io/chroma/config.py b/llama_stack/providers/remote/vector_io/chroma/config.py index a1193905a..89efba1c5 100644 --- a/llama_stack/providers/remote/vector_io/chroma/config.py +++ b/llama_stack/providers/remote/vector_io/chroma/config.py @@ -8,21 +8,20 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class ChromaVectorIOConfig(BaseModel): url: str | None - kvstore: KVStoreConfig = Field(description="Config for KV store backend") + kvstore: KVStoreReference = Field(description="Config for KV store backend") @classmethod def sample_run_config(cls, __distro_dir__: str, url: str = "${env.CHROMADB_URL}", **kwargs: Any) -> dict[str, Any]: return { "url": url, - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="chroma_remote_registry.db", - ), + "kvstore": KVStoreReference( + namespace="vector_io::chroma_remote", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/milvus/config.py b/llama_stack/providers/remote/vector_io/milvus/config.py index 899d3678d..ef21f2f29 100644 --- a/llama_stack/providers/remote/vector_io/milvus/config.py +++ b/llama_stack/providers/remote/vector_io/milvus/config.py @@ -8,7 +8,7 @@ from typing import Any from pydantic import BaseModel, ConfigDict, Field -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @@ -17,7 +17,7 @@ class MilvusVectorIOConfig(BaseModel): uri: str = Field(description="The URI of the Milvus server") token: str | None = Field(description="The token of the Milvus server") consistency_level: str = Field(description="The consistency level of the Milvus server", default="Strong") - kvstore: KVStoreConfig = Field(description="Config for KV store backend") + kvstore: KVStoreReference = Field(description="Config for KV store backend") # This configuration allows additional fields to be passed through to the underlying Milvus client. # See the [Milvus](https://milvus.io/docs/install-overview.md) documentation for more details about Milvus in general. @@ -28,8 +28,7 @@ class MilvusVectorIOConfig(BaseModel): return { "uri": "${env.MILVUS_ENDPOINT}", "token": "${env.MILVUS_TOKEN}", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="milvus_remote_registry.db", - ), + "kvstore": KVStoreReference( + namespace="vector_io::milvus_remote", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/pgvector/config.py b/llama_stack/providers/remote/vector_io/pgvector/config.py index 334cbe5be..46ce55600 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/config.py +++ b/llama_stack/providers/remote/vector_io/pgvector/config.py @@ -8,10 +8,7 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @@ -22,7 +19,9 @@ class PGVectorVectorIOConfig(BaseModel): db: str | None = Field(default="postgres") user: str | None = Field(default="postgres") password: str | None = Field(default="mysecretpassword") - kvstore: KVStoreConfig | None = Field(description="Config for KV store backend (SQLite only for now)", default=None) + kvstore: KVStoreReference | None = Field( + description="Config for KV store backend (SQLite only for now)", default=None + ) @classmethod def sample_run_config( @@ -41,8 +40,7 @@ class PGVectorVectorIOConfig(BaseModel): "db": db, "user": user, "password": password, - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="pgvector_registry.db", - ), + "kvstore": KVStoreReference( + namespace="vector_io::pgvector", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/qdrant/config.py b/llama_stack/providers/remote/vector_io/qdrant/config.py index ff5506236..2dc89386f 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/config.py +++ b/llama_stack/providers/remote/vector_io/qdrant/config.py @@ -8,10 +8,7 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @@ -27,14 +24,13 @@ class QdrantVectorIOConfig(BaseModel): prefix: str | None = None timeout: int | None = None host: str | None = None - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { "api_key": "${env.QDRANT_API_KEY:=}", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="qdrant_registry.db", - ), + "kvstore": KVStoreReference( + namespace="vector_io::qdrant_remote", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/weaviate/config.py b/llama_stack/providers/remote/vector_io/weaviate/config.py index b693e294e..65781c61f 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/config.py +++ b/llama_stack/providers/remote/vector_io/weaviate/config.py @@ -8,10 +8,7 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @@ -19,7 +16,9 @@ from llama_stack.schema_utils import json_schema_type class WeaviateVectorIOConfig(BaseModel): weaviate_api_key: str | None = Field(description="The API key for the Weaviate instance", default=None) weaviate_cluster_url: str | None = Field(description="The URL of the Weaviate cluster", default="localhost:8080") - kvstore: KVStoreConfig | None = Field(description="Config for KV store backend (SQLite only for now)", default=None) + kvstore: KVStoreReference | None = Field( + description="Config for KV store backend (SQLite only for now)", default=None + ) @classmethod def sample_run_config( @@ -30,8 +29,7 @@ class WeaviateVectorIOConfig(BaseModel): return { "weaviate_api_key": None, "weaviate_cluster_url": "${env.WEAVIATE_CLUSTER_URL:=localhost:8080}", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="weaviate_registry.db", - ), + "kvstore": KVStoreReference( + namespace="vector_io::weaviate", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/utils/inference/inference_store.py b/llama_stack/providers/utils/inference/inference_store.py index 901f77c67..13cde66b1 100644 --- a/llama_stack/providers/utils/inference/inference_store.py +++ b/llama_stack/providers/utils/inference/inference_store.py @@ -15,12 +15,13 @@ from llama_stack.apis.inference import ( OpenAIMessageParam, Order, ) -from llama_stack.core.datatypes import AccessRule, InferenceStoreConfig +from llama_stack.core.datatypes import AccessRule +from llama_stack.core.storage.datatypes import InferenceStoreReference, StorageBackendType from llama_stack.log import get_logger from ..sqlstore.api import ColumnDefinition, ColumnType from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore -from ..sqlstore.sqlstore import SqlStoreConfig, SqlStoreType, sqlstore_impl +from ..sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl logger = get_logger(name=__name__, category="inference") @@ -28,33 +29,31 @@ logger = get_logger(name=__name__, category="inference") class InferenceStore: def __init__( self, - config: InferenceStoreConfig | SqlStoreConfig, + reference: InferenceStoreReference, policy: list[AccessRule], ): - # Handle backward compatibility - if not isinstance(config, InferenceStoreConfig): - # Legacy: SqlStoreConfig passed directly as config - config = InferenceStoreConfig( - sql_store_config=config, - ) - - self.config = config - self.sql_store_config = config.sql_store_config + self.reference = reference self.sql_store = None self.policy = policy - # Disable write queue for SQLite to avoid concurrency issues - self.enable_write_queue = self.sql_store_config.type != SqlStoreType.sqlite - # Async write queue and worker control self._queue: asyncio.Queue[tuple[OpenAIChatCompletion, list[OpenAIMessageParam]]] | None = None self._worker_tasks: list[asyncio.Task[Any]] = [] - self._max_write_queue_size: int = config.max_write_queue_size - self._num_writers: int = max(1, config.num_writers) + self._max_write_queue_size: int = reference.max_write_queue_size + self._num_writers: int = max(1, reference.num_writers) async def initialize(self): """Create the necessary tables if they don't exist.""" - self.sql_store = AuthorizedSqlStore(sqlstore_impl(self.sql_store_config), self.policy) + base_store = sqlstore_impl(self.reference) + self.sql_store = AuthorizedSqlStore(base_store, self.policy) + + # Disable write queue for SQLite to avoid concurrency issues + backend_name = self.reference.backend or list(_SQLSTORE_BACKENDS.keys())[0] if _SQLSTORE_BACKENDS else None + if backend_name: + backend_config = _SQLSTORE_BACKENDS.get(backend_name) + self.enable_write_queue = backend_config.type != StorageBackendType.SQL_SQLITE + else: + self.enable_write_queue = True await self.sql_store.create_table( "chat_completions", { diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py index 7b6a79350..c0582abc4 100644 --- a/llama_stack/providers/utils/kvstore/config.py +++ b/llama_stack/providers/utils/kvstore/config.py @@ -4,143 +4,20 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import re -from enum import Enum -from typing import Annotated, Literal +from typing import Annotated -from pydantic import BaseModel, Field, field_validator - -from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR - - -class KVStoreType(Enum): - redis = "redis" - sqlite = "sqlite" - postgres = "postgres" - mongodb = "mongodb" - - -class CommonConfig(BaseModel): - namespace: str | None = Field( - default=None, - description="All keys will be prefixed with this namespace", - ) - - -class RedisKVStoreConfig(CommonConfig): - type: Literal["redis"] = KVStoreType.redis.value - host: str = "localhost" - port: int = 6379 - - @property - def url(self) -> str: - return f"redis://{self.host}:{self.port}" - - @classmethod - def pip_packages(cls) -> list[str]: - return ["redis"] - - @classmethod - def sample_run_config(cls): - return { - "type": "redis", - "host": "${env.REDIS_HOST:=localhost}", - "port": "${env.REDIS_PORT:=6379}", - } - - -class SqliteKVStoreConfig(CommonConfig): - type: Literal["sqlite"] = KVStoreType.sqlite.value - db_path: str = Field( - default=(RUNTIME_BASE_DIR / "kvstore.db").as_posix(), - description="File path for the sqlite database", - ) - - @classmethod - def pip_packages(cls) -> list[str]: - return ["aiosqlite"] - - @classmethod - def sample_run_config(cls, __distro_dir__: str, db_name: str = "kvstore.db"): - return { - "type": "sqlite", - "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + db_name, - } - - -class PostgresKVStoreConfig(CommonConfig): - type: Literal["postgres"] = KVStoreType.postgres.value - host: str = "localhost" - port: int = 5432 - db: str = "llamastack" - user: str - password: str | None = None - ssl_mode: str | None = None - ca_cert_path: str | None = None - table_name: str = "llamastack_kvstore" - - @classmethod - def sample_run_config(cls, table_name: str = "llamastack_kvstore", **kwargs): - return { - "type": "postgres", - "host": "${env.POSTGRES_HOST:=localhost}", - "port": "${env.POSTGRES_PORT:=5432}", - "db": "${env.POSTGRES_DB:=llamastack}", - "user": "${env.POSTGRES_USER:=llamastack}", - "password": "${env.POSTGRES_PASSWORD:=llamastack}", - "table_name": "${env.POSTGRES_TABLE_NAME:=" + table_name + "}", - } - - @classmethod - @field_validator("table_name") - def validate_table_name(cls, v: str) -> str: - # PostgreSQL identifiers rules: - # - Must start with a letter or underscore - # - Can contain letters, numbers, and underscores - # - Maximum length is 63 bytes - pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*$" - if not re.match(pattern, v): - raise ValueError( - "Invalid table name. Must start with letter or underscore and contain only letters, numbers, and underscores" - ) - if len(v) > 63: - raise ValueError("Table name must be less than 63 characters") - return v - - @classmethod - def pip_packages(cls) -> list[str]: - return ["psycopg2-binary"] - - -class MongoDBKVStoreConfig(CommonConfig): - type: Literal["mongodb"] = KVStoreType.mongodb.value - host: str = "localhost" - port: int = 27017 - db: str = "llamastack" - user: str | None = None - password: str | None = None - collection_name: str = "llamastack_kvstore" - - @classmethod - def pip_packages(cls) -> list[str]: - return ["pymongo"] - - @classmethod - def sample_run_config(cls, collection_name: str = "llamastack_kvstore"): - return { - "type": "mongodb", - "host": "${env.MONGODB_HOST:=localhost}", - "port": "${env.MONGODB_PORT:=5432}", - "db": "${env.MONGODB_DB}", - "user": "${env.MONGODB_USER}", - "password": "${env.MONGODB_PASSWORD}", - "collection_name": "${env.MONGODB_COLLECTION_NAME:=" + collection_name + "}", - } +from pydantic import Field +from llama_stack.core.storage.datatypes import ( + MongoDBKVStoreConfig, + PostgresKVStoreConfig, + RedisKVStoreConfig, + SqliteKVStoreConfig, + StorageBackendType, +) KVStoreConfig = Annotated[ - RedisKVStoreConfig | SqliteKVStoreConfig | PostgresKVStoreConfig | MongoDBKVStoreConfig, - Field(discriminator="type", default=KVStoreType.sqlite.value), + RedisKVStoreConfig | SqliteKVStoreConfig | PostgresKVStoreConfig | MongoDBKVStoreConfig, Field(discriminator="type") ] @@ -148,13 +25,13 @@ def get_pip_packages(store_config: dict | KVStoreConfig) -> list[str]: """Get pip packages for KV store config, handling both dict and object cases.""" if isinstance(store_config, dict): store_type = store_config.get("type") - if store_type == "sqlite": + if store_type == StorageBackendType.KV_SQLITE.value: return SqliteKVStoreConfig.pip_packages() - elif store_type == "postgres": + elif store_type == StorageBackendType.KV_POSTGRES.value: return PostgresKVStoreConfig.pip_packages() - elif store_type == "redis": + elif store_type == StorageBackendType.KV_REDIS.value: return RedisKVStoreConfig.pip_packages() - elif store_type == "mongodb": + elif store_type == StorageBackendType.KV_MONGODB.value: return MongoDBKVStoreConfig.pip_packages() else: raise ValueError(f"Unknown KV store type: {store_type}") diff --git a/llama_stack/providers/utils/kvstore/kvstore.py b/llama_stack/providers/utils/kvstore/kvstore.py index 426523d8e..442bbeb41 100644 --- a/llama_stack/providers/utils/kvstore/kvstore.py +++ b/llama_stack/providers/utils/kvstore/kvstore.py @@ -1,12 +1,14 @@ # 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 __future__ import annotations + +from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig, StorageBackendType from .api import KVStore -from .config import KVStoreConfig, KVStoreType +from .config import KVStoreConfig def kvstore_dependencies(): @@ -44,20 +46,58 @@ class InmemoryKVStoreImpl(KVStore): del self._store[key] -async def kvstore_impl(config: KVStoreConfig) -> KVStore: - if config.type == KVStoreType.redis.value: +_KVSTORE_BACKENDS: dict[str, KVStoreConfig] = {} +_KVSTORE_DEFAULT_BACKEND: str | None = None + + +def register_kvstore_backends(backends: dict[str, StorageBackendConfig]) -> None: + """Register the set of available KV store backends for reference resolution.""" + global _KVSTORE_BACKENDS + + def _set_default_backend(name: str) -> None: + global _KVSTORE_DEFAULT_BACKEND + if _KVSTORE_DEFAULT_BACKEND and _KVSTORE_DEFAULT_BACKEND != name: + raise ValueError( + f"Multiple KVStore backends marked as default: '{_KVSTORE_DEFAULT_BACKEND}' and '{name}'. " + "Only one backend can be the default." + ) + _KVSTORE_DEFAULT_BACKEND = name + + _KVSTORE_BACKENDS.clear() + for name, cfg in backends.items(): + if cfg.default: + _set_default_backend(name) + _KVSTORE_BACKENDS[name] = cfg + + +async def kvstore_impl(reference: KVStoreReference) -> KVStore: + backend_name = reference.backend or _KVSTORE_DEFAULT_BACKEND + if not backend_name: + raise ValueError( + "KVStore reference did not specify a backend and no default backend is configured. " + f"Available backends: {sorted(_KVSTORE_BACKENDS)}" + ) + + backend_config = _KVSTORE_BACKENDS.get(backend_name) + if backend_config is None: + raise ValueError(f"Unknown KVStore backend '{backend_name}'. Registered backends: {sorted(_KVSTORE_BACKENDS)}") + + config = backend_config.model_copy() + config.namespace = reference.namespace + + if config.type == StorageBackendType.KV_REDIS.value: from .redis import RedisKVStoreImpl impl = RedisKVStoreImpl(config) - elif config.type == KVStoreType.sqlite.value: + elif config.type == StorageBackendType.KV_SQLITE.value: from .sqlite import SqliteKVStoreImpl impl = SqliteKVStoreImpl(config) - elif config.type == KVStoreType.postgres.value: + elif config.type == StorageBackendType.KV_POSTGRES.value: from .postgres import PostgresKVStoreImpl impl = PostgresKVStoreImpl(config) - elif config.type == KVStoreType.mongodb.value: + elif config.type == StorageBackendType.KV_MONGODB.value: from .mongodb import MongoDBKVStoreImpl impl = MongoDBKVStoreImpl(config) diff --git a/llama_stack/providers/utils/sqlstore/authorized_sqlstore.py b/llama_stack/providers/utils/sqlstore/authorized_sqlstore.py index e1da4db6e..d2a66213e 100644 --- a/llama_stack/providers/utils/sqlstore/authorized_sqlstore.py +++ b/llama_stack/providers/utils/sqlstore/authorized_sqlstore.py @@ -15,7 +15,7 @@ from llama_stack.core.request_headers import get_authenticated_user from llama_stack.log import get_logger from .api import ColumnDefinition, ColumnType, PaginatedResponse, SqlStore -from .sqlstore import SqlStoreType +from .sqlstore import StorageBackendType logger = get_logger(name=__name__, category="providers::utils") @@ -82,8 +82,8 @@ class AuthorizedSqlStore: if not hasattr(self.sql_store, "config"): raise ValueError("SqlStore must have a config attribute to be used with AuthorizedSqlStore") - self.database_type = self.sql_store.config.type - if self.database_type not in [SqlStoreType.postgres, SqlStoreType.sqlite]: + self.database_type = self.sql_store.config.type.value + if self.database_type not in [StorageBackendType.SQL_POSTGRES.value, StorageBackendType.SQL_SQLITE.value]: raise ValueError(f"Unsupported database type: {self.database_type}") def _validate_sql_optimized_policy(self) -> None: @@ -220,9 +220,9 @@ class AuthorizedSqlStore: Returns: SQL expression to extract JSON value """ - if self.database_type == SqlStoreType.postgres: + if self.database_type == StorageBackendType.SQL_POSTGRES.value: return f"{column}->'{path}'" - elif self.database_type == SqlStoreType.sqlite: + elif self.database_type == StorageBackendType.SQL_SQLITE.value: return f"JSON_EXTRACT({column}, '$.{path}')" else: raise ValueError(f"Unsupported database type: {self.database_type}") @@ -237,9 +237,9 @@ class AuthorizedSqlStore: Returns: SQL expression to extract JSON value as text """ - if self.database_type == SqlStoreType.postgres: + if self.database_type == StorageBackendType.SQL_POSTGRES.value: return f"{column}->>'{path}'" - elif self.database_type == SqlStoreType.sqlite: + elif self.database_type == StorageBackendType.SQL_SQLITE.value: return f"JSON_EXTRACT({column}, '$.{path}')" else: raise ValueError(f"Unsupported database type: {self.database_type}") @@ -248,10 +248,10 @@ class AuthorizedSqlStore: """Get the SQL conditions for public access.""" # Public records are records that have no owner_principal or access_attributes conditions = ["owner_principal = ''"] - if self.database_type == SqlStoreType.postgres: + if self.database_type == StorageBackendType.SQL_POSTGRES.value: # Postgres stores JSON null as 'null' conditions.append("access_attributes::text = 'null'") - elif self.database_type == SqlStoreType.sqlite: + elif self.database_type == StorageBackendType.SQL_SQLITE.value: conditions.append("access_attributes = 'null'") else: raise ValueError(f"Unsupported database type: {self.database_type}") diff --git a/llama_stack/providers/utils/sqlstore/sqlstore.py b/llama_stack/providers/utils/sqlstore/sqlstore.py index fc44402ae..06f83116c 100644 --- a/llama_stack/providers/utils/sqlstore/sqlstore.py +++ b/llama_stack/providers/utils/sqlstore/sqlstore.py @@ -4,90 +4,29 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -from abc import abstractmethod -from enum import StrEnum -from pathlib import Path -from typing import Annotated, Literal +from typing import Annotated -from pydantic import BaseModel, Field +from pydantic import Field -from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR +from llama_stack.core.storage.datatypes import ( + PostgresSqlStoreConfig, + SqliteSqlStoreConfig, + SqlStoreReference, + StorageBackendConfig, + StorageBackendType, +) from .api import SqlStore sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"] - -class SqlStoreType(StrEnum): - sqlite = "sqlite" - postgres = "postgres" - - -class SqlAlchemySqlStoreConfig(BaseModel): - @property - @abstractmethod - def engine_str(self) -> str: ... - - # TODO: move this when we have a better way to specify dependencies with internal APIs - @classmethod - def pip_packages(cls) -> list[str]: - return ["sqlalchemy[asyncio]"] - - -class SqliteSqlStoreConfig(SqlAlchemySqlStoreConfig): - type: Literal[SqlStoreType.sqlite] = SqlStoreType.sqlite - db_path: str = Field( - default=(RUNTIME_BASE_DIR / "sqlstore.db").as_posix(), - description="Database path, e.g. ~/.llama/distributions/ollama/sqlstore.db", - ) - - @property - def engine_str(self) -> str: - return "sqlite+aiosqlite:///" + Path(self.db_path).expanduser().as_posix() - - @classmethod - def sample_run_config(cls, __distro_dir__: str, db_name: str = "sqlstore.db"): - return { - "type": "sqlite", - "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + db_name, - } - - @classmethod - def pip_packages(cls) -> list[str]: - return super().pip_packages() + ["aiosqlite"] - - -class PostgresSqlStoreConfig(SqlAlchemySqlStoreConfig): - type: Literal[SqlStoreType.postgres] = SqlStoreType.postgres - host: str = "localhost" - port: int = 5432 - db: str = "llamastack" - user: str - password: str | None = None - - @property - def engine_str(self) -> str: - return f"postgresql+asyncpg://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}" - - @classmethod - def pip_packages(cls) -> list[str]: - return super().pip_packages() + ["asyncpg"] - - @classmethod - def sample_run_config(cls, **kwargs): - return { - "type": "postgres", - "host": "${env.POSTGRES_HOST:=localhost}", - "port": "${env.POSTGRES_PORT:=5432}", - "db": "${env.POSTGRES_DB:=llamastack}", - "user": "${env.POSTGRES_USER:=llamastack}", - "password": "${env.POSTGRES_PASSWORD:=llamastack}", - } +_SQLSTORE_BACKENDS: dict[str, StorageBackendConfig] = {} +_SQLSTORE_DEFAULT_BACKEND: str | None = None SqlStoreConfig = Annotated[ SqliteSqlStoreConfig | PostgresSqlStoreConfig, - Field(discriminator="type", default=SqlStoreType.sqlite.value), + Field(discriminator="type"), ] @@ -95,9 +34,9 @@ def get_pip_packages(store_config: dict | SqlStoreConfig) -> list[str]: """Get pip packages for SQL store config, handling both dict and object cases.""" if isinstance(store_config, dict): store_type = store_config.get("type") - if store_type == "sqlite": + if store_type == StorageBackendType.SQL_SQLITE.value: return SqliteSqlStoreConfig.pip_packages() - elif store_type == "postgres": + elif store_type == StorageBackendType.SQL_POSTGRES.value: return PostgresSqlStoreConfig.pip_packages() else: raise ValueError(f"Unknown SQL store type: {store_type}") @@ -105,12 +44,41 @@ def get_pip_packages(store_config: dict | SqlStoreConfig) -> list[str]: return store_config.pip_packages() -def sqlstore_impl(config: SqlStoreConfig) -> SqlStore: - if config.type in [SqlStoreType.sqlite, SqlStoreType.postgres]: +def sqlstore_impl(reference: SqlStoreReference) -> SqlStore: + backend_name = reference.backend or _SQLSTORE_DEFAULT_BACKEND + if not backend_name: + raise ValueError( + "SQL store reference did not specify a backend and no default backend is configured. " + f"Available backends: {sorted(_SQLSTORE_BACKENDS)}" + ) + + backend_config = _SQLSTORE_BACKENDS.get(backend_name) + if backend_config.type in [StorageBackendType.SQL_SQLITE, StorageBackendType.SQL_POSTGRES]: from .sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl - impl = SqlAlchemySqlStoreImpl(config) + config = backend_config.model_copy() + config.table_name = reference.table_name + return SqlAlchemySqlStoreImpl(config) else: - raise ValueError(f"Unknown sqlstore type {config.type}") + raise ValueError(f"Unknown sqlstore type {backend_config.type}") - return impl + +def register_sqlstore_backends(backends: dict[str, StorageBackendConfig]) -> None: + """Register the set of available SQL store backends for reference resolution.""" + global _SQLSTORE_BACKENDS + + def _set_default_backend(name: str) -> None: + global _SQLSTORE_DEFAULT_BACKEND + if _SQLSTORE_DEFAULT_BACKEND and _SQLSTORE_DEFAULT_BACKEND != name: + raise ValueError( + f"Multiple SQL store backends marked as default: '{_SQLSTORE_DEFAULT_BACKEND}' and '{name}'. " + "Only one backend can be the default." + ) + _SQLSTORE_DEFAULT_BACKEND = name + + _SQLSTORE_BACKENDS.clear() + for name, cfg in backends.items(): + if cfg.default: + _set_default_backend(name) + + _SQLSTORE_BACKENDS[name] = cfg diff --git a/tests/integration/test_persistence_integration.py b/tests/integration/test_persistence_integration.py new file mode 100644 index 000000000..93c44bd6e --- /dev/null +++ b/tests/integration/test_persistence_integration.py @@ -0,0 +1,63 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import yaml + +from llama_stack.core.datatypes import StackRunConfig +from llama_stack.core.storage.datatypes import ( + PostgresSqlStoreConfig, + SqliteKVStoreConfig, + SqliteSqlStoreConfig, +) + + +def test_starter_distribution_config_loads_and_resolves(): + """Integration: Actual starter config should parse and have correct storage structure.""" + with open("llama_stack/distributions/starter/run.yaml") as f: + config_dict = yaml.safe_load(f) + + config = StackRunConfig(**config_dict) + + # Config should have storage with default backend + assert config.storage is not None + assert "default" in config.storage.backends + assert isinstance(config.storage.backends["default"], SqliteSqlStoreConfig) + + # Stores should reference the default backend + assert config.storage.metadata is not None + assert config.storage.metadata.backend == "default" + assert config.storage.metadata.namespace is not None + + assert config.storage.inference is not None + assert config.storage.inference.backend == "default" + assert config.storage.inference.table_name is not None + assert config.storage.inference.max_write_queue_size > 0 + assert config.storage.inference.num_writers > 0 + + +def test_postgres_demo_distribution_config_loads(): + """Integration: Postgres demo should use Postgres backend for all stores.""" + with open("llama_stack/distributions/postgres-demo/run.yaml") as f: + config_dict = yaml.safe_load(f) + + config = StackRunConfig(**config_dict) + + # Should have postgres backend + assert config.storage is not None + assert "default" in config.storage.backends + assert isinstance(config.storage.backends["default"], PostgresSqlStoreConfig) + + # Both stores use same postgres backend + assert config.storage.metadata is not None + assert config.storage.metadata.backend == "default" + + assert config.storage.inference is not None + assert config.storage.inference.backend == "default" + + # Backend config should be Postgres + postgres_backend = config.storage.backends["default"] + assert isinstance(postgres_backend, PostgresSqlStoreConfig) + assert postgres_backend.host == "${env.POSTGRES_HOST:=localhost}" diff --git a/tests/unit/core/test_persistence_config.py b/tests/unit/core/test_persistence_config.py new file mode 100644 index 000000000..0b99b3cab --- /dev/null +++ b/tests/unit/core/test_persistence_config.py @@ -0,0 +1,82 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import pytest +from pydantic import ValidationError + +from llama_stack.core.datatypes import ( + InferenceStoreReference, + PersistenceConfig, + StoreReference, + StoresConfig, +) +from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig +from llama_stack.providers.utils.sqlstore.sqlstore import ( + PostgresSqlStoreConfig, + SqliteSqlStoreConfig, +) + + +def test_backend_reference_validation_catches_missing_backend(): + """Critical: Catch user typos in backend references before runtime.""" + with pytest.raises(ValidationError, match="not defined in persistence.backends"): + PersistenceConfig( + backends={ + "default": SqliteSqlStoreConfig(db_path="/tmp/store.db"), + }, + stores=StoresConfig( + metadata=StoreReference(backend="typo_backend"), # User typo + ), + ) + + +def test_backend_reference_validation_accepts_valid_config(): + """Valid config should parse without errors.""" + config = PersistenceConfig( + backends={ + "default": SqliteSqlStoreConfig(db_path="/tmp/store.db"), + }, + stores=StoresConfig( + metadata=StoreReference(backend="default"), + inference=InferenceStoreReference(backend="default"), + ), + ) + assert config.stores.metadata.backend == "default" + assert config.stores.inference.backend == "default" + + +def test_multiple_stores_can_share_same_backend(): + """Core use case: metadata and inference both use 'default' backend.""" + config = PersistenceConfig( + backends={ + "default": SqliteSqlStoreConfig(db_path="/tmp/shared.db"), + }, + stores=StoresConfig( + metadata=StoreReference(backend="default", namespace="metadata"), + inference=InferenceStoreReference(backend="default"), + conversations=StoreReference(backend="default"), + ), + ) + # All reference the same backend + assert config.stores.metadata.backend == "default" + assert config.stores.inference.backend == "default" + assert config.stores.conversations.backend == "default" + + +def test_mixed_backend_types_allowed(): + """Should support KVStore and SqlStore backends simultaneously.""" + config = PersistenceConfig( + backends={ + "kvstore": SqliteKVStoreConfig(db_path="/tmp/kv.db"), + "sqlstore": PostgresSqlStoreConfig(user="test", password="test", host="localhost", db="test"), + }, + stores=StoresConfig( + metadata=StoreReference(backend="kvstore"), + inference=InferenceStoreReference(backend="sqlstore"), + ), + ) + assert isinstance(config.backends["kvstore"], SqliteKVStoreConfig) + assert isinstance(config.backends["sqlstore"], PostgresSqlStoreConfig)