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:
@@ -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))
|
||||
|
||||
|
||||
|
||||
@@ -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")
|
||||
Reference in New Issue
Block a user