Files
FabledScribe/tests/test_embeddings.py
T
bvandeusen 52d6a8ed53 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>
2026-05-26 21:02:30 -04:00

91 lines
3.2 KiB
Python

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