feat: Use FastEmbed as the default dense embedder

Merged
AW
1mo ago
+1178-179

Proposed Changes

build_embedder("<model name>") now returns a FastEmbed embedder instead of a sentence-transformers embedder. FastEmbed runs ONNX models on the CPU, so the default embedder no longer needs PyTorch.

  • FastEmbedEmbedder in agrag/embedding/fastembed_dense.py implements Embedder the same way SentenceTransformerEmbedder does: lazy model load, blocking calls in a worker thread, a cache key of text, model, and normalize, and the same tracing spans.
  • Always-on embed span. Both dense embedders open agrag.embedding.embed on every embed() call with text_count, cache_hit_count, and model, so fully cached calls stay traceable. The miss-only agrag.embedding.encode span is unchanged, and the new span is listed in the trace-spans reference.
  • fastembed moves into the base dependencies (>=0.8.0). The base install grows by about 38 MB of wheels, mostly onnxruntime. The qdrant extra keeps qdrant-client. The BM25 embedder no longer raises EmbeddingMissingExtraError, because its package is always installed.
  • SentenceTransformerEmbedder stays behind the embed-local extra. Pass an instance to use it. There is no alias for the old behavior.
  • Granite model shim. FastEmbed added the default model in qdrant/fastembed#702, after its 0.8.1 release, so no release has it. The embedder registers the same ONNX export through TextEmbedding.add_custom_model, only while list_supported_models() lacks it. The shim uses CLS pooling (mean pooling gives cosine 0.87). It registers raw vectors, because the embedder applies EmbeddingSettings.normalize itself, as FastEmbed's own entry does.
  • Settings. model, normalize, and batch_size apply to both embedders. cache_folder is FastEmbed's cache directory. FastEmbed ignores device, and the settings docs say so.
  • Docs. Install commands no longer list embed-local where it only supplied the default embedder. The configuration reference, the storage guide, and the README describe the new default.

How did you test it?

  • 20 new unit tests in tests/unit/embedding/test_fastembed_dense.py. They mock FastEmbed at its boundary and cover normalization, batching, the worker thread, lazy and shared loading, the cache, tracing, and when the granite model is registered. The full unit suite passes (2,219 tests). ruff, ruff format, ty, and typos pass.
  • Integration tests with the real model (network): the vectors from FastEmbed and sentence-transformers have cosine above 0.999 for short texts and a 644-token text. In a separate check on fastembed 0.8.0, the minimum cosine was 1.0, the raw vectors differed by at most 2e-6, and a 1,084-token text matched.
  • The first run downloads about 180 MB of ONNX weights, against about 95 MB for the sentence-transformers weights.

Notes for the reviewer

  • Vectors stay compatible: the model is the same, and the cosine against sentence-transformers is 1.0, so existing indexes and benchmark comparisons hold. Cache keys do not change.
  • build_embedder now fails for a model that FastEmbed does not support. Pass a SentenceTransformerEmbedder for those.
  • Rebased onto main past #78 and #79; the branch is mergeable.
  • Follow-ups, not in this PR: move the neo4j and llm extras into the base install, and delete the granite shim once the minimum fastembed version lists the model.

Checklist

  • I have updated the related issue with new insights and changes
  • I added unit tests and updated the docstrings
  • I've used one of the conventional commit types for my PR title: fix:, feat:, build:, chore:, ci:, docs:, style:, refactor:, perf:, test:
  • I documented my code

Summary by cubic

Makes FastEmbed the default dense embedder and hardens LLM retry behavior. build_embedder("<model>") now returns a FastEmbedEmbedder that runs ONNX models on the CPU, so the default embedder no longer needs PyTorch or the embed-local extra. Retry now jitters backoff waits, supports an overall deadline, and retries successful calls whose results are unusable.

  • fastembed moves into the base dependencies, and SentenceTransformerEmbedder stays behind the embed-local extra.
  • Vectors and cache keys are unchanged, so existing indexes and benchmark comparisons hold; build_embedder now fails for models FastEmbed does not support, so pass a SentenceTransformerEmbedder for those.
  • Registers the default granite model with FastEmbed through add_custom_model, only while an installed release lacks it.
  • Every dense embed call, including full cache hits, now opens an agrag.embedding.embed span.
  • The base install grows by about 38 MB of wheels, mostly onnxruntime; the first run downloads about 180 MB of model weights.
  • Retry sleeps now include jitter, RetryConfig.timeout_ms bounds the whole call, sleeps are clamped to the remaining deadline, and no new attempt starts past it; accept_result retries rejected results until the budget runs out.
  • A fallback-selection LLM span test no longer depends on span start order.
  • The README badge now shows total PyPI downloads instead of monthly downloads.