Add persistence resolver for backend resolution

This commit is contained in:
Ashwin Bharambe 2025-10-05 13:09:36 -07:00
parent 2cc2cefe2f
commit bf07e8c44f

View file

@ -0,0 +1,145 @@
# 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 typing import Callable, TypeVar
from llama_stack.core.datatypes import (
InferenceStoreReference,
PersistenceConfig,
StoreReference,
)
from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR, RUNTIME_BASE_DIR
from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
from llama_stack.providers.utils.sqlstore.sqlstore import (
SqlStoreConfig,
SqliteSqlStoreConfig,
)
T = TypeVar("T", KVStoreConfig, SqlStoreConfig)
def resolve_backend(
persistence: PersistenceConfig | None,
store_ref: StoreReference | None,
default_factory: Callable[[], T],
store_name: str,
) -> T:
"""
Resolve a store reference to actual backend config.
Args:
persistence: Global persistence config
store_ref: Store reference (e.g., metadata, inference)
default_factory: Function to create default config if not specified
store_name: Name for error messages
Returns:
Resolved backend config with store-specific overlays applied
"""
if not persistence or not store_ref:
return default_factory()
backend_config = persistence.backends.get(store_ref.backend)
if not backend_config:
raise ValueError(
f"Backend '{store_ref.backend}' referenced by store '{store_name}' "
f"not found in persistence.backends"
)
# Clone backend and apply namespace if KVStore
if isinstance(backend_config, (KVStoreConfig.__args__)): # type: ignore
config_dict = backend_config.model_dump()
if store_ref.namespace:
config_dict["namespace"] = store_ref.namespace
return type(backend_config)(**config_dict) # type: ignore
return backend_config # type: ignore
def resolve_inference_store_config(
persistence: PersistenceConfig | None,
) -> tuple[SqlStoreConfig, int, int]:
"""
Resolve inference store configuration.
Returns:
(sql_config, max_queue_size, num_writers)
"""
if not persistence or not persistence.inference:
# Default SQLite
return (
SqliteSqlStoreConfig(
db_path=(RUNTIME_BASE_DIR / "inference.db").as_posix(),
),
10000,
4,
)
inference_ref = persistence.inference
backend_config = persistence.backends.get(inference_ref.backend)
if not backend_config:
raise ValueError(
f"Backend '{inference_ref.backend}' referenced by inference store "
f"not found in persistence.backends"
)
if not isinstance(backend_config, (SqlStoreConfig.__args__)): # type: ignore
raise ValueError(
f"Inference store requires SqlStore backend, got {type(backend_config).__name__}"
)
return (
backend_config, # type: ignore
inference_ref.max_write_queue_size,
inference_ref.num_writers,
)
def resolve_metadata_store_config(
persistence: PersistenceConfig | None,
image_name: str,
) -> KVStoreConfig:
"""
Resolve metadata store configuration.
Args:
persistence: Global persistence config
image_name: Distribution image name for default path
Returns:
Resolved KVStore config
"""
default_config = SqliteKVStoreConfig(
db_path=(DISTRIBS_BASE_DIR / image_name / "kvstore.db").as_posix()
)
return resolve_backend(
persistence=persistence,
store_ref=persistence.metadata if persistence else None,
default_factory=lambda: default_config,
store_name="metadata",
)
def resolve_conversations_store_config(
persistence: PersistenceConfig | None,
) -> SqlStoreConfig:
"""
Resolve conversations store configuration.
Returns:
Resolved SqlStore config
"""
default_config = SqliteSqlStoreConfig(
db_path=(RUNTIME_BASE_DIR / "conversations.db").as_posix()
)
return resolve_backend(
persistence=persistence,
store_ref=persistence.conversations if persistence else None,
default_factory=lambda: default_config,
store_name="conversations",
)