cores/embedding/cache.py

"""Embedding cache — load-once, reuse-many for embedding models.

When a provider needs to generate an embedding, it asks this cache for
the model.  The first call loads the model; subsequent calls return the
cached instance.  This ensures we never load the same model twice.
"""

from __future__ import annotations

import threading
from typing import Any, Callable, Dict


class EmbeddingCache:
    """Thread-safe cache for embedding models.

    Usage:
        cache = EmbeddingCache()
        model = cache.get_or_load("clip-vit-base", lambda: load_clip_model())
        embedding = model.encode(img)
    """

    def __init__(self) -> None:
        self._cache: Dict[str, Any] = {}
        self._lock = threading.RLock()

    def get_or_load(self, key: str, loader: Callable[[], Any]) -> Any:
        """Return the cached model, or load it via `loader` and cache it."""
        with self._lock:
            if key not in self._cache:
                self._cache[key] = loader()
            return self._cache[key]

    def is_loaded(self, key: str) -> bool:
        with self._lock:
            return key in self._cache

    def evict(self, key: str) -> bool:
        with self._lock:
            return self._cache.pop(key, None) is not None

    def clear(self) -> int:
        with self._lock:
            n = len(self._cache)
            self._cache.clear()
            return n

    def keys(self) -> list[str]:
        with self._lock:
            return list(self._cache.keys())