mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-23 00:27:26 +00:00
145 lines
4.3 KiB
Python
145 lines
4.3 KiB
Python
# 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",
|
|
)
|