diff --git a/src/fabledassistant/services/embeddings.py b/src/fabledassistant/services/embeddings.py index 69643ef..c3c898a 100644 --- a/src/fabledassistant/services/embeddings.py +++ b/src/fabledassistant/services/embeddings.py @@ -1,18 +1,21 @@ -"""Semantic note search via Ollama embedding model (nomic-embed-text). +"""Semantic note search via fastembed (in-process ONNX, no external service). -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. +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 -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 @@ -20,31 +23,57 @@ 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]. +# 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, model: str | None = None) -> list[float]: - """Get an embedding vector from Ollama for the given text. + """Get an embedding vector for the given text. - Raises httpx.HTTPError on failure — callers should handle this. + The ``model`` parameter is preserved for backward compatibility with the + Ollama era but is now ignored — fastembed uses a single fixed model. + + Raises if the fastembed model fails to load. Callers should catch and + degrade to keyword search. """ - 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] + 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 vectors.""" + """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)) @@ -60,7 +89,7 @@ async def upsert_note_embedding(note_id: int, user_id: int, text: str) -> None: try: embedding = await get_embedding(text) except Exception: - logger.debug("Skipping embedding for note %d — model unavailable", note_id) + logger.debug("Skipping embedding for note %d — embedder unavailable", note_id) return try: @@ -89,14 +118,14 @@ async def semantic_search_notes( 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. + 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 — embedding model unavailable") + logger.debug("Semantic search skipped — embedder unavailable") return [] try: @@ -141,7 +170,7 @@ 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. + so a large backfill doesn't peg CPU. """ try: async with async_session() as session: @@ -176,5 +205,3 @@ async def backfill_note_embeddings() -> None: await asyncio.sleep(0.05) # gentle pacing logger.info("Embedding backfill complete: %d/%d notes embedded", success, len(notes_to_embed)) - - diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py new file mode 100644 index 0000000..3706904 --- /dev/null +++ b/tests/test_embeddings.py @@ -0,0 +1,90 @@ +"""Tests for services.embeddings — fastembed backend. + +We don't actually load the fastembed model in tests (heavy download). +Instead, mock _get_model to return a fake that produces deterministic vectors. +""" +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from fabledassistant.services.embeddings import ( + _cosine_similarity, get_embedding, +) + + +# ─── cosine_similarity (pure logic) ────────────────────────────────────────── + + +def test_cosine_similarity_orthogonal_is_zero(): + assert _cosine_similarity([1.0, 0.0], [0.0, 1.0]) == 0.0 + + +def test_cosine_similarity_identical_is_one(): + assert _cosine_similarity([1.0, 0.0], [1.0, 0.0]) == pytest.approx(1.0) + + +def test_cosine_similarity_opposite_is_negative_one(): + assert _cosine_similarity([1.0, 0.0], [-1.0, 0.0]) == pytest.approx(-1.0) + + +def test_cosine_similarity_zero_length_safe(): + """Zero-magnitude vector must not divide-by-zero.""" + assert _cosine_similarity([0.0, 0.0], [1.0, 0.0]) == 0.0 + assert _cosine_similarity([1.0, 0.0], [0.0, 0.0]) == 0.0 + + +def test_cosine_similarity_mismatched_dim_returns_zero(): + """Cross-migration safety: a 768-dim vs 384-dim comparison must not crash.""" + assert _cosine_similarity([1.0] * 5, [1.0] * 3) == 0.0 + + +def test_cosine_similarity_empty_inputs(): + assert _cosine_similarity([], []) == 0.0 + assert _cosine_similarity([], [1.0]) == 0.0 + + +# ─── get_embedding (fastembed path) ────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_get_embedding_returns_list_of_floats(): + """get_embedding wraps the embedder; the result is a Python list of floats.""" + fake_vec = MagicMock() + fake_vec.tolist.return_value = [0.1, 0.2, 0.3, 0.4] + fake_embedder = MagicMock() + fake_embedder.embed = MagicMock(return_value=iter([fake_vec])) + with patch( + "fabledassistant.services.embeddings._get_model", + AsyncMock(return_value=fake_embedder), + ): + out = await get_embedding("hello world") + assert out == [0.1, 0.2, 0.3, 0.4] + fake_embedder.embed.assert_called_once_with(["hello world"]) + + +@pytest.mark.asyncio +async def test_get_embedding_ignores_legacy_model_param(): + """The `model` kwarg is preserved for backward-compat but should not affect + the fastembed call.""" + fake_vec = MagicMock() + fake_vec.tolist.return_value = [0.0] + fake_embedder = MagicMock() + fake_embedder.embed = MagicMock(return_value=iter([fake_vec])) + with patch( + "fabledassistant.services.embeddings._get_model", + AsyncMock(return_value=fake_embedder), + ): + out = await get_embedding("x", model="ignored-model-name") + assert out == [0.0] + + +@pytest.mark.asyncio +async def test_get_embedding_propagates_model_load_failures(): + """If fastembed can't initialize, the error propagates — callers catch + and degrade to keyword search.""" + with patch( + "fabledassistant.services.embeddings._get_model", + AsyncMock(side_effect=RuntimeError("model load failed")), + ): + with pytest.raises(RuntimeError, match="model load failed"): + await get_embedding("x")