Spaces:
Sleeping
Sleeping
| """Persistent disk cache for dense embeddings. | |
| LangChain v1 moved ``CacheBackedEmbeddings`` into the optional | |
| ``langchain-classic`` package. To avoid pulling in a legacy package for a | |
| 30-line feature, we ship our own ``Embeddings`` adapter that disk-caches by | |
| SHA256(model_name + text). Cached values are stored as raw float32 bytes — | |
| ~6 KB per 1536-dim embedding. | |
| """ | |
| from __future__ import annotations | |
| import array | |
| import hashlib | |
| import logging | |
| from pathlib import Path | |
| from langchain_core.embeddings import Embeddings | |
| logger = logging.getLogger(__name__) | |
| class _CachedEmbeddings(Embeddings): | |
| """Wrap any ``Embeddings`` with a per-text disk cache.""" | |
| def __init__( | |
| self, | |
| underlying: Embeddings, | |
| *, | |
| namespace: str, | |
| cache_dir: Path, | |
| ) -> None: | |
| self._underlying = underlying | |
| self._namespace = namespace | |
| self._cache_dir = cache_dir | |
| self._cache_dir.mkdir(parents=True, exist_ok=True) | |
| # ---- Embeddings interface ----------------------------------------- | |
| def embed_documents(self, texts: list[str]) -> list[list[float]]: | |
| results: list[list[float] | None] = [None] * len(texts) | |
| misses: list[int] = [] | |
| miss_texts: list[str] = [] | |
| for i, text in enumerate(texts): | |
| cached = self._read(text) | |
| if cached is not None: | |
| results[i] = cached | |
| else: | |
| misses.append(i) | |
| miss_texts.append(text) | |
| if miss_texts: | |
| logger.info( | |
| f"Embedding cache: {len(texts) - len(misses)} hits, " | |
| f"{len(misses)} misses → calling model" | |
| ) | |
| fresh = self._underlying.embed_documents(miss_texts) | |
| for idx, vec in zip(misses, fresh, strict=True): | |
| results[idx] = vec | |
| self._write(texts[idx], vec) | |
| else: | |
| logger.info(f"Embedding cache: all {len(texts)} hits") | |
| return [r for r in results if r is not None] | |
| def embed_query(self, text: str) -> list[float]: | |
| cached = self._read(text) | |
| if cached is not None: | |
| return cached | |
| vec = self._underlying.embed_query(text) | |
| self._write(text, vec) | |
| return vec | |
| # ---- internals ---------------------------------------------------- | |
| def _key(self, text: str) -> str: | |
| h = hashlib.sha256() | |
| h.update(self._namespace.encode("utf-8")) | |
| h.update(b"\x00") | |
| h.update(text.encode("utf-8")) | |
| return h.hexdigest() | |
| def _path(self, text: str) -> Path: | |
| return self._cache_dir / f"{self._key(text)}.bin" | |
| def _read(self, text: str) -> list[float] | None: | |
| path = self._path(text) | |
| if not path.exists(): | |
| return None | |
| try: | |
| arr = array.array("f") | |
| arr.frombytes(path.read_bytes()) | |
| return arr.tolist() | |
| except OSError as exc: | |
| logger.warning(f"Embedding cache read failed for {path.name}: {exc}") | |
| return None | |
| def _write(self, text: str, vector: list[float]) -> None: | |
| path = self._path(text) | |
| try: | |
| arr = array.array("f", vector) | |
| path.write_bytes(arr.tobytes()) | |
| except OSError as exc: | |
| logger.warning(f"Embedding cache write failed for {path.name}: {exc}") | |
| def cached_embeddings( | |
| underlying: Embeddings, | |
| *, | |
| namespace: str, | |
| cache_dir: str | Path = "./.cache/embeddings", | |
| ) -> Embeddings: | |
| """Return ``underlying`` wrapped in a SHA256-keyed disk cache. | |
| ``namespace`` distinguishes outputs from different models — typically pass | |
| the model name. Same text + different namespace → different cache entry. | |
| """ | |
| return _CachedEmbeddings( | |
| underlying, | |
| namespace=namespace, | |
| cache_dir=Path(cache_dir), | |
| ) | |