feat(prompts): attach prompts to storage stores in run configs

This commit is contained in:
r3v5 2025-10-23 14:40:58 +01:00
parent 658fb2c777
commit 943a1e1d51
No known key found for this signature in database
GPG key ID: C7611ACB4FECAD54
27 changed files with 96 additions and 11 deletions

View file

@ -91,6 +91,9 @@ jobs:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
server: server:
port: 8321 port: 8321
EOF EOF

View file

@ -115,13 +115,21 @@ data:
db: ${env.POSTGRES_DB:=llamastack} db: ${env.POSTGRES_DB:=llamastack}
user: ${env.POSTGRES_USER:=llamastack} user: ${env.POSTGRES_USER:=llamastack}
password: ${env.POSTGRES_PASSWORD:=llamastack} password: ${env.POSTGRES_PASSWORD:=llamastack}
references: stores:
metadata: metadata:
backend: kv_default backend: kv_default
namespace: registry namespace: registry
inference: inference:
backend: sql_default backend: sql_default
table_name: inference_store table_name: inference_store
max_write_queue_size: 10000
num_writers: 4
conversations:
backend: sql_default
table_name: openai_conversations
prompts:
backend: kv_default
namespace: prompts
models: models:
- metadata: - metadata:
embedding_dimension: 768 embedding_dimension: 768

View file

@ -115,6 +115,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: - metadata:

View file

@ -63,13 +63,21 @@ storage:
sql_default: sql_default:
type: sql_sqlite type: sql_sqlite
db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ollama}/sqlstore.db db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/ollama}/sqlstore.db
references: stores:
metadata: metadata:
backend: kv_default backend: kv_default
namespace: registry namespace: registry
inference: inference:
backend: sql_default backend: sql_default
table_name: inference_store table_name: inference_store
max_write_queue_size: 10000
num_writers: 4
conversations:
backend: sql_default
table_name: openai_conversations
prompts:
backend: kv_default
namespace: prompts
models: models:
- metadata: {} - metadata: {}
model_id: ${env.INFERENCE_MODEL} model_id: ${env.INFERENCE_MODEL}

View file

@ -113,13 +113,21 @@ data:
db: ${env.POSTGRES_DB:=llamastack} db: ${env.POSTGRES_DB:=llamastack}
user: ${env.POSTGRES_USER:=llamastack} user: ${env.POSTGRES_USER:=llamastack}
password: ${env.POSTGRES_PASSWORD:=llamastack} password: ${env.POSTGRES_PASSWORD:=llamastack}
references: stores:
metadata: metadata:
backend: kv_default backend: kv_default
namespace: registry namespace: registry
inference: inference:
backend: sql_default backend: sql_default
table_name: inference_store table_name: inference_store
max_write_queue_size: 10000
num_writers: 4
conversations:
backend: sql_default
table_name: openai_conversations
prompts:
backend: kv_default
namespace: prompts
models: models:
- metadata: - metadata:
embedding_dimension: 768 embedding_dimension: 768

View file

@ -113,6 +113,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: - metadata:

View file

@ -582,6 +582,7 @@ can be instantiated multiple times (with different configs) if necessary.
_ensure_backend(stores.inference, sql_backends, "storage.stores.inference") _ensure_backend(stores.inference, sql_backends, "storage.stores.inference")
_ensure_backend(stores.conversations, sql_backends, "storage.stores.conversations") _ensure_backend(stores.conversations, sql_backends, "storage.stores.conversations")
_ensure_backend(stores.responses, sql_backends, "storage.stores.responses") _ensure_backend(stores.responses, sql_backends, "storage.stores.responses")
_ensure_backend(stores.prompts, kv_backends, "storage.stores.prompts")
return self return self

View file

@ -11,7 +11,6 @@ from pydantic import BaseModel
from llama_stack.apis.prompts import ListPromptsResponse, Prompt, Prompts from llama_stack.apis.prompts import ListPromptsResponse, Prompt, Prompts
from llama_stack.core.datatypes import StackRunConfig from llama_stack.core.datatypes import StackRunConfig
from llama_stack.core.storage.datatypes import KVStoreReference
from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
@ -40,11 +39,10 @@ class PromptServiceImpl(Prompts):
self.kvstore: KVStore self.kvstore: KVStore
async def initialize(self) -> None: async def initialize(self) -> None:
# Use metadata store backend with prompts-specific namespace # Use prompts store reference from run config
metadata_ref = self.config.run_config.storage.stores.metadata prompts_ref = self.config.run_config.storage.stores.prompts
if not metadata_ref: if not prompts_ref:
raise ValueError("storage.stores.metadata must be configured in run config") raise ValueError("storage.stores.prompts must be configured in run config")
prompts_ref = KVStoreReference(namespace="prompts", backend=metadata_ref.backend)
self.kvstore = await kvstore_impl(prompts_ref) self.kvstore = await kvstore_impl(prompts_ref)
def _get_default_key(self, prompt_id: str) -> str: def _get_default_key(self, prompt_id: str) -> str:

View file

@ -540,6 +540,7 @@ def run_config_from_adhoc_config_spec(
metadata=KVStoreReference(backend="kv_default", namespace="registry"), metadata=KVStoreReference(backend="kv_default", namespace="registry"),
inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"), inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"),
conversations=SqlStoreReference(backend="sql_default", table_name="openai_conversations"), conversations=SqlStoreReference(backend="sql_default", table_name="openai_conversations"),
prompts=KVStoreReference(backend="kv_default", namespace="prompts"),
), ),
), ),
) )

View file

@ -271,6 +271,10 @@ class ServerStoresConfig(BaseModel):
default=None, default=None,
description="Responses store configuration (uses SQL backend)", description="Responses store configuration (uses SQL backend)",
) )
prompts: KVStoreReference | None = Field(
default=None,
description="Prompts store configuration (uses KV backend)",
)
class StorageConfig(BaseModel): class StorageConfig(BaseModel):

View file

@ -247,6 +247,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: [] models: []
shields: shields:

View file

@ -109,6 +109,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: {} - metadata: {}

View file

@ -105,6 +105,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: {} - metadata: {}

View file

@ -122,6 +122,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: {} - metadata: {}

View file

@ -112,6 +112,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: {} - metadata: {}

View file

@ -111,6 +111,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: {} - metadata: {}

View file

@ -100,6 +100,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: [] models: []
shields: [] shields: []

View file

@ -142,6 +142,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: {} - metadata: {}

View file

@ -87,6 +87,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: models:
- metadata: {} - metadata: {}

View file

@ -250,6 +250,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: [] models: []
shields: shields:

View file

@ -247,6 +247,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: [] models: []
shields: shields:

View file

@ -257,6 +257,10 @@ class RunConfigSettings(BaseModel):
backend="sql_default", backend="sql_default",
table_name="openai_conversations", table_name="openai_conversations",
).model_dump(exclude_none=True), ).model_dump(exclude_none=True),
"prompts": KVStoreReference(
backend="kv_default",
namespace="prompts",
).model_dump(exclude_none=True),
} }
storage_config = dict( storage_config = dict(

View file

@ -115,6 +115,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
registered_resources: registered_resources:
models: [] models: []
shields: [] shields: []

View file

@ -25,6 +25,9 @@ storage:
conversations: conversations:
table_name: openai_conversations table_name: openai_conversations
backend: sql_default backend: sql_default
prompts:
namespace: prompts
backend: kv_default
external_apis_dir: ~/.llama/apis.d external_apis_dir: ~/.llama/apis.d
external_providers_dir: ~/.llama/providers.d external_providers_dir: ~/.llama/providers.d
server: server:

View file

@ -44,6 +44,9 @@ def config_with_image_name_int():
responses: responses:
backend: sql_default backend: sql_default
table_name: responses table_name: responses
prompts:
backend: kv_default
namespace: prompts
providers: providers:
inference: inference:
- provider_id: provider1 - provider_id: provider1

View file

@ -48,6 +48,7 @@ def _default_storage() -> StorageConfig:
metadata=KVStoreReference(backend="kv_default", namespace="registry"), metadata=KVStoreReference(backend="kv_default", namespace="registry"),
inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"), inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"),
conversations=SqlStoreReference(backend="sql_default", table_name="conversations"), conversations=SqlStoreReference(backend="sql_default", table_name="conversations"),
prompts=KVStoreReference(backend="kv_default", namespace="prompts"),
), ),
) )

View file

@ -18,7 +18,7 @@ from llama_stack.core.storage.datatypes import (
SqlStoreReference, SqlStoreReference,
StorageConfig, StorageConfig,
) )
from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends from llama_stack.providers.utils.kvstore import register_kvstore_backends
@pytest.fixture @pytest.fixture
@ -38,6 +38,7 @@ async def temp_prompt_store(tmp_path_factory):
metadata=KVStoreReference(backend="kv_test", namespace="registry"), metadata=KVStoreReference(backend="kv_test", namespace="registry"),
inference=InferenceStoreReference(backend="sql_test", table_name="inference"), inference=InferenceStoreReference(backend="sql_test", table_name="inference"),
conversations=SqlStoreReference(backend="sql_test", table_name="conversations"), conversations=SqlStoreReference(backend="sql_test", table_name="conversations"),
prompts=KVStoreReference(backend="kv_test", namespace="prompts"),
), ),
) )
mock_run_config = StackRunConfig( mock_run_config = StackRunConfig(
@ -50,6 +51,6 @@ async def temp_prompt_store(tmp_path_factory):
store = PromptServiceImpl(config, deps={}) store = PromptServiceImpl(config, deps={})
register_kvstore_backends({"kv_test": storage.backends["kv_test"]}) register_kvstore_backends({"kv_test": storage.backends["kv_test"]})
store.kvstore = await kvstore_impl(KVStoreReference(backend="kv_test", namespace="prompts")) await store.initialize()
yield store yield store