From af7472cdb0f61a47bc5bebfcb61dc4f48c7ea5ba Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Thu, 16 Oct 2025 17:12:09 -0700 Subject: [PATCH] updates --- llama_stack/cli/stack/_build.py | 1 - llama_stack/core/datatypes.py | 2 +- llama_stack/distributions/ci-tests/build.yaml | 1 - llama_stack/distributions/ci-tests/run.yaml | 28 +++++--- llama_stack/distributions/dell/build.yaml | 4 +- .../distributions/dell/run-with-safety.yaml | 28 +++++--- llama_stack/distributions/dell/run.yaml | 28 +++++--- .../meta-reference-gpu/build.yaml | 4 +- .../meta-reference-gpu/run-with-safety.yaml | 28 +++++--- .../distributions/meta-reference-gpu/run.yaml | 28 +++++--- llama_stack/distributions/nvidia/build.yaml | 4 +- .../distributions/nvidia/run-with-safety.yaml | 28 +++++--- llama_stack/distributions/nvidia/run.yaml | 28 +++++--- .../distributions/open-benchmark/build.yaml | 4 +- .../distributions/open-benchmark/run.yaml | 28 +++++--- .../distributions/postgres-demo/build.yaml | 5 +- .../postgres-demo/postgres_demo.py | 16 ++--- .../distributions/postgres-demo/run.yaml | 65 ++++++++++--------- .../distributions/starter-gpu/build.yaml | 1 - .../distributions/starter-gpu/run.yaml | 28 +++++--- llama_stack/distributions/starter/build.yaml | 1 - llama_stack/distributions/starter/run.yaml | 28 +++++--- llama_stack/distributions/template.py | 57 ++++++++++------ llama_stack/distributions/watsonx/build.yaml | 4 +- llama_stack/distributions/watsonx/run.yaml | 28 +++++--- .../test_persistence_integration.py | 1 - 26 files changed, 293 insertions(+), 185 deletions(-) diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py index 1b60e9bf5..55989b05e 100644 --- a/llama_stack/cli/stack/_build.py +++ b/llama_stack/cli/stack/_build.py @@ -45,7 +45,6 @@ from llama_stack.core.utils.dynamic import instantiate_class_type from llama_stack.core.utils.exec import formulate_run_args, run_command from llama_stack.core.utils.image_types import LlamaStackImageType from llama_stack.providers.datatypes import Api -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig DISTRIBS_PATH = Path(__file__).parent.parent.parent / "distributions" diff --git a/llama_stack/core/datatypes.py b/llama_stack/core/datatypes.py index 766e5b550..e1576538f 100644 --- a/llama_stack/core/datatypes.py +++ b/llama_stack/core/datatypes.py @@ -465,7 +465,7 @@ 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 +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). """, ) diff --git a/llama_stack/distributions/ci-tests/build.yaml b/llama_stack/distributions/ci-tests/build.yaml index 191d0ae59..f70935922 100644 --- a/llama_stack/distributions/ci-tests/build.yaml +++ b/llama_stack/distributions/ci-tests/build.yaml @@ -52,6 +52,5 @@ 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 3e592e8d5..b2a3169c1 100644 --- a/llama_stack/distributions/ci-tests/run.yaml +++ b/llama_stack/distributions/ci-tests/run.yaml @@ -198,15 +198,25 @@ providers: config: kvstore: namespace: batches -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/kvstore.db + default_sql_store: + 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 models: [] shields: - shield_id: llama-guard diff --git a/llama_stack/distributions/dell/build.yaml b/llama_stack/distributions/dell/build.yaml index 7bc26ca9e..224a83877 100644 --- a/llama_stack/distributions/dell/build.yaml +++ b/llama_stack/distributions/dell/build.yaml @@ -28,6 +28,4 @@ distribution_spec: - provider_type: remote::tavily-search - provider_type: inline::rag-runtime image_type: venv -additional_pip_packages: -- aiosqlite -- sqlalchemy[asyncio] +additional_pip_packages: [] diff --git a/llama_stack/distributions/dell/run-with-safety.yaml b/llama_stack/distributions/dell/run-with-safety.yaml index 03e42c674..055d71b04 100644 --- a/llama_stack/distributions/dell/run-with-safety.yaml +++ b/llama_stack/distributions/dell/run-with-safety.yaml @@ -81,15 +81,25 @@ providers: max_results: 3 - provider_id: rag-runtime provider_type: inline::rag-runtime -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/kvstore.db + default_sql_store: + 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 models: - metadata: {} model_id: ${env.INFERENCE_MODEL} diff --git a/llama_stack/distributions/dell/run.yaml b/llama_stack/distributions/dell/run.yaml index 9e6c03454..f74cf7802 100644 --- a/llama_stack/distributions/dell/run.yaml +++ b/llama_stack/distributions/dell/run.yaml @@ -77,15 +77,25 @@ providers: max_results: 3 - provider_id: rag-runtime provider_type: inline::rag-runtime -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/kvstore.db + default_sql_store: + 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 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 1513742a7..186be4b42 100644 --- a/llama_stack/distributions/meta-reference-gpu/build.yaml +++ b/llama_stack/distributions/meta-reference-gpu/build.yaml @@ -27,6 +27,4 @@ distribution_spec: - provider_type: inline::rag-runtime - provider_type: remote::model-context-protocol image_type: venv -additional_pip_packages: -- aiosqlite -- sqlalchemy[asyncio] +additional_pip_packages: [] 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 201c0bf7e..a2fe6c027 100644 --- a/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml +++ b/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml @@ -94,15 +94,25 @@ providers: provider_type: inline::rag-runtime - provider_id: model-context-protocol provider_type: remote::model-context-protocol -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db + default_sql_store: + 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 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 088abd02d..13bd08895 100644 --- a/llama_stack/distributions/meta-reference-gpu/run.yaml +++ b/llama_stack/distributions/meta-reference-gpu/run.yaml @@ -84,15 +84,25 @@ providers: provider_type: inline::rag-runtime - provider_id: model-context-protocol provider_type: remote::model-context-protocol -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db + default_sql_store: + 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 models: - metadata: {} model_id: ${env.INFERENCE_MODEL} diff --git a/llama_stack/distributions/nvidia/build.yaml b/llama_stack/distributions/nvidia/build.yaml index 8ddd12439..dc27c994a 100644 --- a/llama_stack/distributions/nvidia/build.yaml +++ b/llama_stack/distributions/nvidia/build.yaml @@ -24,6 +24,4 @@ distribution_spec: files: - provider_type: inline::localfs image_type: venv -additional_pip_packages: -- aiosqlite -- sqlalchemy[asyncio] +additional_pip_packages: [] diff --git a/llama_stack/distributions/nvidia/run-with-safety.yaml b/llama_stack/distributions/nvidia/run-with-safety.yaml index 038d03829..92efdd2ba 100644 --- a/llama_stack/distributions/nvidia/run-with-safety.yaml +++ b/llama_stack/distributions/nvidia/run-with-safety.yaml @@ -84,15 +84,25 @@ providers: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/nvidia/files} metadata_store: table_name: files_metadata -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/kvstore.db + default_sql_store: + 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 models: - metadata: {} model_id: ${env.INFERENCE_MODEL} diff --git a/llama_stack/distributions/nvidia/run.yaml b/llama_stack/distributions/nvidia/run.yaml index 2c08f021a..2bbaf37a1 100644 --- a/llama_stack/distributions/nvidia/run.yaml +++ b/llama_stack/distributions/nvidia/run.yaml @@ -74,15 +74,25 @@ providers: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/nvidia/files} metadata_store: table_name: files_metadata -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/kvstore.db + default_sql_store: + 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 models: [] shields: [] vector_dbs: [] diff --git a/llama_stack/distributions/open-benchmark/build.yaml b/llama_stack/distributions/open-benchmark/build.yaml index 05acd98e3..98e635506 100644 --- a/llama_stack/distributions/open-benchmark/build.yaml +++ b/llama_stack/distributions/open-benchmark/build.yaml @@ -31,6 +31,4 @@ distribution_spec: - provider_type: inline::rag-runtime - provider_type: remote::model-context-protocol image_type: venv -additional_pip_packages: -- aiosqlite -- sqlalchemy[asyncio] +additional_pip_packages: [] diff --git a/llama_stack/distributions/open-benchmark/run.yaml b/llama_stack/distributions/open-benchmark/run.yaml index 84337d935..d6859016b 100644 --- a/llama_stack/distributions/open-benchmark/run.yaml +++ b/llama_stack/distributions/open-benchmark/run.yaml @@ -112,15 +112,25 @@ providers: provider_type: inline::rag-runtime - provider_id: model-context-protocol provider_type: remote::model-context-protocol -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/kvstore.db + default_sql_store: + 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 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 063dc3999..09245f255 100644 --- a/llama_stack/distributions/postgres-demo/build.yaml +++ b/llama_stack/distributions/postgres-demo/build.yaml @@ -17,7 +17,4 @@ distribution_spec: - provider_type: inline::rag-runtime - provider_type: remote::model-context-protocol image_type: venv -additional_pip_packages: -- asyncpg -- psycopg2-binary -- sqlalchemy[asyncio] +additional_pip_packages: [] diff --git a/llama_stack/distributions/postgres-demo/postgres_demo.py b/llama_stack/distributions/postgres-demo/postgres_demo.py index 1f3e88b3b..9c18a305f 100644 --- a/llama_stack/distributions/postgres-demo/postgres_demo.py +++ b/llama_stack/distributions/postgres-demo/postgres_demo.py @@ -105,22 +105,18 @@ def get_distribution_template() -> DistributionTemplate: provider_overrides={ "inference": inference_providers + [embedding_provider], "vector_io": vector_io_providers, - "agents": [ - Provider( - provider_id="meta-reference", - provider_type="inline::meta-reference", - config=dict( - persistence_store=postgres_config, - responses_store=postgres_config, - ), - ) - ], }, 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( + table_name="llamastack_kvstore", + ), + "default_sql_store": PostgresSqlStoreConfig.sample_run_config(), + }, ), }, run_config_env_vars={ diff --git a/llama_stack/distributions/postgres-demo/run.yaml b/llama_stack/distributions/postgres-demo/run.yaml index d592db5ff..634289d5b 100644 --- a/llama_stack/distributions/postgres-demo/run.yaml +++ b/llama_stack/distributions/postgres-demo/run.yaml @@ -33,20 +33,11 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sql_postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} - responses_store: - type: 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} + persistence: + agent_state: + namespace: agents + responses: + table_name: responses tool_runtime: - provider_id: brave-search provider_type: remote::brave-search @@ -62,24 +53,34 @@ providers: provider_type: inline::rag-runtime - provider_id: model-context-protocol provider_type: remote::model-context-protocol -metadata_store: - type: kv_postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} - table_name: ${env.POSTGRES_TABLE_NAME:=llamastack_kvstore} -inference_store: - type: sql_postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/postgres-demo}/conversations.db +storage: + backends: + default_kv_store: + type: kv_postgres + host: ${env.POSTGRES_HOST:=localhost} + port: ${env.POSTGRES_PORT:=5432} + db: ${env.POSTGRES_DB:=llamastack} + user: ${env.POSTGRES_USER:=llamastack} + password: ${env.POSTGRES_PASSWORD:=llamastack} + table_name: ${env.POSTGRES_TABLE_NAME:=llamastack_kvstore} + default_sql_store: + 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 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 943c6134d..138e3b39e 100644 --- a/llama_stack/distributions/starter-gpu/build.yaml +++ b/llama_stack/distributions/starter-gpu/build.yaml @@ -53,6 +53,5 @@ 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 1e3b3b966..375d1fbc9 100644 --- a/llama_stack/distributions/starter-gpu/run.yaml +++ b/llama_stack/distributions/starter-gpu/run.yaml @@ -201,15 +201,25 @@ providers: config: kvstore: namespace: batches -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/kvstore.db + default_sql_store: + 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 models: [] shields: - shield_id: llama-guard diff --git a/llama_stack/distributions/starter/build.yaml b/llama_stack/distributions/starter/build.yaml index c2719d50d..4d149cf05 100644 --- a/llama_stack/distributions/starter/build.yaml +++ b/llama_stack/distributions/starter/build.yaml @@ -53,6 +53,5 @@ 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 9c618e95d..25ca22137 100644 --- a/llama_stack/distributions/starter/run.yaml +++ b/llama_stack/distributions/starter/run.yaml @@ -198,15 +198,25 @@ providers: config: kvstore: namespace: batches -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/kvstore.db + default_sql_store: + 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 models: [] shields: - shield_id: llama-guard diff --git a/llama_stack/distributions/template.py b/llama_stack/distributions/template.py index 807829999..e72744429 100644 --- a/llama_stack/distributions/template.py +++ b/llama_stack/distributions/template.py @@ -29,6 +29,12 @@ from llama_stack.core.datatypes import ( ToolGroupInput, ) from llama_stack.core.distribution import get_provider_registry +from llama_stack.core.storage.datatypes import ( + InferenceStoreReference, + KVStoreReference, + SqlStoreReference, + StorageBackendType, +) from llama_stack.core.utils.dynamic import instantiate_class_type from llama_stack.core.utils.image_types import LlamaStackImageType from llama_stack.providers.utils.inference.model_registry import ProviderModelEntry @@ -180,10 +186,8 @@ class RunConfigSettings(BaseModel): default_tool_groups: list[ToolGroupInput] | None = None default_datasets: list[DatasetInput] | None = None default_benchmarks: list[BenchmarkInput] | None = None - metadata_store: dict | None = None - inference_store: dict | None = None - conversations_store: dict | None = None telemetry: TelemetryConfig = Field(default_factory=lambda: TelemetryConfig(enabled=True)) + storage_backends: dict[str, Any] | None = None def run_config( self, @@ -226,6 +230,33 @@ class RunConfigSettings(BaseModel): # Get unique set of APIs from providers apis = sorted(providers.keys()) + storage_backends = self.storage_backends or { + "default_kv_store": SqliteKVStoreConfig.sample_run_config( + __distro_dir__=f"~/.llama/distributions/{name}", + db_name="kvstore.db", + ), + "default_sql_store": 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), + ) + # Return a dict that matches StackRunConfig structure return { "version": LLAMA_STACK_RUN_CONFIG_VERSION, @@ -233,21 +264,7 @@ class RunConfigSettings(BaseModel): "container_image": container_image, "apis": apis, "providers": provider_configs, - "metadata_store": self.metadata_store - or SqliteKVStoreConfig.sample_run_config( - __distro_dir__=f"~/.llama/distributions/{name}", - db_name="registry.db", - ), - "inference_store": self.inference_store - or SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=f"~/.llama/distributions/{name}", - db_name="inference_store.db", - ), - "conversations_store": self.conversations_store - or SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=f"~/.llama/distributions/{name}", - db_name="conversations.db", - ), + "storage": storage_config, "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": [], @@ -387,11 +404,13 @@ class DistributionTemplate(BaseModel): def enum_representer(dumper, data): return dumper.represent_scalar("tag:yaml.org,2002:str", data.value) - # Register YAML representer for ModelType + # Register YAML representer for enums yaml.add_representer(ModelType, enum_representer) yaml.add_representer(DatasetPurpose, enum_representer) + yaml.add_representer(StorageBackendType, enum_representer) yaml.SafeDumper.add_representer(ModelType, enum_representer) yaml.SafeDumper.add_representer(DatasetPurpose, enum_representer) + yaml.SafeDumper.add_representer(StorageBackendType, enum_representer) for output_dir in [yaml_output_dir, doc_output_dir]: output_dir.mkdir(parents=True, exist_ok=True) diff --git a/llama_stack/distributions/watsonx/build.yaml b/llama_stack/distributions/watsonx/build.yaml index dba1a94e2..f81977347 100644 --- a/llama_stack/distributions/watsonx/build.yaml +++ b/llama_stack/distributions/watsonx/build.yaml @@ -28,6 +28,4 @@ distribution_spec: files: - provider_type: inline::localfs image_type: venv -additional_pip_packages: -- aiosqlite -- sqlalchemy[asyncio] +additional_pip_packages: [] diff --git a/llama_stack/distributions/watsonx/run.yaml b/llama_stack/distributions/watsonx/run.yaml index 8d1bfd58f..5b2a05d3d 100644 --- a/llama_stack/distributions/watsonx/run.yaml +++ b/llama_stack/distributions/watsonx/run.yaml @@ -86,15 +86,25 @@ providers: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/watsonx/files} metadata_store: table_name: files_metadata -metadata_store: - type: kv_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/registry.db -inference_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/inference_store.db -conversations_store: - type: sql_sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/conversations.db +storage: + backends: + default_kv_store: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/kvstore.db + default_sql_store: + 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 models: [] shields: [] vector_dbs: [] diff --git a/tests/integration/test_persistence_integration.py b/tests/integration/test_persistence_integration.py index 93c44bd6e..88b43ecc4 100644 --- a/tests/integration/test_persistence_integration.py +++ b/tests/integration/test_persistence_integration.py @@ -9,7 +9,6 @@ import yaml from llama_stack.core.datatypes import StackRunConfig from llama_stack.core.storage.datatypes import ( PostgresSqlStoreConfig, - SqliteKVStoreConfig, SqliteSqlStoreConfig, )