mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-22 16:23:08 +00:00
Add persistence resolver for backend resolution
This commit is contained in:
parent
2cc2cefe2f
commit
bf07e8c44f
1 changed files with 145 additions and 0 deletions
145
llama_stack/core/persistence_resolver.py
Normal file
145
llama_stack/core/persistence_resolver.py
Normal 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",
|
||||
)
|
Loading…
Add table
Add a link
Reference in a new issue