diff --git a/openspace/cloud/search.py b/openspace/cloud/search.py index bd25436..c92177a 100644 --- a/openspace/cloud/search.py +++ b/openspace/cloud/search.py @@ -21,6 +21,21 @@ logger = logging.getLogger("openspace.cloud") CLOUD_EMBEDDING_SEARCH_MAX_LIMIT = 300 +# Shared SkillRanker singleton. Its pickle cache survives process restarts; +# the singleton itself is per-process and avoids reloading the pickle on every +# search_skills invocation. +_shared_ranker = None + + +def _get_shared_ranker(): + """Lazy-init shared ``SkillRanker`` (with persistent embedding cache).""" + global _shared_ranker + if _shared_ranker is None: + from openspace.skill_engine.skill_ranker import SkillRanker + _shared_ranker = SkillRanker(enable_cache=True) + return _shared_ranker + + def _check_safety(text: str) -> list[str]: """Lazy wrapper — avoids importing skill_engine at module load time.""" from openspace.skill_engine.skill_utils import check_skill_safety @@ -129,9 +144,9 @@ class SkillSearchEngine: limit: int, ) -> List[Dict[str, Any]]: """BM25 rough-rank to keep top candidates for embedding stage.""" - from openspace.skill_engine.skill_ranker import SkillRanker, SkillCandidate + from openspace.skill_engine.skill_ranker import SkillCandidate - ranker = SkillRanker(enable_cache=True) + ranker = _get_shared_ranker() bm25_candidates = [ SkillCandidate( skill_id=c.get("skill_id", ""), @@ -424,15 +439,33 @@ async def hybrid_search_skills( try: query_embedding = await asyncio.to_thread(generate_embedding, normalized_query) if query_embedding: + # Route candidate embedding generation through SkillRanker's persistent + # cache (pickle on disk) instead of re-computing on every query. + # Cloud candidates that already carry ``_embedding`` (from server) are + # left untouched. + from openspace.skill_engine.skill_ranker import SkillCandidate + ranker = _get_shared_ranker() for candidate in candidates: - if not candidate.get("_embedding") and candidate.get("_embedding_text"): - candidate_embedding = await asyncio.to_thread( - generate_embedding, candidate["_embedding_text"], - ) - if candidate_embedding: - candidate["_embedding"] = candidate_embedding - except Exception: - pass + if candidate.get("_embedding") or not candidate.get("_embedding_text"): + continue + sid = candidate.get("skill_id") or "" + if not sid: + # Without a stable skill_id the cache would collide; skip. + continue + cand = SkillCandidate( + skill_id=sid, + name=candidate.get("name", ""), + description=candidate.get("description", ""), + body="", + embedding_text=candidate["_embedding_text"], + ) + candidate_embedding = await asyncio.to_thread( + ranker.get_or_compute_embedding, cand, + ) + if candidate_embedding: + candidate["_embedding"] = candidate_embedding + except Exception as e: + logger.warning(f"hybrid_search_skills: embedding unavailable: {e}") engine = SkillSearchEngine() return engine.search(normalized_query, candidates, query_embedding=query_embedding, limit=limit) diff --git a/openspace/skill_engine/skill_ranker.py b/openspace/skill_engine/skill_ranker.py index 503eda5..6ad21b1 100644 --- a/openspace/skill_engine/skill_ranker.py +++ b/openspace/skill_engine/skill_ranker.py @@ -18,11 +18,15 @@ Reused by: from __future__ import annotations +import base64 +import hashlib import json import math import os import pickle import re +import tempfile +import threading from dataclasses import dataclass, field from datetime import datetime from pathlib import Path @@ -43,8 +47,11 @@ PREFILTER_THRESHOLD = 10 # How many candidates to keep after BM25 rough-rank (before embedding re-rank) BM25_CANDIDATES_MULTIPLIER = 3 # top_k * 3 -# Cache version — increment when format changes -_CACHE_VERSION = 1 +# Cache version — increment when format changes. +# +# v2 intentionally does not migrate v1 skill_id-only entries: old entries did +# not store the source text hash, so they cannot be proven fresh. +_CACHE_VERSION = 2 @dataclass @@ -81,9 +88,10 @@ class SkillRanker: cache_dir: Optional[Path] = None, enable_cache: bool = True, ) -> None: - # Embedding cache: skill_id → List[float] + # Embedding cache: encoded skill_id + content hash → List[float] self._embedding_cache: Dict[str, List[float]] = {} self._enable_cache = enable_cache + self._cache_lock = threading.RLock() if cache_dir is None: try: @@ -146,6 +154,36 @@ class SkillRanker: """Embedding-only ranking.""" return self._embedding_rank(query, candidates, top_k) + @staticmethod + def _skill_key_prefix(skill_id: str) -> str: + """Return a collision-safe prefix for all cache entries of a skill.""" + encoded = base64.urlsafe_b64encode(skill_id.encode("utf-8")).decode("ascii") + return f"{encoded}:" + + @classmethod + def _content_key(cls, skill_id: str, embedding_text: str) -> str: + """Compose a content-addressed cache key. + + The cache stores embeddings under ``"{encoded_skill_id}:{sha256(text)[:16]}"`` + so that any change to the text used for embedding (name, description, + or body) automatically invalidates the cached entry without relying + on an external trigger. + """ + digest = hashlib.sha256(embedding_text.encode("utf-8")).hexdigest()[:16] + return f"{cls._skill_key_prefix(skill_id)}{digest}" + + def _drop_stale_entries_locked(self, skill_id: str, keep_key: str) -> None: + """Drop older content-addressed entries for a skill. + + Caller must hold ``self._cache_lock``. + """ + prefix = self._skill_key_prefix(skill_id) + for stale in [ + k for k in self._embedding_cache + if k != keep_key and k.startswith(prefix) + ]: + self._embedding_cache.pop(stale, None) + def get_or_compute_embedding( self, candidate: SkillCandidate, ) -> Optional[List[float]]: @@ -157,30 +195,65 @@ class SkillRanker: if candidate.embedding: return candidate.embedding - # Check cache - cached = self._embedding_cache.get(candidate.skill_id) - if cached: - candidate.embedding = cached - return cached + text = self._build_embedding_text(candidate) + if not candidate.skill_id: + emb = self._generate_embedding(text) + if emb: + candidate.embedding = emb + return emb + + content_key = self._content_key(candidate.skill_id, text) + dropped_legacy = False + with self._cache_lock: + # Check content-addressed cache. If a legacy in-memory entry exists, + # discard it instead of migrating unverifiable stale data. + cached = self._embedding_cache.get(content_key) + if cached: + candidate.embedding = cached + return cached + dropped_legacy = self._embedding_cache.pop(candidate.skill_id, None) is not None # Compute - text = self._build_embedding_text(candidate) emb = self._generate_embedding(text) - if emb: - candidate.embedding = emb - self._embedding_cache[candidate.skill_id] = emb - self._save_cache() + with self._cache_lock: + cached = self._embedding_cache.get(content_key) + if cached: + candidate.embedding = cached + return cached + if emb: + candidate.embedding = emb + self._embedding_cache[content_key] = emb + # Bound cache growth: previous versions of this skill are now + # obsolete, drop them in the same write. + self._drop_stale_entries_locked(candidate.skill_id, content_key) + self._save_cache() + elif dropped_legacy: + self._save_cache() return emb def invalidate_cache(self, skill_id: str) -> None: - """Remove a skill's cached embedding (e.g. after evolution).""" - self._embedding_cache.pop(skill_id, None) - self._save_cache() + """Remove all cached embeddings for a skill (e.g. after evolution). + + Removes every cache entry whose key matches either the exact + ``skill_id`` (legacy format) or the ``"{skill_id}:*"`` content-addressed + prefix, covering any historical content version that might linger. + """ + prefix = self._skill_key_prefix(skill_id) + with self._cache_lock: + keys_to_drop = [ + k for k in self._embedding_cache + if k == skill_id or k.startswith(prefix) + ] + for k in keys_to_drop: + self._embedding_cache.pop(k, None) + if keys_to_drop: + self._save_cache() def clear_cache(self) -> None: """Clear all cached embeddings.""" - self._embedding_cache.clear() - self._save_cache() + with self._cache_lock: + self._embedding_cache.clear() + self._save_cache() @staticmethod def _tokenize(text: str) -> List[str]: @@ -273,21 +346,41 @@ class SkillRanker: if not query_emb: return [] - # Ensure all candidates have embeddings + # Ensure all candidates have embeddings (content-addressed cache + # with backward-compat fallback for legacy skill_id-only keys). + cache_dirty = False for c in candidates: - if not c.embedding: - cached = self._embedding_cache.get(c.skill_id) + if c.embedding: + continue + text = self._build_embedding_text(c) + if not c.skill_id: + emb = self._generate_embedding(text, api_key=api_key) + if emb: + c.embedding = emb + continue + content_key = self._content_key(c.skill_id, text) + with self._cache_lock: + cached = self._embedding_cache.get(content_key) if cached: c.embedding = cached - else: - text = self._build_embedding_text(c) - emb = self._generate_embedding(text, api_key=api_key) - if emb: - c.embedding = emb - self._embedding_cache[c.skill_id] = emb + continue + # Do not migrate legacy skill_id-only entries; they may be + # stale because the old format did not store a text hash. + if self._embedding_cache.pop(c.skill_id, None) is not None: + cache_dirty = True + emb = self._generate_embedding(text, api_key=api_key) + if emb: + c.embedding = emb + with self._cache_lock: + self._embedding_cache[content_key] = emb + # Bound cache growth: previous versions of this skill are + # obsolete, drop them. + self._drop_stale_entries_locked(c.skill_id, content_key) + cache_dirty = True - # Save newly computed embeddings - self._save_cache() + # Save newly computed / migrated embeddings + if cache_dirty: + self._save_cache() # Score for c in candidates: @@ -361,8 +454,14 @@ class SkillRanker: try: with open(path, "rb") as f: data = pickle.load(f) - if isinstance(data, dict) and data.get("version") == _CACHE_VERSION: - self._embedding_cache = data.get("embeddings", {}) + if ( + isinstance(data, dict) + and data.get("version") == _CACHE_VERSION + and data.get("model") == SKILL_EMBEDDING_MODEL + and isinstance(data.get("embeddings"), dict) + ): + with self._cache_lock: + self._embedding_cache = data.get("embeddings", {}) logger.debug(f"Loaded {len(self._embedding_cache)} skill embeddings from cache") except Exception as e: logger.warning(f"Failed to load skill embedding cache: {e}") @@ -370,18 +469,32 @@ class SkillRanker: def _save_cache(self) -> None: """Persist embedding cache to disk.""" - if not self._enable_cache or not self._embedding_cache: + if not self._enable_cache: return try: self._cache_dir.mkdir(parents=True, exist_ok=True) - data = { - "version": _CACHE_VERSION, - "model": SKILL_EMBEDDING_MODEL, - "last_updated": datetime.now().isoformat(), - "embeddings": self._embedding_cache, - } - with open(self._cache_file(), "wb") as f: - pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL) + with self._cache_lock: + data = { + "version": _CACHE_VERSION, + "model": SKILL_EMBEDDING_MODEL, + "last_updated": datetime.now().isoformat(), + "embeddings": dict(self._embedding_cache), + } + tmp_name = "" + try: + with tempfile.NamedTemporaryFile( + "wb", + dir=self._cache_dir, + prefix=".skill_embeddings_", + suffix=".tmp", + delete=False, + ) as f: + tmp_name = f.name + pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL) + os.replace(tmp_name, self._cache_file()) + finally: + if tmp_name and os.path.exists(tmp_name): + os.unlink(tmp_name) except Exception as e: logger.warning(f"Failed to save skill embedding cache: {e}")