From 46715e7889e2eb6359f4ee46de7723beebfa94c4 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Mon, 10 Nov 2025 14:52:36 -0800 Subject: [PATCH] better --- .../ci-tests/run-with-postgres-store.yaml | 16 +++++----- .../starter-gpu/run-with-postgres-store.yaml | 16 +++++----- .../starter/run-with-postgres-store.yaml | 16 +++++----- .../distributions/starter/starter.py | 31 +++++++------------ 4 files changed, 35 insertions(+), 44 deletions(-) diff --git a/src/llama_stack/distributions/ci-tests/run-with-postgres-store.yaml b/src/llama_stack/distributions/ci-tests/run-with-postgres-store.yaml index c149fd696..bb08e3f9e 100644 --- a/src/llama_stack/distributions/ci-tests/run-with-postgres-store.yaml +++ b/src/llama_stack/distributions/ci-tests/run-with-postgres-store.yaml @@ -153,7 +153,7 @@ providers: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/ci-tests/files} metadata_store: table_name: files_metadata - backend: sql_postgres + backend: sql_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -234,10 +234,10 @@ providers: config: kvstore: namespace: batches - backend: kv_postgres + backend: kv_default storage: backends: - kv_postgres: + kv_default: type: kv_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} @@ -245,7 +245,7 @@ storage: user: ${env.POSTGRES_USER:=llamastack} password: ${env.POSTGRES_PASSWORD:=llamastack} table_name: ${env.POSTGRES_TABLE_NAME:=llamastack_kvstore} - sql_postgres: + sql_default: type: sql_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} @@ -255,18 +255,18 @@ storage: stores: metadata: namespace: registry - backend: kv_postgres + backend: kv_default inference: table_name: inference_store - backend: sql_postgres + backend: sql_default max_write_queue_size: 10000 num_writers: 4 conversations: table_name: openai_conversations - backend: sql_postgres + backend: sql_default prompts: namespace: prompts - backend: kv_postgres + backend: kv_default registered_resources: models: [] shields: [] diff --git a/src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml b/src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml index 07a60b69c..a31d7bf80 100644 --- a/src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml +++ b/src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml @@ -153,7 +153,7 @@ providers: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter-gpu/files} metadata_store: table_name: files_metadata - backend: sql_postgres + backend: sql_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -237,10 +237,10 @@ providers: config: kvstore: namespace: batches - backend: kv_postgres + backend: kv_default storage: backends: - kv_postgres: + kv_default: type: kv_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} @@ -248,7 +248,7 @@ storage: user: ${env.POSTGRES_USER:=llamastack} password: ${env.POSTGRES_PASSWORD:=llamastack} table_name: ${env.POSTGRES_TABLE_NAME:=llamastack_kvstore} - sql_postgres: + sql_default: type: sql_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} @@ -258,18 +258,18 @@ storage: stores: metadata: namespace: registry - backend: kv_postgres + backend: kv_default inference: table_name: inference_store - backend: sql_postgres + backend: sql_default max_write_queue_size: 10000 num_writers: 4 conversations: table_name: openai_conversations - backend: sql_postgres + backend: sql_default prompts: namespace: prompts - backend: kv_postgres + backend: kv_default registered_resources: models: [] shields: [] diff --git a/src/llama_stack/distributions/starter/run-with-postgres-store.yaml b/src/llama_stack/distributions/starter/run-with-postgres-store.yaml index fdbe3ef8e..5d747e1ee 100644 --- a/src/llama_stack/distributions/starter/run-with-postgres-store.yaml +++ b/src/llama_stack/distributions/starter/run-with-postgres-store.yaml @@ -153,7 +153,7 @@ providers: storage_dir: ${env.FILES_STORAGE_DIR:=~/.llama/distributions/starter/files} metadata_store: table_name: files_metadata - backend: sql_postgres + backend: sql_default safety: - provider_id: llama-guard provider_type: inline::llama-guard @@ -234,10 +234,10 @@ providers: config: kvstore: namespace: batches - backend: kv_postgres + backend: kv_default storage: backends: - kv_postgres: + kv_default: type: kv_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} @@ -245,7 +245,7 @@ storage: user: ${env.POSTGRES_USER:=llamastack} password: ${env.POSTGRES_PASSWORD:=llamastack} table_name: ${env.POSTGRES_TABLE_NAME:=llamastack_kvstore} - sql_postgres: + sql_default: type: sql_postgres host: ${env.POSTGRES_HOST:=localhost} port: ${env.POSTGRES_PORT:=5432} @@ -255,18 +255,18 @@ storage: stores: metadata: namespace: registry - backend: kv_postgres + backend: kv_default inference: table_name: inference_store - backend: sql_postgres + backend: sql_default max_write_queue_size: 10000 num_writers: 4 conversations: table_name: openai_conversations - backend: sql_postgres + backend: sql_default prompts: namespace: prompts - backend: kv_postgres + backend: kv_default registered_resources: models: [] shields: [] diff --git a/src/llama_stack/distributions/starter/starter.py b/src/llama_stack/distributions/starter/starter.py index 0e8993f76..29a275e91 100644 --- a/src/llama_stack/distributions/starter/starter.py +++ b/src/llama_stack/distributions/starter/starter.py @@ -160,15 +160,6 @@ def get_distribution_template(name: str = "starter") -> DistributionTemplate: provider_type="inline::localfs", config=files_config, ) - postgres_files_config = dict(files_config) - postgres_metadata_store = dict(postgres_files_config.get("metadata_store", {})) - postgres_metadata_store["backend"] = "sql_postgres" - postgres_files_config["metadata_store"] = postgres_metadata_store - postgres_files_provider = Provider( - provider_id="meta-reference-files", - provider_type="inline::localfs", - config=postgres_files_config, - ) embedding_provider = Provider( provider_id="sentence-transformers", provider_type="inline::sentence-transformers", @@ -197,7 +188,8 @@ def get_distribution_template(name: str = "starter") -> DistributionTemplate: provider_shield_id="${env.CODE_SCANNER_MODEL:=}", ), ] - postgres_config = PostgresSqlStoreConfig.sample_run_config() + postgres_sql_config = PostgresSqlStoreConfig.sample_run_config() + postgres_kv_config = PostgresKVStoreConfig.sample_run_config() default_overrides = { "inference": remote_inference_providers + [embedding_provider], "vector_io": [ @@ -282,14 +274,13 @@ def get_distribution_template(name: str = "starter") -> DistributionTemplate: "run-with-postgres-store.yaml": RunConfigSettings( provider_overrides={ **default_overrides, - "files": [postgres_files_provider], "agents": [ Provider( provider_id="meta-reference", provider_type="inline::meta-reference", config=dict( - persistence_store=postgres_config, - responses_store=postgres_config, + persistence_store=postgres_sql_config, + responses_store=postgres_sql_config, ), ) ], @@ -299,7 +290,7 @@ def get_distribution_template(name: str = "starter") -> DistributionTemplate: provider_type="inline::reference", config=dict( kvstore=KVStoreReference( - backend="kv_postgres", + backend="kv_default", namespace="batches", ).model_dump(exclude_none=True), ), @@ -307,24 +298,24 @@ def get_distribution_template(name: str = "starter") -> DistributionTemplate: ], }, storage_backends={ - "kv_postgres": PostgresKVStoreConfig.sample_run_config(), - "sql_postgres": postgres_config, + "kv_default": postgres_kv_config, + "sql_default": postgres_sql_config, }, storage_stores={ "metadata": KVStoreReference( - backend="kv_postgres", + backend="kv_default", namespace="registry", ).model_dump(exclude_none=True), "inference": InferenceStoreReference( - backend="sql_postgres", + backend="sql_default", table_name="inference_store", ).model_dump(exclude_none=True), "conversations": SqlStoreReference( - backend="sql_postgres", + backend="sql_default", table_name="openai_conversations", ).model_dump(exclude_none=True), "prompts": KVStoreReference( - backend="kv_postgres", + backend="kv_default", namespace="prompts", ).model_dump(exclude_none=True), },