"""Tests for FastEmbedEmbeddingProvider.""" import builtins import sys import pytest from basic_memory.repository.fastembed_provider import FastEmbedEmbeddingProvider from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError class _StubVector: def __init__(self, values): self._values = values def tolist(self): return self._values class _StubTextEmbedding: init_count = 0 last_init_kwargs: dict = {} last_embed_kwargs: dict = {} def __init__(self, model_name: str, cache_dir: str | None = None, threads: int | None = None): self.model_name = model_name self.embed_calls = 0 _StubTextEmbedding.last_init_kwargs = { "model_name": model_name, "cache_dir": cache_dir, "threads": threads, } _StubTextEmbedding.init_count += 1 def embed(self, texts: list[str], batch_size: int = 64, **kwargs): self.embed_calls += 1 _StubTextEmbedding.last_embed_kwargs = {"batch_size": batch_size, **kwargs} for text in texts: if "wide" in text: yield _StubVector([1.0, 0.0, 0.0, 0.0, 0.5]) else: yield _StubVector([1.0, 0.0, 0.0, 0.0]) @pytest.mark.asyncio async def test_fastembed_provider_lazy_loads_and_reuses_model(monkeypatch): """Provider should instantiate FastEmbed lazily and reuse the loaded model.""" module = type(sys)("fastembed") module.TextEmbedding = _StubTextEmbedding monkeypatch.setitem(sys.modules, "fastembed", module) _StubTextEmbedding.init_count = 0 provider = FastEmbedEmbeddingProvider(model_name="stub-model", dimensions=4) assert provider._model is None first = await provider.embed_query("auth query") second = await provider.embed_documents(["database query"]) assert _StubTextEmbedding.init_count == 1 assert provider._model is not None assert len(first) == 4 assert len(second) == 1 assert len(second[0]) == 4 @pytest.mark.asyncio async def test_fastembed_provider_dimension_mismatch_raises_error(monkeypatch): """Provider should fail fast when model output dimensions differ from configured dimensions.""" module = type(sys)("fastembed") module.TextEmbedding = _StubTextEmbedding monkeypatch.setitem(sys.modules, "fastembed", module) provider = FastEmbedEmbeddingProvider(model_name="stub-model", dimensions=4) with pytest.raises(RuntimeError, match="5-dimensional vectors"): await provider.embed_documents(["wide vector"]) @pytest.mark.asyncio async def test_fastembed_provider_missing_dependency_raises_actionable_error(monkeypatch): """Missing fastembed package should raise SemanticDependenciesMissingError.""" monkeypatch.delitem(sys.modules, "fastembed", raising=False) original_import = builtins.__import__ def _raising_import(name, globals=None, locals=None, fromlist=(), level=0): if name == "fastembed": raise ImportError("fastembed not installed") return original_import(name, globals, locals, fromlist, level) monkeypatch.setattr(builtins, "__import__", _raising_import) provider = FastEmbedEmbeddingProvider(model_name="stub-model") with pytest.raises(SemanticDependenciesMissingError) as error: await provider.embed_query("test") assert "pip install -U basic-memory" in str(error.value) @pytest.mark.asyncio async def test_fastembed_provider_passes_runtime_knobs_to_fastembed(monkeypatch): """Provider should pass optional runtime tuning knobs through to FastEmbed.""" module = type(sys)("fastembed") module.TextEmbedding = _StubTextEmbedding monkeypatch.setitem(sys.modules, "fastembed", module) _StubTextEmbedding.last_init_kwargs = {} _StubTextEmbedding.last_embed_kwargs = {} provider = FastEmbedEmbeddingProvider( model_name="stub-model", dimensions=4, batch_size=8, cache_dir="/tmp/fastembed-cache", threads=3, parallel=2, ) await provider.embed_documents(["runtime knobs"]) assert _StubTextEmbedding.last_init_kwargs == { "model_name": "stub-model", "cache_dir": "/tmp/fastembed-cache", "threads": 3, } assert _StubTextEmbedding.last_embed_kwargs == {"batch_size": 8, "parallel": 2} @pytest.mark.asyncio async def test_fastembed_provider_parallel_one_disables_multiprocessing(monkeypatch): """parallel=1 should not pass FastEmbed multiprocessing kwargs.""" module = type(sys)("fastembed") module.TextEmbedding = _StubTextEmbedding monkeypatch.setitem(sys.modules, "fastembed", module) _StubTextEmbedding.last_embed_kwargs = {} provider = FastEmbedEmbeddingProvider(model_name="stub-model", dimensions=4, parallel=1) await provider.embed_documents(["parallel guardrail"]) assert _StubTextEmbedding.last_embed_kwargs == {"batch_size": 64} @pytest.mark.asyncio async def test_fastembed_provider_parallel_two_passes_multiprocessing(monkeypatch): """parallel>1 should keep passing FastEmbed multiprocessing kwargs.""" module = type(sys)("fastembed") module.TextEmbedding = _StubTextEmbedding monkeypatch.setitem(sys.modules, "fastembed", module) _StubTextEmbedding.last_embed_kwargs = {} provider = FastEmbedEmbeddingProvider(model_name="stub-model", dimensions=4, parallel=2) await provider.embed_documents(["parallel enabled"]) assert _StubTextEmbedding.last_embed_kwargs == {"batch_size": 64, "parallel": 2}