Skip to content

Commit 943a1e1

Browse files
committed
feat(prompts): attach prompts to storage stores in run configs
1 parent 658fb2c commit 943a1e1

File tree

27 files changed

+96
-11
lines changed

27 files changed

+96
-11
lines changed

.github/workflows/integration-auth-tests.yml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,9 @@ jobs:
9191
conversations:
9292
table_name: openai_conversations
9393
backend: sql_default
94+
prompts:
95+
namespace: prompts
96+
backend: kv_default
9497
server:
9598
port: 8321
9699
EOF

benchmarking/k8s-benchmark/stack-configmap.yaml

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,13 +115,21 @@ data:
115115
db: ${env.POSTGRES_DB:=llamastack}
116116
user: ${env.POSTGRES_USER:=llamastack}
117117
password: ${env.POSTGRES_PASSWORD:=llamastack}
118-
references:
118+
stores:
119119
metadata:
120120
backend: kv_default
121121
namespace: registry
122122
inference:
123123
backend: sql_default
124124
table_name: inference_store
125+
max_write_queue_size: 10000
126+
num_writers: 4
127+
conversations:
128+
backend: sql_default
129+
table_name: openai_conversations
130+
prompts:
131+
backend: kv_default
132+
namespace: prompts
125133
models:
126134
- metadata:
127135
embedding_dimension: 768

benchmarking/k8s-benchmark/stack_run_config.yaml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,9 @@ storage:
115115
conversations:
116116
table_name: openai_conversations
117117
backend: sql_default
118+
prompts:
119+
namespace: prompts
120+
backend: kv_default
118121
registered_resources:
119122
models:
120123
- metadata:

docs/docs/distributions/configuration.mdx

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,13 +63,21 @@ storage:
6363
sql_default:
6464
type: sql_sqlite
6565
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ollama}/sqlstore.db
66-
references:
66+
stores:
6767
metadata:
6868
backend: kv_default
6969
namespace: registry
7070
inference:
7171
backend: sql_default
7272
table_name: inference_store
73+
max_write_queue_size: 10000
74+
num_writers: 4
75+
conversations:
76+
backend: sql_default
77+
table_name: openai_conversations
78+
prompts:
79+
backend: kv_default
80+
namespace: prompts
7381
models:
7482
- metadata: {}
7583
model_id: ${env.INFERENCE_MODEL}

docs/docs/distributions/k8s/stack-configmap.yaml

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,13 +113,21 @@ data:
113113
db: ${env.POSTGRES_DB:=llamastack}
114114
user: ${env.POSTGRES_USER:=llamastack}
115115
password: ${env.POSTGRES_PASSWORD:=llamastack}
116-
references:
116+
stores:
117117
metadata:
118118
backend: kv_default
119119
namespace: registry
120120
inference:
121121
backend: sql_default
122122
table_name: inference_store
123+
max_write_queue_size: 10000
124+
num_writers: 4
125+
conversations:
126+
backend: sql_default
127+
table_name: openai_conversations
128+
prompts:
129+
backend: kv_default
130+
namespace: prompts
123131
models:
124132
- metadata:
125133
embedding_dimension: 768

docs/docs/distributions/k8s/stack_run_config.yaml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,9 @@ storage:
113113
conversations:
114114
table_name: openai_conversations
115115
backend: sql_default
116+
prompts:
117+
namespace: prompts
118+
backend: kv_default
116119
registered_resources:
117120
models:
118121
- metadata:

llama_stack/core/datatypes.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -582,6 +582,7 @@ def _ensure_backend(reference, expected_set, store_name: str) -> None:
582582
_ensure_backend(stores.inference, sql_backends, "storage.stores.inference")
583583
_ensure_backend(stores.conversations, sql_backends, "storage.stores.conversations")
584584
_ensure_backend(stores.responses, sql_backends, "storage.stores.responses")
585+
_ensure_backend(stores.prompts, kv_backends, "storage.stores.prompts")
585586
return self
586587

587588

llama_stack/core/prompts/prompts.py

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111

1212
from llama_stack.apis.prompts import ListPromptsResponse, Prompt, Prompts
1313
from llama_stack.core.datatypes import StackRunConfig
14-
from llama_stack.core.storage.datatypes import KVStoreReference
1514
from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
1615

1716

@@ -40,11 +39,10 @@ def __init__(self, config: PromptServiceConfig, deps: dict[Any, Any]):
4039
self.kvstore: KVStore
4140

4241
async def initialize(self) -> None:
43-
# Use metadata store backend with prompts-specific namespace
44-
metadata_ref = self.config.run_config.storage.stores.metadata
45-
if not metadata_ref:
46-
raise ValueError("storage.stores.metadata must be configured in run config")
47-
prompts_ref = KVStoreReference(namespace="prompts", backend=metadata_ref.backend)
42+
# Use prompts store reference from run config
43+
prompts_ref = self.config.run_config.storage.stores.prompts
44+
if not prompts_ref:
45+
raise ValueError("storage.stores.prompts must be configured in run config")
4846
self.kvstore = await kvstore_impl(prompts_ref)
4947

5048
def _get_default_key(self, prompt_id: str) -> str:

llama_stack/core/stack.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -540,6 +540,7 @@ def run_config_from_adhoc_config_spec(
540540
metadata=KVStoreReference(backend="kv_default", namespace="registry"),
541541
inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"),
542542
conversations=SqlStoreReference(backend="sql_default", table_name="openai_conversations"),
543+
prompts=KVStoreReference(backend="kv_default", namespace="prompts"),
543544
),
544545
),
545546
)

llama_stack/core/storage/datatypes.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,10 @@ class ServerStoresConfig(BaseModel):
271271
default=None,
272272
description="Responses store configuration (uses SQL backend)",
273273
)
274+
prompts: KVStoreReference | None = Field(
275+
default=None,
276+
description="Prompts store configuration (uses KV backend)",
277+
)
274278

275279

276280
class StorageConfig(BaseModel):

0 commit comments

Comments
 (0)