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