FabledScribe/src/scribe/services/embeddings.py

"""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))