Files
Paul Hernandez 8acdb49a41 fix(core): reuse a single embedding provider per process (#903)
Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-07 19:14:07 -05:00

674 lines
26 KiB
Python

"""Tests for FastEmbedEmbeddingProvider."""
import builtins
import math
import sys
from dataclasses import dataclass
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,
enable_cpu_mem_arena: bool | 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,
"enable_cpu_mem_arena": enable_cpu_mem_arena,
}
_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")
setattr(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")
setattr(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")
setattr(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,
# onnxruntime CPU mem arena is disabled so transient extra loads free memory (#872)
"enable_cpu_mem_arena": False,
}
assert _StubTextEmbedding.last_embed_kwargs == {"batch_size": 8, "parallel": 2}
@pytest.mark.asyncio
async def test_fastembed_provider_disables_cpu_mem_arena_by_default(monkeypatch):
"""Even with no cache_dir/threads set, the ONNX CPU memory arena must be disabled.
onnxruntime's CPU arena never returns memory to the OS, so a duplicate model
load would leak tens of GB in a long-running process (#872). The provider must
always pass enable_cpu_mem_arena=False to FastEmbed regardless of other knobs.
"""
module = type(sys)("fastembed")
setattr(module, "TextEmbedding", _StubTextEmbedding)
monkeypatch.setitem(sys.modules, "fastembed", module)
_StubTextEmbedding.last_init_kwargs = {}
provider = FastEmbedEmbeddingProvider(model_name="stub-model", dimensions=4)
await provider.embed_documents(["arena default"])
assert _StubTextEmbedding.last_init_kwargs == {
"model_name": "stub-model",
"cache_dir": None,
"threads": None,
"enable_cpu_mem_arena": False,
}
@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")
setattr(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")
setattr(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}
class _UnormalizedVector:
"""Stub vector with norm != 1 (simulates multilingual models like paraphrase-multilingual-*)."""
def __init__(self, values):
self._values = values
def tolist(self):
return self._values
class _UnnormalizedTextEmbedding:
def __init__(self, model_name: str, **_kwargs):
self.model_name = model_name
def embed(self, texts: list[str], **_kwargs):
# Return a vector with norm ~= 2.9 (typical for multilingual MiniLM models)
for _ in texts:
yield _UnormalizedVector([1.5, 2.0, 1.0, 0.5])
@pytest.mark.asyncio
async def test_fastembed_provider_l2_normalizes_output_vectors(monkeypatch):
"""Returned vectors must be unit-normalized regardless of the raw model output.
sqlite_search_repository uses a formula that assumes norm == 1. Models such as
paraphrase-multilingual-MiniLM-L12-v2 return vectors with norm ~2.9, which breaks
cosine similarity scoring. The provider must apply L2 normalization before returning.
"""
module = type(sys)("fastembed")
setattr(module, "TextEmbedding", _UnnormalizedTextEmbedding)
monkeypatch.setitem(sys.modules, "fastembed", module)
provider = FastEmbedEmbeddingProvider(model_name="stub-multilingual", dimensions=4)
result = await provider.embed_documents(["some text"])
assert len(result) == 1
norm = math.sqrt(sum(x * x for x in result[0]))
assert abs(norm - 1.0) < 1e-6, f"Expected unit norm, got {norm}"
@pytest.mark.asyncio
async def test_fastembed_provider_zero_vector_does_not_raise(monkeypatch):
"""A zero vector from the model must be returned as-is without a division error."""
class _ZeroEmbedding:
def __init__(self, model_name: str, **_kwargs):
pass
def embed(self, texts: list[str], **_kwargs):
for _ in texts:
yield _UnormalizedVector([0.0, 0.0, 0.0, 0.0])
module = type(sys)("fastembed")
setattr(module, "TextEmbedding", _ZeroEmbedding)
monkeypatch.setitem(sys.modules, "fastembed", module)
provider = FastEmbedEmbeddingProvider(model_name="stub-zero", dimensions=4)
result = await provider.embed_documents(["zero vector"])
assert result == [[0.0, 0.0, 0.0, 0.0]]
# --- Self-heal of corrupt/partial model cache (#895) ---
#
# A real interrupted FastEmbed download is non-deterministic and offline-unfriendly, so we
# stub TextEmbedding to (a) advertise an HF source + model_file via _list_supported_models so
# the provider can compute the exact models--<org>--<repo> cache subdir and the artifact name,
# and (b) raise a NO_SUCHFILE-style ONNX error on the first construction. This is the justified
# mock case called out in the task. The purge is gated on a filesystem confirmation that the
# snapshot dir exists but the artifact is missing, so each test stages the cache accordingly.
@dataclass
class _StubModelSource:
hf: str
@dataclass
class _StubModelDescription:
model: str
sources: _StubModelSource
model_file: str = "model_optimized.onnx"
class _SelfHealStubTextEmbedding:
"""Raises a NO_SUCHFILE-style ONNX error on the first N constructions, then succeeds."""
fail_first_n = 1
construct_count = 0
HF_SOURCE = "stub-org/stub-model-onnx-q"
RESOLVED_MODEL = "stub-model"
MODEL_FILE = "model_optimized.onnx"
def __init__(
self,
model_name: str,
cache_dir: str | None = None,
threads: int | None = None,
**_kwargs,
):
type(self).construct_count += 1
if type(self).construct_count <= type(self).fail_first_n:
raise RuntimeError(
"[ONNXRuntimeError] : 3 : NO_SUCHFILE : Load model from "
f"{cache_dir}/models--stub-org--stub-model-onnx-q/snapshots/abc123/"
"model_optimized.onnx failed. File doesn't exist"
)
self.model_name = model_name
def embed(self, texts: list[str], batch_size: int = 64, **kwargs):
for _ in texts:
yield _StubVector([1.0, 0.0, 0.0, 0.0])
@classmethod
def _list_supported_models(cls):
# Include decoys so the resolver's skip branches are exercised: a model with a
# different name (name-mismatch skip) and one with an empty HF source (no-source skip).
return [
_StubModelDescription(
model="some-other-model",
sources=_StubModelSource(hf="other-org/other-model"),
model_file=cls.MODEL_FILE,
),
_StubModelDescription(
model=cls.RESOLVED_MODEL,
sources=_StubModelSource(hf=""),
model_file=cls.MODEL_FILE,
),
_StubModelDescription(
model=cls.RESOLVED_MODEL,
sources=_StubModelSource(hf=cls.HF_SOURCE),
model_file=cls.MODEL_FILE,
),
]
def _install_self_heal_stub(monkeypatch):
module = type(sys)("fastembed")
setattr(module, "TextEmbedding", _SelfHealStubTextEmbedding)
monkeypatch.setitem(sys.modules, "fastembed", module)
_SelfHealStubTextEmbedding.construct_count = 0
_SelfHealStubTextEmbedding.fail_first_n = 1
_SelfHealStubTextEmbedding.RESOLVED_MODEL = "stub-model"
@pytest.mark.asyncio
async def test_fastembed_provider_self_heals_corrupt_model_cache(monkeypatch, tmp_path):
"""A NO_SUCHFILE load failure should purge the model cache subdir and retry once."""
_install_self_heal_stub(monkeypatch)
# Simulate the partial-download artifact: the model's HF cache subdir exists on disk
# but is incomplete. The provider must remove exactly this subdir, not the whole cache.
cache_dir = tmp_path / "fastembed_cache"
model_subdir = cache_dir / "models--stub-org--stub-model-onnx-q"
model_subdir.mkdir(parents=True)
(model_subdir / "stale.bin").write_text("partial download")
unrelated_subdir = cache_dir / "models--other--keep-me"
unrelated_subdir.mkdir(parents=True)
(unrelated_subdir / "data.bin").write_text("do not delete")
provider = FastEmbedEmbeddingProvider(
model_name="stub-model", dimensions=4, cache_dir=str(cache_dir)
)
vectors = await provider.embed_documents(["recover after corrupt cache"])
# Construction was attempted exactly twice: the failing load, then the post-purge retry.
assert _SelfHealStubTextEmbedding.construct_count == 2
# The corrupt model subdir was removed; the unrelated model cache was untouched.
assert not model_subdir.exists()
assert unrelated_subdir.exists()
assert (unrelated_subdir / "data.bin").read_text() == "do not delete"
# The retry produced real vectors.
assert len(vectors) == 1
assert len(vectors[0]) == 4
@pytest.mark.asyncio
async def test_fastembed_provider_fails_fast_on_persistent_corrupt_cache(monkeypatch, tmp_path):
"""A second consecutive NO_SUCHFILE failure must fail fast (no infinite retry loop)."""
_install_self_heal_stub(monkeypatch)
# Both constructions fail — the retry does not loop.
_SelfHealStubTextEmbedding.fail_first_n = 2
cache_dir = tmp_path / "fastembed_cache"
model_subdir = cache_dir / "models--stub-org--stub-model-onnx-q"
model_subdir.mkdir(parents=True)
provider = FastEmbedEmbeddingProvider(
model_name="stub-model", dimensions=4, cache_dir=str(cache_dir)
)
with pytest.raises(RuntimeError, match="NO_SUCHFILE"):
await provider.embed_documents(["still broken"])
# Exactly one retry: two total construction attempts, then fail fast.
assert _SelfHealStubTextEmbedding.construct_count == 2
@pytest.mark.asyncio
async def test_fastembed_provider_fails_fast_when_purge_silently_noops(monkeypatch, tmp_path):
"""If rmtree silently fails (e.g. Windows locked files), do not claim success or retry.
shutil.rmtree(ignore_errors=True) can no-op when a file is locked. Treating that as a
successful purge would retry against the same broken cache; instead the load must fail
fast with the original error. We inject a no-op rmtree to simulate the locked-file case.
"""
import basic_memory.repository.fastembed_provider as fastembed_provider
_install_self_heal_stub(monkeypatch)
cache_dir = tmp_path / "fastembed_cache"
model_subdir = cache_dir / "models--stub-org--stub-model-onnx-q"
model_subdir.mkdir(parents=True)
(model_subdir / "stale.bin").write_text("partial download")
# Simulate a deletion that silently fails to remove the directory.
monkeypatch.setattr(
fastembed_provider.shutil, "rmtree", lambda *args, **kwargs: None, raising=True
)
provider = FastEmbedEmbeddingProvider(
model_name="stub-model", dimensions=4, cache_dir=str(cache_dir)
)
with pytest.raises(RuntimeError, match="NO_SUCHFILE"):
await provider.embed_documents(["locked cache"])
# rmtree no-oped, so the subdir survives and no retry was attempted.
assert model_subdir.exists()
assert _SelfHealStubTextEmbedding.construct_count == 1
@pytest.mark.asyncio
async def test_fastembed_provider_does_not_purge_on_unrelated_error(monkeypatch, tmp_path):
"""A non-cache load error must propagate without deleting any cache subdir."""
class _ConfigErrorTextEmbedding:
construct_count = 0
def __init__(self, model_name: str, cache_dir: str | None = None, **_kwargs):
type(self).construct_count += 1
raise ValueError("invalid model configuration")
@classmethod
def _list_supported_models(cls):
return [
_StubModelDescription(
model="stub-model",
sources=_StubModelSource(hf="stub-org/stub-model-onnx-q"),
)
]
module = type(sys)("fastembed")
setattr(module, "TextEmbedding", _ConfigErrorTextEmbedding)
monkeypatch.setitem(sys.modules, "fastembed", module)
cache_dir = tmp_path / "fastembed_cache"
model_subdir = cache_dir / "models--stub-org--stub-model-onnx-q"
model_subdir.mkdir(parents=True)
(model_subdir / "keep.bin").write_text("keep")
provider = FastEmbedEmbeddingProvider(
model_name="stub-model", dimensions=4, cache_dir=str(cache_dir)
)
with pytest.raises(ValueError, match="invalid model configuration"):
await provider.embed_documents(["bad config"])
# No retry and no deletion for errors that are not missing-artifact failures.
assert _ConfigErrorTextEmbedding.construct_count == 1
assert model_subdir.exists()
@pytest.mark.asyncio
async def test_fastembed_provider_cold_load_does_not_purge_or_retry(monkeypatch, tmp_path):
"""A cold load (snapshot dir absent) must NOT be misread as corruption.
This is the CI happy-path regression: on a cold model cache the first load can fail
before the model is downloaded, but with no snapshot dir there is nothing corrupt to
purge. The original error must propagate unchanged with no retry, so a normal
not-yet-downloaded model is never deleted.
"""
_install_self_heal_stub(monkeypatch)
cache_dir = tmp_path / "fastembed_cache"
cache_dir.mkdir(parents=True)
# Intentionally do NOT create the model subdir: this is a normal cold load.
provider = FastEmbedEmbeddingProvider(
model_name="stub-model", dimensions=4, cache_dir=str(cache_dir)
)
with pytest.raises(RuntimeError, match="NO_SUCHFILE"):
await provider.embed_documents(["nothing to purge"])
# Only the initial attempt ran — no snapshot dir means no confirmed corruption, no retry.
assert _SelfHealStubTextEmbedding.construct_count == 1
@pytest.mark.asyncio
async def test_fastembed_provider_does_not_purge_when_artifact_present(monkeypatch, tmp_path):
"""A NO_SUCHFILE-shaped error must NOT purge when the artifact is actually on disk.
The error-text gate alone is not enough: if filesystem inspection finds the model
artifact present in the snapshot, the cache is not corrupt and must be left intact.
Re-raise the original error rather than deleting a healthy cache.
"""
_install_self_heal_stub(monkeypatch)
# Construction keeps failing with the NO_SUCHFILE text regardless of cache state.
_SelfHealStubTextEmbedding.fail_first_n = 99
cache_dir = tmp_path / "fastembed_cache"
snapshot_dir = cache_dir / "models--stub-org--stub-model-onnx-q" / "snapshots" / "rev1"
snapshot_dir.mkdir(parents=True)
artifact = snapshot_dir / "model_optimized.onnx"
artifact.write_text("valid model artifact")
provider = FastEmbedEmbeddingProvider(
model_name="stub-model", dimensions=4, cache_dir=str(cache_dir)
)
with pytest.raises(RuntimeError, match="NO_SUCHFILE"):
await provider.embed_documents(["artifact is fine"])
# No purge and no retry: the artifact is present, so the cache is not corrupt.
assert _SelfHealStubTextEmbedding.construct_count == 1
assert artifact.exists()
assert artifact.read_text() == "valid model artifact"
@pytest.mark.asyncio
async def test_fastembed_provider_self_heals_when_current_revision_corrupt(monkeypatch, tmp_path):
"""A corrupt current revision must be detected even when an older revision is complete.
HuggingFace keeps multiple revisions under one models--<repo> tree. Per-revision
inspection is required: a whole-tree rglob would find the OLD revision's artifact and
wrongly conclude the cache is healthy, leaving the broken current snapshot
self-perpetuating (PR #900 review).
"""
_install_self_heal_stub(monkeypatch)
cache_dir = tmp_path / "fastembed_cache"
snapshots = cache_dir / "models--stub-org--stub-model-onnx-q" / "snapshots"
# Old revision: complete (has the artifact).
good_rev = snapshots / "rev_old"
good_rev.mkdir(parents=True)
(good_rev / "model_optimized.onnx").write_text("complete old artifact")
# Current revision: interrupted download — directory present, artifact missing.
bad_rev = snapshots / "rev_current"
bad_rev.mkdir(parents=True)
(bad_rev / "stale.partial").write_text("partial download")
provider = FastEmbedEmbeddingProvider(
model_name="stub-model", dimensions=4, cache_dir=str(cache_dir)
)
vectors = await provider.embed_documents(["recover from mixed-revision cache"])
# The corrupt-current-revision cache was detected (not masked by the old revision),
# purged, and the retry succeeded.
assert _SelfHealStubTextEmbedding.construct_count == 2
assert not (cache_dir / "models--stub-org--stub-model-onnx-q").exists()
assert len(vectors) == 1
@pytest.mark.asyncio
async def test_fastembed_provider_self_heals_with_case_insensitive_model_name(
monkeypatch, tmp_path
):
"""A lower-cased model name must still resolve the HF cache subdir for the purge.
FastEmbed matches model names case-insensitively, so a config like
model="baai/bge-small-en-v1.5" is valid. The purge resolver must mirror that, otherwise
the corrupt subdir resolves to nothing and self-heal silently does nothing.
"""
_install_self_heal_stub(monkeypatch)
# Advertise the model under its canonical mixed-case name.
_SelfHealStubTextEmbedding.RESOLVED_MODEL = "Stub-Model"
cache_dir = tmp_path / "fastembed_cache"
model_subdir = cache_dir / "models--stub-org--stub-model-onnx-q"
model_subdir.mkdir(parents=True)
(model_subdir / "stale.bin").write_text("partial download")
# Configure the provider with the lower-cased spelling.
provider = FastEmbedEmbeddingProvider(
model_name="stub-model", dimensions=4, cache_dir=str(cache_dir)
)
vectors = await provider.embed_documents(["recover with case-insensitive name"])
assert _SelfHealStubTextEmbedding.construct_count == 2
assert not model_subdir.exists()
assert len(vectors) == 1
@pytest.mark.asyncio
async def test_fastembed_provider_fails_fast_without_cache_dir(monkeypatch):
"""Without a configured cache_dir there is nothing to purge, so fail fast."""
_install_self_heal_stub(monkeypatch)
# cache_dir defaults to None — _model_cache_candidates() returns no candidates.
provider = FastEmbedEmbeddingProvider(model_name="stub-model", dimensions=4)
with pytest.raises(RuntimeError, match="NO_SUCHFILE"):
await provider.embed_documents(["no cache dir"])
assert _SelfHealStubTextEmbedding.construct_count == 1
@pytest.mark.asyncio
async def test_factory_loads_native_model_once_across_repo_constructions(monkeypatch):
"""The native ONNX model must load exactly once per process despite reuse (#872).
Counting native model loads requires a stub TextEmbedding that increments a
counter on construction — there is no way to observe real ONNX loads otherwise.
This is the justified mock: the rest of the path (factory cache, repository
injection) is exercised with real implementations.
The test resolves the provider several times with a *drifting* CPU budget — the
exact condition that previously produced a fresh cache key and a second model
load — then builds multiple search repositories and embeds across all of them,
asserting the stub was constructed only once.
"""
from typing import Any, cast
from basic_memory.config import BasicMemoryConfig, DatabaseBackend, ProjectEntry
from basic_memory.repository import embedding_provider_factory as factory_module
from basic_memory.repository.embedding_provider_factory import (
create_embedding_provider,
reset_embedding_provider_cache,
)
from basic_memory.repository.search_repository import create_search_repository
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
module = type(sys)("fastembed")
setattr(module, "TextEmbedding", _StubTextEmbedding)
monkeypatch.setitem(sys.modules, "fastembed", module)
_StubTextEmbedding.init_count = 0
reset_embedding_provider_cache()
config = BasicMemoryConfig(
env="test",
projects={"test-project": ProjectEntry(path="/tmp/basic-memory-test")},
default_project="test-project",
database_backend=DatabaseBackend.SQLITE,
semantic_search_enabled=True,
semantic_embedding_provider="fastembed",
semantic_embedding_model="stub-model",
semantic_embedding_dimensions=4,
semantic_embedding_threads=None,
semantic_embedding_parallel=None,
)
try:
# First resolution under one CPU budget.
monkeypatch.setattr(factory_module.os, "process_cpu_count", lambda: 8)
monkeypatch.setattr(factory_module.os, "cpu_count", lambda: 8)
provider_first = create_embedding_provider(config)
# CPU budget drifts (cgroup throttling) — used to force a second model load.
monkeypatch.setattr(factory_module.os, "process_cpu_count", lambda: 4)
monkeypatch.setattr(factory_module.os, "cpu_count", lambda: 4)
# Build several repositories the way per-request/per-sync code does. Each
# one is injected with the cached provider rather than deriving its own.
# session_maker is unused during construction and during embed_documents,
# so None is sufficient for this provider-identity assertion.
repos = [
cast(
SQLiteSearchRepository,
create_search_repository(cast(Any, None), project_id=project_id, app_config=config),
)
for project_id in (1, 2, 3)
]
for repo in repos:
assert repo._embedding_provider is provider_first
# Embed several times to trigger lazy model loads; because every repo shares
# provider_first, repeated embeds must still construct the stub model once.
for _ in repos:
await provider_first.embed_documents(["auth token session"])
assert _StubTextEmbedding.init_count == 1
finally:
reset_embedding_provider_cache()