"""Semantic note search via fastembed (in-process ONNX, no external service). Embeddings are stored as JSONB lists in the note_embeddings table (one row per note). All search operations degrade gracefully — if the embedder fails to initialize the callers fall back to keyword search. Model: BAAI/bge-small-en-v1.5 (384-dim). The first call downloads the model into `FASTEMBED_CACHE_DIR` (defaults to /data/fastembed-cache, a mounted volume so subsequent boots are instant). """ import asyncio import logging import math import os from sqlalchemy import delete, select from scribe.models import async_session from scribe.models.embedding import NoteEmbedding from scribe.models.note import Note logger = logging.getLogger(__name__) # Minimum cosine similarity to include a note in context results. # bge-small-en-v1.5 produces unit-normalized vectors, so range is [-1, 1]. # 0.45 keeps only genuinely relevant notes; lower values like 0.30 let in # loosely-related results that pad the sidebar without adding real value. _SIMILARITY_THRESHOLD = 0.45 _MODEL_NAME = "BAAI/bge-small-en-v1.5" _CACHE_DIR = os.environ.get("FASTEMBED_CACHE_DIR", "/data/fastembed-cache") _model = None # lazy singleton; first call downloads model files _model_lock = asyncio.Lock() async def _get_model(): """Return the singleton fastembed.TextEmbedding instance, loading on first call.""" global _model if _model is None: async with _model_lock: if _model is None: # Defer the import so module import doesn't pull in onnxruntime # for non-embedding code paths (cheaper cold-start for tests etc.) from fastembed import TextEmbedding _model = await asyncio.to_thread( TextEmbedding, model_name=_MODEL_NAME, cache_dir=_CACHE_DIR, ) logger.info("Loaded fastembed model %s (cache: %s)", _MODEL_NAME, _CACHE_DIR) return _model async def get_embedding(text: str) -> list[float]: """Get an embedding vector for the given text. Raises if the fastembed model fails to load. Callers should catch and degrade to keyword search. """ embedder = await _get_model() # embed() is synchronous CPU work; offload so we don't block the event loop. vecs = await asyncio.to_thread(lambda: list(embedder.embed([text]))) return vecs[0].tolist() def _cosine_similarity(a: list[float], b: list[float]) -> float: """Cosine similarity between two vectors. Returns 0 for zero-length or mismatched-length inputs (defensive — mixed-dim vectors can sneak in across the migration boundary).""" if not a or not b or len(a) != len(b): return 0.0 dot = sum(x * y for x, y in zip(a, b)) mag_a = math.sqrt(sum(x * x for x in a)) mag_b = math.sqrt(sum(x * x for x in b)) if mag_a == 0.0 or mag_b == 0.0: return 0.0 return dot / (mag_a * mag_b) async def upsert_note_embedding(note_id: int, user_id: int, text: str) -> None: """Generate and persist an embedding for a note. Safe to fire-and-forget.""" if not text or not text.strip(): return try: embedding = await get_embedding(text) except Exception: logger.debug("Skipping embedding for note %d — embedder unavailable", note_id) return try: async with async_session() as session: await session.execute( delete(NoteEmbedding).where(NoteEmbedding.note_id == note_id) ) session.add(NoteEmbedding(note_id=note_id, user_id=user_id, embedding=embedding)) await session.commit() logger.debug("Upserted embedding for note %d", note_id) except Exception: logger.warning("Failed to persist embedding for note %d", note_id, exc_info=True) async def semantic_search_notes( user_id: int, query: str, exclude_ids: set[int] | None = None, limit: int = 8, threshold: float = _SIMILARITY_THRESHOLD, project_id: int | None = None, is_task: bool | None = None, orphan_only: bool = False, ) -> list[tuple[float, Note]]: """Return up to *limit* (score, note) pairs most relevant to *query*. Scores are cosine similarities in [-1, 1]; only notes at or above *threshold* are returned, sorted highest-first. Returns an empty list if the embedder is unavailable or on any error. """ if not query or not query.strip(): return [] try: query_vec = await get_embedding(query) except Exception: logger.debug("Semantic search skipped — embedder unavailable") return [] try: async with async_session() as session: stmt = ( select(NoteEmbedding, Note) .join(Note, NoteEmbedding.note_id == Note.id) .where(NoteEmbedding.user_id == user_id, Note.deleted_at.is_(None)) ) if orphan_only: stmt = stmt.where(Note.project_id.is_(None)) elif project_id is not None: stmt = stmt.where(Note.project_id == project_id) if is_task is True: stmt = stmt.where(Note.status.isnot(None)) elif is_task is False: stmt = stmt.where(Note.status.is_(None)) if exclude_ids: stmt = stmt.where(NoteEmbedding.note_id.notin_(exclude_ids)) rows = list((await session.execute(stmt)).all()) except Exception: logger.warning("Failed to query note embeddings", exc_info=True) return [] if not rows: return [] def _score() -> list[tuple[float, Note]]: out: list[tuple[float, Note]] = [] for ne, note in rows: try: sim = _cosine_similarity(query_vec, ne.embedding) except Exception: continue if sim >= threshold: out.append((sim, note)) out.sort(key=lambda x: x[0], reverse=True) return out[:limit] # Offload the O(rows) cosine scoring off the event loop so a large corpus # doesn't stall other requests while ranking. Results are unchanged; the # real scaling fix (ORDER BY / LIMIT in pgvector) is a separate effort. return await asyncio.to_thread(_score) async def backfill_note_embeddings() -> None: """Generate embeddings for all notes that don't have one yet. Runs as a background task at startup. Adds a small sleep between notes so a large backfill doesn't peg CPU. """ try: async with async_session() as session: existing = { row[0] for row in ( await session.execute(select(NoteEmbedding.note_id)) ).fetchall() } result = await session.execute( select(Note.id, Note.user_id, Note.title, Note.body) ) notes_to_embed = [ row for row in result.fetchall() if row[0] not in existing ] except Exception: logger.warning("Embedding backfill: failed to query notes", exc_info=True) return if not notes_to_embed: logger.info("Embedding backfill: all notes already have embeddings") return logger.info("Embedding backfill: generating embeddings for %d notes", len(notes_to_embed)) success = 0 for note_id, user_id, title, body in notes_to_embed: text = f"{title}\n{body}".strip() if body else (title or "") if not text: continue await upsert_note_embedding(note_id, user_id, text) success += 1 await asyncio.sleep(0.05) # gentle pacing logger.info("Embedding backfill complete: %d/%d notes embedded", success, len(notes_to_embed))