feat(embeddings): swap Ollama for fastembed (in-process ONNX)

Replaces the Ollama HTTP get_embedding with a fastembed.TextEmbedding
singleton loaded lazily on first call. Model: BAAI/bge-small-en-v1.5
(384-dim), cached to /data/fastembed-cache.

Public API unchanged:
  - get_embedding(text, model=None) — `model` now silently ignored
  - upsert_note_embedding
  - semantic_search_notes
  - backfill_note_embeddings

_cosine_similarity gains a defensive length-mismatch check so any
stale 768-dim row that survived the migration is treated as 0.0
similarity rather than crashing zip().

The Ollama client dep stays in pyproject for now (other services still
use it); Phase 7 removes it once chat/journal/curator are gone.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-26 21:02:30 -04:00
parent 12d0ebeb84
commit 52d6a8ed53
2 changed files with 143 additions and 26 deletions
+53 -26
View File
@@ -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 %dmodel unavailable", note_id)
logger.debug("Skipping embedding for note %dembedder 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))
+90
View File
@@ -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")