llama-stack 0.3.5__py3-none-any.whl → 0.4.1__py3-none-any.whl

llama_stack/core/stack.py

@@ -6,36 +6,18 @@
 
 import asyncio
 import importlib.resources
+import inspect
 import os
 import re
 import tempfile
-from typing import Any
+from typing import Any, get_type_hints
 
 import yaml
+from pydantic import BaseModel
 
-from llama_stack.apis.agents import Agents
-from llama_stack.apis.benchmarks import Benchmarks
-from llama_stack.apis.conversations import Conversations
-from llama_stack.apis.datasetio import DatasetIO
-from llama_stack.apis.datasets import Datasets
-from llama_stack.apis.eval import Eval
-from llama_stack.apis.files import Files
-from llama_stack.apis.inference import Inference
-from llama_stack.apis.inspect import Inspect
-from llama_stack.apis.models import Models
-from llama_stack.apis.post_training import PostTraining
-from llama_stack.apis.prompts import Prompts
-from llama_stack.apis.providers import Providers
-from llama_stack.apis.safety import Safety
-from llama_stack.apis.scoring import Scoring
-from llama_stack.apis.scoring_functions import ScoringFunctions
-from llama_stack.apis.shields import Shields
-from llama_stack.apis.synthetic_data_generation import SyntheticDataGeneration
-from llama_stack.apis.telemetry import Telemetry
-from llama_stack.apis.tools import RAGToolRuntime, ToolGroups, ToolRuntime
-from llama_stack.apis.vector_io import VectorIO
+from llama_stack.core.admin import AdminImpl, AdminImplConfig
 from llama_stack.core.conversations.conversations import ConversationServiceConfig, ConversationServiceImpl
-from llama_stack.core.datatypes import Provider, StackRunConfig, VectorStoresConfig
+from llama_stack.core.datatypes import Provider, QualifiedModel, SafetyConfig, StackConfig, VectorStoresConfig
 from llama_stack.core.distribution import get_provider_registry
 from llama_stack.core.inspect import DistributionInspectConfig, DistributionInspectImpl
 from llama_stack.core.prompts.prompts import PromptServiceConfig, PromptServiceImpl
@@ -55,7 +37,30 @@ from llama_stack.core.storage.datatypes import (
 from llama_stack.core.store.registry import create_dist_registry
 from llama_stack.core.utils.dynamic import instantiate_class_type
 from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import Api
+from llama_stack_api import (
+    Agents,
+    Api,
+    Batches,
+    Benchmarks,
+    Conversations,
+    DatasetIO,
+    Datasets,
+    Eval,
+    Files,
+    Inference,
+    Inspect,
+    Models,
+    PostTraining,
+    Prompts,
+    Providers,
+    Safety,
+    Scoring,
+    ScoringFunctions,
+    Shields,
+    ToolGroups,
+    ToolRuntime,
+    VectorIO,
+)
 
 logger = get_logger(name=__name__, category="core")
 
@@ -64,10 +69,9 @@ class LlamaStack(
     Providers,
     Inference,
     Agents,
+    Batches,
     Safety,
-    SyntheticDataGeneration,
     Datasets,
-    Telemetry,
     PostTraining,
     VectorIO,
     Eval,
@@ -80,7 +84,6 @@ class LlamaStack(
     Inspect,
     ToolGroups,
     ToolRuntime,
-    RAGToolRuntime,
     Files,
     Prompts,
     Conversations,
@@ -107,8 +110,95 @@ REGISTRY_REFRESH_INTERVAL_SECONDS = 300
 REGISTRY_REFRESH_TASK = None
 TEST_RECORDING_CONTEXT = None
 
+# ID fields for registered resources that should trigger skipping
+# when they resolve to empty/None (from conditional env vars like :+)
+RESOURCE_ID_FIELDS = [
+    "vector_store_id",
+    "model_id",
+    "shield_id",
+    "dataset_id",
+    "scoring_fn_id",
+    "benchmark_id",
+    "toolgroup_id",
+]
+
+
+def is_request_model(t: Any) -> bool:
+    """Check if a type is a request model (Pydantic BaseModel).
+
+    Args:
+        t: The type to check
+
+    Returns:
+        True if the type is a Pydantic BaseModel subclass, False otherwise
+    """
+
+    return inspect.isclass(t) and issubclass(t, BaseModel)
+
+
+async def invoke_with_optional_request(method: Any) -> Any:
+    """Invoke a method, automatically creating a request instance if needed.
+
+    For APIs that use request models, this will create an empty request object.
+    For backward compatibility, falls back to calling without arguments.
+
+    Uses get_type_hints() to resolve forward references (e.g., "ListBenchmarksRequest" -> actual class).
+
+    Handles methods with:
+    - No parameters: calls without arguments
+    - One or more request model parameters: creates empty instances for each
+    - Mixed parameters: creates request models, uses defaults for others
+    - Required non-request-model parameters without defaults: falls back to calling without arguments
+
+    Args:
+        method: The method to invoke
+
+    Returns:
+        The result of calling the method
+    """
+    try:
+        hints = get_type_hints(method)
+    except Exception:
+        # Forward references can't be resolved, fall back to calling without request
+        return await method()
 
-async def register_resources(run_config: StackRunConfig, impls: dict[Api, Any]):
+    params = list(inspect.signature(method).parameters.values())
+    params = [p for p in params if p.name != "self"]
+
+    if not params:
+        return await method()
+
+    # Build arguments for the method call
+    args: dict[str, Any] = {}
+    can_call = True
+
+    for param in params:
+        param_type = hints.get(param.name)
+
+        # If it's a request model, try to create an empty instance
+        if param_type and is_request_model(param_type):
+            try:
+                args[param.name] = param_type()
+            except Exception:
+                # Request model requires arguments, can't create empty instance
+                can_call = False
+                break
+        # If it has a default value, we can skip it (will use default)
+        elif param.default != inspect.Parameter.empty:
+            continue
+        # Required parameter that's not a request model - can't provide it
+        else:
+            can_call = False
+            break
+
+    if can_call and args:
+        return await method(**args)
+
+    # Fall back to calling without arguments for backward compatibility
+    return await method()
+
+
+async def register_resources(run_config: StackConfig, impls: dict[Api, Any]):
     for rsrc, api, register_method, list_method in RESOURCES:
         objects = getattr(run_config.registered_resources, rsrc)
         if api not in impls:
@@ -129,7 +219,7 @@ async def register_resources(run_config: StackRunConfig, impls: dict[Api, Any]):
             await method(**{k: getattr(obj, k) for k in obj.model_dump().keys()})
 
         method = getattr(impls[api], list_method)
-        response = await method()
+        response = await invoke_with_optional_request(method)
 
         objects_to_process = response.data if hasattr(response, "data") else response
 
@@ -144,35 +234,93 @@ async def validate_vector_stores_config(vector_stores_config: VectorStoresConfig
     if vector_stores_config is None:
         return
 
-    default_embedding_model = vector_stores_config.default_embedding_model
-    if default_embedding_model is None:
-        return
+    # Validate default embedding model
+    if vector_stores_config.default_embedding_model is not None:
+        await _validate_embedding_model(vector_stores_config.default_embedding_model, impls)
+
+    # Validate rewrite query params
+    if vector_stores_config.rewrite_query_params:
+        if vector_stores_config.rewrite_query_params.model:
+            await _validate_rewrite_query_model(vector_stores_config.rewrite_query_params.model, impls)
 
-    provider_id = default_embedding_model.provider_id
-    model_id = default_embedding_model.model_id
-    default_model_id = f"{provider_id}/{model_id}"
+
+async def _validate_embedding_model(embedding_model: QualifiedModel, impls: dict[Api, Any]) -> None:
+    """Validate that an embedding model exists and has required metadata."""
+    provider_id = embedding_model.provider_id
+    model_id = embedding_model.model_id
+    model_identifier = f"{provider_id}/{model_id}"
 
     if Api.models not in impls:
-        raise ValueError(f"Models API is not available but vector_stores config requires model '{default_model_id}'")
+        raise ValueError(f"Models API is not available but vector_stores config requires model '{model_identifier}'")
 
     models_impl = impls[Api.models]
     response = await models_impl.list_models()
     models_list = {m.identifier: m for m in response.data if m.model_type == "embedding"}
 
-    default_model = models_list.get(default_model_id)
-    if default_model is None:
-        raise ValueError(f"Embedding model '{default_model_id}' not found. Available embedding models: {models_list}")
+    model = models_list.get(model_identifier)
+    if model is None:
+        raise ValueError(
+            f"Embedding model '{model_identifier}' not found. Available embedding models: {list(models_list.keys())}"
+        )
 
-    embedding_dimension = default_model.metadata.get("embedding_dimension")
+    embedding_dimension = model.metadata.get("embedding_dimension")
     if embedding_dimension is None:
-        raise ValueError(f"Embedding model '{default_model_id}' is missing 'embedding_dimension' in metadata")
+        raise ValueError(f"Embedding model '{model_identifier}' is missing 'embedding_dimension' in metadata")
 
     try:
         int(embedding_dimension)
     except ValueError as err:
         raise ValueError(f"Embedding dimension '{embedding_dimension}' cannot be converted to an integer") from err
 
-    logger.debug(f"Validated default embedding model: {default_model_id} (dimension: {embedding_dimension})")
+    logger.debug(f"Validated embedding model: {model_identifier} (dimension: {embedding_dimension})")
+
+
+async def _validate_rewrite_query_model(rewrite_query_model: QualifiedModel, impls: dict[Api, Any]) -> None:
+    """Validate that a rewrite query model exists and is accessible."""
+    provider_id = rewrite_query_model.provider_id
+    model_id = rewrite_query_model.model_id
+    model_identifier = f"{provider_id}/{model_id}"
+
+    if Api.models not in impls:
+        raise ValueError(
+            f"Models API is not available but vector_stores config requires rewrite query model '{model_identifier}'"
+        )
+
+    models_impl = impls[Api.models]
+    response = await models_impl.list_models()
+    llm_models_list = {m.identifier: m for m in response.data if m.model_type == "llm"}
+
+    model = llm_models_list.get(model_identifier)
+    if model is None:
+        raise ValueError(
+            f"Rewrite query model '{model_identifier}' not found. Available LLM models: {list(llm_models_list.keys())}"
+        )
+
+    logger.debug(f"Validated rewrite query model: {model_identifier}")
+
+
+async def validate_safety_config(safety_config: SafetyConfig | None, impls: dict[Api, Any]):
+    if safety_config is None or safety_config.default_shield_id is None:
+        return
+
+    if Api.shields not in impls:
+        raise ValueError("Safety configuration requires the shields API to be enabled")
+
+    if Api.safety not in impls:
+        raise ValueError("Safety configuration requires the safety API to be enabled")
+
+    shields_impl = impls[Api.shields]
+    response = await shields_impl.list_shields()
+    shields_by_id = {shield.identifier: shield for shield in response.data}
+
+    default_shield_id = safety_config.default_shield_id
+    # don't validate if there are no shields registered
+    if shields_by_id and default_shield_id not in shields_by_id:
+        available = sorted(shields_by_id)
+        raise ValueError(
+            f"Configured default_shield_id '{default_shield_id}' not found among registered shields."
+            f" Available shields: {available}"
+        )
 
 
 class EnvVarError(Exception):
@@ -210,15 +358,33 @@ def replace_env_vars(config: Any, path: str = "") -> Any:
                             logger.debug(
                                 f"Skipping config env variable expansion for disabled provider: {v.get('provider_id', '')}"
                             )
-                            # Create a copy with resolved provider_id but original config
-                            disabled_provider = v.copy()
-                            disabled_provider["provider_id"] = resolved_provider_id
                             continue
                     except EnvVarError:
                         # If we can't resolve the provider_id, continue with normal processing
                         pass
 
-                # Normal processing for non-disabled providers
+                # Special handling for registered resources: check if ID field resolves to empty/None
+                # from conditional env vars (e.g., ${env.VAR:+value}) and skip the entry if so
+                if isinstance(v, dict):
+                    should_skip = False
+                    for id_field in RESOURCE_ID_FIELDS:
+                        if id_field in v:
+                            try:
+                                resolved_id = replace_env_vars(v[id_field], f"{path}[{i}].{id_field}")
+                                if resolved_id is None or resolved_id == "":
+                                    logger.debug(
+                                        f"Skipping {path}[{i}] with empty {id_field} (conditional env var not set)"
+                                    )
+                                    should_skip = True
+                                    break
+                            except EnvVarError as e:
+                                logger.warning(
+                                    f"Could not resolve {id_field} in {path}[{i}], env var '{e.var_name}': {e}"
+                                )
+                    if should_skip:
+                        continue
+
+                # Normal processing
                 result.append(replace_env_vars(v, f"{path}[{i}]"))
             except EnvVarError as e:
                 raise EnvVarError(e.var_name, e.path) from None
@@ -317,39 +483,44 @@ def cast_image_name_to_string(config_dict: dict[str, Any]) -> dict[str, Any]:
     return config_dict
 
 
-def add_internal_implementations(impls: dict[Api, Any], run_config: StackRunConfig) -> None:
-    """Add internal implementations (inspect and providers) to the implementations dictionary.
-
+def add_internal_implementations(impls: dict[Api, Any], config: StackConfig) -> None:
+    """Add internal implementations (inspect, providers, and admin) to the implementations dictionary.
     Args:
         impls: Dictionary of API implementations
         run_config: Stack run configuration
     """
     inspect_impl = DistributionInspectImpl(
-        DistributionInspectConfig(run_config=run_config),
+        DistributionInspectConfig(config=config),
         deps=impls,
     )
     impls[Api.inspect] = inspect_impl
 
     providers_impl = ProviderImpl(
-        ProviderImplConfig(run_config=run_config),
+        ProviderImplConfig(config=config),
         deps=impls,
     )
     impls[Api.providers] = providers_impl
 
+    admin_impl = AdminImpl(
+        AdminImplConfig(config=config),
+        deps=impls,
+    )
+    impls[Api.admin] = admin_impl
+
     prompts_impl = PromptServiceImpl(
-        PromptServiceConfig(run_config=run_config),
+        PromptServiceConfig(config=config),
         deps=impls,
     )
     impls[Api.prompts] = prompts_impl
 
     conversations_impl = ConversationServiceImpl(
-        ConversationServiceConfig(run_config=run_config),
+        ConversationServiceConfig(config=config),
         deps=impls,
     )
     impls[Api.conversations] = conversations_impl
 
 
-def _initialize_storage(run_config: StackRunConfig):
+def _initialize_storage(run_config: StackConfig):
     kv_backends: dict[str, StorageBackendConfig] = {}
     sql_backends: dict[str, StorageBackendConfig] = {}
     for backend_name, backend_config in run_config.storage.backends.items():
@@ -361,15 +532,15 @@ def _initialize_storage(run_config: StackRunConfig):
         else:
             raise ValueError(f"Unknown storage backend type: {type}")
 
-    from llama_stack.providers.utils.kvstore.kvstore import register_kvstore_backends
-    from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+    from llama_stack.core.storage.kvstore.kvstore import register_kvstore_backends
+    from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 
     register_kvstore_backends(kv_backends)
     register_sqlstore_backends(sql_backends)
 
 
 class Stack:
-    def __init__(self, run_config: StackRunConfig, provider_registry: ProviderRegistry | None = None):
+    def __init__(self, run_config: StackConfig, provider_registry: ProviderRegistry | None = None):
         self.run_config = run_config
         self.provider_registry = provider_registry
         self.impls = None
@@ -412,6 +583,7 @@ class Stack:
         await register_resources(self.run_config, impls)
         await refresh_registry_once(impls)
         await validate_vector_stores_config(self.run_config.vector_stores, impls)
+        await validate_safety_config(self.run_config.safety, impls)
         self.impls = impls
 
     def create_registry_refresh_task(self):
@@ -474,20 +646,20 @@ async def refresh_registry_task(impls: dict[Api, Any]):
         await asyncio.sleep(REGISTRY_REFRESH_INTERVAL_SECONDS)
 
 
-def get_stack_run_config_from_distro(distro: str) -> StackRunConfig:
-    distro_path = importlib.resources.files("llama_stack") / f"distributions/{distro}/run.yaml"
+def get_stack_run_config_from_distro(distro: str) -> StackConfig:
+    distro_path = importlib.resources.files("llama_stack") / f"distributions/{distro}/config.yaml"
 
     with importlib.resources.as_file(distro_path) as path:
         if not path.exists():
             raise ValueError(f"Distribution '{distro}' not found at {distro_path}")
         run_config = yaml.safe_load(path.open())
 
-    return StackRunConfig(**replace_env_vars(run_config))
+    return StackConfig(**replace_env_vars(run_config))
 
 
 def run_config_from_adhoc_config_spec(
     adhoc_config_spec: str, provider_registry: ProviderRegistry | None = None
-) -> StackRunConfig:
+) -> StackConfig:
     """
     Create an adhoc distribution from a list of API providers.
 
@@ -527,7 +699,7 @@ def run_config_from_adhoc_config_spec(
                 config=provider_config,
             )
         ]
-    config = StackRunConfig(
+    config = StackConfig(
         image_name="distro-test",
         apis=list(provider_configs_by_api.keys()),
         providers=provider_configs_by_api,
@@ -540,6 +712,7 @@ def run_config_from_adhoc_config_spec(
                 metadata=KVStoreReference(backend="kv_default", namespace="registry"),
                 inference=InferenceStoreReference(backend="sql_default", table_name="inference_store"),
                 conversations=SqlStoreReference(backend="sql_default", table_name="openai_conversations"),
+                prompts=KVStoreReference(backend="kv_default", namespace="prompts"),
             ),
         ),
     )

llama_stack/core/storage/datatypes.py

@@ -12,6 +12,8 @@ from typing import Annotated, Literal
 
 from pydantic import BaseModel, Field, field_validator
 
+from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR
+
 
 class StorageBackendType(StrEnum):
     KV_REDIS = "kv_redis"
@@ -256,25 +258,46 @@ class ResponsesStoreReference(InferenceStoreReference):
 
 class ServerStoresConfig(BaseModel):
     metadata: KVStoreReference | None = Field(
-        default=None,
+        default=KVStoreReference(
+            backend="kv_default",
+            namespace="registry",
+        ),
         description="Metadata store configuration (uses KV backend)",
     )
     inference: InferenceStoreReference | None = Field(
-        default=None,
+        default=InferenceStoreReference(
+            backend="sql_default",
+            table_name="inference_store",
+        ),
         description="Inference store configuration (uses SQL backend)",
     )
     conversations: SqlStoreReference | None = Field(
-        default=None,
+        default=SqlStoreReference(
+            backend="sql_default",
+            table_name="openai_conversations",
+        ),
         description="Conversations store configuration (uses SQL backend)",
     )
     responses: ResponsesStoreReference | None = Field(
         default=None,
         description="Responses store configuration (uses SQL backend)",
     )
+    prompts: KVStoreReference | None = Field(
+        default=KVStoreReference(backend="kv_default", namespace="prompts"),
+        description="Prompts store configuration (uses KV backend)",
+    )
 
 
 class StorageConfig(BaseModel):
     backends: dict[str, StorageBackendConfig] = Field(
+        default={
+            "kv_default": SqliteKVStoreConfig(
+                db_path=f"${{env.SQLITE_STORE_DIR:={DISTRIBS_BASE_DIR}}}/kvstore.db",
+            ),
+            "sql_default": SqliteSqlStoreConfig(
+                db_path=f"${{env.SQLITE_STORE_DIR:={DISTRIBS_BASE_DIR}}}/sql_store.db",
+            ),
+        },
         description="Named backend configurations (e.g., 'default', 'cache')",
     )
     stores: ServerStoresConfig = Field(

llama_stack/{providers/utils → core/storage}/kvstore/__init__.py

@@ -4,4 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from llama_stack_api.internal.kvstore import KVStore as KVStore
+
 from .kvstore import *  # noqa: F401, F403

llama_stack/{providers/utils → core/storage}/kvstore/kvstore.py

@@ -11,10 +11,21 @@
 
 from __future__ import annotations
 
-from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig, StorageBackendType
+import asyncio
+from collections import defaultdict
+from datetime import datetime
+from typing import cast
 
-from .api import KVStore
-from .config import KVStoreConfig
+from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig
+from llama_stack_api.internal.kvstore import KVStore
+
+from .config import (
+    KVStoreConfig,
+    MongoDBKVStoreConfig,
+    PostgresKVStoreConfig,
+    RedisKVStoreConfig,
+    SqliteKVStoreConfig,
+)
 
 
 def kvstore_dependencies():
@@ -30,7 +41,7 @@ def kvstore_dependencies():
 
 class InmemoryKVStoreImpl(KVStore):
     def __init__(self):
-        self._store = {}
+        self._store: dict[str, str] = {}
 
     async def initialize(self) -> None:
         pass
@@ -38,7 +49,7 @@ class InmemoryKVStoreImpl(KVStore):
     async def get(self, key: str) -> str | None:
         return self._store.get(key)
 
-    async def set(self, key: str, value: str) -> None:
+    async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         self._store[key] = value
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
@@ -53,45 +64,65 @@ class InmemoryKVStoreImpl(KVStore):
 
 
 _KVSTORE_BACKENDS: dict[str, KVStoreConfig] = {}
+_KVSTORE_INSTANCES: dict[tuple[str, str], KVStore] = {}
+_KVSTORE_LOCKS: defaultdict[tuple[str, str], asyncio.Lock] = defaultdict(asyncio.Lock)
 
 
 def register_kvstore_backends(backends: dict[str, StorageBackendConfig]) -> None:
     """Register the set of available KV store backends for reference resolution."""
     global _KVSTORE_BACKENDS
+    global _KVSTORE_INSTANCES
+    global _KVSTORE_LOCKS
 
     _KVSTORE_BACKENDS.clear()
+    _KVSTORE_INSTANCES.clear()
+    _KVSTORE_LOCKS.clear()
     for name, cfg in backends.items():
-        _KVSTORE_BACKENDS[name] = cfg
+        typed_cfg = cast(KVStoreConfig, cfg)
+        _KVSTORE_BACKENDS[name] = typed_cfg
 
 
 async def kvstore_impl(reference: KVStoreReference) -> KVStore:
     backend_name = reference.backend
+    cache_key = (backend_name, reference.namespace)
+
+    existing = _KVSTORE_INSTANCES.get(cache_key)
+    if existing:
+        return existing
 
     backend_config = _KVSTORE_BACKENDS.get(backend_name)
     if backend_config is None:
         raise ValueError(f"Unknown KVStore backend '{backend_name}'. Registered backends: {sorted(_KVSTORE_BACKENDS)}")
 
-    config = backend_config.model_copy()
-    config.namespace = reference.namespace
+    lock = _KVSTORE_LOCKS[cache_key]
+    async with lock:
+        existing = _KVSTORE_INSTANCES.get(cache_key)
+        if existing:
+            return existing
+
+        config = backend_config.model_copy()
+        config.namespace = reference.namespace
 
-    if config.type == StorageBackendType.KV_REDIS.value:
-        from .redis import RedisKVStoreImpl
+        impl: KVStore
+        if isinstance(config, RedisKVStoreConfig):
+            from .redis import RedisKVStoreImpl
 
-        impl = RedisKVStoreImpl(config)
-    elif config.type == StorageBackendType.KV_SQLITE.value:
-        from .sqlite import SqliteKVStoreImpl
+            impl = RedisKVStoreImpl(config)
+        elif isinstance(config, SqliteKVStoreConfig):
+            from .sqlite import SqliteKVStoreImpl
 
-        impl = SqliteKVStoreImpl(config)
-    elif config.type == StorageBackendType.KV_POSTGRES.value:
-        from .postgres import PostgresKVStoreImpl
+            impl = SqliteKVStoreImpl(config)
+        elif isinstance(config, PostgresKVStoreConfig):
+            from .postgres import PostgresKVStoreImpl
 
-        impl = PostgresKVStoreImpl(config)
-    elif config.type == StorageBackendType.KV_MONGODB.value:
-        from .mongodb import MongoDBKVStoreImpl
+            impl = PostgresKVStoreImpl(config)
+        elif isinstance(config, MongoDBKVStoreConfig):
+            from .mongodb import MongoDBKVStoreImpl
 
-        impl = MongoDBKVStoreImpl(config)
-    else:
-        raise ValueError(f"Unknown kvstore type {config.type}")
+            impl = MongoDBKVStoreImpl(config)
+        else:
+            raise ValueError(f"Unknown kvstore type {config.type}")
 
-    await impl.initialize()
-    return impl
+        await impl.initialize()
+        _KVSTORE_INSTANCES[cache_key] = impl
+        return impl