diff --git a/.github/actions/run-and-record-tests/action.yml b/.github/actions/run-and-record-tests/action.yml index a5aa31af4..3929df09c 100644 --- a/.github/actions/run-and-record-tests/action.yml +++ b/.github/actions/run-and-record-tests/action.yml @@ -82,11 +82,14 @@ runs: echo "No recording changes" fi - - name: Write inference logs to file + - name: Write docker logs to file if: ${{ always() }} shell: bash run: | sudo docker logs ollama > ollama-${{ inputs.inference-mode }}.log || true + distro_name=$(echo "${{ inputs.stack-config }}" | sed 's/^docker://' | sed 's/^server://') + stack_container_name="llama-stack-test-$distro_name" + sudo docker logs $stack_container_name > docker-${distro_name}-${{ inputs.inference-mode }}.log || true - name: Upload logs if: ${{ always() }} diff --git a/.github/workflows/integration-auth-tests.yml b/.github/workflows/integration-auth-tests.yml index ea3ff2b64..30a8063ea 100644 --- a/.github/workflows/integration-auth-tests.yml +++ b/.github/workflows/integration-auth-tests.yml @@ -73,6 +73,24 @@ jobs: image_name: kube apis: [] providers: {} + storage: + backends: + kv_default: + type: kv_sqlite + db_path: $run_dir/kvstore.db + sql_default: + type: sql_sqlite + db_path: $run_dir/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + conversations: + table_name: openai_conversations + backend: sql_default server: port: 8321 EOF diff --git a/benchmarking/k8s-benchmark/stack-configmap.yaml b/benchmarking/k8s-benchmark/stack-configmap.yaml index bb8a48d65..e1ca170f5 100644 --- a/benchmarking/k8s-benchmark/stack-configmap.yaml +++ b/benchmarking/k8s-benchmark/stack-configmap.yaml @@ -98,21 +98,30 @@ data: - provider_id: model-context-protocol provider_type: remote::model-context-protocol config: {} - metadata_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} - table_name: llamastack_kvstore - inference_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} + storage: + backends: + kv_default: + 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} + 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} + references: + metadata: + backend: kv_default + namespace: registry + inference: + backend: sql_default + table_name: inference_store models: - metadata: embedding_dimension: 768 @@ -137,5 +146,4 @@ data: port: 8323 kind: ConfigMap metadata: - creationTimestamp: null name: llama-stack-config diff --git a/benchmarking/k8s-benchmark/stack_run_config.yaml b/benchmarking/k8s-benchmark/stack_run_config.yaml index e2fbfd7a4..2ccaa21aa 100644 --- a/benchmarking/k8s-benchmark/stack_run_config.yaml +++ b/benchmarking/k8s-benchmark/stack_run_config.yaml @@ -95,21 +95,30 @@ providers: - provider_id: model-context-protocol provider_type: remote::model-context-protocol config: {} -metadata_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} - table_name: llamastack_kvstore -inference_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} +storage: + backends: + kv_default: + 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} + 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} + references: + metadata: + backend: kv_default + namespace: registry + inference: + backend: sql_default + table_name: inference_store models: - metadata: embedding_dimension: 768 diff --git a/docs/docs/distributions/configuration.mdx b/docs/docs/distributions/configuration.mdx index 81243c97b..bf3156865 100644 --- a/docs/docs/distributions/configuration.mdx +++ b/docs/docs/distributions/configuration.mdx @@ -44,18 +44,32 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - namespace: null - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ollama}/agents_store.db + persistence: + agent_state: + backend: kv_default + namespace: agents + responses: + backend: sql_default + table_name: responses telemetry: - provider_id: meta-reference provider_type: inline::meta-reference config: {} -metadata_store: - namespace: null - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ollama}/registry.db +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ollama}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ollama}/sqlstore.db + references: + metadata: + backend: kv_default + namespace: registry + inference: + backend: sql_default + table_name: inference_store models: - metadata: {} model_id: ${env.INFERENCE_MODEL} diff --git a/docs/docs/distributions/k8s/stack-configmap.yaml b/docs/docs/distributions/k8s/stack-configmap.yaml index 3dbb0da97..c71ab05d8 100644 --- a/docs/docs/distributions/k8s/stack-configmap.yaml +++ b/docs/docs/distributions/k8s/stack-configmap.yaml @@ -1,56 +1,155 @@ apiVersion: v1 data: - stack_run_config.yaml: "version: '2'\nimage_name: kubernetes-demo\napis:\n- agents\n- - inference\n- files\n- safety\n- telemetry\n- tool_runtime\n- vector_io\nproviders:\n - \ inference:\n - provider_id: vllm-inference\n provider_type: remote::vllm\n - \ config:\n url: ${env.VLLM_URL:=http://localhost:8000/v1}\n max_tokens: - ${env.VLLM_MAX_TOKENS:=4096}\n api_token: ${env.VLLM_API_TOKEN:=fake}\n tls_verify: - ${env.VLLM_TLS_VERIFY:=true}\n - provider_id: vllm-safety\n provider_type: - remote::vllm\n config:\n url: ${env.VLLM_SAFETY_URL:=http://localhost:8000/v1}\n - \ max_tokens: ${env.VLLM_MAX_TOKENS:=4096}\n api_token: ${env.VLLM_API_TOKEN:=fake}\n - \ tls_verify: ${env.VLLM_TLS_VERIFY:=true}\n - provider_id: sentence-transformers\n - \ provider_type: inline::sentence-transformers\n config: {}\n vector_io:\n - \ - provider_id: ${env.ENABLE_CHROMADB:+chromadb}\n provider_type: remote::chromadb\n - \ config:\n url: ${env.CHROMADB_URL:=}\n kvstore:\n type: postgres\n - \ host: ${env.POSTGRES_HOST:=localhost}\n port: ${env.POSTGRES_PORT:=5432}\n - \ db: ${env.POSTGRES_DB:=llamastack}\n user: ${env.POSTGRES_USER:=llamastack}\n - \ password: ${env.POSTGRES_PASSWORD:=llamastack}\n files:\n - provider_id: - meta-reference-files\n provider_type: inline::localfs\n config:\n storage_dir: - ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter/files}\n metadata_store:\n - \ type: sqlite\n db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/files_metadata.db - \ \n safety:\n - provider_id: llama-guard\n provider_type: inline::llama-guard\n - \ config:\n excluded_categories: []\n agents:\n - provider_id: meta-reference\n - \ provider_type: inline::meta-reference\n config:\n persistence_store:\n - \ type: postgres\n host: ${env.POSTGRES_HOST:=localhost}\n port: - ${env.POSTGRES_PORT:=5432}\n db: ${env.POSTGRES_DB:=llamastack}\n user: - ${env.POSTGRES_USER:=llamastack}\n password: ${env.POSTGRES_PASSWORD:=llamastack}\n - \ responses_store:\n type: postgres\n host: ${env.POSTGRES_HOST:=localhost}\n - \ port: ${env.POSTGRES_PORT:=5432}\n db: ${env.POSTGRES_DB:=llamastack}\n - \ user: ${env.POSTGRES_USER:=llamastack}\n password: ${env.POSTGRES_PASSWORD:=llamastack}\n - \ telemetry:\n - provider_id: meta-reference\n provider_type: inline::meta-reference\n - \ config:\n service_name: \"${env.OTEL_SERVICE_NAME:=\\u200B}\"\n sinks: - ${env.TELEMETRY_SINKS:=console}\n tool_runtime:\n - provider_id: brave-search\n - \ provider_type: remote::brave-search\n config:\n api_key: ${env.BRAVE_SEARCH_API_KEY:+}\n - \ max_results: 3\n - provider_id: tavily-search\n provider_type: remote::tavily-search\n - \ config:\n api_key: ${env.TAVILY_SEARCH_API_KEY:+}\n max_results: - 3\n - provider_id: rag-runtime\n provider_type: inline::rag-runtime\n config: - {}\n - provider_id: model-context-protocol\n provider_type: remote::model-context-protocol\n - \ config: {}\nmetadata_store:\n type: postgres\n host: ${env.POSTGRES_HOST:=localhost}\n - \ port: ${env.POSTGRES_PORT:=5432}\n db: ${env.POSTGRES_DB:=llamastack}\n user: - ${env.POSTGRES_USER:=llamastack}\n password: ${env.POSTGRES_PASSWORD:=llamastack}\n - \ table_name: llamastack_kvstore\ninference_store:\n type: postgres\n host: - ${env.POSTGRES_HOST:=localhost}\n port: ${env.POSTGRES_PORT:=5432}\n db: ${env.POSTGRES_DB:=llamastack}\n - \ user: ${env.POSTGRES_USER:=llamastack}\n password: ${env.POSTGRES_PASSWORD:=llamastack}\nmodels:\n- - metadata:\n embedding_dimension: 384\n model_id: all-MiniLM-L6-v2\n provider_id: - sentence-transformers\n model_type: embedding\n- metadata: {}\n model_id: ${env.INFERENCE_MODEL}\n - \ provider_id: vllm-inference\n model_type: llm\n- metadata: {}\n model_id: - ${env.SAFETY_MODEL:=meta-llama/Llama-Guard-3-1B}\n provider_id: vllm-safety\n - \ model_type: llm\nshields:\n- shield_id: ${env.SAFETY_MODEL:=meta-llama/Llama-Guard-3-1B}\nvector_dbs: - []\ndatasets: []\nscoring_fns: []\nbenchmarks: []\ntool_groups:\n- toolgroup_id: - builtin::websearch\n provider_id: tavily-search\n- toolgroup_id: builtin::rag\n - \ provider_id: rag-runtime\nserver:\n port: 8321\n auth:\n provider_config:\n - \ type: github_token\n" + stack_run_config.yaml: | + version: '2' + image_name: kubernetes-demo + apis: + - agents + - inference + - files + - safety + - telemetry + - tool_runtime + - vector_io + providers: + inference: + - provider_id: vllm-inference + provider_type: remote::vllm + config: + url: ${env.VLLM_URL:=http://localhost:8000/v1} + max_tokens: ${env.VLLM_MAX_TOKENS:=4096} + api_token: ${env.VLLM_API_TOKEN:=fake} + tls_verify: ${env.VLLM_TLS_VERIFY:=true} + - provider_id: vllm-safety + provider_type: remote::vllm + config: + url: ${env.VLLM_SAFETY_URL:=http://localhost:8000/v1} + max_tokens: ${env.VLLM_MAX_TOKENS:=4096} + api_token: ${env.VLLM_API_TOKEN:=fake} + tls_verify: ${env.VLLM_TLS_VERIFY:=true} + - provider_id: sentence-transformers + provider_type: inline::sentence-transformers + config: {} + vector_io: + - provider_id: ${env.ENABLE_CHROMADB:+chromadb} + provider_type: remote::chromadb + config: + url: ${env.CHROMADB_URL:=} + kvstore: + type: postgres + host: ${env.POSTGRES_HOST:=localhost} + port: ${env.POSTGRES_PORT:=5432} + db: ${env.POSTGRES_DB:=llamastack} + user: ${env.POSTGRES_USER:=llamastack} + password: ${env.POSTGRES_PASSWORD:=llamastack} + files: + - provider_id: meta-reference-files + provider_type: inline::localfs + config: + storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter/files} + metadata_store: + type: sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/files_metadata.db + safety: + - provider_id: llama-guard + provider_type: inline::llama-guard + config: + excluded_categories: [] + agents: + - provider_id: meta-reference + provider_type: inline::meta-reference + config: + persistence_store: + type: postgres + host: ${env.POSTGRES_HOST:=localhost} + port: ${env.POSTGRES_PORT:=5432} + db: ${env.POSTGRES_DB:=llamastack} + user: ${env.POSTGRES_USER:=llamastack} + password: ${env.POSTGRES_PASSWORD:=llamastack} + responses_store: + type: postgres + host: ${env.POSTGRES_HOST:=localhost} + port: ${env.POSTGRES_PORT:=5432} + db: ${env.POSTGRES_DB:=llamastack} + user: ${env.POSTGRES_USER:=llamastack} + password: ${env.POSTGRES_PASSWORD:=llamastack} + telemetry: + - provider_id: meta-reference + provider_type: inline::meta-reference + config: + service_name: "${env.OTEL_SERVICE_NAME:=\u200B}" + sinks: ${env.TELEMETRY_SINKS:=console} + tool_runtime: + - provider_id: brave-search + provider_type: remote::brave-search + config: + api_key: ${env.BRAVE_SEARCH_API_KEY:+} + max_results: 3 + - provider_id: tavily-search + provider_type: remote::tavily-search + config: + api_key: ${env.TAVILY_SEARCH_API_KEY:+} + max_results: 3 + - provider_id: rag-runtime + provider_type: inline::rag-runtime + config: {} + - provider_id: model-context-protocol + provider_type: remote::model-context-protocol + config: {} + storage: + backends: + kv_default: + 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} + 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} + references: + metadata: + backend: kv_default + namespace: registry + inference: + backend: sql_default + table_name: inference_store + models: + - metadata: + embedding_dimension: 768 + model_id: nomic-embed-text-v1.5 + provider_id: sentence-transformers + model_type: embedding + - metadata: {} + model_id: ${env.INFERENCE_MODEL} + provider_id: vllm-inference + model_type: llm + - metadata: {} + model_id: ${env.SAFETY_MODEL:=meta-llama/Llama-Guard-3-1B} + provider_id: vllm-safety + model_type: llm + shields: + - shield_id: ${env.SAFETY_MODEL:=meta-llama/Llama-Guard-3-1B} + vector_dbs: [] + datasets: [] + scoring_fns: [] + benchmarks: [] + tool_groups: + - toolgroup_id: builtin::websearch + provider_id: tavily-search + - toolgroup_id: builtin::rag + provider_id: rag-runtime + server: + port: 8321 + auth: + provider_config: + type: github_token kind: ConfigMap metadata: - creationTimestamp: null name: llama-stack-config diff --git a/docs/docs/distributions/k8s/stack_run_config.yaml b/docs/docs/distributions/k8s/stack_run_config.yaml index ee28a1ea8..863565fdf 100644 --- a/docs/docs/distributions/k8s/stack_run_config.yaml +++ b/docs/docs/distributions/k8s/stack_run_config.yaml @@ -93,21 +93,30 @@ providers: - provider_id: model-context-protocol provider_type: remote::model-context-protocol config: {} -metadata_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} - table_name: llamastack_kvstore -inference_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} +storage: + backends: + kv_default: + 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} + 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} + references: + metadata: + backend: kv_default + namespace: registry + inference: + backend: sql_default + table_name: inference_store models: - metadata: embedding_dimension: 768 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..0be5cd5b3 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 | +| `persistence` | `` | 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 +persistence: + 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..3a1fba055 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 | | +| `persistence` | `` | No | | | ## Sample Configuration ```yaml -kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/faiss_store.db +persistence: + 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..17fd40cf5 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 | | +| `persistence` | `` | No | | | ## Sample Configuration ```yaml -kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/faiss_store.db +persistence: + 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..6063edab1 100644 --- a/docs/docs/providers/vector_io/inline_milvus.mdx +++ b/docs/docs/providers/vector_io/inline_milvus.mdx @@ -17,14 +17,14 @@ 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) | +| `persistence` | `` | No | | Config for KV store backend (SQLite only for now) | | `consistency_level` | `` | No | Strong | The consistency level of the Milvus server | ## Sample Configuration ```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 +persistence: + 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..057d96761 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 | | +| `persistence` | `` | 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 +persistence: + 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..98a372250 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) | +| `persistence` | `` | 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 +persistence: + 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..67cbd0021 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) | +| `persistence` | `` | 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 +persistence: + 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..2aee3eeca 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 | +| `persistence` | `` | 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 +persistence: + 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..bf9935d61 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 | +| `persistence` | `` | 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 @@ -420,7 +420,7 @@ This configuration class accepts additional fields beyond those listed above. Yo ```yaml uri: ${env.MILVUS_ENDPOINT} token: ${env.MILVUS_TOKEN} -kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/milvus_remote_registry.db +persistence: + 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..cb70f35d1 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) | +| `persistence` | `llama_stack.core.storage.datatypes.KVStoreReference \| None` | No | | Config for KV store backend (SQLite only for now) | ## Sample Configuration @@ -228,7 +228,7 @@ port: ${env.PGVECTOR_PORT:=5432} 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 +persistence: + 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..dff9642b5 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 | | +| `persistence` | `` | No | | | ## Sample Configuration ```yaml api_key: ${env.QDRANT_API_KEY:=} -kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/qdrant_registry.db +persistence: + 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..b809bed2e 100644 --- a/docs/docs/providers/vector_io/remote_weaviate.mdx +++ b/docs/docs/providers/vector_io/remote_weaviate.mdx @@ -75,14 +75,14 @@ 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) | +| `persistence` | `llama_stack.core.storage.datatypes.KVStoreReference \| None` | No | | Config for KV store backend (SQLite only for now) | ## Sample Configuration ```yaml 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 +persistence: + namespace: vector_io::weaviate + backend: kv_default ``` diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py index 471d5cb66..2a30ff394 100644 --- a/llama_stack/cli/stack/_build.py +++ b/llama_stack/cli/stack/_build.py @@ -40,12 +40,20 @@ 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, + ServerStoresConfig, + 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 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" @@ -286,21 +294,42 @@ 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 = DISTRIBS_BASE_DIR / 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", + ), + }, + stores=ServerStoresConfig( + metadata=KVStoreReference( + backend="kv_default", + namespace="registry", + ), + inference=InferenceStoreReference( + backend="sql_default", + table_name="inference_store", + ), + conversations=SqlStoreReference( + backend="sql_default", + table_name="openai_conversations", + ), + ), + ) + 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, external_providers_dir=build_config.external_providers_dir if build_config.external_providers_dir else EXTERNAL_PROVIDERS_DIR, ) - if not run_config.inference_store: - run_config.inference_store = SqliteSqlStoreConfig( - **SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=(DISTRIBS_BASE_DIR / image_name).as_posix(), db_name="inference_store.db" - ) - ) # build providers dict provider_registry = get_provider_registry(build_config) for api in apis: diff --git a/llama_stack/cli/stack/utils.py b/llama_stack/cli/stack/utils.py index 4d4c1b538..cc1ca051b 100644 --- a/llama_stack/cli/stack/utils.py +++ b/llama_stack/cli/stack/utils.py @@ -17,10 +17,19 @@ from llama_stack.core.datatypes import ( BuildConfig, Provider, StackRunConfig, + StorageConfig, ) from llama_stack.core.distribution import get_provider_registry from llama_stack.core.resolver import InvalidProviderError -from llama_stack.core.utils.config_dirs import EXTERNAL_PROVIDERS_DIR +from llama_stack.core.storage.datatypes import ( + InferenceStoreReference, + KVStoreReference, + ServerStoresConfig, + SqliteKVStoreConfig, + SqliteSqlStoreConfig, + SqlStoreReference, +) +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.image_types import LlamaStackImageType from llama_stack.providers.datatypes import Api @@ -51,11 +60,23 @@ 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 = DISTRIBS_BASE_DIR / image_name 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=StorageConfig( + backends={ + "kv_default": SqliteKVStoreConfig(db_path=str(distro_dir / "kvstore.db")), + "sql_default": SqliteSqlStoreConfig(db_path=str(distro_dir / "sql_store.db")), + }, + stores=ServerStoresConfig( + metadata=KVStoreReference(backend="kv_default", namespace="registry"), + inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"), + conversations=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, diff --git a/llama_stack/core/configure.py b/llama_stack/core/configure.py index bfa2c6d71..734839ea9 100644 --- a/llama_stack/core/configure.py +++ b/llama_stack/core/configure.py @@ -159,6 +159,37 @@ def upgrade_from_routing_table( config_dict["apis"] = config_dict["apis_to_serve"] config_dict.pop("apis_to_serve", None) + # Add default storage config if not present + if "storage" not in config_dict: + config_dict["storage"] = { + "backends": { + "kv_default": { + "type": "kv_sqlite", + "db_path": "~/.llama/kvstore.db", + }, + "sql_default": { + "type": "sql_sqlite", + "db_path": "~/.llama/sql_store.db", + }, + }, + "stores": { + "metadata": { + "namespace": "registry", + "backend": "kv_default", + }, + "inference": { + "table_name": "inference_store", + "backend": "sql_default", + "max_write_queue_size": 10000, + "num_writers": 4, + }, + "conversations": { + "table_name": "openai_conversations", + "backend": "sql_default", + }, + }, + } + return config_dict diff --git a/llama_stack/core/conversations/conversations.py b/llama_stack/core/conversations/conversations.py index d2537c7ee..66880ca36 100644 --- a/llama_stack/core/conversations/conversations.py +++ b/llama_stack/core/conversations/conversations.py @@ -4,7 +4,6 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import os import secrets import time from typing import Any @@ -21,16 +20,11 @@ from llama_stack.apis.conversations.conversations import ( Conversations, Metadata, ) -from llama_stack.core.datatypes import AccessRule -from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR +from llama_stack.core.datatypes import AccessRule, StackRunConfig from llama_stack.log import get_logger from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore -from llama_stack.providers.utils.sqlstore.sqlstore import ( - SqliteSqlStoreConfig, - SqlStoreConfig, - sqlstore_impl, -) +from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl logger = get_logger(name=__name__, category="openai_conversations") @@ -38,13 +32,11 @@ logger = get_logger(name=__name__, category="openai_conversations") class ConversationServiceConfig(BaseModel): """Configuration for the built-in conversation service. - :param conversations_store: SQL store configuration for conversations (defaults to SQLite) + :param run_config: Stack run configuration for resolving persistence :param policy: Access control rules """ - conversations_store: SqlStoreConfig = SqliteSqlStoreConfig( - db_path=(DISTRIBS_BASE_DIR / "conversations.db").as_posix() - ) + run_config: StackRunConfig policy: list[AccessRule] = [] @@ -63,14 +55,16 @@ class ConversationServiceImpl(Conversations): self.deps = deps self.policy = config.policy - base_sql_store = sqlstore_impl(config.conversations_store) + # Use conversations store reference from run config + conversations_ref = config.run_config.storage.stores.conversations + if not conversations_ref: + raise ValueError("storage.stores.conversations must be configured in run config") + + base_sql_store = sqlstore_impl(conversations_ref) self.sql_store = AuthorizedSqlStore(base_sql_store, self.policy) async def initialize(self) -> None: """Initialize the store and create tables.""" - if isinstance(self.config.conversations_store, SqliteSqlStoreConfig): - os.makedirs(os.path.dirname(self.config.conversations_store.db_path), exist_ok=True) - await self.sql_store.create_table( "openai_conversations", { diff --git a/llama_stack/core/datatypes.py b/llama_stack/core/datatypes.py index 94222d49e..d692da3b3 100644 --- a/llama_stack/core/datatypes.py +++ b/llama_stack/core/datatypes.py @@ -26,9 +26,12 @@ 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, + StorageBackendType, + StorageConfig, +) from llama_stack.providers.datatypes import Api, ProviderSpec -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig -from llama_stack.providers.utils.sqlstore.sqlstore import SqlStoreConfig LLAMA_STACK_BUILD_CONFIG_VERSION = 2 LLAMA_STACK_RUN_CONFIG_VERSION = 2 @@ -356,7 +359,7 @@ class QuotaPeriod(StrEnum): class QuotaConfig(BaseModel): - kvstore: SqliteKVStoreConfig = Field(description="Config for KV store backend (SQLite only for now)") + kvstore: KVStoreReference = Field(description="Config for KV store backend (SQLite only for now)") anonymous_max_requests: int = Field(default=100, description="Max requests for unauthenticated clients per period") authenticated_max_requests: int = Field( default=1000, description="Max requests for authenticated clients per period" @@ -438,18 +441,6 @@ class ServerConfig(BaseModel): ) -class InferenceStoreConfig(BaseModel): - sql_store_config: SqlStoreConfig - max_write_queue_size: int = Field(default=10000, description="Max queued writes for inference store") - num_writers: int = Field(default=4, description="Number of concurrent background writers") - - -class ResponsesStoreConfig(BaseModel): - sql_store_config: SqlStoreConfig - max_write_queue_size: int = Field(default=10000, description="Max queued writes for responses store") - num_writers: int = Field(default=4, description="Number of concurrent background writers") - - class StackRunConfig(BaseModel): version: int = LLAMA_STACK_RUN_CONFIG_VERSION @@ -476,26 +467,8 @@ One or more providers to use for each API. The same provider_type (e.g., meta-re can be instantiated multiple times (with different configs) if necessary. """, ) - metadata_store: KVStoreConfig | None = Field( - default=None, - description=""" -Configuration for the persistence store used by the distribution registry. If not specified, -a default SQLite store will be used.""", - ) - - inference_store: InferenceStoreConfig | SqlStoreConfig | None = Field( - default=None, - description=""" -Configuration for the persistence store used by the inference API. Can be either a -InferenceStoreConfig (with queue tuning parameters) or a SqlStoreConfig (deprecated). -If not specified, a default SQLite store will be used.""", - ) - - conversations_store: SqlStoreConfig | None = Field( - default=None, - description=""" -Configuration for the persistence store used by the conversations API. -If not specified, a default SQLite store will be used.""", + storage: StorageConfig = Field( + description="Catalog of named storage backends and references available to the stack", ) # registry of "resources" in the distribution @@ -535,6 +508,49 @@ If not specified, a default SQLite store will be used.""", return Path(v) return v + @model_validator(mode="after") + def validate_server_stores(self) -> "StackRunConfig": + backend_map = self.storage.backends + stores = self.storage.stores + 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(stores.metadata, kv_backends, "storage.stores.metadata") + _ensure_backend(stores.inference, sql_backends, "storage.stores.inference") + _ensure_backend(stores.conversations, sql_backends, "storage.stores.conversations") + _ensure_backend(stores.responses, sql_backends, "storage.stores.responses") + 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 26e8f5cef..856397ca5 100644 --- a/llama_stack/core/prompts/prompts.py +++ b/llama_stack/core/prompts/prompts.py @@ -11,9 +11,8 @@ from pydantic import BaseModel from llama_stack.apis.prompts import ListPromptsResponse, Prompt, Prompts from llama_stack.core.datatypes import StackRunConfig -from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl -from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig class PromptServiceConfig(BaseModel): @@ -41,10 +40,12 @@ class PromptServiceImpl(Prompts): self.kvstore: KVStore async def initialize(self) -> None: - kvstore_config = SqliteKVStoreConfig( - db_path=(DISTRIBS_BASE_DIR / self.config.run_config.image_name / "prompts.db").as_posix() - ) - self.kvstore = await kvstore_impl(kvstore_config) + # Use metadata store backend with prompts-specific namespace + metadata_ref = self.config.run_config.storage.stores.metadata + if not metadata_ref: + raise ValueError("storage.stores.metadata 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: """Get the KVStore key that stores the default version number.""" diff --git a/llama_stack/core/routers/__init__.py b/llama_stack/core/routers/__init__.py index 4463d2460..0573fc2c7 100644 --- a/llama_stack/core/routers/__init__.py +++ b/llama_stack/core/routers/__init__.py @@ -6,7 +6,10 @@ from typing import Any -from llama_stack.core.datatypes import AccessRule, RoutedProtocol +from llama_stack.core.datatypes import ( + AccessRule, + RoutedProtocol, +) from llama_stack.core.stack import StackRunConfig from llama_stack.core.store import DistributionRegistry from llama_stack.providers.datatypes import Api, RoutingTable @@ -76,9 +79,13 @@ async def get_auto_router_impl( api_to_dep_impl[dep_name] = deps[dep_api] # TODO: move pass configs to routers instead - if api == Api.inference and run_config.inference_store: + if api == Api.inference: + inference_ref = run_config.storage.stores.inference + if not inference_ref: + raise ValueError("storage.stores.inference must be configured in run config") + inference_store = InferenceStore( - config=run_config.inference_store, + reference=inference_ref, policy=policy, ) await inference_store.initialize() diff --git a/llama_stack/core/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 733b55262..15d0198b1 100644 --- a/llama_stack/core/stack.py +++ b/llama_stack/core/stack.py @@ -42,6 +42,16 @@ from llama_stack.core.prompts.prompts import PromptServiceConfig, PromptServiceI from llama_stack.core.providers import ProviderImpl, ProviderImplConfig from llama_stack.core.resolver import ProviderRegistry, resolve_impls from llama_stack.core.routing_tables.common import CommonRoutingTableImpl +from llama_stack.core.storage.datatypes import ( + InferenceStoreReference, + KVStoreReference, + ServerStoresConfig, + SqliteKVStoreConfig, + SqliteSqlStoreConfig, + SqlStoreReference, + StorageBackendConfig, + StorageConfig, +) from llama_stack.core.store.registry import create_dist_registry from llama_stack.core.utils.dynamic import instantiate_class_type from llama_stack.log import get_logger @@ -329,6 +339,25 @@ def add_internal_implementations(impls: dict[Api, Any], run_config: StackRunConf impls[Api.conversations] = conversations_impl +def _initialize_storage(run_config: StackRunConfig): + kv_backends: dict[str, StorageBackendConfig] = {} + sql_backends: dict[str, StorageBackendConfig] = {} + for backend_name, backend_config in run_config.storage.backends.items(): + type = backend_config.type.value + if type.startswith("kv_"): + kv_backends[backend_name] = backend_config + elif type.startswith("sql_"): + sql_backends[backend_name] = backend_config + else: + raise ValueError(f"Unknown storage backend type: {type}") + + from llama_stack.providers.utils.kvstore.kvstore import register_kvstore_backends + from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends + + register_kvstore_backends(kv_backends) + register_sqlstore_backends(sql_backends) + + class Stack: def __init__(self, run_config: StackRunConfig, provider_registry: ProviderRegistry | None = None): self.run_config = run_config @@ -347,7 +376,11 @@ class Stack: TEST_RECORDING_CONTEXT.__enter__() logger.info(f"API recording enabled: mode={os.environ.get('LLAMA_STACK_TEST_INFERENCE_MODE')}") - dist_registry, _ = await create_dist_registry(self.run_config.metadata_store, self.run_config.image_name) + _initialize_storage(self.run_config) + stores = self.run_config.storage.stores + if not stores.metadata: + raise ValueError("storage.stores.metadata must be configured with a kv_* backend") + dist_registry, _ = await create_dist_registry(stores.metadata, self.run_config.image_name) policy = self.run_config.server.auth.access_policy if self.run_config.server.auth else [] internal_impls = {} @@ -488,5 +521,16 @@ def run_config_from_adhoc_config_spec( image_name="distro-test", apis=list(provider_configs_by_api.keys()), providers=provider_configs_by_api, + storage=StorageConfig( + backends={ + "kv_default": SqliteKVStoreConfig(db_path=f"{distro_dir}/kvstore.db"), + "sql_default": SqliteSqlStoreConfig(db_path=f"{distro_dir}/sql_store.db"), + }, + stores=ServerStoresConfig( + metadata=KVStoreReference(backend="kv_default", namespace="registry"), + inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"), + conversations=SqlStoreReference(backend="sql_default", table_name="openai_conversations"), + ), + ), ) return config diff --git a/llama_stack/core/storage/__init__.py b/llama_stack/core/storage/__init__.py new file mode 100644 index 000000000..756f351d8 --- /dev/null +++ b/llama_stack/core/storage/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. diff --git a/llama_stack/core/storage/datatypes.py b/llama_stack/core/storage/datatypes.py new file mode 100644 index 000000000..9df170e10 --- /dev/null +++ b/llama_stack/core/storage/datatypes.py @@ -0,0 +1,283 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import re +from abc import abstractmethod +from enum import StrEnum +from pathlib import Path +from typing import Annotated, Literal + +from pydantic import BaseModel, Field, field_validator + + +class StorageBackendType(StrEnum): + KV_REDIS = "kv_redis" + KV_SQLITE = "kv_sqlite" + KV_POSTGRES = "kv_postgres" + KV_MONGODB = "kv_mongodb" + SQL_SQLITE = "sql_sqlite" + SQL_POSTGRES = "sql_postgres" + + +class CommonConfig(BaseModel): + namespace: str | None = Field( + default=None, + description="All keys will be prefixed with this namespace", + ) + + +class RedisKVStoreConfig(CommonConfig): + type: Literal[StorageBackendType.KV_REDIS] = StorageBackendType.KV_REDIS + host: str = "localhost" + port: int = 6379 + + @property + def url(self) -> str: + return f"redis://{self.host}:{self.port}" + + @classmethod + def pip_packages(cls) -> list[str]: + return ["redis"] + + @classmethod + def sample_run_config(cls): + return { + "type": StorageBackendType.KV_REDIS.value, + "host": "${env.REDIS_HOST:=localhost}", + "port": "${env.REDIS_PORT:=6379}", + } + + +class SqliteKVStoreConfig(CommonConfig): + type: Literal[StorageBackendType.KV_SQLITE] = StorageBackendType.KV_SQLITE + db_path: str = Field( + description="File path for the sqlite database", + ) + + @classmethod + def pip_packages(cls) -> list[str]: + return ["aiosqlite"] + + @classmethod + def sample_run_config(cls, __distro_dir__: str, db_name: str = "kvstore.db"): + return { + "type": StorageBackendType.KV_SQLITE.value, + "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + db_name, + } + + +class PostgresKVStoreConfig(CommonConfig): + type: Literal[StorageBackendType.KV_POSTGRES] = StorageBackendType.KV_POSTGRES + host: str = "localhost" + port: int | str = 5432 + db: str = "llamastack" + user: str + password: str | None = None + ssl_mode: str | None = None + ca_cert_path: str | None = None + table_name: str = "llamastack_kvstore" + + @classmethod + def sample_run_config(cls, table_name: str = "llamastack_kvstore", **kwargs): + return { + "type": StorageBackendType.KV_POSTGRES.value, + "host": "${env.POSTGRES_HOST:=localhost}", + "port": "${env.POSTGRES_PORT:=5432}", + "db": "${env.POSTGRES_DB:=llamastack}", + "user": "${env.POSTGRES_USER:=llamastack}", + "password": "${env.POSTGRES_PASSWORD:=llamastack}", + "table_name": "${env.POSTGRES_TABLE_NAME:=" + table_name + "}", + } + + @classmethod + @field_validator("table_name") + def validate_table_name(cls, v: str) -> str: + # PostgreSQL identifiers rules: + # - Must start with a letter or underscore + # - Can contain letters, numbers, and underscores + # - Maximum length is 63 bytes + pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*$" + if not re.match(pattern, v): + raise ValueError( + "Invalid table name. Must start with letter or underscore and contain only letters, numbers, and underscores" + ) + if len(v) > 63: + raise ValueError("Table name must be less than 63 characters") + return v + + @classmethod + def pip_packages(cls) -> list[str]: + return ["psycopg2-binary"] + + +class MongoDBKVStoreConfig(CommonConfig): + type: Literal[StorageBackendType.KV_MONGODB] = StorageBackendType.KV_MONGODB + host: str = "localhost" + port: int = 27017 + db: str = "llamastack" + user: str | None = None + password: str | None = None + collection_name: str = "llamastack_kvstore" + + @classmethod + def pip_packages(cls) -> list[str]: + return ["pymongo"] + + @classmethod + def sample_run_config(cls, collection_name: str = "llamastack_kvstore"): + return { + "type": StorageBackendType.KV_MONGODB.value, + "host": "${env.MONGODB_HOST:=localhost}", + "port": "${env.MONGODB_PORT:=5432}", + "db": "${env.MONGODB_DB}", + "user": "${env.MONGODB_USER}", + "password": "${env.MONGODB_PASSWORD}", + "collection_name": "${env.MONGODB_COLLECTION_NAME:=" + collection_name + "}", + } + + +class SqlAlchemySqlStoreConfig(BaseModel): + @property + @abstractmethod + def engine_str(self) -> str: ... + + # TODO: move this when we have a better way to specify dependencies with internal APIs + @classmethod + def pip_packages(cls) -> list[str]: + return ["sqlalchemy[asyncio]"] + + +class SqliteSqlStoreConfig(SqlAlchemySqlStoreConfig): + type: Literal[StorageBackendType.SQL_SQLITE] = StorageBackendType.SQL_SQLITE + db_path: str = Field( + description="Database path, e.g. ~/.llama/distributions/ollama/sqlstore.db", + ) + + @property + def engine_str(self) -> str: + return "sqlite+aiosqlite:///" + Path(self.db_path).expanduser().as_posix() + + @classmethod + def sample_run_config(cls, __distro_dir__: str, db_name: str = "sqlstore.db"): + return { + "type": StorageBackendType.SQL_SQLITE.value, + "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + db_name, + } + + @classmethod + def pip_packages(cls) -> list[str]: + return super().pip_packages() + ["aiosqlite"] + + +class PostgresSqlStoreConfig(SqlAlchemySqlStoreConfig): + type: Literal[StorageBackendType.SQL_POSTGRES] = StorageBackendType.SQL_POSTGRES + host: str = "localhost" + port: int | str = 5432 + db: str = "llamastack" + user: str + password: str | None = None + + @property + def engine_str(self) -> str: + return f"postgresql+asyncpg://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}" + + @classmethod + def pip_packages(cls) -> list[str]: + return super().pip_packages() + ["asyncpg"] + + @classmethod + def sample_run_config(cls, **kwargs): + return { + "type": StorageBackendType.SQL_POSTGRES.value, + "host": "${env.POSTGRES_HOST:=localhost}", + "port": "${env.POSTGRES_PORT:=5432}", + "db": "${env.POSTGRES_DB:=llamastack}", + "user": "${env.POSTGRES_USER:=llamastack}", + "password": "${env.POSTGRES_PASSWORD:=llamastack}", + } + + +# reference = (backend_name, table_name) +class SqlStoreReference(BaseModel): + """A reference to a 'SQL-like' persistent store. A table name must be provided.""" + + table_name: str = Field( + description="Name of the table to use for the SqlStore", + ) + + backend: str = Field( + description="Name of backend from storage.backends", + ) + + +# reference = (backend_name, namespace) +class KVStoreReference(BaseModel): + """A reference to a 'key-value' persistent store. A namespace must be provided.""" + + namespace: str = Field( + description="Key prefix for KVStore backends", + ) + + backend: str = Field( + description="Name of backend from storage.backends", + ) + + +StorageBackendConfig = Annotated[ + RedisKVStoreConfig + | SqliteKVStoreConfig + | PostgresKVStoreConfig + | MongoDBKVStoreConfig + | SqliteSqlStoreConfig + | PostgresSqlStoreConfig, + Field(discriminator="type"), +] + + +class InferenceStoreReference(SqlStoreReference): + """Inference store configuration with queue tuning.""" + + max_write_queue_size: int = Field( + default=10000, + description="Max queued writes for inference store", + ) + num_writers: int = Field( + default=4, + description="Number of concurrent background writers", + ) + + +class ResponsesStoreReference(InferenceStoreReference): + """Responses store configuration with queue tuning.""" + + +class ServerStoresConfig(BaseModel): + metadata: KVStoreReference | None = Field( + default=None, + description="Metadata store configuration (uses KV backend)", + ) + inference: InferenceStoreReference | None = Field( + default=None, + description="Inference store configuration (uses SQL backend)", + ) + conversations: SqlStoreReference | None = Field( + default=None, + description="Conversations store configuration (uses SQL backend)", + ) + responses: ResponsesStoreReference | None = Field( + default=None, + description="Responses store configuration (uses SQL backend)", + ) + + +class StorageConfig(BaseModel): + backends: dict[str, StorageBackendConfig] = Field( + description="Named backend configurations (e.g., 'default', 'cache')", + ) + stores: ServerStoresConfig = Field( + default_factory=lambda: ServerStoresConfig(), + description="Named references to storage backends used by the stack core", + ) diff --git a/llama_stack/core/store/registry.py b/llama_stack/core/store/registry.py index 04581bab5..6ff9e575b 100644 --- a/llama_stack/core/store/registry.py +++ b/llama_stack/core/store/registry.py @@ -11,10 +11,9 @@ from typing import Protocol import pydantic from llama_stack.core.datatypes import RoutableObjectWithProvider -from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.log import get_logger from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig logger = get_logger(__name__, category="core::registry") @@ -191,16 +190,10 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry): async def create_dist_registry( - metadata_store: KVStoreConfig | None, - image_name: str, + metadata_store: KVStoreReference, image_name: str ) -> tuple[CachedDiskDistributionRegistry, KVStore]: # instantiate kvstore for storing and retrieving distribution metadata - if metadata_store: - dist_kvstore = await kvstore_impl(metadata_store) - else: - dist_kvstore = await kvstore_impl( - SqliteKVStoreConfig(db_path=(DISTRIBS_BASE_DIR / image_name / "kvstore.db").as_posix()) - ) + 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/run.yaml b/llama_stack/distributions/ci-tests/run.yaml index a6a6b7c0d..f9e741474 100644 --- a/llama_stack/distributions/ci-tests/run.yaml +++ b/llama_stack/distributions/ci-tests/run.yaml @@ -93,30 +93,30 @@ providers: - provider_id: faiss provider_type: inline::faiss config: - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/faiss_store.db + persistence: + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/sqlite_vec_registry.db + persistence: + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/milvus_registry.db + persistence: + namespace: vector_io::milvus + backend: kv_default - provider_id: ${env.CHROMADB_URL:+chromadb} provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests/}/chroma_remote_registry.db + persistence: + namespace: vector_io::chroma_remote + backend: kv_default - provider_id: ${env.PGVECTOR_DB:+pgvector} provider_type: remote::pgvector config: @@ -125,17 +125,17 @@ providers: db: ${env.PGVECTOR_DB:=} user: ${env.PGVECTOR_USER:=} password: ${env.PGVECTOR_PASSWORD:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/pgvector_registry.db + persistence: + namespace: vector_io::pgvector + backend: kv_default files: - provider_id: meta-reference-files provider_type: inline::localfs config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/ci-tests/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/files_metadata.db + table_name: files_metadata + backend: sql_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -147,12 +147,15 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/responses_store.db + persistence: + agent_state: + namespace: agents + 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 @@ -163,21 +166,21 @@ providers: provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -207,17 +210,28 @@ providers: provider_type: inline::reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/batches.db -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/conversations.db + namespace: batches + backend: kv_default +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ci-tests}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default models: [] shields: - shield_id: llama-guard diff --git a/llama_stack/distributions/dell/run-with-safety.yaml b/llama_stack/distributions/dell/run-with-safety.yaml index 5da3cf511..3130285b9 100644 --- a/llama_stack/distributions/dell/run-with-safety.yaml +++ b/llama_stack/distributions/dell/run-with-safety.yaml @@ -26,9 +26,9 @@ providers: provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell/}/chroma_remote_registry.db + persistence: + namespace: vector_io::chroma_remote + backend: kv_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -38,32 +38,35 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/responses_store.db + persistence: + agent_state: + namespace: agents + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -86,15 +89,26 @@ providers: max_results: 3 - provider_id: rag-runtime provider_type: inline::rag-runtime -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/conversations.db +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + 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 ac0fdc0fa..af1a96a21 100644 --- a/llama_stack/distributions/dell/run.yaml +++ b/llama_stack/distributions/dell/run.yaml @@ -22,9 +22,9 @@ providers: provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell/}/chroma_remote_registry.db + persistence: + namespace: vector_io::chroma_remote + backend: kv_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -34,32 +34,35 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/responses_store.db + persistence: + agent_state: + namespace: agents + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -82,15 +85,26 @@ providers: max_results: 3 - provider_id: rag-runtime provider_type: inline::rag-runtime -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/conversations.db +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default models: - metadata: {} model_id: ${env.INFERENCE_MODEL} diff --git a/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml b/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml index 874c5050f..b43d1ff19 100644 --- a/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml +++ b/llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml @@ -37,9 +37,9 @@ providers: - provider_id: faiss provider_type: inline::faiss config: - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/faiss_store.db + persistence: + namespace: vector_io::faiss + backend: kv_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -49,32 +49,35 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/responses_store.db + persistence: + agent_state: + namespace: agents + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -99,15 +102,26 @@ providers: provider_type: inline::rag-runtime - provider_id: model-context-protocol provider_type: remote::model-context-protocol -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/conversations.db +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + 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 50553d2c7..59e2d8129 100644 --- a/llama_stack/distributions/meta-reference-gpu/run.yaml +++ b/llama_stack/distributions/meta-reference-gpu/run.yaml @@ -27,9 +27,9 @@ providers: - provider_id: faiss provider_type: inline::faiss config: - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/faiss_store.db + persistence: + namespace: vector_io::faiss + backend: kv_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -39,32 +39,35 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/responses_store.db + persistence: + agent_state: + namespace: agents + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -89,15 +92,26 @@ providers: provider_type: inline::rag-runtime - provider_id: model-context-protocol provider_type: remote::model-context-protocol -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/conversations.db +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default models: - metadata: {} model_id: ${env.INFERENCE_MODEL} diff --git a/llama_stack/distributions/nvidia/run-with-safety.yaml b/llama_stack/distributions/nvidia/run-with-safety.yaml index e0482f67d..e06787d0b 100644 --- a/llama_stack/distributions/nvidia/run-with-safety.yaml +++ b/llama_stack/distributions/nvidia/run-with-safety.yaml @@ -28,9 +28,9 @@ providers: - provider_id: faiss provider_type: inline::faiss config: - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/faiss_store.db + persistence: + namespace: vector_io::faiss + backend: kv_default safety: - provider_id: nvidia provider_type: remote::nvidia @@ -41,12 +41,15 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/responses_store.db + persistence: + agent_state: + namespace: agents + 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 @@ -65,8 +68,8 @@ providers: provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default - provider_id: nvidia provider_type: remote::nvidia config: @@ -86,17 +89,28 @@ providers: config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/nvidia/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/files_metadata.db -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/conversations.db + table_name: files_metadata + backend: sql_default +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + 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 950782eed..85e0743e4 100644 --- a/llama_stack/distributions/nvidia/run.yaml +++ b/llama_stack/distributions/nvidia/run.yaml @@ -23,9 +23,9 @@ providers: - provider_id: faiss provider_type: inline::faiss config: - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/faiss_store.db + persistence: + namespace: vector_io::faiss + backend: kv_default safety: - provider_id: nvidia provider_type: remote::nvidia @@ -36,12 +36,15 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/responses_store.db + persistence: + agent_state: + namespace: agents + 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 @@ -75,17 +78,28 @@ providers: config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/nvidia/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/files_metadata.db -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/conversations.db + table_name: files_metadata + backend: sql_default +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/nvidia}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default models: [] shields: [] vector_dbs: [] diff --git a/llama_stack/distributions/open-benchmark/run.yaml b/llama_stack/distributions/open-benchmark/run.yaml index a738887b4..2c6936bfc 100644 --- a/llama_stack/distributions/open-benchmark/run.yaml +++ b/llama_stack/distributions/open-benchmark/run.yaml @@ -39,16 +39,16 @@ providers: provider_type: inline::sqlite-vec config: db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/sqlite_vec.db - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/sqlite_vec_registry.db + persistence: + namespace: vector_io::sqlite_vec + backend: kv_default - provider_id: ${env.ENABLE_CHROMADB:+chromadb} provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/chroma_remote_registry.db + persistence: + namespace: vector_io::chroma_remote + backend: kv_default - provider_id: ${env.ENABLE_PGVECTOR:+pgvector} provider_type: remote::pgvector config: @@ -57,9 +57,9 @@ providers: db: ${env.PGVECTOR_DB:=} user: ${env.PGVECTOR_USER:=} password: ${env.PGVECTOR_PASSWORD:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/pgvector_registry.db + persistence: + namespace: vector_io::pgvector + backend: kv_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -69,32 +69,35 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/responses_store.db + persistence: + agent_state: + namespace: agents + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -119,15 +122,26 @@ providers: provider_type: inline::rag-runtime - provider_id: model-context-protocol provider_type: remote::model-context-protocol -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/conversations.db +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/open-benchmark}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default models: - metadata: {} model_id: gpt-4o diff --git a/llama_stack/distributions/postgres-demo/postgres_demo.py b/llama_stack/distributions/postgres-demo/postgres_demo.py index 1f3e88b3b..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", @@ -105,22 +104,16 @@ 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={ + "kv_default": PostgresKVStoreConfig.sample_run_config( + table_name="llamastack_kvstore", + ), + "sql_default": 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 62faf3f62..9556b1287 100644 --- a/llama_stack/distributions/postgres-demo/run.yaml +++ b/llama_stack/distributions/postgres-demo/run.yaml @@ -22,9 +22,9 @@ providers: provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/postgres-demo}/chroma_remote_registry.db + persistence: + namespace: vector_io::chroma_remote + backend: kv_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -34,20 +34,15 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} - responses_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} + 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 @@ -63,24 +58,35 @@ providers: provider_type: inline::rag-runtime - provider_id: model-context-protocol provider_type: remote::model-context-protocol -metadata_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} - table_name: ${env.POSTGRES_TABLE_NAME:=llamastack_kvstore} -inference_store: - type: postgres - host: ${env.POSTGRES_HOST:=localhost} - port: ${env.POSTGRES_PORT:=5432} - db: ${env.POSTGRES_DB:=llamastack} - user: ${env.POSTGRES_USER:=llamastack} - password: ${env.POSTGRES_PASSWORD:=llamastack} -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/postgres-demo}/conversations.db +storage: + backends: + kv_default: + 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} + 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} + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default models: - metadata: {} model_id: ${env.INFERENCE_MODEL} diff --git a/llama_stack/distributions/starter-gpu/run.yaml b/llama_stack/distributions/starter-gpu/run.yaml index 370d4b516..abfa579a7 100644 --- a/llama_stack/distributions/starter-gpu/run.yaml +++ b/llama_stack/distributions/starter-gpu/run.yaml @@ -93,30 +93,30 @@ providers: - provider_id: faiss provider_type: inline::faiss config: - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/faiss_store.db + persistence: + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/sqlite_vec_registry.db + persistence: + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/milvus_registry.db + persistence: + namespace: vector_io::milvus + backend: kv_default - provider_id: ${env.CHROMADB_URL:+chromadb} provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu/}/chroma_remote_registry.db + persistence: + namespace: vector_io::chroma_remote + backend: kv_default - provider_id: ${env.PGVECTOR_DB:+pgvector} provider_type: remote::pgvector config: @@ -125,17 +125,17 @@ providers: db: ${env.PGVECTOR_DB:=} user: ${env.PGVECTOR_USER:=} password: ${env.PGVECTOR_PASSWORD:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/pgvector_registry.db + persistence: + namespace: vector_io::pgvector + backend: kv_default files: - provider_id: meta-reference-files provider_type: inline::localfs config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter-gpu/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/files_metadata.db + table_name: files_metadata + backend: sql_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -147,12 +147,15 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/responses_store.db + persistence: + agent_state: + namespace: agents + 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 @@ -166,21 +169,21 @@ providers: provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -210,17 +213,28 @@ providers: provider_type: inline::reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/batches.db -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/conversations.db + namespace: batches + backend: kv_default +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter-gpu}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default models: [] shields: - shield_id: llama-guard diff --git a/llama_stack/distributions/starter/run.yaml b/llama_stack/distributions/starter/run.yaml index 2f4e7f350..fc58a4afe 100644 --- a/llama_stack/distributions/starter/run.yaml +++ b/llama_stack/distributions/starter/run.yaml @@ -93,30 +93,30 @@ providers: - provider_id: faiss provider_type: inline::faiss config: - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/faiss_store.db + persistence: + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/sqlite_vec_registry.db + persistence: + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/milvus_registry.db + persistence: + namespace: vector_io::milvus + backend: kv_default - provider_id: ${env.CHROMADB_URL:+chromadb} provider_type: remote::chromadb config: url: ${env.CHROMADB_URL:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter/}/chroma_remote_registry.db + persistence: + namespace: vector_io::chroma_remote + backend: kv_default - provider_id: ${env.PGVECTOR_DB:+pgvector} provider_type: remote::pgvector config: @@ -125,17 +125,17 @@ providers: db: ${env.PGVECTOR_DB:=} user: ${env.PGVECTOR_USER:=} password: ${env.PGVECTOR_PASSWORD:=} - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/pgvector_registry.db + persistence: + namespace: vector_io::pgvector + backend: kv_default files: - provider_id: meta-reference-files provider_type: inline::localfs config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/files_metadata.db + table_name: files_metadata + backend: sql_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -147,12 +147,15 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/responses_store.db + persistence: + agent_state: + namespace: agents + 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 @@ -163,21 +166,21 @@ providers: provider_type: inline::meta-reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -207,17 +210,28 @@ providers: provider_type: inline::reference config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/batches.db -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/conversations.db + namespace: batches + backend: kv_default +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/starter}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + 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 807829999..542c7bea9 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,9 @@ 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 + storage_stores: dict[str, Any] | None = None def run_config( self, @@ -226,6 +231,37 @@ class RunConfigSettings(BaseModel): # Get unique set of APIs from providers apis = sorted(providers.keys()) + storage_backends = self.storage_backends or { + "kv_default": SqliteKVStoreConfig.sample_run_config( + __distro_dir__=f"~/.llama/distributions/{name}", + db_name="kvstore.db", + ), + "sql_default": SqliteSqlStoreConfig.sample_run_config( + __distro_dir__=f"~/.llama/distributions/{name}", + db_name="sql_store.db", + ), + } + + storage_stores = self.storage_stores or { + "metadata": KVStoreReference( + backend="kv_default", + namespace="registry", + ).model_dump(exclude_none=True), + "inference": InferenceStoreReference( + backend="sql_default", + table_name="inference_store", + ).model_dump(exclude_none=True), + "conversations": SqlStoreReference( + backend="sql_default", + table_name="openai_conversations", + ).model_dump(exclude_none=True), + } + + storage_config = dict( + backends=storage_backends, + stores=storage_stores, + ) + # Return a dict that matches StackRunConfig structure return { "version": LLAMA_STACK_RUN_CONFIG_VERSION, @@ -233,21 +269,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": [], @@ -297,11 +319,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) @@ -387,11 +413,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/run.yaml b/llama_stack/distributions/watsonx/run.yaml index c3db4eeb8..37866cb32 100644 --- a/llama_stack/distributions/watsonx/run.yaml +++ b/llama_stack/distributions/watsonx/run.yaml @@ -22,9 +22,9 @@ providers: - provider_id: faiss provider_type: inline::faiss config: - kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/faiss_store.db + persistence: + namespace: vector_io::faiss + backend: kv_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -34,32 +34,35 @@ providers: - provider_id: meta-reference provider_type: inline::meta-reference config: - persistence_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/agents_store.db - responses_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/responses_store.db + persistence: + agent_state: + namespace: agents + 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: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/meta_reference_eval.db + namespace: eval + backend: kv_default datasetio: - provider_id: huggingface provider_type: remote::huggingface config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/huggingface_datasetio.db + namespace: datasetio::huggingface + backend: kv_default - provider_id: localfs provider_type: inline::localfs config: kvstore: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/localfs_datasetio.db + namespace: datasetio::localfs + backend: kv_default scoring: - provider_id: basic provider_type: inline::basic @@ -90,17 +93,28 @@ providers: config: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/watsonx/files} metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/files_metadata.db -metadata_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/registry.db -inference_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/inference_store.db -conversations_store: - type: sqlite - db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/conversations.db + table_name: files_metadata + backend: sql_default +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/watsonx}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default models: [] shields: [] vector_dbs: [] diff --git a/llama_stack/providers/inline/agents/meta_reference/agents.py b/llama_stack/providers/inline/agents/meta_reference/agents.py index 810c063e6..c2f6ea640 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agents.py +++ b/llama_stack/providers/inline/agents/meta_reference/agents.py @@ -83,8 +83,8 @@ class MetaReferenceAgentsImpl(Agents): self.policy = policy async def initialize(self) -> None: - self.persistence_store = await kvstore_impl(self.config.persistence_store) - self.responses_store = ResponsesStore(self.config.responses_store, self.policy) + self.persistence_store = await kvstore_impl(self.config.persistence.agent_state) + self.responses_store = ResponsesStore(self.config.persistence.responses, self.policy) await self.responses_store.initialize() self.openai_responses_impl = OpenAIResponsesImpl( inference_api=self.inference_api, diff --git a/llama_stack/providers/inline/agents/meta_reference/config.py b/llama_stack/providers/inline/agents/meta_reference/config.py index 1c392f29c..a800b426b 100644 --- a/llama_stack/providers/inline/agents/meta_reference/config.py +++ b/llama_stack/providers/inline/agents/meta_reference/config.py @@ -8,24 +8,30 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore import KVStoreConfig -from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference, ResponsesStoreReference + + +class AgentPersistenceConfig(BaseModel): + """Nested persistence configuration for agents.""" + + agent_state: KVStoreReference + responses: ResponsesStoreReference class MetaReferenceAgentsImplConfig(BaseModel): - persistence_store: KVStoreConfig - responses_store: SqlStoreConfig + persistence: AgentPersistenceConfig @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { - "persistence_store": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="agents_store.db", - ), - "responses_store": SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="responses_store.db", - ), + "persistence": { + "agent_state": KVStoreReference( + backend="kv_default", + namespace="agents", + ).model_dump(exclude_none=True), + "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 d8d06868b..f896a897d 100644 --- a/llama_stack/providers/inline/batches/reference/config.py +++ b/llama_stack/providers/inline/batches/reference/config.py @@ -6,13 +6,13 @@ from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference class ReferenceBatchesImplConfig(BaseModel): """Configuration for the Reference Batches implementation.""" - kvstore: KVStoreConfig = Field( + kvstore: KVStoreReference = Field( description="Configuration for the key-value store backend.", ) @@ -33,8 +33,8 @@ class ReferenceBatchesImplConfig(BaseModel): @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="batches.db", - ), + "kvstore": KVStoreReference( + 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 b450e8777..6e878df62 100644 --- a/llama_stack/providers/inline/datasetio/localfs/config.py +++ b/llama_stack/providers/inline/datasetio/localfs/config.py @@ -7,20 +7,17 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference class LocalFSDatasetIOConfig(BaseModel): - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="localfs_datasetio.db", - ) + "kvstore": KVStoreReference( + 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 2a4a29998..b496c855e 100644 --- a/llama_stack/providers/inline/eval/meta_reference/config.py +++ b/llama_stack/providers/inline/eval/meta_reference/config.py @@ -7,20 +7,17 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference class MetaReferenceEvalConfig(BaseModel): - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="meta_reference_eval.db", - ) + "kvstore": KVStoreReference( + 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 6c767af8f..0c2dd3b21 100644 --- a/llama_stack/providers/inline/files/localfs/config.py +++ b/llama_stack/providers/inline/files/localfs/config.py @@ -8,14 +8,14 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig +from llama_stack.core.storage.datatypes import SqlStoreReference class LocalfsFilesImplConfig(BaseModel): storage_dir: str = Field( description="Directory to store uploaded files", ) - metadata_store: SqlStoreConfig = Field( + metadata_store: SqlStoreReference = Field( description="SQL store configuration for file metadata", ) ttl_secs: int = 365 * 24 * 60 * 60 # 1 year @@ -24,8 +24,8 @@ class LocalfsFilesImplConfig(BaseModel): def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { "storage_dir": "${env.FILES_STORAGE_DIR:=" + __distro_dir__ + "/files}", - "metadata_store": SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="files_metadata.db", - ), + "metadata_store": SqlStoreReference( + 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 a9566f7ff..1798f10de 100644 --- a/llama_stack/providers/inline/vector_io/chroma/config.py +++ b/llama_stack/providers/inline/vector_io/chroma/config.py @@ -8,14 +8,14 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class ChromaVectorIOConfig(BaseModel): db_path: str - kvstore: KVStoreConfig = Field(description="Config for KV store backend") + persistence: KVStoreReference = Field(description="Config for KV store backend") @classmethod def sample_run_config( @@ -23,8 +23,8 @@ class ChromaVectorIOConfig(BaseModel): ) -> dict[str, Any]: return { "db_path": db_path, - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="chroma_inline_registry.db", - ), + "persistence": 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 cbcbb1762..dd7a7aeca 100644 --- a/llama_stack/providers/inline/vector_io/faiss/config.py +++ b/llama_stack/providers/inline/vector_io/faiss/config.py @@ -8,22 +8,19 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class FaissVectorIOConfig(BaseModel): - kvstore: KVStoreConfig + persistence: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="faiss_store.db", - ) + "persistence": KVStoreReference( + backend="kv_default", + namespace="vector_io::faiss", + ).model_dump(exclude_none=True) } diff --git a/llama_stack/providers/inline/vector_io/faiss/faiss.py b/llama_stack/providers/inline/vector_io/faiss/faiss.py index df0864db8..ff1a6aa4c 100644 --- a/llama_stack/providers/inline/vector_io/faiss/faiss.py +++ b/llama_stack/providers/inline/vector_io/faiss/faiss.py @@ -214,7 +214,7 @@ class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPr self.cache: dict[str, VectorDBWithIndex] = {} async def initialize(self) -> None: - self.kvstore = await kvstore_impl(self.config.kvstore) + self.kvstore = await kvstore_impl(self.config.persistence) # Load existing banks from kvstore start_key = VECTOR_DBS_PREFIX end_key = f"{VECTOR_DBS_PREFIX}\xff" diff --git a/llama_stack/providers/inline/vector_io/milvus/config.py b/llama_stack/providers/inline/vector_io/milvus/config.py index 8cbd056be..b333b04ea 100644 --- a/llama_stack/providers/inline/vector_io/milvus/config.py +++ b/llama_stack/providers/inline/vector_io/milvus/config.py @@ -8,25 +8,22 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class MilvusVectorIOConfig(BaseModel): db_path: str - kvstore: KVStoreConfig = Field(description="Config for KV store backend (SQLite only for now)") + persistence: KVStoreReference = Field(description="Config for KV store backend (SQLite only for now)") consistency_level: str = Field(description="The consistency level of the Milvus server", default="Strong") @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { "db_path": "${env.MILVUS_DB_PATH:=" + __distro_dir__ + "}/" + "milvus.db", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="milvus_registry.db", - ), + "persistence": 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 e15c27ea1..e7ecde7b7 100644 --- a/llama_stack/providers/inline/vector_io/qdrant/config.py +++ b/llama_stack/providers/inline/vector_io/qdrant/config.py @@ -9,23 +9,21 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class QdrantVectorIOConfig(BaseModel): path: str - kvstore: KVStoreConfig + persistence: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { "path": "${env.QDRANT_PATH:=~/.llama/" + __distro_dir__ + "}/" + "qdrant.db", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, db_name="qdrant_registry.db" - ), + "persistence": 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 525ed4b1f..596f8fc95 100644 --- a/llama_stack/providers/inline/vector_io/sqlite_vec/config.py +++ b/llama_stack/providers/inline/vector_io/sqlite_vec/config.py @@ -8,22 +8,19 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference class SQLiteVectorIOConfig(BaseModel): db_path: str = Field(description="Path to the SQLite database file") - kvstore: KVStoreConfig = Field(description="Config for KV store backend (SQLite only for now)") + persistence: KVStoreReference = Field(description="Config for KV store backend (SQLite only for now)") @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + "sqlite_vec.db", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="sqlite_vec_registry.db", - ), + "persistence": KVStoreReference( + backend="kv_default", + namespace="vector_io::sqlite_vec", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py b/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py index 8bc3b04cb..a58aa05b8 100644 --- a/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py +++ b/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py @@ -425,7 +425,7 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc self.vector_db_store = None async def initialize(self) -> None: - self.kvstore = await kvstore_impl(self.config.kvstore) + self.kvstore = await kvstore_impl(self.config.persistence) start_key = VECTOR_DBS_PREFIX end_key = f"{VECTOR_DBS_PREFIX}\xff" diff --git a/llama_stack/providers/remote/datasetio/huggingface/config.py b/llama_stack/providers/remote/datasetio/huggingface/config.py index 38f933728..35297cb58 100644 --- a/llama_stack/providers/remote/datasetio/huggingface/config.py +++ b/llama_stack/providers/remote/datasetio/huggingface/config.py @@ -7,20 +7,17 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference class HuggingfaceDatasetIOConfig(BaseModel): - kvstore: KVStoreConfig + kvstore: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="huggingface_datasetio.db", - ) + "kvstore": KVStoreReference( + 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 da20d8668..cd4b1adda 100644 --- a/llama_stack/providers/remote/files/s3/config.py +++ b/llama_stack/providers/remote/files/s3/config.py @@ -8,7 +8,7 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig +from llama_stack.core.storage.datatypes import SqlStoreReference class S3FilesImplConfig(BaseModel): @@ -24,7 +24,7 @@ class S3FilesImplConfig(BaseModel): auto_create_bucket: bool = Field( default=False, description="Automatically create the S3 bucket if it doesn't exist" ) - metadata_store: SqlStoreConfig = Field(description="SQL store configuration for file metadata") + metadata_store: SqlStoreReference = Field(description="SQL store configuration for file metadata") @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: @@ -35,8 +35,8 @@ class S3FilesImplConfig(BaseModel): "aws_secret_access_key": "${env.AWS_SECRET_ACCESS_KEY:=}", "endpoint_url": "${env.S3_ENDPOINT_URL:=}", "auto_create_bucket": "${env.S3_AUTO_CREATE_BUCKET:=false}", - "metadata_store": SqliteSqlStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="s3_files_metadata.db", - ), + "metadata_store": SqlStoreReference( + backend="sql_default", + table_name="s3_files_metadata", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/chroma/chroma.py b/llama_stack/providers/remote/vector_io/chroma/chroma.py index 5792a83c6..b07207cc6 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/llama_stack/providers/remote/vector_io/chroma/chroma.py @@ -151,7 +151,7 @@ class ChromaVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP self.vector_db_store = None async def initialize(self) -> None: - self.kvstore = await kvstore_impl(self.config.kvstore) + self.kvstore = await kvstore_impl(self.config.persistence) self.vector_db_store = self.kvstore if isinstance(self.config, RemoteChromaVectorIOConfig): diff --git a/llama_stack/providers/remote/vector_io/chroma/config.py b/llama_stack/providers/remote/vector_io/chroma/config.py index a1193905a..209ba90bb 100644 --- a/llama_stack/providers/remote/vector_io/chroma/config.py +++ b/llama_stack/providers/remote/vector_io/chroma/config.py @@ -8,21 +8,21 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @json_schema_type class ChromaVectorIOConfig(BaseModel): url: str | None - kvstore: KVStoreConfig = Field(description="Config for KV store backend") + persistence: KVStoreReference = Field(description="Config for KV store backend") @classmethod def sample_run_config(cls, __distro_dir__: str, url: str = "${env.CHROMADB_URL}", **kwargs: Any) -> dict[str, Any]: return { "url": url, - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="chroma_remote_registry.db", - ), + "persistence": 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 899d3678d..8ff9e1328 100644 --- a/llama_stack/providers/remote/vector_io/milvus/config.py +++ b/llama_stack/providers/remote/vector_io/milvus/config.py @@ -8,7 +8,7 @@ from typing import Any from pydantic import BaseModel, ConfigDict, Field -from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @@ -17,7 +17,7 @@ class MilvusVectorIOConfig(BaseModel): uri: str = Field(description="The URI of the Milvus server") token: str | None = Field(description="The token of the Milvus server") consistency_level: str = Field(description="The consistency level of the Milvus server", default="Strong") - kvstore: KVStoreConfig = Field(description="Config for KV store backend") + persistence: KVStoreReference = Field(description="Config for KV store backend") # This configuration allows additional fields to be passed through to the underlying Milvus client. # See the [Milvus](https://milvus.io/docs/install-overview.md) documentation for more details about Milvus in general. @@ -28,8 +28,8 @@ class MilvusVectorIOConfig(BaseModel): return { "uri": "${env.MILVUS_ENDPOINT}", "token": "${env.MILVUS_TOKEN}", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="milvus_remote_registry.db", - ), + "persistence": KVStoreReference( + backend="kv_default", + namespace="vector_io::milvus_remote", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index d7147a7f0..1f689d1a9 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -321,7 +321,7 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP self.metadata_collection_name = "openai_vector_stores_metadata" async def initialize(self) -> None: - self.kvstore = await kvstore_impl(self.config.kvstore) + self.kvstore = await kvstore_impl(self.config.persistence) start_key = VECTOR_DBS_PREFIX end_key = f"{VECTOR_DBS_PREFIX}\xff" stored_vector_dbs = await self.kvstore.values_in_range(start_key, end_key) diff --git a/llama_stack/providers/remote/vector_io/pgvector/config.py b/llama_stack/providers/remote/vector_io/pgvector/config.py index 334cbe5be..d81e524e4 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/config.py +++ b/llama_stack/providers/remote/vector_io/pgvector/config.py @@ -8,10 +8,7 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @@ -22,7 +19,9 @@ class PGVectorVectorIOConfig(BaseModel): db: str | None = Field(default="postgres") user: str | None = Field(default="postgres") password: str | None = Field(default="mysecretpassword") - kvstore: KVStoreConfig | None = Field(description="Config for KV store backend (SQLite only for now)", default=None) + persistence: KVStoreReference | None = Field( + description="Config for KV store backend (SQLite only for now)", default=None + ) @classmethod def sample_run_config( @@ -41,8 +40,8 @@ class PGVectorVectorIOConfig(BaseModel): "db": db, "user": user, "password": password, - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="pgvector_registry.db", - ), + "persistence": KVStoreReference( + backend="kv_default", + namespace="vector_io::pgvector", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index d55c13103..691cf965c 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -358,7 +358,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco async def initialize(self) -> None: log.info(f"Initializing PGVector memory adapter with config: {self.config}") - self.kvstore = await kvstore_impl(self.config.kvstore) + self.kvstore = await kvstore_impl(self.config.persistence) await self.initialize_openai_vector_stores() try: diff --git a/llama_stack/providers/remote/vector_io/qdrant/config.py b/llama_stack/providers/remote/vector_io/qdrant/config.py index ff5506236..01fbcc5cb 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/config.py +++ b/llama_stack/providers/remote/vector_io/qdrant/config.py @@ -8,10 +8,7 @@ from typing import Any from pydantic import BaseModel -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @@ -27,14 +24,14 @@ class QdrantVectorIOConfig(BaseModel): prefix: str | None = None timeout: int | None = None host: str | None = None - kvstore: KVStoreConfig + persistence: KVStoreReference @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { "api_key": "${env.QDRANT_API_KEY:=}", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="qdrant_registry.db", - ), + "persistence": KVStoreReference( + backend="kv_default", + namespace="vector_io::qdrant_remote", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py index 8b90935cd..eba8333e4 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py +++ b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py @@ -174,9 +174,9 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP self._qdrant_lock = asyncio.Lock() async def initialize(self) -> None: - client_config = self.config.model_dump(exclude_none=True, exclude={"kvstore"}) + client_config = self.config.model_dump(exclude_none=True, exclude={"persistence"}) self.client = AsyncQdrantClient(**client_config) - self.kvstore = await kvstore_impl(self.config.kvstore) + self.kvstore = await kvstore_impl(self.config.persistence) start_key = VECTOR_DBS_PREFIX end_key = f"{VECTOR_DBS_PREFIX}\xff" diff --git a/llama_stack/providers/remote/vector_io/weaviate/config.py b/llama_stack/providers/remote/vector_io/weaviate/config.py index b693e294e..06242c6b4 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/config.py +++ b/llama_stack/providers/remote/vector_io/weaviate/config.py @@ -8,10 +8,7 @@ from typing import Any from pydantic import BaseModel, Field -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.core.storage.datatypes import KVStoreReference from llama_stack.schema_utils import json_schema_type @@ -19,7 +16,9 @@ from llama_stack.schema_utils import json_schema_type class WeaviateVectorIOConfig(BaseModel): weaviate_api_key: str | None = Field(description="The API key for the Weaviate instance", default=None) weaviate_cluster_url: str | None = Field(description="The URL of the Weaviate cluster", default="localhost:8080") - kvstore: KVStoreConfig | None = Field(description="Config for KV store backend (SQLite only for now)", default=None) + persistence: KVStoreReference | None = Field( + description="Config for KV store backend (SQLite only for now)", default=None + ) @classmethod def sample_run_config( @@ -30,8 +29,8 @@ class WeaviateVectorIOConfig(BaseModel): return { "weaviate_api_key": None, "weaviate_cluster_url": "${env.WEAVIATE_CLUSTER_URL:=localhost:8080}", - "kvstore": SqliteKVStoreConfig.sample_run_config( - __distro_dir__=__distro_dir__, - db_name="weaviate_registry.db", - ), + "persistence": KVStoreReference( + backend="kv_default", + namespace="vector_io::weaviate", + ).model_dump(exclude_none=True), } diff --git a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py index d8b11c441..06ffc8706 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py +++ b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py @@ -320,8 +320,8 @@ class WeaviateVectorIOAdapter( async def initialize(self) -> None: """Set up KV store and load existing vector DBs and OpenAI vector stores.""" # Initialize KV store for metadata if configured - if self.config.kvstore is not None: - self.kvstore = await kvstore_impl(self.config.kvstore) + if self.config.persistence is not None: + self.kvstore = await kvstore_impl(self.config.persistence) else: self.kvstore = None log.info("No kvstore configured, registry will not persist across restarts") diff --git a/llama_stack/providers/utils/inference/inference_store.py b/llama_stack/providers/utils/inference/inference_store.py index 901f77c67..8e20bca6b 100644 --- a/llama_stack/providers/utils/inference/inference_store.py +++ b/llama_stack/providers/utils/inference/inference_store.py @@ -15,12 +15,13 @@ from llama_stack.apis.inference import ( OpenAIMessageParam, Order, ) -from llama_stack.core.datatypes import AccessRule, InferenceStoreConfig +from llama_stack.core.datatypes import AccessRule +from llama_stack.core.storage.datatypes import InferenceStoreReference, StorageBackendType from llama_stack.log import get_logger from ..sqlstore.api import ColumnDefinition, ColumnType from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore -from ..sqlstore.sqlstore import SqlStoreConfig, SqlStoreType, sqlstore_impl +from ..sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl logger = get_logger(name=__name__, category="inference") @@ -28,33 +29,32 @@ logger = get_logger(name=__name__, category="inference") class InferenceStore: def __init__( self, - config: InferenceStoreConfig | SqlStoreConfig, + reference: InferenceStoreReference, policy: list[AccessRule], ): - # Handle backward compatibility - if not isinstance(config, InferenceStoreConfig): - # Legacy: SqlStoreConfig passed directly as config - config = InferenceStoreConfig( - sql_store_config=config, - ) - - self.config = config - self.sql_store_config = config.sql_store_config + self.reference = reference self.sql_store = None self.policy = policy - # Disable write queue for SQLite to avoid concurrency issues - self.enable_write_queue = self.sql_store_config.type != SqlStoreType.sqlite - # Async write queue and worker control self._queue: asyncio.Queue[tuple[OpenAIChatCompletion, list[OpenAIMessageParam]]] | None = None self._worker_tasks: list[asyncio.Task[Any]] = [] - self._max_write_queue_size: int = config.max_write_queue_size - self._num_writers: int = max(1, config.num_writers) + self._max_write_queue_size: int = reference.max_write_queue_size + self._num_writers: int = max(1, reference.num_writers) async def initialize(self): """Create the necessary tables if they don't exist.""" - self.sql_store = AuthorizedSqlStore(sqlstore_impl(self.sql_store_config), self.policy) + base_store = sqlstore_impl(self.reference) + self.sql_store = AuthorizedSqlStore(base_store, self.policy) + + # Disable write queue for SQLite to avoid concurrency issues + backend_name = self.reference.backend + 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/config.py b/llama_stack/providers/utils/kvstore/config.py index 7b6a79350..c0582abc4 100644 --- a/llama_stack/providers/utils/kvstore/config.py +++ b/llama_stack/providers/utils/kvstore/config.py @@ -4,143 +4,20 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import re -from enum import Enum -from typing import Annotated, Literal +from typing import Annotated -from pydantic import BaseModel, Field, field_validator - -from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR - - -class KVStoreType(Enum): - redis = "redis" - sqlite = "sqlite" - postgres = "postgres" - mongodb = "mongodb" - - -class CommonConfig(BaseModel): - namespace: str | None = Field( - default=None, - description="All keys will be prefixed with this namespace", - ) - - -class RedisKVStoreConfig(CommonConfig): - type: Literal["redis"] = KVStoreType.redis.value - host: str = "localhost" - port: int = 6379 - - @property - def url(self) -> str: - return f"redis://{self.host}:{self.port}" - - @classmethod - def pip_packages(cls) -> list[str]: - return ["redis"] - - @classmethod - def sample_run_config(cls): - return { - "type": "redis", - "host": "${env.REDIS_HOST:=localhost}", - "port": "${env.REDIS_PORT:=6379}", - } - - -class SqliteKVStoreConfig(CommonConfig): - type: Literal["sqlite"] = KVStoreType.sqlite.value - db_path: str = Field( - default=(RUNTIME_BASE_DIR / "kvstore.db").as_posix(), - description="File path for the sqlite database", - ) - - @classmethod - def pip_packages(cls) -> list[str]: - return ["aiosqlite"] - - @classmethod - def sample_run_config(cls, __distro_dir__: str, db_name: str = "kvstore.db"): - return { - "type": "sqlite", - "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + db_name, - } - - -class PostgresKVStoreConfig(CommonConfig): - type: Literal["postgres"] = KVStoreType.postgres.value - host: str = "localhost" - port: int = 5432 - db: str = "llamastack" - user: str - password: str | None = None - ssl_mode: str | None = None - ca_cert_path: str | None = None - table_name: str = "llamastack_kvstore" - - @classmethod - def sample_run_config(cls, table_name: str = "llamastack_kvstore", **kwargs): - return { - "type": "postgres", - "host": "${env.POSTGRES_HOST:=localhost}", - "port": "${env.POSTGRES_PORT:=5432}", - "db": "${env.POSTGRES_DB:=llamastack}", - "user": "${env.POSTGRES_USER:=llamastack}", - "password": "${env.POSTGRES_PASSWORD:=llamastack}", - "table_name": "${env.POSTGRES_TABLE_NAME:=" + table_name + "}", - } - - @classmethod - @field_validator("table_name") - def validate_table_name(cls, v: str) -> str: - # PostgreSQL identifiers rules: - # - Must start with a letter or underscore - # - Can contain letters, numbers, and underscores - # - Maximum length is 63 bytes - pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*$" - if not re.match(pattern, v): - raise ValueError( - "Invalid table name. Must start with letter or underscore and contain only letters, numbers, and underscores" - ) - if len(v) > 63: - raise ValueError("Table name must be less than 63 characters") - return v - - @classmethod - def pip_packages(cls) -> list[str]: - return ["psycopg2-binary"] - - -class MongoDBKVStoreConfig(CommonConfig): - type: Literal["mongodb"] = KVStoreType.mongodb.value - host: str = "localhost" - port: int = 27017 - db: str = "llamastack" - user: str | None = None - password: str | None = None - collection_name: str = "llamastack_kvstore" - - @classmethod - def pip_packages(cls) -> list[str]: - return ["pymongo"] - - @classmethod - def sample_run_config(cls, collection_name: str = "llamastack_kvstore"): - return { - "type": "mongodb", - "host": "${env.MONGODB_HOST:=localhost}", - "port": "${env.MONGODB_PORT:=5432}", - "db": "${env.MONGODB_DB}", - "user": "${env.MONGODB_USER}", - "password": "${env.MONGODB_PASSWORD}", - "collection_name": "${env.MONGODB_COLLECTION_NAME:=" + collection_name + "}", - } +from pydantic import Field +from llama_stack.core.storage.datatypes import ( + MongoDBKVStoreConfig, + PostgresKVStoreConfig, + RedisKVStoreConfig, + SqliteKVStoreConfig, + StorageBackendType, +) KVStoreConfig = Annotated[ - RedisKVStoreConfig | SqliteKVStoreConfig | PostgresKVStoreConfig | MongoDBKVStoreConfig, - Field(discriminator="type", default=KVStoreType.sqlite.value), + RedisKVStoreConfig | SqliteKVStoreConfig | PostgresKVStoreConfig | MongoDBKVStoreConfig, Field(discriminator="type") ] @@ -148,13 +25,13 @@ def get_pip_packages(store_config: dict | KVStoreConfig) -> list[str]: """Get pip packages for KV store config, handling both dict and object cases.""" if isinstance(store_config, dict): store_type = store_config.get("type") - if store_type == "sqlite": + if store_type == StorageBackendType.KV_SQLITE.value: return SqliteKVStoreConfig.pip_packages() - elif store_type == "postgres": + elif store_type == StorageBackendType.KV_POSTGRES.value: return PostgresKVStoreConfig.pip_packages() - elif store_type == "redis": + elif store_type == StorageBackendType.KV_REDIS.value: return RedisKVStoreConfig.pip_packages() - elif store_type == "mongodb": + elif store_type == StorageBackendType.KV_MONGODB.value: return MongoDBKVStoreConfig.pip_packages() else: raise ValueError(f"Unknown KV store type: {store_type}") diff --git a/llama_stack/providers/utils/kvstore/kvstore.py b/llama_stack/providers/utils/kvstore/kvstore.py index 426523d8e..eee51e5d9 100644 --- a/llama_stack/providers/utils/kvstore/kvstore.py +++ b/llama_stack/providers/utils/kvstore/kvstore.py @@ -4,9 +4,17 @@ # 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 +# the root directory of this source tree. + +from __future__ import annotations + +from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig, StorageBackendType from .api import KVStore -from .config import KVStoreConfig, KVStoreType +from .config import KVStoreConfig def kvstore_dependencies(): @@ -44,20 +52,41 @@ class InmemoryKVStoreImpl(KVStore): del self._store[key] -async def kvstore_impl(config: KVStoreConfig) -> KVStore: - if config.type == KVStoreType.redis.value: +_KVSTORE_BACKENDS: dict[str, KVStoreConfig] = {} + + +def register_kvstore_backends(backends: dict[str, StorageBackendConfig]) -> None: + """Register the set of available KV store backends for reference resolution.""" + global _KVSTORE_BACKENDS + + _KVSTORE_BACKENDS.clear() + for name, cfg in backends.items(): + _KVSTORE_BACKENDS[name] = cfg + + +async def kvstore_impl(reference: KVStoreReference) -> KVStore: + backend_name = reference.backend + + backend_config = _KVSTORE_BACKENDS.get(backend_name) + if backend_config is None: + raise ValueError(f"Unknown KVStore backend '{backend_name}'. Registered backends: {sorted(_KVSTORE_BACKENDS)}") + + config = backend_config.model_copy() + config.namespace = reference.namespace + + if config.type == StorageBackendType.KV_REDIS.value: from .redis import RedisKVStoreImpl impl = RedisKVStoreImpl(config) - elif config.type == KVStoreType.sqlite.value: + elif config.type == StorageBackendType.KV_SQLITE.value: from .sqlite import SqliteKVStoreImpl impl = SqliteKVStoreImpl(config) - elif config.type == KVStoreType.postgres.value: + elif config.type == StorageBackendType.KV_POSTGRES.value: from .postgres import PostgresKVStoreImpl impl = PostgresKVStoreImpl(config) - elif config.type == KVStoreType.mongodb.value: + elif config.type == StorageBackendType.KV_MONGODB.value: from .mongodb import MongoDBKVStoreImpl impl = MongoDBKVStoreImpl(config) diff --git a/llama_stack/providers/utils/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 e1da4db6e..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 SqlStoreType logger = get_logger(name=__name__, category="providers::utils") @@ -82,8 +82,8 @@ class AuthorizedSqlStore: if not hasattr(self.sql_store, "config"): raise ValueError("SqlStore must have a config attribute to be used with AuthorizedSqlStore") - self.database_type = self.sql_store.config.type - if self.database_type not in [SqlStoreType.postgres, SqlStoreType.sqlite]: + self.database_type = self.sql_store.config.type.value + if self.database_type not in [StorageBackendType.SQL_POSTGRES.value, StorageBackendType.SQL_SQLITE.value]: raise ValueError(f"Unsupported database type: {self.database_type}") def _validate_sql_optimized_policy(self) -> None: @@ -220,9 +220,9 @@ class AuthorizedSqlStore: Returns: SQL expression to extract JSON value """ - if self.database_type == SqlStoreType.postgres: + if self.database_type == StorageBackendType.SQL_POSTGRES.value: return f"{column}->'{path}'" - elif self.database_type == SqlStoreType.sqlite: + elif self.database_type == StorageBackendType.SQL_SQLITE.value: return f"JSON_EXTRACT({column}, '$.{path}')" else: raise ValueError(f"Unsupported database type: {self.database_type}") @@ -237,9 +237,9 @@ class AuthorizedSqlStore: Returns: SQL expression to extract JSON value as text """ - if self.database_type == SqlStoreType.postgres: + if self.database_type == StorageBackendType.SQL_POSTGRES.value: return f"{column}->>'{path}'" - elif self.database_type == SqlStoreType.sqlite: + elif self.database_type == StorageBackendType.SQL_SQLITE.value: return f"JSON_EXTRACT({column}, '$.{path}')" else: raise ValueError(f"Unsupported database type: {self.database_type}") @@ -248,10 +248,10 @@ class AuthorizedSqlStore: """Get the SQL conditions for public access.""" # Public records are records that have no owner_principal or access_attributes conditions = ["owner_principal = ''"] - if self.database_type == SqlStoreType.postgres: + if self.database_type == StorageBackendType.SQL_POSTGRES.value: # Postgres stores JSON null as 'null' conditions.append("access_attributes::text = 'null'") - elif self.database_type == SqlStoreType.sqlite: + elif self.database_type == StorageBackendType.SQL_SQLITE.value: conditions.append("access_attributes = 'null'") else: raise ValueError(f"Unsupported database type: {self.database_type}") diff --git a/llama_stack/providers/utils/sqlstore/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 fc44402ae..31801c4ca 100644 --- a/llama_stack/providers/utils/sqlstore/sqlstore.py +++ b/llama_stack/providers/utils/sqlstore/sqlstore.py @@ -4,90 +4,28 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -from abc import abstractmethod -from enum import StrEnum -from pathlib import Path -from typing import Annotated, Literal +from typing import Annotated, cast -from pydantic import BaseModel, Field +from pydantic import Field -from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR +from llama_stack.core.storage.datatypes import ( + PostgresSqlStoreConfig, + SqliteSqlStoreConfig, + SqlStoreReference, + StorageBackendConfig, + StorageBackendType, +) from .api import SqlStore sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"] - -class SqlStoreType(StrEnum): - sqlite = "sqlite" - postgres = "postgres" - - -class SqlAlchemySqlStoreConfig(BaseModel): - @property - @abstractmethod - def engine_str(self) -> str: ... - - # TODO: move this when we have a better way to specify dependencies with internal APIs - @classmethod - def pip_packages(cls) -> list[str]: - return ["sqlalchemy[asyncio]"] - - -class SqliteSqlStoreConfig(SqlAlchemySqlStoreConfig): - type: Literal[SqlStoreType.sqlite] = SqlStoreType.sqlite - db_path: str = Field( - default=(RUNTIME_BASE_DIR / "sqlstore.db").as_posix(), - description="Database path, e.g. ~/.llama/distributions/ollama/sqlstore.db", - ) - - @property - def engine_str(self) -> str: - return "sqlite+aiosqlite:///" + Path(self.db_path).expanduser().as_posix() - - @classmethod - def sample_run_config(cls, __distro_dir__: str, db_name: str = "sqlstore.db"): - return { - "type": "sqlite", - "db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + db_name, - } - - @classmethod - def pip_packages(cls) -> list[str]: - return super().pip_packages() + ["aiosqlite"] - - -class PostgresSqlStoreConfig(SqlAlchemySqlStoreConfig): - type: Literal[SqlStoreType.postgres] = SqlStoreType.postgres - host: str = "localhost" - port: int = 5432 - db: str = "llamastack" - user: str - password: str | None = None - - @property - def engine_str(self) -> str: - return f"postgresql+asyncpg://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}" - - @classmethod - def pip_packages(cls) -> list[str]: - return super().pip_packages() + ["asyncpg"] - - @classmethod - def sample_run_config(cls, **kwargs): - return { - "type": "postgres", - "host": "${env.POSTGRES_HOST:=localhost}", - "port": "${env.POSTGRES_PORT:=5432}", - "db": "${env.POSTGRES_DB:=llamastack}", - "user": "${env.POSTGRES_USER:=llamastack}", - "password": "${env.POSTGRES_PASSWORD:=llamastack}", - } +_SQLSTORE_BACKENDS: dict[str, StorageBackendConfig] = {} SqlStoreConfig = Annotated[ SqliteSqlStoreConfig | PostgresSqlStoreConfig, - Field(discriminator="type", default=SqlStoreType.sqlite.value), + Field(discriminator="type"), ] @@ -95,9 +33,9 @@ def get_pip_packages(store_config: dict | SqlStoreConfig) -> list[str]: """Get pip packages for SQL store config, handling both dict and object cases.""" if isinstance(store_config, dict): store_type = store_config.get("type") - if store_type == "sqlite": + if store_type == StorageBackendType.SQL_SQLITE.value: return SqliteSqlStoreConfig.pip_packages() - elif store_type == "postgres": + elif store_type == StorageBackendType.SQL_POSTGRES.value: return PostgresSqlStoreConfig.pip_packages() else: raise ValueError(f"Unknown SQL store type: {store_type}") @@ -105,12 +43,28 @@ def get_pip_packages(store_config: dict | SqlStoreConfig) -> list[str]: return store_config.pip_packages() -def sqlstore_impl(config: SqlStoreConfig) -> SqlStore: - if config.type in [SqlStoreType.sqlite, SqlStoreType.postgres]: +def sqlstore_impl(reference: SqlStoreReference) -> SqlStore: + backend_name = reference.backend + + backend_config = _SQLSTORE_BACKENDS.get(backend_name) + 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 - impl = SqlAlchemySqlStoreImpl(config) + config = cast(SqliteSqlStoreConfig | PostgresSqlStoreConfig, backend_config).model_copy() + return SqlAlchemySqlStoreImpl(config) else: - raise ValueError(f"Unknown sqlstore type {config.type}") + raise ValueError(f"Unknown sqlstore type {backend_config.type}") - return impl + +def register_sqlstore_backends(backends: dict[str, StorageBackendConfig]) -> None: + """Register the set of available SQL store backends for reference resolution.""" + global _SQLSTORE_BACKENDS + + _SQLSTORE_BACKENDS.clear() + for name, cfg in backends.items(): + _SQLSTORE_BACKENDS[name] = cfg diff --git a/scripts/docker.sh b/scripts/docker.sh index 1ba1d9adf..7a5c3e6e0 100755 --- a/scripts/docker.sh +++ b/scripts/docker.sh @@ -236,7 +236,7 @@ start_container() { echo "=== Starting Docker Container ===" # Get the repo root for volume mount - SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) + SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd) REPO_ROOT=$(cd "$SCRIPT_DIR/.." && pwd) # Determine the actual image name (may have localhost/ prefix) diff --git a/tests/external/run-byoa.yaml b/tests/external/run-byoa.yaml index 5774ae9da..4d63046c6 100644 --- a/tests/external/run-byoa.yaml +++ b/tests/external/run-byoa.yaml @@ -7,6 +7,24 @@ providers: - provider_id: kaze provider_type: remote::kaze config: {} +storage: + backends: + kv_default: + type: kv_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/external}/kvstore.db + sql_default: + type: sql_sqlite + db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/external}/sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + conversations: + table_name: openai_conversations + backend: sql_default external_apis_dir: ~/.llama/apis.d external_providers_dir: ~/.llama/providers.d server: diff --git a/tests/integration/fixtures/common.py b/tests/integration/fixtures/common.py index 68a30fc69..eb6840e60 100644 --- a/tests/integration/fixtures/common.py +++ b/tests/integration/fixtures/common.py @@ -238,7 +238,7 @@ def instantiate_llama_stack_client(session): run_config = run_config_from_adhoc_config_spec(config) run_config_file = tempfile.NamedTemporaryFile(delete=False, suffix=".yaml") with open(run_config_file.name, "w") as f: - yaml.dump(run_config.model_dump(), f) + yaml.dump(run_config.model_dump(mode="json"), f) config = run_config_file.name client = LlamaStackAsLibraryClient( 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 new file mode 100644 index 000000000..e9b80dc0c --- /dev/null +++ b/tests/integration/test_persistence_integration.py @@ -0,0 +1,71 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import yaml + +from llama_stack.core.datatypes import StackRunConfig +from llama_stack.core.storage.datatypes import ( + PostgresKVStoreConfig, + PostgresSqlStoreConfig, + SqliteKVStoreConfig, + SqliteSqlStoreConfig, +) + + +def test_starter_distribution_config_loads_and_resolves(): + """Integration: Actual starter config should parse and have correct storage structure.""" + with open("llama_stack/distributions/starter/run.yaml") as f: + config_dict = yaml.safe_load(f) + + config = StackRunConfig(**config_dict) + + # Config should have named backends and explicit store references + assert config.storage is not None + 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 = config.storage.stores + assert stores.metadata is not None + assert stores.metadata.backend == "kv_default" + assert stores.metadata.namespace == "registry" + + assert stores.inference is not None + assert stores.inference.backend == "sql_default" + assert stores.inference.table_name == "inference_store" + assert stores.inference.max_write_queue_size > 0 + assert stores.inference.num_writers > 0 + + assert stores.conversations is not None + assert stores.conversations.backend == "sql_default" + assert stores.conversations.table_name == "openai_conversations" + + +def test_postgres_demo_distribution_config_loads(): + """Integration: Postgres demo should use Postgres backend for all stores.""" + with open("llama_stack/distributions/postgres-demo/run.yaml") as f: + config_dict = yaml.safe_load(f) + + config = StackRunConfig(**config_dict) + + # Should have postgres backend + assert config.storage is not None + assert "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 = config.storage.stores + # Stores target the Postgres backends explicitly + assert stores.metadata is not None + assert stores.metadata.backend == "kv_default" + assert stores.inference is not None + assert stores.inference.backend == "sql_default" diff --git a/tests/unit/cli/test_stack_config.py b/tests/unit/cli/test_stack_config.py index daaf229e5..7b9f3ca0c 100644 --- a/tests/unit/cli/test_stack_config.py +++ b/tests/unit/cli/test_stack_config.py @@ -23,6 +23,27 @@ def config_with_image_name_int(): image_name: 1234 apis_to_serve: [] built_at: {datetime.now().isoformat()} + storage: + backends: + kv_default: + type: kv_sqlite + db_path: /tmp/test_kv.db + sql_default: + type: sql_sqlite + db_path: /tmp/test_sql.db + stores: + metadata: + backend: kv_default + namespace: metadata + inference: + backend: sql_default + table_name: inference + conversations: + backend: sql_default + table_name: conversations + responses: + backend: sql_default + table_name: responses providers: inference: - provider_id: provider1 @@ -54,6 +75,27 @@ def up_to_date_config(): image_name: foo apis_to_serve: [] built_at: {datetime.now().isoformat()} + storage: + backends: + kv_default: + type: kv_sqlite + db_path: /tmp/test_kv.db + sql_default: + type: sql_sqlite + db_path: /tmp/test_sql.db + stores: + metadata: + backend: kv_default + namespace: metadata + inference: + backend: sql_default + table_name: inference + conversations: + backend: sql_default + table_name: conversations + responses: + backend: sql_default + table_name: responses providers: inference: - provider_id: provider1 diff --git a/tests/unit/conversations/test_conversations.py b/tests/unit/conversations/test_conversations.py index 65c3e2333..ff6dd243d 100644 --- a/tests/unit/conversations/test_conversations.py +++ b/tests/unit/conversations/test_conversations.py @@ -20,7 +20,14 @@ from llama_stack.core.conversations.conversations import ( ConversationServiceConfig, ConversationServiceImpl, ) -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig +from llama_stack.core.datatypes import StackRunConfig +from llama_stack.core.storage.datatypes import ( + ServerStoresConfig, + SqliteSqlStoreConfig, + SqlStoreReference, + StorageConfig, +) +from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends @pytest.fixture @@ -28,7 +35,18 @@ async def service(): with tempfile.TemporaryDirectory() as tmpdir: db_path = Path(tmpdir) / "test_conversations.db" - config = ConversationServiceConfig(conversations_store=SqliteSqlStoreConfig(db_path=str(db_path)), policy=[]) + storage = StorageConfig( + backends={ + "sql_test": SqliteSqlStoreConfig(db_path=str(db_path)), + }, + stores=ServerStoresConfig( + conversations=SqlStoreReference(backend="sql_test", table_name="openai_conversations"), + ), + ) + register_sqlstore_backends({"sql_test": storage.backends["sql_test"]}) + run_config = StackRunConfig(image_name="test", apis=[], providers={}, storage=storage) + + config = ConversationServiceConfig(run_config=run_config, policy=[]) service = ConversationServiceImpl(config, {}) await service.initialize() yield service @@ -121,9 +139,18 @@ async def test_policy_configuration(): AccessRule(forbid=Scope(principal="test_user", actions=[Action.CREATE, Action.READ], resource="*")) ] - config = ConversationServiceConfig( - conversations_store=SqliteSqlStoreConfig(db_path=str(db_path)), policy=restrictive_policy + storage = StorageConfig( + backends={ + "sql_test": SqliteSqlStoreConfig(db_path=str(db_path)), + }, + stores=ServerStoresConfig( + conversations=SqlStoreReference(backend="sql_test", table_name="openai_conversations"), + ), ) + register_sqlstore_backends({"sql_test": storage.backends["sql_test"]}) + run_config = StackRunConfig(image_name="test", apis=[], providers={}, storage=storage) + + config = ConversationServiceConfig(run_config=run_config, policy=restrictive_policy) service = ConversationServiceImpl(config, {}) await service.initialize() diff --git a/tests/unit/core/test_storage_references.py b/tests/unit/core/test_storage_references.py new file mode 100644 index 000000000..7bceba74d --- /dev/null +++ b/tests/unit/core/test_storage_references.py @@ -0,0 +1,84 @@ +# 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, + ServerStoresConfig, + SqliteKVStoreConfig, + SqliteSqlStoreConfig, + SqlStoreReference, + StorageConfig, +) + + +def _base_run_config(**overrides): + metadata_reference = overrides.pop( + "metadata_reference", + KVStoreReference(backend="kv_default", namespace="registry"), + ) + inference_reference = overrides.pop( + "inference_reference", + InferenceStoreReference(backend="sql_default", table_name="inference"), + ) + conversations_reference = overrides.pop( + "conversations_reference", + SqlStoreReference(backend="sql_default", table_name="conversations"), + ) + storage = overrides.pop( + "storage", + StorageConfig( + backends={ + "kv_default": SqliteKVStoreConfig(db_path="/tmp/kv.db"), + "sql_default": SqliteSqlStoreConfig(db_path="/tmp/sql.db"), + }, + stores=ServerStoresConfig( + metadata=metadata_reference, + inference=inference_reference, + conversations=conversations_reference, + ), + ), + ) + return StackRunConfig( + version=LLAMA_STACK_RUN_CONFIG_VERSION, + image_name="test-distro", + apis=[], + providers={}, + storage=storage, + **overrides, + ) + + +def test_references_require_known_backend(): + with pytest.raises(ValidationError, match="unknown backend 'missing'"): + _base_run_config(metadata_reference=KVStoreReference(backend="missing", namespace="registry")) + + +def test_references_must_match_backend_family(): + with pytest.raises(ValidationError, match="kv_.* is required"): + _base_run_config(metadata_reference=KVStoreReference(backend="sql_default", namespace="registry")) + + with pytest.raises(ValidationError, match="sql_.* is required"): + _base_run_config( + inference_reference=InferenceStoreReference(backend="kv_default", table_name="inference"), + ) + + +def test_valid_configuration_passes_validation(): + config = _base_run_config() + stores = config.storage.stores + assert stores.metadata is not None and stores.metadata.backend == "kv_default" + assert stores.inference is not None and stores.inference.backend == "sql_default" + assert stores.conversations is not None and stores.conversations.backend == "sql_default" diff --git a/tests/unit/distribution/test_distribution.py b/tests/unit/distribution/test_distribution.py index 08a376008..3b0643a13 100644 --- a/tests/unit/distribution/test_distribution.py +++ b/tests/unit/distribution/test_distribution.py @@ -13,6 +13,15 @@ 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, + ServerStoresConfig, + SqliteKVStoreConfig, + SqliteSqlStoreConfig, + SqlStoreReference, + StorageConfig, +) from llama_stack.providers.datatypes import ProviderSpec @@ -29,6 +38,32 @@ class SampleConfig(BaseModel): } +def _default_storage() -> StorageConfig: + return StorageConfig( + backends={ + "kv_default": SqliteKVStoreConfig(db_path=":memory:"), + "sql_default": SqliteSqlStoreConfig(db_path=":memory:"), + }, + stores=ServerStoresConfig( + metadata=KVStoreReference(backend="kv_default", namespace="registry"), + inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"), + conversations=SqlStoreReference(backend="sql_default", table_name="conversations"), + ), + ) + + +def make_stack_config(**overrides) -> StackRunConfig: + storage = overrides.pop("storage", _default_storage()) + defaults = dict( + image_name="test_image", + apis=[], + providers={}, + storage=storage, + ) + defaults.update(overrides) + return StackRunConfig(**defaults) + + @pytest.fixture def mock_providers(): """Mock the available_providers function to return test providers.""" @@ -47,8 +82,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 +257,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 +313,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 +327,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 +351,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 +374,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 +431,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 +457,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 +474,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 +594,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 +618,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 +642,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 +666,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 +690,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 +715,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 +745,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 +753,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 +777,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 +790,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 +812,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 +829,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 +853,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..fe30e1a77 100644 --- a/tests/unit/prompts/prompts/conftest.py +++ b/tests/unit/prompts/prompts/conftest.py @@ -9,7 +9,16 @@ 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, + ServerStoresConfig, + SqliteKVStoreConfig, + SqliteSqlStoreConfig, + SqlStoreReference, + StorageConfig, +) +from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends @pytest.fixture @@ -19,12 +28,28 @@ 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")), + }, + stores=ServerStoresConfig( + metadata=KVStoreReference(backend="kv_test", namespace="registry"), + inference=InferenceStoreReference(backend="sql_test", table_name="inference"), + conversations=SqlStoreReference(backend="sql_test", table_name="conversations"), + ), + ) + mock_run_config = StackRunConfig( + image_name="test-distribution", + apis=[], + providers={}, + storage=storage, + ) 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/agent/test_meta_reference_agent.py b/tests/unit/providers/agent/test_meta_reference_agent.py index cfb3e1327..dfd9b6d52 100644 --- a/tests/unit/providers/agent/test_meta_reference_agent.py +++ b/tests/unit/providers/agent/test_meta_reference_agent.py @@ -26,6 +26,20 @@ from llama_stack.providers.inline.agents.meta_reference.config import MetaRefere from llama_stack.providers.inline.agents.meta_reference.persistence import AgentInfo +@pytest.fixture(autouse=True) +def setup_backends(tmp_path): + """Register KV and SQL store backends for testing.""" + from llama_stack.core.storage.datatypes import SqliteKVStoreConfig, SqliteSqlStoreConfig + from llama_stack.providers.utils.kvstore.kvstore import register_kvstore_backends + from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends + + kv_path = str(tmp_path / "test_kv.db") + sql_path = str(tmp_path / "test_sql.db") + + register_kvstore_backends({"kv_default": SqliteKVStoreConfig(db_path=kv_path)}) + register_sqlstore_backends({"sql_default": SqliteSqlStoreConfig(db_path=sql_path)}) + + @pytest.fixture def mock_apis(): return { @@ -40,15 +54,20 @@ def mock_apis(): @pytest.fixture def config(tmp_path): + from llama_stack.core.storage.datatypes import KVStoreReference, ResponsesStoreReference + from llama_stack.providers.inline.agents.meta_reference.config import AgentPersistenceConfig + return MetaReferenceAgentsImplConfig( - persistence_store={ - "type": "sqlite", - "db_path": str(tmp_path / "test.db"), - }, - responses_store={ - "type": "sqlite", - "db_path": str(tmp_path / "test.db"), - }, + persistence=AgentPersistenceConfig( + agent_state=KVStoreReference( + backend="kv_default", + namespace="agents", + ), + responses=ResponsesStoreReference( + backend="sql_default", + table_name="responses", + ), + ) ) 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 54c1820fb..f31ec0c28 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 @@ -917,8 +917,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..6d0367beb 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") @@ -138,7 +140,7 @@ async def sqlite_vec_vec_index(embedding_dimension, tmp_path_factory): async def sqlite_vec_adapter(sqlite_vec_db_path, unique_kvstore_config, mock_inference_api, embedding_dimension): config = SQLiteVectorIOConfig( db_path=sqlite_vec_db_path, - kvstore=unique_kvstore_config, + persistence=unique_kvstore_config, ) adapter = SQLiteVecVectorIOAdapter( config=config, @@ -177,7 +179,7 @@ async def faiss_vec_index(embedding_dimension): @pytest.fixture async def faiss_vec_adapter(unique_kvstore_config, mock_inference_api, embedding_dimension): config = FaissVectorIOConfig( - kvstore=unique_kvstore_config, + persistence=unique_kvstore_config, ) adapter = FaissVectorIOAdapter( config=config, @@ -253,7 +255,7 @@ async def pgvector_vec_adapter(unique_kvstore_config, mock_inference_api, embedd db="test_db", user="test_user", password="test_password", - kvstore=unique_kvstore_config, + persistence=unique_kvstore_config, ) adapter = PGVectorVectorIOAdapter(config, mock_inference_api, None) 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..b44f12f7e 100644 --- a/tests/unit/server/test_resolver.py +++ b/tests/unit/server/test_resolver.py @@ -12,15 +12,22 @@ 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, + ServerStoresConfig, + SqliteKVStoreConfig, + SqliteSqlStoreConfig, + SqlStoreReference, + StorageConfig, +) from llama_stack.providers.datatypes import InlineProviderSpec, ProviderSpec +from llama_stack.providers.utils.kvstore import register_kvstore_backends +from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends def add_protocol_methods(cls: type, protocol: type[Protocol]) -> None: @@ -65,6 +72,35 @@ 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:"), + }, + stores=ServerStoresConfig( + metadata=KVStoreReference(backend="kv_default", namespace="registry"), + inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"), + conversations=SqlStoreReference(backend="sql_default", table_name="conversations"), + ), + ), + ) + register_kvstore_backends({name: cfg for name, cfg in storage.backends.items() if cfg.type.value.startswith("kv_")}) + register_sqlstore_backends( + {name: cfg for name, cfg in storage.backends.items() if cfg.type.value.startswith("sql_")} + ) + defaults = dict( + image_name="test_image", + apis=[], + providers={}, + storage=storage, + ) + 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/inference/test_inference_store.py b/tests/unit/utils/inference/test_inference_store.py index f6d63490a..d2de1c759 100644 --- a/tests/unit/utils/inference/test_inference_store.py +++ b/tests/unit/utils/inference/test_inference_store.py @@ -5,7 +5,6 @@ # the root directory of this source tree. import time -from tempfile import TemporaryDirectory import pytest @@ -16,8 +15,16 @@ from llama_stack.apis.inference import ( OpenAIUserMessageParam, Order, ) +from llama_stack.core.storage.datatypes import InferenceStoreReference, SqliteSqlStoreConfig from llama_stack.providers.utils.inference.inference_store import InferenceStore -from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig +from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends + + +@pytest.fixture(autouse=True) +def setup_backends(tmp_path): + """Register SQL store backends for testing.""" + db_path = str(tmp_path / "test.db") + register_sqlstore_backends({"sql_default": SqliteSqlStoreConfig(db_path=db_path)}) def create_test_chat_completion( @@ -44,167 +51,162 @@ def create_test_chat_completion( async def test_inference_store_pagination_basic(): """Test basic pagination functionality.""" - with TemporaryDirectory() as tmp_dir: - db_path = tmp_dir + "/test.db" - store = InferenceStore(SqliteSqlStoreConfig(db_path=db_path), policy=[]) - await store.initialize() + reference = InferenceStoreReference(backend="sql_default", table_name="chat_completions") + store = InferenceStore(reference, policy=[]) + await store.initialize() - # Create test data with different timestamps - base_time = int(time.time()) - test_data = [ - ("zebra-task", base_time + 1), - ("apple-job", base_time + 2), - ("moon-work", base_time + 3), - ("banana-run", base_time + 4), - ("car-exec", base_time + 5), - ] + # Create test data with different timestamps + base_time = int(time.time()) + test_data = [ + ("zebra-task", base_time + 1), + ("apple-job", base_time + 2), + ("moon-work", base_time + 3), + ("banana-run", base_time + 4), + ("car-exec", base_time + 5), + ] - # Store test chat completions - for completion_id, timestamp in test_data: - completion = create_test_chat_completion(completion_id, timestamp) - input_messages = [OpenAIUserMessageParam(role="user", content=f"Test message for {completion_id}")] - await store.store_chat_completion(completion, input_messages) + # Store test chat completions + for completion_id, timestamp in test_data: + completion = create_test_chat_completion(completion_id, timestamp) + input_messages = [OpenAIUserMessageParam(role="user", content=f"Test message for {completion_id}")] + await store.store_chat_completion(completion, input_messages) - # Wait for all queued writes to complete - await store.flush() + # Wait for all queued writes to complete + await store.flush() - # Test 1: First page with limit=2, descending order (default) - result = await store.list_chat_completions(limit=2, order=Order.desc) - assert len(result.data) == 2 - assert result.data[0].id == "car-exec" # Most recent first - assert result.data[1].id == "banana-run" - assert result.has_more is True - assert result.last_id == "banana-run" + # Test 1: First page with limit=2, descending order (default) + result = await store.list_chat_completions(limit=2, order=Order.desc) + assert len(result.data) == 2 + assert result.data[0].id == "car-exec" # Most recent first + assert result.data[1].id == "banana-run" + assert result.has_more is True + assert result.last_id == "banana-run" - # Test 2: Second page using 'after' parameter - result2 = await store.list_chat_completions(after="banana-run", limit=2, order=Order.desc) - assert len(result2.data) == 2 - assert result2.data[0].id == "moon-work" - assert result2.data[1].id == "apple-job" - assert result2.has_more is True + # Test 2: Second page using 'after' parameter + result2 = await store.list_chat_completions(after="banana-run", limit=2, order=Order.desc) + assert len(result2.data) == 2 + assert result2.data[0].id == "moon-work" + assert result2.data[1].id == "apple-job" + assert result2.has_more is True - # Test 3: Final page - result3 = await store.list_chat_completions(after="apple-job", limit=2, order=Order.desc) - assert len(result3.data) == 1 - assert result3.data[0].id == "zebra-task" - assert result3.has_more is False + # Test 3: Final page + result3 = await store.list_chat_completions(after="apple-job", limit=2, order=Order.desc) + assert len(result3.data) == 1 + assert result3.data[0].id == "zebra-task" + assert result3.has_more is False async def test_inference_store_pagination_ascending(): """Test pagination with ascending order.""" - with TemporaryDirectory() as tmp_dir: - db_path = tmp_dir + "/test.db" - store = InferenceStore(SqliteSqlStoreConfig(db_path=db_path), policy=[]) - await store.initialize() + reference = InferenceStoreReference(backend="sql_default", table_name="chat_completions") + store = InferenceStore(reference, policy=[]) + await store.initialize() - # Create test data - base_time = int(time.time()) - test_data = [ - ("delta-item", base_time + 1), - ("charlie-task", base_time + 2), - ("alpha-work", base_time + 3), - ] + # Create test data + base_time = int(time.time()) + test_data = [ + ("delta-item", base_time + 1), + ("charlie-task", base_time + 2), + ("alpha-work", base_time + 3), + ] - # Store test chat completions - for completion_id, timestamp in test_data: - completion = create_test_chat_completion(completion_id, timestamp) - input_messages = [OpenAIUserMessageParam(role="user", content=f"Test message for {completion_id}")] - await store.store_chat_completion(completion, input_messages) + # Store test chat completions + for completion_id, timestamp in test_data: + completion = create_test_chat_completion(completion_id, timestamp) + input_messages = [OpenAIUserMessageParam(role="user", content=f"Test message for {completion_id}")] + await store.store_chat_completion(completion, input_messages) - # Wait for all queued writes to complete - await store.flush() + # Wait for all queued writes to complete + await store.flush() - # Test ascending order pagination - result = await store.list_chat_completions(limit=1, order=Order.asc) - assert len(result.data) == 1 - assert result.data[0].id == "delta-item" # Oldest first - assert result.has_more is True + # Test ascending order pagination + result = await store.list_chat_completions(limit=1, order=Order.asc) + assert len(result.data) == 1 + assert result.data[0].id == "delta-item" # Oldest first + assert result.has_more is True - # Second page with ascending order - result2 = await store.list_chat_completions(after="delta-item", limit=1, order=Order.asc) - assert len(result2.data) == 1 - assert result2.data[0].id == "charlie-task" - assert result2.has_more is True + # Second page with ascending order + result2 = await store.list_chat_completions(after="delta-item", limit=1, order=Order.asc) + assert len(result2.data) == 1 + assert result2.data[0].id == "charlie-task" + assert result2.has_more is True async def test_inference_store_pagination_with_model_filter(): """Test pagination combined with model filtering.""" - with TemporaryDirectory() as tmp_dir: - db_path = tmp_dir + "/test.db" - store = InferenceStore(SqliteSqlStoreConfig(db_path=db_path), policy=[]) - await store.initialize() + reference = InferenceStoreReference(backend="sql_default", table_name="chat_completions") + store = InferenceStore(reference, policy=[]) + await store.initialize() - # Create test data with different models - base_time = int(time.time()) - test_data = [ - ("xyz-task", base_time + 1, "model-a"), - ("def-work", base_time + 2, "model-b"), - ("pqr-job", base_time + 3, "model-a"), - ("abc-run", base_time + 4, "model-b"), - ] + # Create test data with different models + base_time = int(time.time()) + test_data = [ + ("xyz-task", base_time + 1, "model-a"), + ("def-work", base_time + 2, "model-b"), + ("pqr-job", base_time + 3, "model-a"), + ("abc-run", base_time + 4, "model-b"), + ] - # Store test chat completions - for completion_id, timestamp, model in test_data: - completion = create_test_chat_completion(completion_id, timestamp, model) - input_messages = [OpenAIUserMessageParam(role="user", content=f"Test message for {completion_id}")] - await store.store_chat_completion(completion, input_messages) + # Store test chat completions + for completion_id, timestamp, model in test_data: + completion = create_test_chat_completion(completion_id, timestamp, model) + input_messages = [OpenAIUserMessageParam(role="user", content=f"Test message for {completion_id}")] + await store.store_chat_completion(completion, input_messages) - # Wait for all queued writes to complete - await store.flush() + # Wait for all queued writes to complete + await store.flush() - # Test pagination with model filter - result = await store.list_chat_completions(limit=1, model="model-a", order=Order.desc) - assert len(result.data) == 1 - assert result.data[0].id == "pqr-job" # Most recent model-a - assert result.data[0].model == "model-a" - assert result.has_more is True + # Test pagination with model filter + result = await store.list_chat_completions(limit=1, model="model-a", order=Order.desc) + assert len(result.data) == 1 + assert result.data[0].id == "pqr-job" # Most recent model-a + assert result.data[0].model == "model-a" + assert result.has_more is True - # Second page with model filter - result2 = await store.list_chat_completions(after="pqr-job", limit=1, model="model-a", order=Order.desc) - assert len(result2.data) == 1 - assert result2.data[0].id == "xyz-task" - assert result2.data[0].model == "model-a" - assert result2.has_more is False + # Second page with model filter + result2 = await store.list_chat_completions(after="pqr-job", limit=1, model="model-a", order=Order.desc) + assert len(result2.data) == 1 + assert result2.data[0].id == "xyz-task" + assert result2.data[0].model == "model-a" + assert result2.has_more is False async def test_inference_store_pagination_invalid_after(): """Test error handling for invalid 'after' parameter.""" - with TemporaryDirectory() as tmp_dir: - db_path = tmp_dir + "/test.db" - store = InferenceStore(SqliteSqlStoreConfig(db_path=db_path), policy=[]) - await store.initialize() + reference = InferenceStoreReference(backend="sql_default", table_name="chat_completions") + store = InferenceStore(reference, policy=[]) + await store.initialize() - # Try to paginate with non-existent ID - with pytest.raises(ValueError, match="Record with id='non-existent' not found in table 'chat_completions'"): - await store.list_chat_completions(after="non-existent", limit=2) + # Try to paginate with non-existent ID + with pytest.raises(ValueError, match="Record with id='non-existent' not found in table 'chat_completions'"): + await store.list_chat_completions(after="non-existent", limit=2) async def test_inference_store_pagination_no_limit(): """Test pagination behavior when no limit is specified.""" - with TemporaryDirectory() as tmp_dir: - db_path = tmp_dir + "/test.db" - store = InferenceStore(SqliteSqlStoreConfig(db_path=db_path), policy=[]) - await store.initialize() + reference = InferenceStoreReference(backend="sql_default", table_name="chat_completions") + store = InferenceStore(reference, policy=[]) + await store.initialize() - # Create test data - base_time = int(time.time()) - test_data = [ - ("omega-first", base_time + 1), - ("beta-second", base_time + 2), - ] + # Create test data + base_time = int(time.time()) + test_data = [ + ("omega-first", base_time + 1), + ("beta-second", base_time + 2), + ] - # Store test chat completions - for completion_id, timestamp in test_data: - completion = create_test_chat_completion(completion_id, timestamp) - input_messages = [OpenAIUserMessageParam(role="user", content=f"Test message for {completion_id}")] - await store.store_chat_completion(completion, input_messages) + # Store test chat completions + for completion_id, timestamp in test_data: + completion = create_test_chat_completion(completion_id, timestamp) + input_messages = [OpenAIUserMessageParam(role="user", content=f"Test message for {completion_id}")] + await store.store_chat_completion(completion, input_messages) - # Wait for all queued writes to complete - await store.flush() + # Wait for all queued writes to complete + await store.flush() - # Test without limit - result = await store.list_chat_completions(order=Order.desc) - assert len(result.data) == 2 - assert result.data[0].id == "beta-second" # Most recent first - assert result.data[1].id == "omega-first" - assert result.has_more is False + # Test without limit + result = await store.list_chat_completions(order=Order.desc) + assert len(result.data) == 2 + assert result.data[0].id == "beta-second" # Most recent first + assert result.data[1].id == "omega-first" + assert result.has_more is False 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