"""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 collections.abc import Sequence from sqlalchemy import delete, or_, 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 # The floor for INTERACTIVE, human-facing feeds — REST /api/search, Browse # search, and the list views' semantic `q`. Deliberately looser than the agent # default above: a human scanning a result list gets value from loosely-related # hits an agent would be misled by. One constant, because this number was # written as `0.3` in two files with the reasoning attached to only one of # them — the exact shape where a later tuner moves one and not the other # (#2463 finding 3). INTERACTIVE_SEARCH_THRESHOLD = 0.3 # 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) # How much a superseded record is pushed down the ranking (#278). # # Chosen against a measurement, not by feel. On 2026-08-07 dev-log #2420 sat at # 0.6120 on a query made of its own title phrase, 8th, behind #1759 at 0.6506 — # a deficit of 0.039 to the top and ~0.014 to its nearest neighbours. A penalty # of 0.05 clears that whole band, so demoting a cluster's stale members actually # reorders it rather than shuffling within a tie. # # It is deliberately NOT large. Supersession is a claim about SOME of a record's # content, so a superseded note that strongly answers a question nothing else # answers should still surface — just behind anything comparable that is # current. A penalty big enough to bury it outright would be hiding by another # name, which is the thing the operator ruled out. _SUPERSESSION_PENALTY = 0.05 # Candidates fetched per requested result when a re-rank follows. Three ranks of # headroom is far more than a 0.05 penalty can move anything through in a corpus # whose neighbours sit ~0.01-0.02 apart. _SUPERSESSION_OVERFETCH = 3 async def _apply_supersession_penalty( scored: list[tuple[float, "Note"]], limit: int ) -> list[tuple[float, "Note"]]: """Push superseded records below their equals, then take the top `limit`. The penalty is applied to the RANKING score and the returned score, so downstream gates see the adjusted value — the auto-inject margin band in particular, which exists to stop near-ties dragging in neighbours and would otherwise re-tie exactly what this just separated. It is NOT applied to the relevance threshold: the floor decides whether a record is relevant at all, the penalty decides which relevant record comes first. Applying it to the floor would drop a superseded record out of the results entirely — hiding, which is the one thing this must not do. Stable within a tie: Python's sort preserves the distance order the database already established, so equal-scoring records keep their original sequence rather than reshuffling per call. """ if not scored: return [] from scribe.services.supersession import superseded_ids try: stale = await superseded_ids([int(note.id) for _score, note in scored]) except Exception: # Fail OPEN, and the direction matters: ranking without the penalty is # the behaviour that shipped for months. Returning nothing, or raising, # would turn a supersession-lookup hiccup into a broken search. logger.warning("Supersession lookup failed — ranking unpenalised", exc_info=True) return scored[:limit] if not stale: return scored[:limit] adjusted = [ (score - _SUPERSESSION_PENALTY if int(note.id) in stale else score, note) for score, note in scored ] adjusted.sort(key=lambda pair: pair[0], reverse=True) return adjusted[:limit] def embedding_text(title: str | None, body: str | None) -> str: """The document a record is embedded AS. One definition, deliberately. This was written out three times — the write path (`notes.embed_note`), the recurring-task spawn, and the startup backfill — and identical copies of a formatting rule are three chances to change one and not the others. The spawn path is the dangerous one: a recurring task embedded to a different shape than everything else would be ranked against a corpus it doesn't match, and nothing would report it. It is also a PRECONDITION for changing the shape at all (#2486). Measured, a dev-log's vector separates from five unrelated dev-logs by 0.023 while a snippet's separates by 0.153 — the difference being that a snippet states its purpose twice in a short document, so the purpose dominates. Testing an alternative shape against three copies would mean testing a shape that is not the one in production. Whether `title\\n{body}` is the RIGHT shape is the open question. That it is one shape is what makes the question answerable. """ title = title or "" body = body or "" return f"{title}\n{body}".strip() if body else title 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, note_type: str | Sequence[str] | None = None, task_kind: str | Sequence[str] | None = None, orphan_only: bool = False, scope: str = "own", demote_superseded: bool = True, system_id: int | None = None, ) -> 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. `note_type` narrows to a record kind, or several (e.g. "snippet", or ("snippet", "note")), for callers that want prior art rather than everything embedded. `task_kind` restricts TASKS to the given kinds while leaving non-task notes untouched. That asymmetry is the point: "recorded experience" is issues plus dev-logs, and those differ on `is_task`, so neither `note_type` nor `is_task` alone can express it. With `note_type="note", task_kind="issue"` a caller gets fixed problems and durable notes without the open to-do list. `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``. `demote_superseded` applies the supersession penalty (#278): a record a later note claims to have overtaken ranks below its equals. Callers asking "what is the current answer" want it; the near-duplicate gate does NOT, and passes False — a superseded record is still a duplicate of what you are about to write, and demoting it there would let the same note be recorded twice, the second time invisibly. 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) # Narrow to records tagged to one System (subsystem/area). An # association filter, not a ranking signal — membership in the # candidate set, decided before scoring, like project_id above. if system_id is not None: from scribe.models.system import RecordSystem stmt = stmt.where( select(RecordSystem.id) .where( RecordSystem.note_id == Note.id, RecordSystem.system_id == system_id, ) .exists() ) if is_task is True: stmt = stmt.where(Note.status.isnot(None)) elif is_task is False: stmt = stmt.where(Note.status.is_(None)) # Narrow to one kind of record, or several. Composes with is_task # rather than replacing it — 'snippet' is a non-task note_type, so a # caller asking for prior art gets snippets and not the dev-log that # mentions them. if note_type: kinds = [note_type] if isinstance(note_type, str) else list(note_type) stmt = stmt.where(Note.note_type.in_(kinds)) # Restrict TASKS to certain kinds while leaving notes alone. A note # has no task_kind that means anything, so a plain `.in_()` would # drop every dev-log — which is exactly the record a caller asking # for prior experience wants most. if task_kind: tkinds = [task_kind] if isinstance(task_kind, str) else list(task_kind) stmt = stmt.where( or_(Note.status.is_(None), Note.task_kind.in_(tkinds)) ) if exclude_ids: stmt = stmt.where(NoteEmbedding.note_id.notin_(exclude_ids)) # OVER-FETCH when a re-rank follows, so the demotion can actually # move something. Demoting after a LIMIT k would be theatre: the cut # already happened, so a superseded record pushed down still sits in # the results and the live record that should have replaced it was # never fetched. # # Ordering stays on RAW distance so pgvector's HNSW index still # serves it (migration 0067). Ordering by `distance + penalty` # instead would be exact, and would turn an indexed top-k into a # scan-and-sort of every embedded note. # # The cost of that trade, stated plainly: a live record outside the # over-fetch window cannot be promoted into the results. With a # penalty far smaller than the window's score spread, that case # needs the true answer to be more than _SUPERSESSION_OVERFETCH # ranks down, which no observed query comes close to. fetch = limit * _SUPERSESSION_OVERFETCH if demote_superseded else limit stmt = ( stmt.where(distance <= max_distance) .order_by(distance.asc()) .limit(fetch) ) 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. scored = [(1.0 - float(dist), note) for note, dist in rows] if not demote_superseded: return scored[:limit] return await _apply_supersession_penalty(scored, limit) 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 = embedding_text(title, body) 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))