From 636764c2a1f30147d560eb58a90a0484cfcd5e7d Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Fri, 17 Oct 2025 10:05:07 -0700 Subject: [PATCH] simplified some, walked back some decisions --- .../agents/inline_meta-reference.mdx | 18 ++-- .../providers/batches/inline_reference.mdx | 6 +- .../providers/datasetio/inline_localfs.mdx | 6 +- .../datasetio/remote_huggingface.mdx | 6 +- .../providers/eval/inline_meta-reference.mdx | 6 +- docs/docs/providers/files/inline_localfs.mdx | 6 +- docs/docs/providers/files/remote_s3.mdx | 6 +- .../providers/vector_io/inline_chromadb.mdx | 6 +- .../docs/providers/vector_io/inline_faiss.mdx | 6 +- .../vector_io/inline_meta-reference.mdx | 6 +- .../providers/vector_io/inline_milvus.mdx | 6 +- .../providers/vector_io/inline_qdrant.mdx | 6 +- .../providers/vector_io/inline_sqlite-vec.mdx | 6 +- .../providers/vector_io/inline_sqlite_vec.mdx | 6 +- .../providers/vector_io/remote_chromadb.mdx | 6 +- .../providers/vector_io/remote_milvus.mdx | 6 +- .../providers/vector_io/remote_pgvector.mdx | 6 +- .../providers/vector_io/remote_qdrant.mdx | 6 +- .../providers/vector_io/remote_weaviate.mdx | 6 +- llama_stack/cli/stack/_build.py | 34 ++++++- .../core/conversations/conversations.py | 6 +- llama_stack/core/datatypes.py | 66 ++++++++++++-- llama_stack/core/prompts/prompts.py | 9 +- llama_stack/core/routers/__init__.py | 4 +- llama_stack/core/server/quota.py | 18 ++-- llama_stack/core/stack.py | 4 +- llama_stack/core/storage/datatypes.py | 41 ++------- llama_stack/core/store/registry.py | 13 +-- llama_stack/distributions/ci-tests/build.yaml | 1 + llama_stack/distributions/ci-tests/run.yaml | 40 ++++++--- llama_stack/distributions/dell/build.yaml | 4 +- .../distributions/dell/run-with-safety.yaml | 34 ++++--- llama_stack/distributions/dell/run.yaml | 34 ++++--- .../meta-reference-gpu/build.yaml | 4 +- .../meta-reference-gpu/run-with-safety.yaml | 34 ++++--- .../distributions/meta-reference-gpu/run.yaml | 34 ++++--- llama_stack/distributions/nvidia/build.yaml | 4 +- .../distributions/nvidia/run-with-safety.yaml | 33 ++++--- llama_stack/distributions/nvidia/run.yaml | 32 ++++--- .../distributions/open-benchmark/build.yaml | 4 +- .../distributions/open-benchmark/run.yaml | 36 +++++--- .../distributions/postgres-demo/build.yaml | 5 +- .../postgres-demo/postgres_demo.py | 7 +- .../distributions/postgres-demo/run.yaml | 31 ++++--- .../distributions/starter-gpu/build.yaml | 1 + .../distributions/starter-gpu/run.yaml | 40 ++++++--- llama_stack/distributions/starter/build.yaml | 1 + llama_stack/distributions/starter/run.yaml | 40 ++++++--- llama_stack/distributions/template.py | 49 ++++++----- llama_stack/distributions/watsonx/build.yaml | 4 +- llama_stack/distributions/watsonx/run.yaml | 35 +++++--- .../inline/agents/meta_reference/config.py | 8 +- .../inline/batches/reference/config.py | 1 + .../inline/datasetio/localfs/config.py | 1 + .../inline/eval/meta_reference/config.py | 1 + .../providers/inline/files/localfs/config.py | 1 + .../inline/vector_io/chroma/config.py | 1 + .../inline/vector_io/faiss/config.py | 1 + .../inline/vector_io/milvus/config.py | 1 + .../inline/vector_io/qdrant/config.py | 1 + .../inline/vector_io/sqlite_vec/config.py | 1 + .../remote/datasetio/huggingface/config.py | 1 + .../providers/remote/files/s3/config.py | 1 + .../providers/remote/files/s3/files.py | 2 +- .../remote/vector_io/chroma/config.py | 1 + .../remote/vector_io/milvus/config.py | 1 + .../remote/vector_io/pgvector/config.py | 1 + .../remote/vector_io/qdrant/config.py | 1 + .../remote/vector_io/weaviate/config.py | 1 + .../utils/inference/inference_store.py | 13 +-- .../providers/utils/kvstore/kvstore.py | 25 ++---- .../utils/responses/responses_store.py | 45 +++++----- .../utils/sqlstore/authorized_sqlstore.py | 2 +- .../utils/sqlstore/sqlalchemy_sqlstore.py | 2 +- .../providers/utils/sqlstore/sqlstore.py | 32 ++----- .../sqlstore/test_authorized_sqlstore.py | 13 ++- .../test_persistence_integration.py | 53 ++++++----- tests/unit/core/test_persistence_config.py | 82 ----------------- tests/unit/core/test_storage_references.py | 77 ++++++++++++++++ tests/unit/distribution/test_distribution.py | 88 +++++++++++++------ tests/unit/files/test_files.py | 8 +- tests/unit/prompts/prompts/conftest.py | 30 ++++++- .../meta_reference/test_openai_responses.py | 8 +- tests/unit/providers/batches/conftest.py | 8 +- tests/unit/providers/files/conftest.py | 7 +- tests/unit/providers/vector_io/conftest.py | 8 +- tests/unit/registry/test_registry.py | 16 +++- tests/unit/server/test_quota.py | 9 +- tests/unit/server/test_resolver.py | 48 ++++++++-- .../utils/responses/test_responses_store.py | 29 ++++-- 90 files changed, 887 insertions(+), 570 deletions(-) delete mode 100644 tests/unit/core/test_persistence_config.py create mode 100644 tests/unit/core/test_storage_references.py diff --git a/docs/docs/providers/agents/inline_meta-reference.mdx b/docs/docs/providers/agents/inline_meta-reference.mdx index fd961745f..fac9b8406 100644 --- a/docs/docs/providers/agents/inline_meta-reference.mdx +++ b/docs/docs/providers/agents/inline_meta-reference.mdx @@ -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` | `` | 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 ``` diff --git a/docs/docs/providers/batches/inline_reference.mdx b/docs/docs/providers/batches/inline_reference.mdx index f43800555..45304fbb1 100644 --- a/docs/docs/providers/batches/inline_reference.mdx +++ b/docs/docs/providers/batches/inline_reference.mdx @@ -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` | `` | No | | Configuration for the key-value store backend. | | `max_concurrent_batches` | `` | No | 1 | Maximum number of concurrent batches to process simultaneously. | | `max_concurrent_requests_per_batch` | `` | 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 ``` diff --git a/docs/docs/providers/datasetio/inline_localfs.mdx b/docs/docs/providers/datasetio/inline_localfs.mdx index b02a3a3bd..a9363376c 100644 --- a/docs/docs/providers/datasetio/inline_localfs.mdx +++ b/docs/docs/providers/datasetio/inline_localfs.mdx @@ -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` | `` | No | | | ## Sample Configuration ```yaml kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default ``` diff --git a/docs/docs/providers/datasetio/remote_huggingface.mdx b/docs/docs/providers/datasetio/remote_huggingface.mdx index 82597d999..de3ffaaa6 100644 --- a/docs/docs/providers/datasetio/remote_huggingface.mdx +++ b/docs/docs/providers/datasetio/remote_huggingface.mdx @@ -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` | `` | No | | | ## Sample Configuration ```yaml kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default ``` diff --git a/docs/docs/providers/eval/inline_meta-reference.mdx b/docs/docs/providers/eval/inline_meta-reference.mdx index b0eb589e0..2c86c18c9 100644 --- a/docs/docs/providers/eval/inline_meta-reference.mdx +++ b/docs/docs/providers/eval/inline_meta-reference.mdx @@ -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` | `` | No | | | ## Sample Configuration ```yaml kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/meta_reference_eval.db + namespace: eval + backend: kv_default ``` diff --git a/docs/docs/providers/files/inline_localfs.mdx b/docs/docs/providers/files/inline_localfs.mdx index 86d141f93..bff0c4eb9 100644 --- a/docs/docs/providers/files/inline_localfs.mdx +++ b/docs/docs/providers/files/inline_localfs.mdx @@ -15,7 +15,7 @@ Local filesystem-based file storage provider for managing files and documents lo | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| | `storage_dir` | `` | 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` | `` | No | | SQL store configuration for file metadata | | `ttl_secs` | `` | 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 ``` diff --git a/docs/docs/providers/files/remote_s3.mdx b/docs/docs/providers/files/remote_s3.mdx index 353cedbfb..65cd545c5 100644 --- a/docs/docs/providers/files/remote_s3.mdx +++ b/docs/docs/providers/files/remote_s3.mdx @@ -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` | `` | 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` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/inline_chromadb.mdx b/docs/docs/providers/vector_io/inline_chromadb.mdx index a1858eacc..97f2ee3d7 100644 --- a/docs/docs/providers/vector_io/inline_chromadb.mdx +++ b/docs/docs/providers/vector_io/inline_chromadb.mdx @@ -79,13 +79,13 @@ See [Chroma's documentation](https://docs.trychroma.com/docs/overview/introducti | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| | `db_path` | `` | 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` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/inline_faiss.mdx b/docs/docs/providers/vector_io/inline_faiss.mdx index 03bc2a928..c7471f319 100644 --- a/docs/docs/providers/vector_io/inline_faiss.mdx +++ b/docs/docs/providers/vector_io/inline_faiss.mdx @@ -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` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/inline_meta-reference.mdx b/docs/docs/providers/vector_io/inline_meta-reference.mdx index bcad86750..968240c0a 100644 --- a/docs/docs/providers/vector_io/inline_meta-reference.mdx +++ b/docs/docs/providers/vector_io/inline_meta-reference.mdx @@ -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` | `` | 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 diff --git a/docs/docs/providers/vector_io/inline_milvus.mdx b/docs/docs/providers/vector_io/inline_milvus.mdx index 7e6f15c81..e7b8a3e20 100644 --- a/docs/docs/providers/vector_io/inline_milvus.mdx +++ b/docs/docs/providers/vector_io/inline_milvus.mdx @@ -17,7 +17,7 @@ Please refer to the remote provider documentation. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| | `db_path` | `` | 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` | `` | No | | Config for KV store backend (SQLite only for now) | | `consistency_level` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/inline_qdrant.mdx b/docs/docs/providers/vector_io/inline_qdrant.mdx index 5c9ab10f2..399690f75 100644 --- a/docs/docs/providers/vector_io/inline_qdrant.mdx +++ b/docs/docs/providers/vector_io/inline_qdrant.mdx @@ -98,13 +98,13 @@ See the [Qdrant documentation](https://qdrant.tech/documentation/) for more deta | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| | `path` | `` | No | | | -| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | | +| `kvstore` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/inline_sqlite-vec.mdx b/docs/docs/providers/vector_io/inline_sqlite-vec.mdx index aa6992a56..0c5a5b827 100644 --- a/docs/docs/providers/vector_io/inline_sqlite-vec.mdx +++ b/docs/docs/providers/vector_io/inline_sqlite-vec.mdx @@ -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` | `` | 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` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/inline_sqlite_vec.mdx b/docs/docs/providers/vector_io/inline_sqlite_vec.mdx index 7f69f617d..604443aa9 100644 --- a/docs/docs/providers/vector_io/inline_sqlite_vec.mdx +++ b/docs/docs/providers/vector_io/inline_sqlite_vec.mdx @@ -17,15 +17,15 @@ Please refer to the sqlite-vec provider documentation. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| | `db_path` | `` | 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` | `` | 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 diff --git a/docs/docs/providers/vector_io/remote_chromadb.mdx b/docs/docs/providers/vector_io/remote_chromadb.mdx index 807771003..3d62a1de6 100644 --- a/docs/docs/providers/vector_io/remote_chromadb.mdx +++ b/docs/docs/providers/vector_io/remote_chromadb.mdx @@ -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` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/remote_milvus.mdx b/docs/docs/providers/vector_io/remote_milvus.mdx index 7f7c08122..0f56d7903 100644 --- a/docs/docs/providers/vector_io/remote_milvus.mdx +++ b/docs/docs/providers/vector_io/remote_milvus.mdx @@ -408,7 +408,7 @@ For more details on TLS configuration, refer to the [TLS setup guide](https://mi | `uri` | `` | No | | The URI of the Milvus server | | `token` | `str \| None` | No | | The token of the Milvus server | | `consistency_level` | `` | 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` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/remote_pgvector.mdx b/docs/docs/providers/vector_io/remote_pgvector.mdx index d21810c68..22fde059c 100644 --- a/docs/docs/providers/vector_io/remote_pgvector.mdx +++ b/docs/docs/providers/vector_io/remote_pgvector.mdx @@ -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 ``` diff --git a/docs/docs/providers/vector_io/remote_qdrant.mdx b/docs/docs/providers/vector_io/remote_qdrant.mdx index c44a2b937..9a58533da 100644 --- a/docs/docs/providers/vector_io/remote_qdrant.mdx +++ b/docs/docs/providers/vector_io/remote_qdrant.mdx @@ -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` | `` | 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 ``` diff --git a/docs/docs/providers/vector_io/remote_weaviate.mdx b/docs/docs/providers/vector_io/remote_weaviate.mdx index 3f1e36422..ffa665687 100644 --- a/docs/docs/providers/vector_io/remote_weaviate.mdx +++ b/docs/docs/providers/vector_io/remote_weaviate.mdx @@ -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 ``` diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py index 55989b05e..cf5ed55ae 100644 --- a/llama_stack/cli/stack/_build.py +++ b/llama_stack/cli/stack/_build.py @@ -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: diff --git a/llama_stack/core/conversations/conversations.py b/llama_stack/core/conversations/conversations.py index eb9a3c399..ef2bca7e3 100644 --- a/llama_stack/core/conversations/conversations.py +++ b/llama_stack/core/conversations/conversations.py @@ -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) diff --git a/llama_stack/core/datatypes.py b/llama_stack/core/datatypes.py index e1576538f..a813b6084 100644 --- a/llama_stack/core/datatypes.py +++ b/llama_stack/core/datatypes.py @@ -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 diff --git a/llama_stack/core/prompts/prompts.py b/llama_stack/core/prompts/prompts.py index a819a2c57..ee8c42596 100644 --- a/llama_stack/core/prompts/prompts.py +++ b/llama_stack/core/prompts/prompts.py @@ -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: diff --git a/llama_stack/core/routers/__init__.py b/llama_stack/core/routers/__init__.py index a0ab8fc3b..62da054de 100644 --- a/llama_stack/core/routers/__init__.py +++ b/llama_stack/core/routers/__init__.py @@ -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, diff --git a/llama_stack/core/server/quota.py b/llama_stack/core/server/quota.py index 693f224c3..689f0e4c3 100644 --- a/llama_stack/core/server/quota.py +++ b/llama_stack/core/server/quota.py @@ -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): diff --git a/llama_stack/core/stack.py b/llama_stack/core/stack.py index 6a68dec7a..95e9a28b0 100644 --- a/llama_stack/core/stack.py +++ b/llama_stack/core/stack.py @@ -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 = {} diff --git a/llama_stack/core/storage/datatypes.py b/llama_stack/core/storage/datatypes.py index 10046c248..4c3f1b99e 100644 --- a/llama_stack/core/storage/datatypes.py +++ b/llama_stack/core/storage/datatypes.py @@ -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)", - ) diff --git a/llama_stack/core/store/registry.py b/llama_stack/core/store/registry.py index ac5d1f1d7..6ff9e575b 100644 --- a/llama_stack/core/store/registry.py +++ b/llama_stack/core/store/registry.py @@ -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 diff --git a/llama_stack/distributions/ci-tests/build.yaml b/llama_stack/distributions/ci-tests/build.yaml index f70935922..191d0ae59 100644 --- a/llama_stack/distributions/ci-tests/build.yaml +++ b/llama_stack/distributions/ci-tests/build.yaml @@ -52,5 +52,6 @@ distribution_spec: - provider_type: inline::reference image_type: venv additional_pip_packages: +- aiosqlite - asyncpg - sqlalchemy[asyncio] diff --git a/llama_stack/distributions/ci-tests/run.yaml b/llama_stack/distributions/ci-tests/run.yaml index b2a3169c1..de950d1b3 100644 --- a/llama_stack/distributions/ci-tests/run.yaml +++ b/llama_stack/distributions/ci-tests/run.yaml @@ -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 diff --git a/llama_stack/distributions/dell/build.yaml b/llama_stack/distributions/dell/build.yaml index 224a83877..7bc26ca9e 100644 --- a/llama_stack/distributions/dell/build.yaml +++ b/llama_stack/distributions/dell/build.yaml @@ -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] diff --git a/llama_stack/distributions/dell/run-with-safety.yaml b/llama_stack/distributions/dell/run-with-safety.yaml index 055d71b04..a705afe94 100644 --- a/llama_stack/distributions/dell/run-with-safety.yaml +++ b/llama_stack/distributions/dell/run-with-safety.yaml @@ -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} diff --git a/llama_stack/distributions/dell/run.yaml b/llama_stack/distributions/dell/run.yaml index f74cf7802..e1feb0c75 100644 --- a/llama_stack/distributions/dell/run.yaml +++ b/llama_stack/distributions/dell/run.yaml @@ -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} diff --git a/llama_stack/distributions/meta-reference-gpu/build.yaml b/llama_stack/distributions/meta-reference-gpu/build.yaml index 186be4b42..1513742a7 100644 --- a/llama_stack/distributions/meta-reference-gpu/build.yaml +++ b/llama_stack/distributions/meta-reference-gpu/build.yaml @@ -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] diff --git a/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml b/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml index a2fe6c027..517886dba 100644 --- a/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml +++ b/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml @@ -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} diff --git a/llama_stack/distributions/meta-reference-gpu/run.yaml b/llama_stack/distributions/meta-reference-gpu/run.yaml index 13bd08895..0133d5cd3 100644 --- a/llama_stack/distributions/meta-reference-gpu/run.yaml +++ b/llama_stack/distributions/meta-reference-gpu/run.yaml @@ -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} diff --git a/llama_stack/distributions/nvidia/build.yaml b/llama_stack/distributions/nvidia/build.yaml index dc27c994a..8ddd12439 100644 --- a/llama_stack/distributions/nvidia/build.yaml +++ b/llama_stack/distributions/nvidia/build.yaml @@ -24,4 +24,6 @@ distribution_spec: files: - provider_type: inline::localfs image_type: venv -additional_pip_packages: [] +additional_pip_packages: +- aiosqlite +- sqlalchemy[asyncio] diff --git a/llama_stack/distributions/nvidia/run-with-safety.yaml b/llama_stack/distributions/nvidia/run-with-safety.yaml index 92efdd2ba..ae705977b 100644 --- a/llama_stack/distributions/nvidia/run-with-safety.yaml +++ b/llama_stack/distributions/nvidia/run-with-safety.yaml @@ -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} diff --git a/llama_stack/distributions/nvidia/run.yaml b/llama_stack/distributions/nvidia/run.yaml index 2bbaf37a1..65bc0ce9a 100644 --- a/llama_stack/distributions/nvidia/run.yaml +++ b/llama_stack/distributions/nvidia/run.yaml @@ -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: [] diff --git a/llama_stack/distributions/open-benchmark/build.yaml b/llama_stack/distributions/open-benchmark/build.yaml index 98e635506..05acd98e3 100644 --- a/llama_stack/distributions/open-benchmark/build.yaml +++ b/llama_stack/distributions/open-benchmark/build.yaml @@ -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] diff --git a/llama_stack/distributions/open-benchmark/run.yaml b/llama_stack/distributions/open-benchmark/run.yaml index d6859016b..e15cbedd8 100644 --- a/llama_stack/distributions/open-benchmark/run.yaml +++ b/llama_stack/distributions/open-benchmark/run.yaml @@ -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 diff --git a/llama_stack/distributions/postgres-demo/build.yaml b/llama_stack/distributions/postgres-demo/build.yaml index 09245f255..063dc3999 100644 --- a/llama_stack/distributions/postgres-demo/build.yaml +++ b/llama_stack/distributions/postgres-demo/build.yaml @@ -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] diff --git a/llama_stack/distributions/postgres-demo/postgres_demo.py b/llama_stack/distributions/postgres-demo/postgres_demo.py index 9c18a305f..876370ef3 100644 --- a/llama_stack/distributions/postgres-demo/postgres_demo.py +++ b/llama_stack/distributions/postgres-demo/postgres_demo.py @@ -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(), }, ), }, diff --git a/llama_stack/distributions/postgres-demo/run.yaml b/llama_stack/distributions/postgres-demo/run.yaml index 634289d5b..df979a8fe 100644 --- a/llama_stack/distributions/postgres-demo/run.yaml +++ b/llama_stack/distributions/postgres-demo/run.yaml @@ -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} diff --git a/llama_stack/distributions/starter-gpu/build.yaml b/llama_stack/distributions/starter-gpu/build.yaml index 138e3b39e..943c6134d 100644 --- a/llama_stack/distributions/starter-gpu/build.yaml +++ b/llama_stack/distributions/starter-gpu/build.yaml @@ -53,5 +53,6 @@ distribution_spec: - provider_type: inline::reference image_type: venv additional_pip_packages: +- aiosqlite - asyncpg - sqlalchemy[asyncio] diff --git a/llama_stack/distributions/starter-gpu/run.yaml b/llama_stack/distributions/starter-gpu/run.yaml index 375d1fbc9..357495368 100644 --- a/llama_stack/distributions/starter-gpu/run.yaml +++ b/llama_stack/distributions/starter-gpu/run.yaml @@ -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 diff --git a/llama_stack/distributions/starter/build.yaml b/llama_stack/distributions/starter/build.yaml index 4d149cf05..c2719d50d 100644 --- a/llama_stack/distributions/starter/build.yaml +++ b/llama_stack/distributions/starter/build.yaml @@ -53,5 +53,6 @@ distribution_spec: - provider_type: inline::reference image_type: venv additional_pip_packages: +- aiosqlite - asyncpg - sqlalchemy[asyncio] diff --git a/llama_stack/distributions/starter/run.yaml b/llama_stack/distributions/starter/run.yaml index 25ca22137..cc5fc92c4 100644 --- a/llama_stack/distributions/starter/run.yaml +++ b/llama_stack/distributions/starter/run.yaml @@ -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 diff --git a/llama_stack/distributions/template.py b/llama_stack/distributions/template.py index e72744429..f7671719c 100644 --- a/llama_stack/distributions/template.py +++ b/llama_stack/distributions/template.py @@ -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) diff --git a/llama_stack/distributions/watsonx/build.yaml b/llama_stack/distributions/watsonx/build.yaml index f81977347..dba1a94e2 100644 --- a/llama_stack/distributions/watsonx/build.yaml +++ b/llama_stack/distributions/watsonx/build.yaml @@ -28,4 +28,6 @@ distribution_spec: files: - provider_type: inline::localfs image_type: venv -additional_pip_packages: [] +additional_pip_packages: +- aiosqlite +- sqlalchemy[asyncio] diff --git a/llama_stack/distributions/watsonx/run.yaml b/llama_stack/distributions/watsonx/run.yaml index 5b2a05d3d..5a5343e88 100644 --- a/llama_stack/distributions/watsonx/run.yaml +++ b/llama_stack/distributions/watsonx/run.yaml @@ -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: [] diff --git a/llama_stack/providers/inline/agents/meta_reference/config.py b/llama_stack/providers/inline/agents/meta_reference/config.py index acda5f463..a800b426b 100644 --- a/llama_stack/providers/inline/agents/meta_reference/config.py +++ b/llama_stack/providers/inline/agents/meta_reference/config.py @@ -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), } diff --git a/llama_stack/providers/inline/batches/reference/config.py b/llama_stack/providers/inline/batches/reference/config.py index 19078b4f5..f896a897d 100644 --- a/llama_stack/providers/inline/batches/reference/config.py +++ b/llama_stack/providers/inline/batches/reference/config.py @@ -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), } diff --git a/llama_stack/providers/inline/datasetio/localfs/config.py b/llama_stack/providers/inline/datasetio/localfs/config.py index f70fb79ec..6e878df62 100644 --- a/llama_stack/providers/inline/datasetio/localfs/config.py +++ b/llama_stack/providers/inline/datasetio/localfs/config.py @@ -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) } diff --git a/llama_stack/providers/inline/eval/meta_reference/config.py b/llama_stack/providers/inline/eval/meta_reference/config.py index 210e23201..b496c855e 100644 --- a/llama_stack/providers/inline/eval/meta_reference/config.py +++ b/llama_stack/providers/inline/eval/meta_reference/config.py @@ -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) } diff --git a/llama_stack/providers/inline/files/localfs/config.py b/llama_stack/providers/inline/files/localfs/config.py index 5b437f41c..0c2dd3b21 100644 --- a/llama_stack/providers/inline/files/localfs/config.py +++ b/llama_stack/providers/inline/files/localfs/config.py @@ -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), } diff --git a/llama_stack/providers/inline/vector_io/chroma/config.py b/llama_stack/providers/inline/vector_io/chroma/config.py index c373c2241..9d61e3ba7 100644 --- a/llama_stack/providers/inline/vector_io/chroma/config.py +++ b/llama_stack/providers/inline/vector_io/chroma/config.py @@ -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), } diff --git a/llama_stack/providers/inline/vector_io/faiss/config.py b/llama_stack/providers/inline/vector_io/faiss/config.py index c156eee37..231382c83 100644 --- a/llama_stack/providers/inline/vector_io/faiss/config.py +++ b/llama_stack/providers/inline/vector_io/faiss/config.py @@ -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) } diff --git a/llama_stack/providers/inline/vector_io/milvus/config.py b/llama_stack/providers/inline/vector_io/milvus/config.py index ba37febaa..f6fcd1db0 100644 --- a/llama_stack/providers/inline/vector_io/milvus/config.py +++ b/llama_stack/providers/inline/vector_io/milvus/config.py @@ -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), } diff --git a/llama_stack/providers/inline/vector_io/qdrant/config.py b/llama_stack/providers/inline/vector_io/qdrant/config.py index d462c7544..892670a0d 100644 --- a/llama_stack/providers/inline/vector_io/qdrant/config.py +++ b/llama_stack/providers/inline/vector_io/qdrant/config.py @@ -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), } diff --git a/llama_stack/providers/inline/vector_io/sqlite_vec/config.py b/llama_stack/providers/inline/vector_io/sqlite_vec/config.py index e478efb37..c2d39b300 100644 --- a/llama_stack/providers/inline/vector_io/sqlite_vec/config.py +++ b/llama_stack/providers/inline/vector_io/sqlite_vec/config.py @@ -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), } diff --git a/llama_stack/providers/remote/datasetio/huggingface/config.py b/llama_stack/providers/remote/datasetio/huggingface/config.py index 1123a103f..35297cb58 100644 --- a/llama_stack/providers/remote/datasetio/huggingface/config.py +++ b/llama_stack/providers/remote/datasetio/huggingface/config.py @@ -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) } diff --git a/llama_stack/providers/remote/files/s3/config.py b/llama_stack/providers/remote/files/s3/config.py index 66e4f6044..cd4b1adda 100644 --- a/llama_stack/providers/remote/files/s3/config.py +++ b/llama_stack/providers/remote/files/s3/config.py @@ -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), } diff --git a/llama_stack/providers/remote/files/s3/files.py b/llama_stack/providers/remote/files/s3/files.py index 7d76f1721..c0e9f81d6 100644 --- a/llama_stack/providers/remote/files/s3/files.py +++ b/llama_stack/providers/remote/files/s3/files.py @@ -166,7 +166,7 @@ class S3FilesImpl(Files): self._client = _create_s3_client(self._config) await _create_bucket_if_not_exists(self._client, self._config) - self._sql_store = AuthorizedSqlStore(sqlstore_impl(self._config.persistence), self.policy) + self._sql_store = AuthorizedSqlStore(sqlstore_impl(self._config.metadata_store), self.policy) await self._sql_store.create_table( "openai_files", { diff --git a/llama_stack/providers/remote/vector_io/chroma/config.py b/llama_stack/providers/remote/vector_io/chroma/config.py index 89efba1c5..f356232ad 100644 --- a/llama_stack/providers/remote/vector_io/chroma/config.py +++ b/llama_stack/providers/remote/vector_io/chroma/config.py @@ -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), } diff --git a/llama_stack/providers/remote/vector_io/milvus/config.py b/llama_stack/providers/remote/vector_io/milvus/config.py index ef21f2f29..0ed7a6b00 100644 --- a/llama_stack/providers/remote/vector_io/milvus/config.py +++ b/llama_stack/providers/remote/vector_io/milvus/config.py @@ -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), } diff --git a/llama_stack/providers/remote/vector_io/pgvector/config.py b/llama_stack/providers/remote/vector_io/pgvector/config.py index 46ce55600..d59fb8fe2 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/config.py +++ b/llama_stack/providers/remote/vector_io/pgvector/config.py @@ -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), } diff --git a/llama_stack/providers/remote/vector_io/qdrant/config.py b/llama_stack/providers/remote/vector_io/qdrant/config.py index 2dc89386f..d2f5686ee 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/config.py +++ b/llama_stack/providers/remote/vector_io/qdrant/config.py @@ -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), } diff --git a/llama_stack/providers/remote/vector_io/weaviate/config.py b/llama_stack/providers/remote/vector_io/weaviate/config.py index 65781c61f..87b127a56 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/config.py +++ b/llama_stack/providers/remote/vector_io/weaviate/config.py @@ -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), } diff --git a/llama_stack/providers/utils/inference/inference_store.py b/llama_stack/providers/utils/inference/inference_store.py index 13cde66b1..8e20bca6b 100644 --- a/llama_stack/providers/utils/inference/inference_store.py +++ b/llama_stack/providers/utils/inference/inference_store.py @@ -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", { diff --git a/llama_stack/providers/utils/kvstore/kvstore.py b/llama_stack/providers/utils/kvstore/kvstore.py index 442bbeb41..eee51e5d9 100644 --- a/llama_stack/providers/utils/kvstore/kvstore.py +++ b/llama_stack/providers/utils/kvstore/kvstore.py @@ -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: diff --git a/llama_stack/providers/utils/responses/responses_store.py b/llama_stack/providers/utils/responses/responses_store.py index 36370b492..d5c243252 100644 --- a/llama_stack/providers/utils/responses/responses_store.py +++ b/llama_stack/providers/utils/responses/responses_store.py @@ -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", { diff --git a/llama_stack/providers/utils/sqlstore/authorized_sqlstore.py b/llama_stack/providers/utils/sqlstore/authorized_sqlstore.py index d2a66213e..3dfc82677 100644 --- a/llama_stack/providers/utils/sqlstore/authorized_sqlstore.py +++ b/llama_stack/providers/utils/sqlstore/authorized_sqlstore.py @@ -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") diff --git a/llama_stack/providers/utils/sqlstore/sqlalchemy_sqlstore.py b/llama_stack/providers/utils/sqlstore/sqlalchemy_sqlstore.py index 23cd6444e..c1ccd73dd 100644 --- a/llama_stack/providers/utils/sqlstore/sqlalchemy_sqlstore.py +++ b/llama_stack/providers/utils/sqlstore/sqlalchemy_sqlstore.py @@ -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") diff --git a/llama_stack/providers/utils/sqlstore/sqlstore.py b/llama_stack/providers/utils/sqlstore/sqlstore.py index 06f83116c..31801c4ca 100644 --- a/llama_stack/providers/utils/sqlstore/sqlstore.py +++ b/llama_stack/providers/utils/sqlstore/sqlstore.py @@ -4,7 +4,7 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -from 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 diff --git a/tests/integration/providers/utils/sqlstore/test_authorized_sqlstore.py b/tests/integration/providers/utils/sqlstore/test_authorized_sqlstore.py index 98bef0f2c..ad9115756 100644 --- a/tests/integration/providers/utils/sqlstore/test_authorized_sqlstore.py +++ b/tests/integration/providers/utils/sqlstore/test_authorized_sqlstore.py @@ -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 diff --git a/tests/integration/test_persistence_integration.py b/tests/integration/test_persistence_integration.py index 88b43ecc4..dcc4cb1af 100644 --- a/tests/integration/test_persistence_integration.py +++ b/tests/integration/test_persistence_integration.py @@ -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" diff --git a/tests/unit/core/test_persistence_config.py b/tests/unit/core/test_persistence_config.py deleted file mode 100644 index 0b99b3cab..000000000 --- a/tests/unit/core/test_persistence_config.py +++ /dev/null @@ -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) diff --git a/tests/unit/core/test_storage_references.py b/tests/unit/core/test_storage_references.py new file mode 100644 index 000000000..206b90304 --- /dev/null +++ b/tests/unit/core/test_storage_references.py @@ -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" diff --git a/tests/unit/distribution/test_distribution.py b/tests/unit/distribution/test_distribution.py index 08a376008..788585328 100644 --- a/tests/unit/distribution/test_distribution.py +++ b/tests/unit/distribution/test_distribution.py @@ -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": [ diff --git a/tests/unit/files/test_files.py b/tests/unit/files/test_files.py index e14e033b9..426e2cf64 100644 --- a/tests/unit/files/test_files.py +++ b/tests/unit/files/test_files.py @@ -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()) diff --git a/tests/unit/prompts/prompts/conftest.py b/tests/unit/prompts/prompts/conftest.py index b2c619e49..0fa52193b 100644 --- a/tests/unit/prompts/prompts/conftest.py +++ b/tests/unit/prompts/prompts/conftest.py @@ -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 diff --git a/tests/unit/providers/agents/meta_reference/test_openai_responses.py b/tests/unit/providers/agents/meta_reference/test_openai_responses.py index e93668a62..59e9df5c0 100644 --- a/tests/unit/providers/agents/meta_reference/test_openai_responses.py +++ b/tests/unit/providers/agents/meta_reference/test_openai_responses.py @@ -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 diff --git a/tests/unit/providers/batches/conftest.py b/tests/unit/providers/batches/conftest.py index df37141b5..d161bf976 100644 --- a/tests/unit/providers/batches/conftest.py +++ b/tests/unit/providers/batches/conftest.py @@ -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) diff --git a/tests/unit/providers/files/conftest.py b/tests/unit/providers/files/conftest.py index 46282e3dc..c64ecc3a3 100644 --- a/tests/unit/providers/files/conftest.py +++ b/tests/unit/providers/files/conftest.py @@ -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"), ) diff --git a/tests/unit/providers/vector_io/conftest.py b/tests/unit/providers/vector_io/conftest.py index 8e5c85cf1..5e2acdc35 100644 --- a/tests/unit/providers/vector_io/conftest.py +++ b/tests/unit/providers/vector_io/conftest.py @@ -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") diff --git a/tests/unit/registry/test_registry.py b/tests/unit/registry/test_registry.py index e49c9dc77..95022ad33 100644 --- a/tests/unit/registry/test_registry.py +++ b/tests/unit/registry/test_registry.py @@ -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 diff --git a/tests/unit/server/test_quota.py b/tests/unit/server/test_quota.py index 85acbc66a..16b1772ce 100644 --- a/tests/unit/server/test_quota.py +++ b/tests/unit/server/test_quota.py @@ -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, diff --git a/tests/unit/server/test_resolver.py b/tests/unit/server/test_resolver.py index 1ee1b2f47..5620c7fa5 100644 --- a/tests/unit/server/test_resolver.py +++ b/tests/unit/server/test_resolver.py @@ -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": [ diff --git a/tests/unit/utils/responses/test_responses_store.py b/tests/unit/utils/responses/test_responses_store.py index c27b5a8e5..34cff3d3f 100644 --- a/tests/unit/utils/responses/test_responses_store.py +++ b/tests/unit/utils/responses/test_responses_store.py @@ -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