From bf07e8c44f7f18664dc7f4bdf9390c3a627325e9 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Sun, 5 Oct 2025 13:09:36 -0700
Subject: [PATCH] Add persistence resolver for backend resolution

---
 llama_stack/core/persistence_resolver.py | 145 +++++++++++++++++++++++
 1 file changed, 145 insertions(+)
 create mode 100644 llama_stack/core/persistence_resolver.py

diff --git a/llama_stack/core/persistence_resolver.py b/llama_stack/core/persistence_resolver.py
new file mode 100644
index 000000000..aa0c469bc
--- /dev/null
+++ b/llama_stack/core/persistence_resolver.py
@@ -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",
+    )