simplified some, walked back some decisions

This commit is contained in:
Ashwin Bharambe 2025-10-17 10:05:07 -07:00
parent af7472cdb0
commit 636764c2a1
90 changed files with 887 additions and 570 deletions

View file

@ -14,16 +14,18 @@ Meta's reference implementation of an agent system that can use tools, access ve
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `persistence_store` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | |
| `responses_store` | `utils.sqlstore.sqlstore.SqliteSqlStoreConfig \| utils.sqlstore.sqlstore.PostgresSqlStoreConfig` | No | sqlite | |
| `persistence` | `<class 'inline.agents.meta_reference.config.AgentPersistenceConfig'>` | No | | |
## Sample Configuration
```yaml
persistence_store:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/agents_store.db
responses_store:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/responses_store.db
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
```

View file

@ -14,7 +14,7 @@ Reference implementation of batches API with KVStore persistence.
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Configuration for the key-value store backend. |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | Configuration for the key-value store backend. |
| `max_concurrent_batches` | `<class 'int'>` | No | 1 | Maximum number of concurrent batches to process simultaneously. |
| `max_concurrent_requests_per_batch` | `<class 'int'>` | No | 10 | Maximum number of concurrent requests to process per batch. |
@ -22,6 +22,6 @@ Reference implementation of batches API with KVStore persistence.
```yaml
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/batches.db
namespace: batches
backend: kv_default
```

View file

@ -14,12 +14,12 @@ Local filesystem-based dataset I/O provider for reading and writing datasets to
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | |
## Sample Configuration
```yaml
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/localfs_datasetio.db
namespace: datasetio::localfs
backend: kv_default
```

View file

@ -14,12 +14,12 @@ HuggingFace datasets provider for accessing and managing datasets from the Huggi
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | |
## Sample Configuration
```yaml
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/huggingface_datasetio.db
namespace: datasetio::huggingface
backend: kv_default
```

View file

@ -14,12 +14,12 @@ Meta's reference implementation of evaluation tasks with support for multiple la
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | |
## Sample Configuration
```yaml
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/meta_reference_eval.db
namespace: eval
backend: kv_default
```

View file

@ -15,7 +15,7 @@ Local filesystem-based file storage provider for managing files and documents lo
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `storage_dir` | `<class 'str'>` | No | | Directory to store uploaded files |
| `metadata_store` | `utils.sqlstore.sqlstore.SqliteSqlStoreConfig \| utils.sqlstore.sqlstore.PostgresSqlStoreConfig` | No | sqlite | SQL store configuration for file metadata |
| `metadata_store` | `<class 'llama_stack.core.storage.datatypes.SqlStoreReference'>` | No | | SQL store configuration for file metadata |
| `ttl_secs` | `<class 'int'>` | No | 31536000 | |
## Sample Configuration
@ -23,6 +23,6 @@ Local filesystem-based file storage provider for managing files and documents lo
```yaml
storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/dummy/files}
metadata_store:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/files_metadata.db
table_name: files_metadata
backend: sql_default
```

View file

@ -20,7 +20,7 @@ AWS S3-based file storage provider for scalable cloud file management with metad
| `aws_secret_access_key` | `str \| None` | No | | AWS secret access key (optional if using IAM roles) |
| `endpoint_url` | `str \| None` | No | | Custom S3 endpoint URL (for MinIO, LocalStack, etc.) |
| `auto_create_bucket` | `<class 'bool'>` | No | False | Automatically create the S3 bucket if it doesn't exist |
| `metadata_store` | `utils.sqlstore.sqlstore.SqliteSqlStoreConfig \| utils.sqlstore.sqlstore.PostgresSqlStoreConfig` | No | sqlite | SQL store configuration for file metadata |
| `metadata_store` | `<class 'llama_stack.core.storage.datatypes.SqlStoreReference'>` | No | | SQL store configuration for file metadata |
## Sample Configuration
@ -32,6 +32,6 @@ 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:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/s3_files_metadata.db
table_name: s3_files_metadata
backend: sql_default
```

View file

@ -79,13 +79,13 @@ See [Chroma's documentation](https://docs.trychroma.com/docs/overview/introducti
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `db_path` | `<class 'str'>` | No | | |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Config for KV store backend |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | Config for KV store backend |
## Sample Configuration
```yaml
db_path: ${env.CHROMADB_PATH}
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/chroma_inline_registry.db
namespace: vector_io::chroma
backend: kv_default
```

View file

@ -95,12 +95,12 @@ more details about Faiss in general.
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | |
## Sample Configuration
```yaml
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/faiss_store.db
namespace: vector_io::faiss
backend: kv_default
```

View file

@ -14,14 +14,14 @@ Meta's reference implementation of a vector database.
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | |
## Sample Configuration
```yaml
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/faiss_store.db
namespace: vector_io::faiss
backend: kv_default
```
## Deprecation Notice

View file

@ -17,7 +17,7 @@ Please refer to the remote provider documentation.
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `db_path` | `<class 'str'>` | No | | |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Config for KV store backend (SQLite only for now) |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | Config for KV store backend (SQLite only for now) |
| `consistency_level` | `<class 'str'>` | No | Strong | The consistency level of the Milvus server |
## Sample Configuration
@ -25,6 +25,6 @@ Please refer to the remote provider documentation.
```yaml
db_path: ${env.MILVUS_DB_PATH:=~/.llama/dummy}/milvus.db
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/milvus_registry.db
namespace: vector_io::milvus
backend: kv_default
```

View file

@ -98,13 +98,13 @@ See the [Qdrant documentation](https://qdrant.tech/documentation/) for more deta
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `path` | `<class 'str'>` | No | | |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | |
## Sample Configuration
```yaml
path: ${env.QDRANT_PATH:=~/.llama/~/.llama/dummy}/qdrant.db
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/qdrant_registry.db
namespace: vector_io::qdrant
backend: kv_default
```

View file

@ -408,13 +408,13 @@ See [sqlite-vec's GitHub repo](https://github.com/asg017/sqlite-vec/tree/main) f
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `db_path` | `<class 'str'>` | No | | Path to the SQLite database file |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Config for KV store backend (SQLite only for now) |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | Config for KV store backend (SQLite only for now) |
## Sample Configuration
```yaml
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/sqlite_vec.db
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/sqlite_vec_registry.db
namespace: vector_io::sqlite_vec
backend: kv_default
```

View file

@ -17,15 +17,15 @@ Please refer to the sqlite-vec provider documentation.
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `db_path` | `<class 'str'>` | No | | Path to the SQLite database file |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Config for KV store backend (SQLite only for now) |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | Config for KV store backend (SQLite only for now) |
## Sample Configuration
```yaml
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/sqlite_vec.db
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/sqlite_vec_registry.db
namespace: vector_io::sqlite_vec
backend: kv_default
```
## Deprecation Notice

View file

@ -78,13 +78,13 @@ See [Chroma's documentation](https://docs.trychroma.com/docs/overview/introducti
| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `url` | `str \| None` | No | | |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Config for KV store backend |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | Config for KV store backend |
## Sample Configuration
```yaml
url: ${env.CHROMADB_URL}
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/chroma_remote_registry.db
namespace: vector_io::chroma_remote
backend: kv_default
```

View file

@ -408,7 +408,7 @@ For more details on TLS configuration, refer to the [TLS setup guide](https://mi
| `uri` | `<class 'str'>` | No | | The URI of the Milvus server |
| `token` | `str \| None` | No | | The token of the Milvus server |
| `consistency_level` | `<class 'str'>` | No | Strong | The consistency level of the Milvus server |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Config for KV store backend |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | Config for KV store backend |
| `config` | `dict` | No | `{}` | 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. |
:::note
@ -421,6 +421,6 @@ This configuration class accepts additional fields beyond those listed above. Yo
uri: ${env.MILVUS_ENDPOINT}
token: ${env.MILVUS_TOKEN}
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/milvus_remote_registry.db
namespace: vector_io::milvus_remote
backend: kv_default
```

View file

@ -218,7 +218,7 @@ See [PGVector's documentation](https://github.com/pgvector/pgvector) for more de
| `db` | `str \| None` | No | postgres | |
| `user` | `str \| None` | No | postgres | |
| `password` | `str \| None` | No | mysecretpassword | |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig, annotation=NoneType, required=False, default='sqlite', discriminator='type'` | No | | Config for KV store backend (SQLite only for now) |
| `kvstore` | `llama_stack.core.storage.datatypes.KVStoreReference \| None` | No | | Config for KV store backend (SQLite only for now) |
## Sample Configuration
@ -229,6 +229,6 @@ db: ${env.PGVECTOR_DB}
user: ${env.PGVECTOR_USER}
password: ${env.PGVECTOR_PASSWORD}
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/pgvector_registry.db
namespace: vector_io::pgvector
backend: kv_default
```

View file

@ -26,13 +26,13 @@ Please refer to the inline provider documentation.
| `prefix` | `str \| None` | No | | |
| `timeout` | `int \| None` | No | | |
| `host` | `str \| None` | No | | |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | |
| `kvstore` | `<class 'llama_stack.core.storage.datatypes.KVStoreReference'>` | No | | |
## Sample Configuration
```yaml
api_key: ${env.QDRANT_API_KEY:=}
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/qdrant_registry.db
namespace: vector_io::qdrant_remote
backend: kv_default
```

View file

@ -75,7 +75,7 @@ See [Weaviate's documentation](https://weaviate.io/developers/weaviate) for more
|-------|------|----------|---------|-------------|
| `weaviate_api_key` | `str \| None` | No | | The API key for the Weaviate instance |
| `weaviate_cluster_url` | `str \| None` | No | localhost:8080 | The URL of the Weaviate cluster |
| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig, annotation=NoneType, required=False, default='sqlite', discriminator='type'` | No | | Config for KV store backend (SQLite only for now) |
| `kvstore` | `llama_stack.core.storage.datatypes.KVStoreReference \| None` | No | | Config for KV store backend (SQLite only for now) |
## Sample Configuration
@ -83,6 +83,6 @@ See [Weaviate's documentation](https://weaviate.io/developers/weaviate) for more
weaviate_api_key: null
weaviate_cluster_url: ${env.WEAVIATE_CLUSTER_URL:=localhost:8080}
kvstore:
type: sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/weaviate_registry.db
namespace: vector_io::weaviate
backend: kv_default
```

View file

@ -40,6 +40,14 @@ from llama_stack.core.distribution import get_provider_registry
from llama_stack.core.external import load_external_apis
from llama_stack.core.resolver import InvalidProviderError
from llama_stack.core.stack import replace_env_vars
from llama_stack.core.storage.datatypes import (
InferenceStoreReference,
KVStoreReference,
SqliteKVStoreConfig,
SqliteSqlStoreConfig,
SqlStoreReference,
StorageConfig,
)
from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR, EXTERNAL_PROVIDERS_DIR
from llama_stack.core.utils.dynamic import instantiate_class_type
from llama_stack.core.utils.exec import formulate_run_args, run_command
@ -285,16 +293,40 @@ def _generate_run_config(
Generate a run.yaml template file for user to edit from a build.yaml file
"""
apis = list(build_config.distribution_spec.providers.keys())
distro_dir = f"~/.llama/distributions/{image_name}"
storage = StorageConfig(
backends={
"kv_default": SqliteKVStoreConfig(
db_path=f"${{env.SQLITE_STORE_DIR:={distro_dir}}}/kvstore.db",
),
"sql_default": SqliteSqlStoreConfig(
db_path=f"${{env.SQLITE_STORE_DIR:={distro_dir}}}/sql_store.db",
),
}
)
run_config = StackRunConfig(
container_image=(image_name if build_config.image_type == LlamaStackImageType.CONTAINER.value else None),
image_name=image_name,
apis=apis,
providers={},
storage=storage,
metadata_store=KVStoreReference(
backend="kv_default",
namespace="registry",
),
inference_store=InferenceStoreReference(
backend="sql_default",
table_name="inference_store",
),
conversations_store=SqlStoreReference(
backend="sql_default",
table_name="openai_conversations",
),
external_providers_dir=build_config.external_providers_dir
if build_config.external_providers_dir
else EXTERNAL_PROVIDERS_DIR,
)
# Persistence config defaults are handled by PersistenceConfig model validators
# build providers dict
provider_registry = get_provider_registry(build_config)
for api in apis:

View file

@ -55,10 +55,10 @@ class ConversationServiceImpl(Conversations):
self.deps = deps
self.policy = config.policy
# Use conversations store reference from storage config
conversations_ref = config.run_config.storage.conversations
# Use conversations store reference from run config
conversations_ref = config.run_config.conversations_store
if not conversations_ref:
raise ValueError("storage.conversations must be configured in run config")
raise ValueError("conversations_store must be configured in run config")
base_sql_store = sqlstore_impl(conversations_ref)
self.sql_store = AuthorizedSqlStore(base_sql_store, self.policy)

View file

@ -26,7 +26,13 @@ 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.core.storage.datatypes import (
InferenceStoreReference,
KVStoreReference,
SqlStoreReference,
StorageBackendType,
StorageConfig,
)
from llama_stack.providers.datatypes import Api, ProviderSpec
LLAMA_STACK_BUILD_CONFIG_VERSION = 2
@ -464,10 +470,19 @@ can be instantiated multiple times (with different configs) if necessary.
""",
)
storage: StorageConfig = Field(
description="""
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).
""",
description="Catalog of named storage backends available to the stack",
)
metadata_store: KVStoreReference | None = Field(
default=None,
description="Reference to the KV store backend used by the distribution registry (kv_* backend).",
)
inference_store: InferenceStoreReference | None = Field(
default=None,
description="Reference to the SQL store backend used by the inference API (sql_* backend).",
)
conversations_store: SqlStoreReference | None = Field(
default=None,
description="Reference to the SQL store backend used by the conversations API (sql_* backend).",
)
# registry of "resources" in the distribution
@ -507,6 +522,47 @@ throughout the Stack (both by its core as well as providers).
return Path(v)
return v
@model_validator(mode="after")
def validate_storage_references(self) -> "StackRunConfig":
backend_map = self.storage.backends if self.storage else {}
kv_backends = {
name
for name, cfg in backend_map.items()
if cfg.type
in {
StorageBackendType.KV_REDIS,
StorageBackendType.KV_SQLITE,
StorageBackendType.KV_POSTGRES,
StorageBackendType.KV_MONGODB,
}
}
sql_backends = {
name
for name, cfg in backend_map.items()
if cfg.type in {StorageBackendType.SQL_SQLITE, StorageBackendType.SQL_POSTGRES}
}
def _ensure_backend(reference, expected_set, store_name: str) -> None:
if reference is None:
return
backend_name = reference.backend
if backend_name not in backend_map:
raise ValueError(
f"{store_name} references unknown backend '{backend_name}'. "
f"Available backends: {sorted(backend_map)}"
)
if backend_name not in expected_set:
raise ValueError(
f"{store_name} references backend '{backend_name}' of type "
f"'{backend_map[backend_name].type.value}', but a backend of type "
f"{'kv_*' if expected_set is kv_backends else 'sql_*'} is required."
)
_ensure_backend(self.metadata_store, kv_backends, "metadata_store")
_ensure_backend(self.inference_store, sql_backends, "inference_store")
_ensure_backend(self.conversations_store, sql_backends, "conversations_store")
return self
class BuildConfig(BaseModel):
version: int = LLAMA_STACK_BUILD_CONFIG_VERSION

View file

@ -41,11 +41,10 @@ class PromptServiceImpl(Prompts):
async def initialize(self) -> None:
# 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,
)
metadata_ref = self.config.run_config.metadata_store
if not metadata_ref:
raise ValueError("metadata_store must be configured in run config")
prompts_ref = KVStoreReference(namespace="prompts", backend=metadata_ref.backend)
self.kvstore = await kvstore_impl(prompts_ref)
def _get_default_key(self, prompt_id: str) -> str:

View file

@ -80,9 +80,9 @@ async def get_auto_router_impl(
# TODO: move pass configs to routers instead
if api == Api.inference:
inference_ref = run_config.storage.inference
inference_ref = run_config.inference_store
if not inference_ref:
raise ValueError("storage.inference must be configured in run config")
raise ValueError("inference_store must be configured in run config")
inference_store = InferenceStore(
reference=inference_ref,

View file

@ -10,10 +10,10 @@ from datetime import UTC, datetime, timedelta
from starlette.types import ASGIApp, Receive, Scope, Send
from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendType
from llama_stack.log import get_logger
from llama_stack.providers.utils.kvstore.api import KVStore
from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
from llama_stack.providers.utils.kvstore.kvstore import kvstore_impl
from llama_stack.providers.utils.kvstore.kvstore import _KVSTORE_BACKENDS, kvstore_impl
logger = get_logger(name=__name__, category="core::server")
@ -33,7 +33,7 @@ class QuotaMiddleware:
def __init__(
self,
app: ASGIApp,
kv_config: KVStoreConfig,
kv_config: KVStoreReference,
anonymous_max_requests: int,
authenticated_max_requests: int,
window_seconds: int = 86400,
@ -45,15 +45,15 @@ class QuotaMiddleware:
self.authenticated_max_requests = authenticated_max_requests
self.window_seconds = window_seconds
if isinstance(self.kv_config, SqliteKVStoreConfig):
logger.warning(
"QuotaMiddleware: Using SQLite backend. Expiry/TTL is not enforced; cleanup is manual. "
f"window_seconds={self.window_seconds}"
)
async def _get_kv(self) -> KVStore:
if self.kv is None:
self.kv = await kvstore_impl(self.kv_config)
backend_config = _KVSTORE_BACKENDS.get(self.kv_config.backend)
if backend_config and backend_config.type == StorageBackendType.KV_SQLITE:
logger.warning(
"QuotaMiddleware: Using SQLite backend. Expiry/TTL is not enforced; cleanup is manual. "
f"window_seconds={self.window_seconds}"
)
return self.kv
async def __call__(self, scope: Scope, receive: Receive, send: Send):

View file

@ -368,7 +368,9 @@ class Stack:
logger.info(f"API recording enabled: mode={os.environ.get('LLAMA_STACK_TEST_INFERENCE_MODE')}")
_initialize_storage(self.run_config)
dist_registry, _ = await create_dist_registry(self.run_config.storage, self.run_config.image_name)
if not self.run_config.metadata_store:
raise ValueError("metadata_store must be configured with a kv_* backend")
dist_registry, _ = await create_dist_registry(self.run_config.metadata_store, self.run_config.image_name)
policy = self.run_config.server.auth.access_policy if self.run_config.server.auth else []
internal_impls = {}

View file

@ -27,10 +27,6 @@ class CommonConfig(BaseModel):
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):
@ -143,13 +139,6 @@ class MongoDBKVStoreConfig(CommonConfig):
}
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
@ -161,7 +150,7 @@ class SqlAlchemySqlStoreConfig(BaseModel):
return ["sqlalchemy[asyncio]"]
class SqliteSqlStoreConfig(SqlAlchemySqlStoreConfig, CommonSqlStoreConfig):
class SqliteSqlStoreConfig(SqlAlchemySqlStoreConfig):
type: Literal[StorageBackendType.SQL_SQLITE] = StorageBackendType.SQL_SQLITE
db_path: str = Field(
description="Database path, e.g. ~/.llama/distributions/ollama/sqlstore.db",
@ -219,9 +208,8 @@ class SqlStoreReference(BaseModel):
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,
backend: str = Field(
description="Name of backend from storage.backends",
)
@ -233,9 +221,8 @@ class KVStoreReference(BaseModel):
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,
backend: str = Field(
description="Name of backend from storage.backends",
)
@ -263,21 +250,11 @@ class InferenceStoreReference(SqlStoreReference):
)
class ResponsesStoreReference(InferenceStoreReference):
"""Responses store configuration with queue tuning."""
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)",
)

View file

@ -11,7 +11,7 @@ from typing import Protocol
import pydantic
from llama_stack.core.datatypes import RoutableObjectWithProvider
from llama_stack.core.storage.datatypes import KVStoreReference, StorageConfig
from llama_stack.core.storage.datatypes import KVStoreReference
from llama_stack.log import get_logger
from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
@ -190,17 +190,10 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry):
async def create_dist_registry(
storage: StorageConfig,
image_name: str,
metadata_store: KVStoreReference, image_name: str
) -> tuple[CachedDiskDistributionRegistry, KVStore]:
# instantiate kvstore for storing and retrieving distribution metadata
# 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_kvstore = await kvstore_impl(metadata_store)
dist_registry = CachedDiskDistributionRegistry(dist_kvstore)
await dist_registry.initialize()
return dist_registry, dist_kvstore

View file

@ -52,5 +52,6 @@ distribution_spec:
- provider_type: inline::reference
image_type: venv
additional_pip_packages:
- aiosqlite
- asyncpg
- sqlalchemy[asyncio]

View file

@ -95,24 +95,28 @@ providers:
config:
kvstore:
namespace: vector_io::faiss
backend: kv_default
- provider_id: sqlite-vec
provider_type: inline::sqlite-vec
config:
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/sqlite_vec.db
kvstore:
namespace: vector_io::sqlite_vec
backend: kv_default
- provider_id: ${env.MILVUS_URL:+milvus}
provider_type: inline::milvus
config:
db_path: ${env.MILVUS_DB_PATH:=~/.llama/distributions/ci-tests}/milvus.db
kvstore:
namespace: vector_io::milvus
backend: kv_default
- provider_id: ${env.CHROMADB_URL:+chromadb}
provider_type: remote::chromadb
config:
url: ${env.CHROMADB_URL:=}
kvstore:
namespace: vector_io::chroma_remote
backend: kv_default
- provider_id: ${env.PGVECTOR_DB:+pgvector}
provider_type: remote::pgvector
config:
@ -123,6 +127,7 @@ providers:
password: ${env.PGVECTOR_PASSWORD:=}
kvstore:
namespace: vector_io::pgvector
backend: kv_default
files:
- provider_id: meta-reference-files
provider_type: inline::localfs
@ -130,6 +135,7 @@ providers:
storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/ci-tests/files}
metadata_store:
table_name: files_metadata
backend: sql_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -144,8 +150,12 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
post_training:
- provider_id: torchtune-cpu
provider_type: inline::torchtune-cpu
@ -157,17 +167,20 @@ providers:
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -198,25 +211,26 @@ providers:
config:
kvstore:
namespace: batches
backend: kv_default
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models: []
shields:
- shield_id: llama-guard

View file

@ -28,4 +28,6 @@ distribution_spec:
- provider_type: remote::tavily-search
- provider_type: inline::rag-runtime
image_type: venv
additional_pip_packages: []
additional_pip_packages:
- aiosqlite
- sqlalchemy[asyncio]

View file

@ -28,6 +28,7 @@ providers:
url: ${env.CHROMADB_URL:=}
kvstore:
namespace: vector_io::chroma_remote
backend: kv_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -40,25 +41,32 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
eval:
- provider_id: meta-reference
provider_type: inline::meta-reference
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -83,23 +91,23 @@ providers:
provider_type: inline::rag-runtime
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models:
- metadata: {}
model_id: ${env.INFERENCE_MODEL}

View file

@ -24,6 +24,7 @@ providers:
url: ${env.CHROMADB_URL:=}
kvstore:
namespace: vector_io::chroma_remote
backend: kv_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -36,25 +37,32 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
eval:
- provider_id: meta-reference
provider_type: inline::meta-reference
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -79,23 +87,23 @@ providers:
provider_type: inline::rag-runtime
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models:
- metadata: {}
model_id: ${env.INFERENCE_MODEL}

View file

@ -27,4 +27,6 @@ distribution_spec:
- provider_type: inline::rag-runtime
- provider_type: remote::model-context-protocol
image_type: venv
additional_pip_packages: []
additional_pip_packages:
- aiosqlite
- sqlalchemy[asyncio]

View file

@ -39,6 +39,7 @@ providers:
config:
kvstore:
namespace: vector_io::faiss
backend: kv_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -51,25 +52,32 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
eval:
- provider_id: meta-reference
provider_type: inline::meta-reference
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -96,23 +104,23 @@ providers:
provider_type: remote::model-context-protocol
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models:
- metadata: {}
model_id: ${env.INFERENCE_MODEL}

View file

@ -29,6 +29,7 @@ providers:
config:
kvstore:
namespace: vector_io::faiss
backend: kv_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -41,25 +42,32 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
eval:
- provider_id: meta-reference
provider_type: inline::meta-reference
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -86,23 +94,23 @@ providers:
provider_type: remote::model-context-protocol
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models:
- metadata: {}
model_id: ${env.INFERENCE_MODEL}

View file

@ -24,4 +24,6 @@ distribution_spec:
files:
- provider_type: inline::localfs
image_type: venv
additional_pip_packages: []
additional_pip_packages:
- aiosqlite
- sqlalchemy[asyncio]

View file

@ -30,6 +30,7 @@ providers:
config:
kvstore:
namespace: vector_io::faiss
backend: kv_default
safety:
- provider_id: nvidia
provider_type: remote::nvidia
@ -43,8 +44,12 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
eval:
- provider_id: nvidia
provider_type: remote::nvidia
@ -64,6 +69,7 @@ providers:
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
- provider_id: nvidia
provider_type: remote::nvidia
config:
@ -84,25 +90,26 @@ providers:
storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/nvidia/files}
metadata_store:
table_name: files_metadata
backend: sql_default
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models:
- metadata: {}
model_id: ${env.INFERENCE_MODEL}

View file

@ -25,6 +25,7 @@ providers:
config:
kvstore:
namespace: vector_io::faiss
backend: kv_default
safety:
- provider_id: nvidia
provider_type: remote::nvidia
@ -38,8 +39,12 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
eval:
- provider_id: nvidia
provider_type: remote::nvidia
@ -74,25 +79,26 @@ providers:
storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/nvidia/files}
metadata_store:
table_name: files_metadata
backend: sql_default
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models: []
shields: []
vector_dbs: []

View file

@ -31,4 +31,6 @@ distribution_spec:
- provider_type: inline::rag-runtime
- provider_type: remote::model-context-protocol
image_type: venv
additional_pip_packages: []
additional_pip_packages:
- aiosqlite
- sqlalchemy[asyncio]

View file

@ -41,12 +41,14 @@ providers:
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/sqlite_vec.db
kvstore:
namespace: vector_io::sqlite_vec
backend: kv_default
- provider_id: ${env.ENABLE_CHROMADB:+chromadb}
provider_type: remote::chromadb
config:
url: ${env.CHROMADB_URL:=}
kvstore:
namespace: vector_io::chroma_remote
backend: kv_default
- provider_id: ${env.ENABLE_PGVECTOR:+pgvector}
provider_type: remote::pgvector
config:
@ -57,6 +59,7 @@ providers:
password: ${env.PGVECTOR_PASSWORD:=}
kvstore:
namespace: vector_io::pgvector
backend: kv_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -69,25 +72,32 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
eval:
- provider_id: meta-reference
provider_type: inline::meta-reference
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -114,23 +124,23 @@ providers:
provider_type: remote::model-context-protocol
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models:
- metadata: {}
model_id: gpt-4o

View file

@ -17,4 +17,7 @@ distribution_spec:
- provider_type: inline::rag-runtime
- provider_type: remote::model-context-protocol
image_type: venv
additional_pip_packages: []
additional_pip_packages:
- asyncpg
- psycopg2-binary
- sqlalchemy[asyncio]

View file

@ -91,7 +91,6 @@ def get_distribution_template() -> DistributionTemplate:
"embedding_dimension": 768,
},
)
postgres_config = PostgresSqlStoreConfig.sample_run_config()
return DistributionTemplate(
name=name,
distro_type="self_hosted",
@ -109,13 +108,11 @@ def get_distribution_template() -> DistributionTemplate:
default_models=default_models + [embedding_model],
default_tool_groups=default_tool_groups,
default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")],
metadata_store=PostgresKVStoreConfig.sample_run_config(),
inference_store=postgres_config,
storage_backends={
"default_kv_store": PostgresKVStoreConfig.sample_run_config(
"kv_default": PostgresKVStoreConfig.sample_run_config(
table_name="llamastack_kvstore",
),
"default_sql_store": PostgresSqlStoreConfig.sample_run_config(),
"sql_default": PostgresSqlStoreConfig.sample_run_config(),
},
),
},

View file

@ -24,6 +24,7 @@ providers:
url: ${env.CHROMADB_URL:=}
kvstore:
namespace: vector_io::chroma_remote
backend: kv_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -36,8 +37,12 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
tool_runtime:
- provider_id: brave-search
provider_type: remote::brave-search
@ -55,7 +60,7 @@ providers:
provider_type: remote::model-context-protocol
storage:
backends:
default_kv_store:
kv_default:
type: kv_postgres
host: ${env.POSTGRES_HOST:=localhost}
port: ${env.POSTGRES_PORT:=5432}
@ -63,24 +68,24 @@ storage:
user: ${env.POSTGRES_USER:=llamastack}
password: ${env.POSTGRES_PASSWORD:=llamastack}
table_name: ${env.POSTGRES_TABLE_NAME:=llamastack_kvstore}
default_sql_store:
sql_default:
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}
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models:
- metadata: {}
model_id: ${env.INFERENCE_MODEL}

View file

@ -53,5 +53,6 @@ distribution_spec:
- provider_type: inline::reference
image_type: venv
additional_pip_packages:
- aiosqlite
- asyncpg
- sqlalchemy[asyncio]

View file

@ -95,24 +95,28 @@ providers:
config:
kvstore:
namespace: vector_io::faiss
backend: kv_default
- provider_id: sqlite-vec
provider_type: inline::sqlite-vec
config:
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/sqlite_vec.db
kvstore:
namespace: vector_io::sqlite_vec
backend: kv_default
- provider_id: ${env.MILVUS_URL:+milvus}
provider_type: inline::milvus
config:
db_path: ${env.MILVUS_DB_PATH:=~/.llama/distributions/starter-gpu}/milvus.db
kvstore:
namespace: vector_io::milvus
backend: kv_default
- provider_id: ${env.CHROMADB_URL:+chromadb}
provider_type: remote::chromadb
config:
url: ${env.CHROMADB_URL:=}
kvstore:
namespace: vector_io::chroma_remote
backend: kv_default
- provider_id: ${env.PGVECTOR_DB:+pgvector}
provider_type: remote::pgvector
config:
@ -123,6 +127,7 @@ providers:
password: ${env.PGVECTOR_PASSWORD:=}
kvstore:
namespace: vector_io::pgvector
backend: kv_default
files:
- provider_id: meta-reference-files
provider_type: inline::localfs
@ -130,6 +135,7 @@ providers:
storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter-gpu/files}
metadata_store:
table_name: files_metadata
backend: sql_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -144,8 +150,12 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
post_training:
- provider_id: huggingface-gpu
provider_type: inline::huggingface-gpu
@ -160,17 +170,20 @@ providers:
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -201,25 +214,26 @@ providers:
config:
kvstore:
namespace: batches
backend: kv_default
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models: []
shields:
- shield_id: llama-guard

View file

@ -53,5 +53,6 @@ distribution_spec:
- provider_type: inline::reference
image_type: venv
additional_pip_packages:
- aiosqlite
- asyncpg
- sqlalchemy[asyncio]

View file

@ -95,24 +95,28 @@ providers:
config:
kvstore:
namespace: vector_io::faiss
backend: kv_default
- provider_id: sqlite-vec
provider_type: inline::sqlite-vec
config:
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/sqlite_vec.db
kvstore:
namespace: vector_io::sqlite_vec
backend: kv_default
- provider_id: ${env.MILVUS_URL:+milvus}
provider_type: inline::milvus
config:
db_path: ${env.MILVUS_DB_PATH:=~/.llama/distributions/starter}/milvus.db
kvstore:
namespace: vector_io::milvus
backend: kv_default
- provider_id: ${env.CHROMADB_URL:+chromadb}
provider_type: remote::chromadb
config:
url: ${env.CHROMADB_URL:=}
kvstore:
namespace: vector_io::chroma_remote
backend: kv_default
- provider_id: ${env.PGVECTOR_DB:+pgvector}
provider_type: remote::pgvector
config:
@ -123,6 +127,7 @@ providers:
password: ${env.PGVECTOR_PASSWORD:=}
kvstore:
namespace: vector_io::pgvector
backend: kv_default
files:
- provider_id: meta-reference-files
provider_type: inline::localfs
@ -130,6 +135,7 @@ providers:
storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter/files}
metadata_store:
table_name: files_metadata
backend: sql_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -144,8 +150,12 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
post_training:
- provider_id: torchtune-cpu
provider_type: inline::torchtune-cpu
@ -157,17 +167,20 @@ providers:
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -198,25 +211,26 @@ providers:
config:
kvstore:
namespace: batches
backend: kv_default
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models: []
shields:
- shield_id: llama-guard

View file

@ -231,31 +231,29 @@ class RunConfigSettings(BaseModel):
apis = sorted(providers.keys())
storage_backends = self.storage_backends or {
"default_kv_store": SqliteKVStoreConfig.sample_run_config(
"kv_default": SqliteKVStoreConfig.sample_run_config(
__distro_dir__=f"~/.llama/distributions/{name}",
db_name="kvstore.db",
),
"default_sql_store": SqliteSqlStoreConfig.sample_run_config(
"sql_default": SqliteSqlStoreConfig.sample_run_config(
__distro_dir__=f"~/.llama/distributions/{name}",
db_name="sql_store.db",
),
}
storage_config = dict(
backends=storage_backends,
metadata=KVStoreReference(
backend="default_kv_store",
namespace="registry",
).model_dump(exclude_none=True),
inference=InferenceStoreReference(
backend="default_sql_store",
table_name="inference_store",
).model_dump(exclude_none=True),
conversations=SqlStoreReference(
backend="default_sql_store",
table_name="openai_conversations",
).model_dump(exclude_none=True),
)
storage_config = dict(backends=storage_backends)
metadata_store = KVStoreReference(
backend="kv_default",
namespace="registry",
).model_dump(exclude_none=True)
inference_store = InferenceStoreReference(
backend="sql_default",
table_name="inference_store",
).model_dump(exclude_none=True)
conversations_store = SqlStoreReference(
backend="sql_default",
table_name="openai_conversations",
).model_dump(exclude_none=True)
# Return a dict that matches StackRunConfig structure
return {
@ -265,6 +263,9 @@ class RunConfigSettings(BaseModel):
"apis": apis,
"providers": provider_configs,
"storage": storage_config,
"metadata_store": metadata_store,
"inference_store": inference_store,
"conversations_store": conversations_store,
"models": [m.model_dump(exclude_none=True) for m in (self.default_models or [])],
"shields": [s.model_dump(exclude_none=True) for s in (self.default_shields or [])],
"vector_dbs": [],
@ -314,11 +315,15 @@ class DistributionTemplate(BaseModel):
# We should have a better way to do this by formalizing the concept of "internal" APIs
# and providers, with a way to specify dependencies for them.
if run_config_.get("inference_store"):
additional_pip_packages.extend(get_sql_pip_packages(run_config_["inference_store"]))
if run_config_.get("metadata_store"):
additional_pip_packages.extend(get_kv_pip_packages(run_config_["metadata_store"]))
storage_cfg = run_config_.get("storage", {})
for backend_cfg in storage_cfg.get("backends", {}).values():
store_type = backend_cfg.get("type")
if not store_type:
continue
if str(store_type).startswith("kv_"):
additional_pip_packages.extend(get_kv_pip_packages(backend_cfg))
elif str(store_type).startswith("sql_"):
additional_pip_packages.extend(get_sql_pip_packages(backend_cfg))
if self.additional_pip_packages:
additional_pip_packages.extend(self.additional_pip_packages)

View file

@ -28,4 +28,6 @@ distribution_spec:
files:
- provider_type: inline::localfs
image_type: venv
additional_pip_packages: []
additional_pip_packages:
- aiosqlite
- sqlalchemy[asyncio]

View file

@ -24,6 +24,7 @@ providers:
config:
kvstore:
namespace: vector_io::faiss
backend: kv_default
safety:
- provider_id: llama-guard
provider_type: inline::llama-guard
@ -36,25 +37,32 @@ providers:
persistence:
agent_state:
namespace: agents
backend: kv_default
responses:
table_name: responses
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
eval:
- provider_id: meta-reference
provider_type: inline::meta-reference
config:
kvstore:
namespace: eval
backend: kv_default
datasetio:
- provider_id: huggingface
provider_type: remote::huggingface
config:
kvstore:
namespace: datasetio::huggingface
backend: kv_default
- provider_id: localfs
provider_type: inline::localfs
config:
kvstore:
namespace: datasetio::localfs
backend: kv_default
scoring:
- provider_id: basic
provider_type: inline::basic
@ -86,25 +94,26 @@ providers:
storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/watsonx/files}
metadata_store:
table_name: files_metadata
backend: sql_default
storage:
backends:
default_kv_store:
kv_default:
type: kv_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/kvstore.db
default_sql_store:
sql_default:
type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/sql_store.db
metadata:
namespace: registry
backend: default_kv_store
inference:
table_name: inference_store
backend: default_sql_store
max_write_queue_size: 10000
num_writers: 4
conversations:
table_name: openai_conversations
backend: default_sql_store
metadata_store:
namespace: registry
backend: kv_default
inference_store:
table_name: inference_store
backend: sql_default
max_write_queue_size: 10000
num_writers: 4
conversations_store:
table_name: openai_conversations
backend: sql_default
models: []
shields: []
vector_dbs: []

View file

@ -8,14 +8,14 @@ from typing import Any
from pydantic import BaseModel
from llama_stack.core.storage.datatypes import KVStoreReference, SqlStoreReference
from llama_stack.core.storage.datatypes import KVStoreReference, ResponsesStoreReference
class AgentPersistenceConfig(BaseModel):
"""Nested persistence configuration for agents."""
agent_state: KVStoreReference
responses: SqlStoreReference
responses: ResponsesStoreReference
class MetaReferenceAgentsImplConfig(BaseModel):
@ -26,9 +26,11 @@ class MetaReferenceAgentsImplConfig(BaseModel):
return {
"persistence": {
"agent_state": KVStoreReference(
backend="kv_default",
namespace="agents",
).model_dump(exclude_none=True),
"responses": SqlStoreReference(
"responses": ResponsesStoreReference(
backend="sql_default",
table_name="responses",
).model_dump(exclude_none=True),
}

View file

@ -34,6 +34,7 @@ class ReferenceBatchesImplConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str) -> dict:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="batches",
).model_dump(exclude_none=True),
}

View file

@ -17,6 +17,7 @@ class LocalFSDatasetIOConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="datasetio::localfs",
).model_dump(exclude_none=True)
}

View file

@ -17,6 +17,7 @@ class MetaReferenceEvalConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="eval",
).model_dump(exclude_none=True)
}

View file

@ -25,6 +25,7 @@ class LocalfsFilesImplConfig(BaseModel):
return {
"storage_dir": "${env.FILES_STORAGE_DIR:=" + __distro_dir__ + "/files}",
"metadata_store": SqlStoreReference(
backend="sql_default",
table_name="files_metadata",
).model_dump(exclude_none=True),
}

View file

@ -24,6 +24,7 @@ class ChromaVectorIOConfig(BaseModel):
return {
"db_path": db_path,
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::chroma",
).model_dump(exclude_none=True),
}

View file

@ -20,6 +20,7 @@ class FaissVectorIOConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::faiss",
).model_dump(exclude_none=True)
}

View file

@ -23,6 +23,7 @@ class MilvusVectorIOConfig(BaseModel):
return {
"db_path": "${env.MILVUS_DB_PATH:=" + __distro_dir__ + "}/" + "milvus.db",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::milvus",
).model_dump(exclude_none=True),
}

View file

@ -23,6 +23,7 @@ class QdrantVectorIOConfig(BaseModel):
return {
"path": "${env.QDRANT_PATH:=~/.llama/" + __distro_dir__ + "}/" + "qdrant.db",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::qdrant",
).model_dump(exclude_none=True),
}

View file

@ -20,6 +20,7 @@ class SQLiteVectorIOConfig(BaseModel):
return {
"db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + "sqlite_vec.db",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::sqlite_vec",
).model_dump(exclude_none=True),
}

View file

@ -17,6 +17,7 @@ class HuggingfaceDatasetIOConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="datasetio::huggingface",
).model_dump(exclude_none=True)
}

View file

@ -36,6 +36,7 @@ class S3FilesImplConfig(BaseModel):
"endpoint_url": "${env.S3_ENDPOINT_URL:=}",
"auto_create_bucket": "${env.S3_AUTO_CREATE_BUCKET:=false}",
"metadata_store": SqlStoreReference(
backend="sql_default",
table_name="s3_files_metadata",
).model_dump(exclude_none=True),
}

View file

@ -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.persistence), self.policy)
self._sql_store = AuthorizedSqlStore(sqlstore_impl(self._config.metadata_store), self.policy)
await self._sql_store.create_table(
"openai_files",
{

View file

@ -22,6 +22,7 @@ class ChromaVectorIOConfig(BaseModel):
return {
"url": url,
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::chroma_remote",
).model_dump(exclude_none=True),
}

View file

@ -29,6 +29,7 @@ class MilvusVectorIOConfig(BaseModel):
"uri": "${env.MILVUS_ENDPOINT}",
"token": "${env.MILVUS_TOKEN}",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::milvus_remote",
).model_dump(exclude_none=True),
}

View file

@ -41,6 +41,7 @@ class PGVectorVectorIOConfig(BaseModel):
"user": user,
"password": password,
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::pgvector",
).model_dump(exclude_none=True),
}

View file

@ -31,6 +31,7 @@ class QdrantVectorIOConfig(BaseModel):
return {
"api_key": "${env.QDRANT_API_KEY:=}",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::qdrant_remote",
).model_dump(exclude_none=True),
}

View file

@ -30,6 +30,7 @@ class WeaviateVectorIOConfig(BaseModel):
"weaviate_api_key": None,
"weaviate_cluster_url": "${env.WEAVIATE_CLUSTER_URL:=localhost:8080}",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::weaviate",
).model_dump(exclude_none=True),
}

View file

@ -48,12 +48,13 @@ class InferenceStore:
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
backend_name = self.reference.backend
backend_config = _SQLSTORE_BACKENDS.get(backend_name)
if backend_config is None:
raise ValueError(
f"Unregistered SQL backend '{backend_name}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
)
self.enable_write_queue = backend_config.type != StorageBackendType.SQL_SQLITE
await self.sql_store.create_table(
"chat_completions",
{

View file

@ -1,3 +1,9 @@
# 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.
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the terms described in the LICENSE file in
@ -47,36 +53,19 @@ class InmemoryKVStoreImpl(KVStore):
_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_name = reference.backend
backend_config = _KVSTORE_BACKENDS.get(backend_name)
if backend_config is None:

View file

@ -18,13 +18,13 @@ from llama_stack.apis.agents.openai_responses import (
OpenAIResponseObjectWithInput,
)
from llama_stack.apis.inference import OpenAIMessageParam
from llama_stack.core.datatypes import AccessRule, ResponsesStoreConfig
from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR
from llama_stack.core.datatypes import AccessRule
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqlStoreReference, StorageBackendType
from llama_stack.log import get_logger
from ..sqlstore.api import ColumnDefinition, ColumnType
from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
from ..sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig, SqlStoreType, sqlstore_impl
from ..sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
logger = get_logger(name=__name__, category="openai_responses")
@ -45,39 +45,38 @@ class _OpenAIResponseObjectWithInputAndMessages(OpenAIResponseObjectWithInput):
class ResponsesStore:
def __init__(
self,
config: ResponsesStoreConfig | SqlStoreConfig,
reference: ResponsesStoreReference | SqlStoreReference,
policy: list[AccessRule],
):
# Handle backward compatibility
if not isinstance(config, ResponsesStoreConfig):
# Legacy: SqlStoreConfig passed directly as config
config = ResponsesStoreConfig(
sql_store_config=config,
)
if isinstance(reference, ResponsesStoreReference):
self.reference = reference
else:
self.reference = ResponsesStoreReference(**reference.model_dump())
self.config = config
self.sql_store_config = config.sql_store_config
if not self.sql_store_config:
self.sql_store_config = SqliteSqlStoreConfig(
db_path=(RUNTIME_BASE_DIR / "sqlstore.db").as_posix(),
)
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
self.sql_store = None
self.enable_write_queue = True
# Async write queue and worker control
self._queue: (
asyncio.Queue[tuple[OpenAIResponseObject, list[OpenAIResponseInput], 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 = self.reference.max_write_queue_size
self._num_writers: int = max(1, self.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)
backend_config = _SQLSTORE_BACKENDS.get(self.reference.backend)
if backend_config is None:
raise ValueError(
f"Unregistered SQL backend '{self.reference.backend}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
)
if backend_config.type == StorageBackendType.SQL_SQLITE:
self.enable_write_queue = False
await self.sql_store.create_table(
"openai_responses",
{

View file

@ -12,10 +12,10 @@ from llama_stack.core.access_control.conditions import ProtectedResource
from llama_stack.core.access_control.datatypes import AccessRule, Action, Scope
from llama_stack.core.datatypes import User
from llama_stack.core.request_headers import get_authenticated_user
from llama_stack.core.storage.datatypes import StorageBackendType
from llama_stack.log import get_logger
from .api import ColumnDefinition, ColumnType, PaginatedResponse, SqlStore
from .sqlstore import StorageBackendType
logger = get_logger(name=__name__, category="providers::utils")

View file

@ -26,10 +26,10 @@ from sqlalchemy.ext.asyncio.engine import AsyncEngine
from sqlalchemy.sql.elements import ColumnElement
from llama_stack.apis.common.responses import PaginatedResponse
from llama_stack.core.storage.datatypes import SqlAlchemySqlStoreConfig
from llama_stack.log import get_logger
from .api import ColumnDefinition, ColumnType, SqlStore
from .sqlstore import SqlAlchemySqlStoreConfig
logger = get_logger(name=__name__, category="providers::utils")

View file

@ -4,7 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Annotated
from typing import Annotated, cast
from pydantic import Field
@ -21,7 +21,6 @@ from .api import SqlStore
sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"]
_SQLSTORE_BACKENDS: dict[str, StorageBackendConfig] = {}
_SQLSTORE_DEFAULT_BACKEND: str | None = None
SqlStoreConfig = Annotated[
@ -45,19 +44,18 @@ def get_pip_packages(store_config: dict | SqlStoreConfig) -> list[str]:
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_name = reference.backend
backend_config = _SQLSTORE_BACKENDS.get(backend_name)
if backend_config.type in [StorageBackendType.SQL_SQLITE, StorageBackendType.SQL_POSTGRES]:
if backend_config is None:
raise ValueError(
f"Unknown SQL store backend '{backend_name}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
)
if isinstance(backend_config, SqliteSqlStoreConfig | PostgresSqlStoreConfig):
from .sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
config = backend_config.model_copy()
config.table_name = reference.table_name
config = cast(SqliteSqlStoreConfig | PostgresSqlStoreConfig, backend_config).model_copy()
return SqlAlchemySqlStoreImpl(config)
else:
raise ValueError(f"Unknown sqlstore type {backend_config.type}")
@ -67,18 +65,6 @@ def register_sqlstore_backends(backends: dict[str, StorageBackendConfig]) -> Non
"""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

View file

@ -12,9 +12,15 @@ import pytest
from llama_stack.core.access_control.access_control import default_policy
from llama_stack.core.datatypes import User
from llama_stack.core.storage.datatypes import SqlStoreReference
from llama_stack.providers.utils.sqlstore.api import ColumnType
from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
from llama_stack.providers.utils.sqlstore.sqlstore import PostgresSqlStoreConfig, SqliteSqlStoreConfig, sqlstore_impl
from llama_stack.providers.utils.sqlstore.sqlstore import (
PostgresSqlStoreConfig,
SqliteSqlStoreConfig,
register_sqlstore_backends,
sqlstore_impl,
)
def get_postgres_config():
@ -55,8 +61,9 @@ def authorized_store(backend_config):
config_func = backend_config
config = config_func()
base_sqlstore = sqlstore_impl(config)
backend_name = f"sql_{type(config).__name__.lower()}"
register_sqlstore_backends({backend_name: config})
base_sqlstore = sqlstore_impl(SqlStoreReference(backend=backend_name, table_name="authorized_store"))
authorized_store = AuthorizedSqlStore(base_sqlstore, default_policy())
yield authorized_store

View file

@ -8,7 +8,9 @@ import yaml
from llama_stack.core.datatypes import StackRunConfig
from llama_stack.core.storage.datatypes import (
PostgresKVStoreConfig,
PostgresSqlStoreConfig,
SqliteKVStoreConfig,
SqliteSqlStoreConfig,
)
@ -20,21 +22,26 @@ def test_starter_distribution_config_loads_and_resolves():
config = StackRunConfig(**config_dict)
# Config should have storage with default backend
# Config should have named backends and explicit store references
assert config.storage is not None
assert "default" in config.storage.backends
assert isinstance(config.storage.backends["default"], SqliteSqlStoreConfig)
assert "kv_default" in config.storage.backends
assert "sql_default" in config.storage.backends
assert isinstance(config.storage.backends["kv_default"], SqliteKVStoreConfig)
assert isinstance(config.storage.backends["sql_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.metadata_store is not None
assert config.metadata_store.backend == "kv_default"
assert config.metadata_store.namespace == "registry"
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
assert config.inference_store is not None
assert config.inference_store.backend == "sql_default"
assert config.inference_store.table_name == "inference_store"
assert config.inference_store.max_write_queue_size > 0
assert config.inference_store.num_writers > 0
assert config.conversations_store is not None
assert config.conversations_store.backend == "sql_default"
assert config.conversations_store.table_name == "openai_conversations"
def test_postgres_demo_distribution_config_loads():
@ -46,17 +53,15 @@ def test_postgres_demo_distribution_config_loads():
# 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 "kv_default" in config.storage.backends
assert "sql_default" in config.storage.backends
postgres_backend = config.storage.backends["sql_default"]
assert isinstance(postgres_backend, PostgresSqlStoreConfig)
assert postgres_backend.host == "${env.POSTGRES_HOST:=localhost}"
kv_backend = config.storage.backends["kv_default"]
assert isinstance(kv_backend, PostgresKVStoreConfig)
# Stores target the Postgres backends explicitly
assert config.metadata_store.backend == "kv_default"
assert config.inference_store.backend == "sql_default"

View file

@ -1,82 +0,0 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
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)

View file

@ -0,0 +1,77 @@
# 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.
"""Unit tests for storage backend/reference validation."""
import pytest
from pydantic import ValidationError
from llama_stack.core.datatypes import (
LLAMA_STACK_RUN_CONFIG_VERSION,
StackRunConfig,
)
from llama_stack.core.storage.datatypes import (
InferenceStoreReference,
KVStoreReference,
SqliteKVStoreConfig,
SqliteSqlStoreConfig,
SqlStoreReference,
StorageConfig,
)
def _base_run_config(**overrides):
storage = overrides.pop(
"storage",
StorageConfig(
backends={
"kv_default": SqliteKVStoreConfig(db_path="/tmp/kv.db"),
"sql_default": SqliteSqlStoreConfig(db_path="/tmp/sql.db"),
}
),
)
return StackRunConfig(
version=LLAMA_STACK_RUN_CONFIG_VERSION,
image_name="test-distro",
apis=[],
providers={},
storage=storage,
metadata_store=overrides.pop(
"metadata_store",
KVStoreReference(backend="kv_default", namespace="registry"),
),
inference_store=overrides.pop(
"inference_store",
InferenceStoreReference(backend="sql_default", table_name="inference"),
),
conversations_store=overrides.pop(
"conversations_store",
SqlStoreReference(backend="sql_default", table_name="conversations"),
),
**overrides,
)
def test_references_require_known_backend():
with pytest.raises(ValidationError, match="unknown backend 'missing'"):
_base_run_config(metadata_store=KVStoreReference(backend="missing", namespace="registry"))
def test_references_must_match_backend_family():
with pytest.raises(ValidationError, match="kv_.* is required"):
_base_run_config(metadata_store=KVStoreReference(backend="sql_default", namespace="registry"))
with pytest.raises(ValidationError, match="sql_.* is required"):
_base_run_config(
inference_store=InferenceStoreReference(backend="kv_default", table_name="inference"),
)
def test_valid_configuration_passes_validation():
config = _base_run_config()
assert config.metadata_store.backend == "kv_default"
assert config.inference_store.backend == "sql_default"
assert config.conversations_store.backend == "sql_default"

View file

@ -13,6 +13,14 @@ from pydantic import BaseModel, Field, ValidationError
from llama_stack.core.datatypes import Api, Provider, StackRunConfig
from llama_stack.core.distribution import INTERNAL_APIS, get_provider_registry, providable_apis
from llama_stack.core.storage.datatypes import (
InferenceStoreReference,
KVStoreReference,
SqliteKVStoreConfig,
SqliteSqlStoreConfig,
SqlStoreReference,
StorageConfig,
)
from llama_stack.providers.datatypes import ProviderSpec
@ -29,6 +37,42 @@ class SampleConfig(BaseModel):
}
def _default_storage() -> StorageConfig:
return StorageConfig(
backends={
"kv_default": SqliteKVStoreConfig(db_path=":memory:"),
"sql_default": SqliteSqlStoreConfig(db_path=":memory:"),
}
)
def make_stack_config(**overrides) -> StackRunConfig:
storage = overrides.pop("storage", _default_storage())
metadata_store = overrides.pop(
"metadata_store",
KVStoreReference(backend="kv_default", namespace="registry"),
)
inference_store = overrides.pop(
"inference_store",
InferenceStoreReference(backend="sql_default", table_name="inference_store"),
)
conversations_store = overrides.pop(
"conversations_store",
SqlStoreReference(backend="sql_default", table_name="conversations"),
)
defaults = dict(
image_name="test_image",
apis=[],
providers={},
storage=storage,
metadata_store=metadata_store,
inference_store=inference_store,
conversations_store=conversations_store,
)
defaults.update(overrides)
return make_stack_config(**defaults)
@pytest.fixture
def mock_providers():
"""Mock the available_providers function to return test providers."""
@ -47,8 +91,8 @@ def mock_providers():
@pytest.fixture
def base_config(tmp_path):
"""Create a base StackRunConfig with common settings."""
return StackRunConfig(
image_name="test_image",
return make_stack_config(
apis=["inference"],
providers={
"inference": [
Provider(
@ -222,8 +266,8 @@ class TestProviderRegistry:
def test_missing_directory(self, mock_providers):
"""Test handling of missing external providers directory."""
config = StackRunConfig(
image_name="test_image",
config = make_stack_config(
apis=["inference"],
providers={
"inference": [
Provider(
@ -278,7 +322,6 @@ pip_packages:
"""Test loading an external provider from a module (success path)."""
from types import SimpleNamespace
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.providers.datatypes import Api, ProviderSpec
# Simulate a provider module with get_provider_spec
@ -293,7 +336,7 @@ pip_packages:
import_module_side_effect = make_import_module_side_effect(external_module=fake_module)
with patch("importlib.import_module", side_effect=import_module_side_effect) as mock_import:
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -317,12 +360,11 @@ pip_packages:
def test_external_provider_from_module_not_found(self, mock_providers):
"""Test handling ModuleNotFoundError for missing provider module."""
from llama_stack.core.datatypes import Provider, StackRunConfig
import_module_side_effect = make_import_module_side_effect(raise_for_external=True)
with patch("importlib.import_module", side_effect=import_module_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -341,12 +383,11 @@ pip_packages:
def test_external_provider_from_module_missing_get_provider_spec(self, mock_providers):
"""Test handling missing get_provider_spec in provider module (should raise ValueError)."""
from llama_stack.core.datatypes import Provider, StackRunConfig
import_module_side_effect = make_import_module_side_effect(missing_get_provider_spec=True)
with patch("importlib.import_module", side_effect=import_module_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -399,13 +440,12 @@ class TestGetExternalProvidersFromModule:
def test_stackrunconfig_provider_without_module(self, mock_providers):
"""Test that providers without module attribute are skipped."""
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
import_module_side_effect = make_import_module_side_effect()
with patch("importlib.import_module", side_effect=import_module_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -426,7 +466,6 @@ class TestGetExternalProvidersFromModule:
"""Test provider with module containing version spec (e.g., package==1.0.0)."""
from types import SimpleNamespace
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
from llama_stack.providers.datatypes import ProviderSpec
@ -444,7 +483,7 @@ class TestGetExternalProvidersFromModule:
raise ModuleNotFoundError(name)
with patch("importlib.import_module", side_effect=import_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -564,7 +603,6 @@ class TestGetExternalProvidersFromModule:
"""Test when get_provider_spec returns a list of specs."""
from types import SimpleNamespace
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
from llama_stack.providers.datatypes import ProviderSpec
@ -589,7 +627,7 @@ class TestGetExternalProvidersFromModule:
raise ModuleNotFoundError(name)
with patch("importlib.import_module", side_effect=import_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -613,7 +651,6 @@ class TestGetExternalProvidersFromModule:
"""Test that list return filters specs by provider_type."""
from types import SimpleNamespace
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
from llama_stack.providers.datatypes import ProviderSpec
@ -638,7 +675,7 @@ class TestGetExternalProvidersFromModule:
raise ModuleNotFoundError(name)
with patch("importlib.import_module", side_effect=import_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -662,7 +699,6 @@ class TestGetExternalProvidersFromModule:
"""Test that list return adds multiple different provider_types when config requests them."""
from types import SimpleNamespace
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
from llama_stack.providers.datatypes import ProviderSpec
@ -688,7 +724,7 @@ class TestGetExternalProvidersFromModule:
raise ModuleNotFoundError(name)
with patch("importlib.import_module", side_effect=import_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -718,7 +754,6 @@ class TestGetExternalProvidersFromModule:
def test_module_not_found_raises_value_error(self, mock_providers):
"""Test that ModuleNotFoundError raises ValueError with helpful message."""
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
def import_side_effect(name):
@ -727,7 +762,7 @@ class TestGetExternalProvidersFromModule:
raise ModuleNotFoundError(name)
with patch("importlib.import_module", side_effect=import_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -751,7 +786,6 @@ class TestGetExternalProvidersFromModule:
"""Test that generic exceptions are properly raised."""
from types import SimpleNamespace
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
def bad_spec():
@ -765,7 +799,7 @@ class TestGetExternalProvidersFromModule:
raise ModuleNotFoundError(name)
with patch("importlib.import_module", side_effect=import_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [
@ -787,10 +821,9 @@ class TestGetExternalProvidersFromModule:
def test_empty_provider_list(self, mock_providers):
"""Test with empty provider list."""
from llama_stack.core.datatypes import StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={},
)
@ -805,7 +838,6 @@ class TestGetExternalProvidersFromModule:
"""Test multiple APIs with providers."""
from types import SimpleNamespace
from llama_stack.core.datatypes import Provider, StackRunConfig
from llama_stack.core.distribution import get_external_providers_from_module
from llama_stack.providers.datatypes import ProviderSpec
@ -830,7 +862,7 @@ class TestGetExternalProvidersFromModule:
raise ModuleNotFoundError(name)
with patch("importlib.import_module", side_effect=import_side_effect):
config = StackRunConfig(
config = make_stack_config(
image_name="test_image",
providers={
"inference": [

View file

@ -11,11 +11,12 @@ from llama_stack.apis.common.errors import ResourceNotFoundError
from llama_stack.apis.common.responses import Order
from llama_stack.apis.files import OpenAIFilePurpose
from llama_stack.core.access_control.access_control import default_policy
from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
from llama_stack.providers.inline.files.localfs import (
LocalfsFilesImpl,
LocalfsFilesImplConfig,
)
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
class MockUploadFile:
@ -36,8 +37,11 @@ async def files_provider(tmp_path):
storage_dir = tmp_path / "files"
db_path = tmp_path / "files_metadata.db"
backend_name = "sql_localfs_test"
register_sqlstore_backends({backend_name: SqliteSqlStoreConfig(db_path=db_path.as_posix())})
config = LocalfsFilesImplConfig(
storage_dir=storage_dir.as_posix(), metadata_store=SqliteSqlStoreConfig(db_path=db_path.as_posix())
storage_dir=storage_dir.as_posix(),
metadata_store=SqlStoreReference(backend=backend_name, table_name="files_metadata"),
)
provider = LocalfsFilesImpl(config, default_policy())

View file

@ -9,7 +9,15 @@ import random
import pytest
from llama_stack.core.prompts.prompts import PromptServiceConfig, PromptServiceImpl
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
from llama_stack.core.storage.datatypes import (
InferenceStoreReference,
KVStoreReference,
SqliteKVStoreConfig,
SqliteSqlStoreConfig,
SqlStoreReference,
StorageConfig,
)
from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends
@pytest.fixture
@ -19,12 +27,26 @@ async def temp_prompt_store(tmp_path_factory):
db_path = str(temp_dir / f"{unique_id}.db")
from llama_stack.core.datatypes import StackRunConfig
from llama_stack.providers.utils.kvstore import kvstore_impl
mock_run_config = StackRunConfig(image_name="test-distribution", apis=[], providers={})
storage = StorageConfig(
backends={
"kv_test": SqliteKVStoreConfig(db_path=db_path),
"sql_test": SqliteSqlStoreConfig(db_path=str(temp_dir / f"{unique_id}_sql.db")),
}
)
mock_run_config = StackRunConfig(
image_name="test-distribution",
apis=[],
providers={},
storage=storage,
metadata_store=KVStoreReference(backend="kv_test", namespace="registry"),
inference_store=InferenceStoreReference(backend="sql_test", table_name="inference"),
conversations_store=SqlStoreReference(backend="sql_test", table_name="conversations"),
)
config = PromptServiceConfig(run_config=mock_run_config)
store = PromptServiceImpl(config, deps={})
store.kvstore = await kvstore_impl(SqliteKVStoreConfig(db_path=db_path))
register_kvstore_backends({"kv_test": storage.backends["kv_test"]})
store.kvstore = await kvstore_impl(KVStoreReference(backend="kv_test", namespace="prompts"))
yield store

View file

@ -42,7 +42,7 @@ from llama_stack.apis.inference import (
)
from llama_stack.apis.tools.tools import ListToolDefsResponse, ToolDef, ToolGroups, ToolInvocationResult, ToolRuntime
from llama_stack.core.access_control.access_control import default_policy
from llama_stack.core.datatypes import ResponsesStoreConfig
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
from llama_stack.providers.inline.agents.meta_reference.responses.openai_responses import (
OpenAIResponsesImpl,
)
@ -50,7 +50,7 @@ from llama_stack.providers.utils.responses.responses_store import (
ResponsesStore,
_OpenAIResponseObjectWithInputAndMessages,
)
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
from tests.unit.providers.agents.meta_reference.fixtures import load_chat_completion_fixture
@ -854,8 +854,10 @@ async def test_responses_store_list_input_items_logic():
# Create mock store and response store
mock_sql_store = AsyncMock()
backend_name = "sql_responses_test"
register_sqlstore_backends({backend_name: SqliteSqlStoreConfig(db_path="mock_db_path")})
responses_store = ResponsesStore(
ResponsesStoreConfig(sql_store_config=SqliteSqlStoreConfig(db_path="mock_db_path")), policy=default_policy()
ResponsesStoreReference(backend=backend_name, table_name="responses"), policy=default_policy()
)
responses_store.sql_store = mock_sql_store

View file

@ -12,10 +12,10 @@ from unittest.mock import AsyncMock
import pytest
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
from llama_stack.providers.inline.batches.reference.batches import ReferenceBatchesImpl
from llama_stack.providers.inline.batches.reference.config import ReferenceBatchesImplConfig
from llama_stack.providers.utils.kvstore import kvstore_impl
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends
@pytest.fixture
@ -23,8 +23,10 @@ async def provider():
"""Create a test provider instance with temporary database."""
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test_batches.db"
backend_name = "kv_batches_test"
kvstore_config = SqliteKVStoreConfig(db_path=str(db_path))
config = ReferenceBatchesImplConfig(kvstore=kvstore_config)
register_kvstore_backends({backend_name: kvstore_config})
config = ReferenceBatchesImplConfig(kvstore=KVStoreReference(backend=backend_name, namespace="batches"))
# Create kvstore and mock APIs
kvstore = await kvstore_impl(config.kvstore)

View file

@ -8,8 +8,9 @@ import boto3
import pytest
from moto import mock_aws
from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
from llama_stack.providers.remote.files.s3 import S3FilesImplConfig, get_adapter_impl
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
class MockUploadFile:
@ -38,11 +39,13 @@ def sample_text_file2():
def s3_config(tmp_path):
db_path = tmp_path / "s3_files_metadata.db"
backend_name = f"sql_s3_{tmp_path.name}"
register_sqlstore_backends({backend_name: SqliteSqlStoreConfig(db_path=db_path.as_posix())})
return S3FilesImplConfig(
bucket_name=f"test-bucket-{tmp_path.name}",
region="not-a-region",
auto_create_bucket=True,
metadata_store=SqliteSqlStoreConfig(db_path=db_path.as_posix()),
metadata_store=SqlStoreReference(backend=backend_name, table_name="s3_files_metadata"),
)

View file

@ -12,13 +12,14 @@ import pytest
from llama_stack.apis.vector_dbs import VectorDB
from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
from llama_stack.providers.inline.vector_io.faiss.faiss import FaissIndex, FaissVectorIOAdapter
from llama_stack.providers.inline.vector_io.sqlite_vec import SQLiteVectorIOConfig
from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import SQLiteVecIndex, SQLiteVecVectorIOAdapter
from llama_stack.providers.remote.vector_io.pgvector.config import PGVectorVectorIOConfig
from llama_stack.providers.remote.vector_io.pgvector.pgvector import PGVectorIndex, PGVectorVectorIOAdapter
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
from llama_stack.providers.utils.kvstore import register_kvstore_backends
EMBEDDING_DIMENSION = 768
COLLECTION_PREFIX = "test_collection"
@ -112,8 +113,9 @@ async def unique_kvstore_config(tmp_path_factory):
unique_id = f"test_kv_{np.random.randint(1e6)}"
temp_dir = tmp_path_factory.getbasetemp()
db_path = str(temp_dir / f"{unique_id}.db")
return SqliteKVStoreConfig(db_path=db_path)
backend_name = f"kv_vector_{unique_id}"
register_kvstore_backends({backend_name: SqliteKVStoreConfig(db_path=db_path)})
return KVStoreReference(backend=backend_name, namespace=f"vector_io::{unique_id}")
@pytest.fixture(scope="session")

View file

@ -10,13 +10,13 @@ import pytest
from llama_stack.apis.inference import Model
from llama_stack.apis.vector_dbs import VectorDB
from llama_stack.core.datatypes import VectorDBWithOwner
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
from llama_stack.core.store.registry import (
KEY_FORMAT,
CachedDiskDistributionRegistry,
DiskDistributionRegistry,
)
from llama_stack.providers.utils.kvstore import kvstore_impl
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends
@pytest.fixture
@ -72,7 +72,11 @@ async def test_cached_registry_initialization(sqlite_kvstore, sample_vector_db,
# Test cached version loads from disk
db_path = sqlite_kvstore.db_path
cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(SqliteKVStoreConfig(db_path=db_path)))
backend_name = "kv_cached_test"
register_kvstore_backends({backend_name: SqliteKVStoreConfig(db_path=db_path)})
cached_registry = CachedDiskDistributionRegistry(
await kvstore_impl(KVStoreReference(backend=backend_name, namespace="registry"))
)
await cached_registry.initialize()
result_vector_db = await cached_registry.get("vector_db", "test_vector_db")
@ -101,7 +105,11 @@ async def test_cached_registry_updates(cached_disk_dist_registry):
# Verify persisted to disk
db_path = cached_disk_dist_registry.kvstore.db_path
new_registry = DiskDistributionRegistry(await kvstore_impl(SqliteKVStoreConfig(db_path=db_path)))
backend_name = "kv_cached_new"
register_kvstore_backends({backend_name: SqliteKVStoreConfig(db_path=db_path)})
new_registry = DiskDistributionRegistry(
await kvstore_impl(KVStoreReference(backend=backend_name, namespace="registry"))
)
await new_registry.initialize()
result_vector_db = await new_registry.get("vector_db", "test_vector_db_2")
assert result_vector_db is not None

View file

@ -4,6 +4,8 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from uuid import uuid4
import pytest
from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
@ -11,7 +13,8 @@ from starlette.middleware.base import BaseHTTPMiddleware
from llama_stack.core.datatypes import QuotaConfig, QuotaPeriod
from llama_stack.core.server.quota import QuotaMiddleware
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
from llama_stack.providers.utils.kvstore import register_kvstore_backends
class InjectClientIDMiddleware(BaseHTTPMiddleware):
@ -29,8 +32,10 @@ class InjectClientIDMiddleware(BaseHTTPMiddleware):
def build_quota_config(db_path) -> QuotaConfig:
backend_name = f"kv_quota_{uuid4().hex}"
register_kvstore_backends({backend_name: SqliteKVStoreConfig(db_path=str(db_path))})
return QuotaConfig(
kvstore=SqliteKVStoreConfig(db_path=str(db_path)),
kvstore=KVStoreReference(backend=backend_name, namespace="quota"),
anonymous_max_requests=1,
authenticated_max_requests=2,
period=QuotaPeriod.DAY,

View file

@ -12,14 +12,18 @@ from unittest.mock import AsyncMock, MagicMock
from pydantic import BaseModel, Field
from llama_stack.apis.inference import Inference
from llama_stack.core.datatypes import (
Api,
Provider,
StackRunConfig,
)
from llama_stack.core.datatypes import Api, Provider, StackRunConfig
from llama_stack.core.resolver import resolve_impls
from llama_stack.core.routers.inference import InferenceRouter
from llama_stack.core.routing_tables.models import ModelsRoutingTable
from llama_stack.core.storage.datatypes import (
InferenceStoreReference,
KVStoreReference,
SqliteKVStoreConfig,
SqliteSqlStoreConfig,
SqlStoreReference,
StorageConfig,
)
from llama_stack.providers.datatypes import InlineProviderSpec, ProviderSpec
@ -65,6 +69,38 @@ class SampleImpl:
pass
def make_run_config(**overrides) -> StackRunConfig:
storage = overrides.pop(
"storage",
StorageConfig(
backends={
"kv_default": SqliteKVStoreConfig(db_path=":memory:"),
"sql_default": SqliteSqlStoreConfig(db_path=":memory:"),
}
),
)
defaults = dict(
image_name="test_image",
apis=[],
providers={},
storage=storage,
metadata_store=overrides.pop(
"metadata_store",
KVStoreReference(backend="kv_default", namespace="registry"),
),
inference_store=overrides.pop(
"inference_store",
InferenceStoreReference(backend="sql_default", table_name="inference_store"),
),
conversations_store=overrides.pop(
"conversations_store",
SqlStoreReference(backend="sql_default", table_name="conversations"),
),
)
defaults.update(overrides)
return StackRunConfig(**defaults)
async def test_resolve_impls_basic():
# Create a real provider spec
provider_spec = InlineProviderSpec(
@ -78,7 +114,7 @@ async def test_resolve_impls_basic():
# Create provider registry with our provider
provider_registry = {Api.inference: {provider_spec.provider_type: provider_spec}}
run_config = StackRunConfig(
run_config = make_run_config(
image_name="test_image",
providers={
"inference": [

View file

@ -6,6 +6,7 @@
import time
from tempfile import TemporaryDirectory
from uuid import uuid4
import pytest
@ -15,8 +16,18 @@ from llama_stack.apis.agents.openai_responses import (
OpenAIResponseObject,
)
from llama_stack.apis.inference import OpenAIMessageParam, OpenAIUserMessageParam
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
from llama_stack.providers.utils.responses.responses_store import ResponsesStore
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
def build_store(db_path: str, policy: list | None = None) -> ResponsesStore:
backend_name = f"sql_responses_{uuid4().hex}"
register_sqlstore_backends({backend_name: SqliteSqlStoreConfig(db_path=db_path)})
return ResponsesStore(
ResponsesStoreReference(backend=backend_name, table_name="responses"),
policy=policy or [],
)
def create_test_response_object(
@ -54,7 +65,7 @@ async def test_responses_store_pagination_basic():
"""Test basic pagination functionality for responses store."""
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/test.db"
store = ResponsesStore(SqliteSqlStoreConfig(db_path=db_path), policy=[])
store = build_store(db_path)
await store.initialize()
# Create test data with different timestamps
@ -103,7 +114,7 @@ async def test_responses_store_pagination_ascending():
"""Test pagination with ascending order."""
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/test.db"
store = ResponsesStore(SqliteSqlStoreConfig(db_path=db_path), policy=[])
store = build_store(db_path)
await store.initialize()
# Create test data
@ -141,7 +152,7 @@ async def test_responses_store_pagination_with_model_filter():
"""Test pagination combined with model filtering."""
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/test.db"
store = ResponsesStore(SqliteSqlStoreConfig(db_path=db_path), policy=[])
store = build_store(db_path)
await store.initialize()
# Create test data with different models
@ -182,7 +193,7 @@ async def test_responses_store_pagination_invalid_after():
"""Test error handling for invalid 'after' parameter."""
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/test.db"
store = ResponsesStore(SqliteSqlStoreConfig(db_path=db_path), policy=[])
store = build_store(db_path)
await store.initialize()
# Try to paginate with non-existent ID
@ -194,7 +205,7 @@ async def test_responses_store_pagination_no_limit():
"""Test pagination behavior when no limit is specified."""
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/test.db"
store = ResponsesStore(SqliteSqlStoreConfig(db_path=db_path), policy=[])
store = build_store(db_path)
await store.initialize()
# Create test data
@ -226,7 +237,7 @@ async def test_responses_store_get_response_object():
"""Test retrieving a single response object."""
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/test.db"
store = ResponsesStore(SqliteSqlStoreConfig(db_path=db_path), policy=[])
store = build_store(db_path)
await store.initialize()
# Store a test response
@ -254,7 +265,7 @@ async def test_responses_store_input_items_pagination():
"""Test pagination functionality for input items."""
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/test.db"
store = ResponsesStore(SqliteSqlStoreConfig(db_path=db_path), policy=[])
store = build_store(db_path)
await store.initialize()
# Store a test response with many inputs with explicit IDs
@ -335,7 +346,7 @@ async def test_responses_store_input_items_before_pagination():
"""Test before pagination functionality for input items."""
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/test.db"
store = ResponsesStore(SqliteSqlStoreConfig(db_path=db_path), policy=[])
store = build_store(db_path)
await store.initialize()
# Store a test response with many inputs with explicit IDs