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

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