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).
RSS item embeddings are stored in rss_item_embeddings (one row per item).
All search operations degrade gracefully — if the embedding model is
unavailable the callers fall back to keyword search.
"""

import asyncio
import logging
import math
from datetime import datetime, timedelta, timezone

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
from fabledassistant.models.rss_feed import RssItem
from fabledassistant.models.rss_item_embedding import RssItemEmbedding

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
_RSS_SIMILARITY_THRESHOLD = 0.55
_RSS_SEARCH_LIMIT = 3
_RSS_SEARCH_DAYS = 30
_RSS_SNIPPET_CHARS = 500


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


# ── RSS item embeddings ───────────────────────────────────────────────────────

async def upsert_rss_item_embedding(item_id: int, user_id: int, title: str, content: str) -> None:
    """Generate and persist an embedding for an RSS item. Safe to fire-and-forget."""
    text = f"{title}\n{content}".strip()
    if not text:
        return
    try:
        embedding = await get_embedding(text)
    except Exception:
        logger.debug("Skipping embedding for RSS item %d — model unavailable", item_id)
        return

    try:
        async with async_session() as session:
            await session.execute(
                delete(RssItemEmbedding).where(RssItemEmbedding.rss_item_id == item_id)
            )
            session.add(RssItemEmbedding(rss_item_id=item_id, user_id=user_id, embedding=embedding))
            await session.commit()
        logger.debug("Upserted embedding for RSS item %d", item_id)
    except Exception:
        logger.warning("Failed to persist embedding for RSS item %d", item_id, exc_info=True)


async def semantic_search_rss_items(
    user_id: int,
    query_vector: list[float],
    limit: int = _RSS_SEARCH_LIMIT,
    days: int = _RSS_SEARCH_DAYS,
) -> list[tuple[float, RssItem]]:
    """Return up to *limit* (score, RssItem) pairs most relevant to *query_vector*.

    Only considers items fetched within the last *days* days.
    Returns an empty list on any error.
    """
    since = datetime.now(timezone.utc) - timedelta(days=days)
    try:
        async with async_session() as session:
            stmt = (
                select(RssItemEmbedding, RssItem)
                .join(RssItem, RssItemEmbedding.rss_item_id == RssItem.id)
                .where(
                    RssItemEmbedding.user_id == user_id,
                    RssItem.fetched_at >= since,
                )
            )
            rows = list((await session.execute(stmt)).all())
    except Exception:
        logger.warning("Failed to query RSS item embeddings", exc_info=True)
        return []

    if not rows:
        return []

    scored: list[tuple[float, RssItem]] = []
    for rie, item in rows:
        try:
            sim = _cosine_similarity(query_vector, rie.embedding)
        except Exception:
            continue
        if sim >= _RSS_SIMILARITY_THRESHOLD:
            scored.append((sim, item))

    scored.sort(key=lambda x: x[0], reverse=True)
    return scored[:limit]


async def backfill_rss_item_embeddings() -> None:
    """Generate embeddings for all RSS items that don't have one yet.

    Runs as a background task at startup. Adds a small sleep between items
    to avoid overwhelming Ollama.
    """
    try:
        async with async_session() as session:
            existing = {
                row[0]
                for row in (
                    await session.execute(select(RssItemEmbedding.rss_item_id))
                ).fetchall()
            }
            result = await session.execute(
                select(RssItem.id, RssItem.feed_id, RssItem.title, RssItem.content)
            )
            items_to_embed = [row for row in result.fetchall() if row[0] not in existing]
    except Exception:
        logger.warning("RSS embedding backfill: failed to query items", exc_info=True)
        return

    if not items_to_embed:
        logger.info("RSS embedding backfill: all items already have embeddings")
        return

    # Resolve user_id per feed_id
    try:
        from fabledassistant.models.rss_feed import RssFeed
        async with async_session() as session:
            result = await session.execute(select(RssFeed.id, RssFeed.user_id))
            feed_user_map = {fid: uid for fid, uid in result.fetchall()}
    except Exception:
        logger.warning("RSS embedding backfill: failed to load feed user map", exc_info=True)
        return

    logger.info("RSS embedding backfill: generating embeddings for %d items", len(items_to_embed))
    success = 0
    for item_id, feed_id, title, content in items_to_embed:
        user_id = feed_user_map.get(feed_id)
        if user_id is None:
            continue
        await upsert_rss_item_embedding(item_id, user_id, title or "", content or "")
        success += 1
        await asyncio.sleep(0.05)

    logger.info("RSS embedding backfill complete: %d/%d items embedded", success, len(items_to_embed))


async def backfill_rss_article_content() -> None:
    """Fetch full article text for RSS items that only have short feed-provided content.

    An item is considered unenriched if its content is shorter than 1000 chars —
    typical of feed summaries/teasers rather than full articles.
    Runs at startup after the embedding backfill.
    """
    from fabledassistant.services.rss import _fetch_full_article, CONTENT_MAX_CHARS
    from fabledassistant.models.rss_feed import RssFeed

    SHORT_THRESHOLD = 1000

    try:
        async with async_session() as session:
            feed_result = await session.execute(select(RssFeed.id, RssFeed.user_id))
            feed_user_map = {fid: uid for fid, uid in feed_result.fetchall()}

            item_result = await session.execute(
                select(RssItem.id, RssItem.feed_id, RssItem.url, RssItem.title, RssItem.content)
                .where(RssItem.url != "")
            )
            candidates = [
                row for row in item_result.fetchall()
                if len(row[4] or "") < SHORT_THRESHOLD
            ]
    except Exception:
        logger.warning("Article content backfill: failed to query items", exc_info=True)
        return

    if not candidates:
        logger.info("Article content backfill: no unenriched items found")
        return

    logger.info("Article content backfill: enriching %d items", len(candidates))
    enriched = 0
    for item_id, feed_id, url, title, _ in candidates:
        user_id = feed_user_map.get(feed_id)
        if user_id is None:
            continue
        full_text = await _fetch_full_article(url)
        if full_text and len(full_text) > SHORT_THRESHOLD:
            try:
                async with async_session() as session:
                    item = await session.get(RssItem, item_id)
                    if item:
                        item.content = full_text[:CONTENT_MAX_CHARS]
                        await session.commit()
                await upsert_rss_item_embedding(item_id, user_id, title or "", full_text[:CONTENT_MAX_CHARS])
                enriched += 1
            except Exception:
                logger.debug("Failed to store enriched content for item %d", item_id, exc_info=True)
        await asyncio.sleep(0.5)

    logger.info("Article content backfill complete: %d/%d items enriched", enriched, len(candidates))