mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-20 21:18:13 +02:00
91d947ff79
The cached payload is the indexing pipeline's embeddings (markdown is chunked then embedded), so "embedding cache" names the expensive output directly and removes the "index" ambiguity (DB index vs vector index vs indexing phase). Renames the service, settings, eligibility, eviction task, metrics, config flags (INDEX_CACHE_* -> EMBEDDING_CACHE_*), object prefix, and the table (index_cache_embedding_sets -> embedding_cache_sets) with its constraint and indexes. Migration 161 renamed accordingly.
51 lines
1.9 KiB
Python
51 lines
1.9 KiB
Python
"""Recall and remember embedding sets, coordinating the index and blob store."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.indexing_pipeline.cache.persistence import CachedEmbeddingSetRepository
|
|
from app.indexing_pipeline.cache.schemas import EmbeddingKey, EmbeddingSet
|
|
from app.indexing_pipeline.cache.storage import EmbeddingCacheStore
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EmbeddingCacheService:
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
self._index = CachedEmbeddingSetRepository(session)
|
|
self._store = EmbeddingCacheStore()
|
|
|
|
async def recall(self, key: EmbeddingKey) -> EmbeddingSet | None:
|
|
"""Return the cached embedding set, or None on a miss."""
|
|
row = await self._index.get(key)
|
|
if row is None:
|
|
return None
|
|
|
|
try:
|
|
embedding_set = await self._store.load(row.storage_key)
|
|
except Exception:
|
|
# Index points at a blob that is gone; treat as a miss and re-embed.
|
|
logger.warning("Cache blob missing: %s", row.storage_key, exc_info=True)
|
|
return None
|
|
|
|
if int(embedding_set.summary_embedding.shape[0]) != key.embedding_dim:
|
|
# A model swapped its dimension under a reused name; never serve it.
|
|
logger.warning("Cached embedding dimension mismatch: %s", row.storage_key)
|
|
return None
|
|
|
|
await self._index.mark_used(row.id)
|
|
return embedding_set
|
|
|
|
async def remember(self, key: EmbeddingKey, embedding_set: EmbeddingSet) -> None:
|
|
"""Store a freshly embedded set for future reuse."""
|
|
storage_key, size_bytes = await self._store.save(key, embedding_set)
|
|
await self._index.insert(
|
|
key=key,
|
|
storage_backend=self._store.backend_name,
|
|
storage_key=storage_key,
|
|
size_bytes=size_bytes,
|
|
chunk_count=embedding_set.chunk_count,
|
|
)
|