"""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 from scribe.services.access import notes_visibility_clause 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 # Public alias so callers (and telemetry) can record the effective default # threshold without reaching for the underscored name. DEFAULT_SIMILARITY_THRESHOLD = _SIMILARITY_THRESHOLD _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, scope: str = "own", ) -> 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. `scope` ("own" | "browse" | "read", see access.notes_visibility_clause) decides how far this may see. It exists because this one function serves three different kinds of act: an explicit search, which should reach everything the caller may read; passive auto-injection, which must not pull an unrequested record into their context; and the near-duplicate gate, whose verdict must not depend on other people's notes at all. Defaults to "own" so a caller that forgets is wrong in the safe direction. Ranking and the top-k cut happen in Postgres via pgvector's cosine-distance operator (`<=>`, exposed as ``Vector.cosine_distance``) backed by the HNSW index from migration 0067 — so this is an indexed ``ORDER BY ... LIMIT k`` rather than a full-table scan. Cosine distance is ``1 - cosine_similarity``, so a similarity floor of *threshold* is a distance ceiling of ``1 - threshold`` and similarity is recovered as ``1 - distance``. 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 [] # Distance ceiling equivalent to the similarity floor. Clamp to the valid # cosine-distance range [0, 2] so a threshold of, say, -1 doesn't produce a # nonsensical ceiling. max_distance = min(2.0, max(0.0, 1.0 - threshold)) distance = NoteEmbedding.embedding.cosine_distance(query_vec) try: async with async_session() as session: # Scope on Note, not NoteEmbedding.user_id: the embedding row belongs # to the note's owner, so filtering it would pin every scope to "own" # and leave shared records unreachable by meaning. stmt = ( select(Note, distance.label("distance")) .select_from(NoteEmbedding) .join(Note, NoteEmbedding.note_id == Note.id) .where( notes_visibility_clause(user_id, scope), 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)) stmt = stmt.where(distance <= max_distance).order_by(distance.asc()).limit(limit) rows = list((await session.execute(stmt)).all()) except Exception: logger.warning("Failed to query note embeddings", exc_info=True) return [] # Recover similarity (1 - distance) and preserve the highest-first contract. return [(1.0 - float(dist), note) for note, dist in rows] 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))