FabledScribe/src/fabledassistant/services/embeddings.py

"""Semantic note search via Ollama embedding model (nomic-embed-text).

Embeddings are stored in the note_embeddings table (one row per note).
All search operations degrade gracefully — if the embedding model is
unavailable the callers fall back to keyword search.
"""

import asyncio
import logging
import math

import httpx
from sqlalchemy import delete, select

from fabledassistant.config import Config
from fabledassistant.models import async_session
from fabledassistant.models.embedding import NoteEmbedding
from fabledassistant.models.note import Note

logger = logging.getLogger(__name__)

# Minimum cosine similarity to include a note in context results.
# nomic-embed-text 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


async def get_embedding(text: str, model: str | None = None) -> list[float]:
    """Get an embedding vector from Ollama for the given text.

    Raises httpx.HTTPError on failure — callers should handle this.
    """
    m = model or Config.EMBEDDING_MODEL
    async with httpx.AsyncClient(timeout=30.0) as client:
        resp = await client.post(
            f"{Config.OLLAMA_URL}/api/embed",
            json={"model": m, "input": text},
        )
        resp.raise_for_status()
        data = resp.json()
        # Ollama /api/embed → {"embeddings": [[float, ...]]}
        return data["embeddings"][0]


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    """Cosine similarity between two vectors. Returns 0 for zero-length vectors."""
    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."""
    try:
        embedding = await get_embedding(text)
    except Exception:
        logger.debug("Skipping embedding for note %d — model 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 embedding model is unavailable or on any error.
    """
    try:
        query_vec = await get_embedding(query)
    except Exception:
        logger.debug("Semantic search skipped — embedding model 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)
            )
            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 []

    scored: list[tuple[float, Note]] = []
    for ne, note in rows:
        try:
            sim = _cosine_similarity(query_vec, ne.embedding)
        except Exception:
            continue
        if sim >= threshold:
            scored.append((sim, note))

    scored.sort(key=lambda x: x[0], reverse=True)
    return 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
    to avoid overwhelming Ollama.
    """
    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))