Skip to content

feat: Use FastEmbed as the default dense embedder - #80

Merged
awinml merged 9 commits into
mainfrom
fastembed_default_embedder
Oct 7, 2026
Merged

awinml merged 9 commits into
mainfrom
fastembed_default_embedder

Conversation

@awinml

@awinml awinml commented Oct 7, 2026 •

Copy link
Copy Markdown
Member

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 feat: Add ibm-granite/granite-embedding-small-english-r2 model 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 build: Rework docs, docs-site and speedup CI #78 and ci: Run agent integration tests in their own job #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.

Written for commit 25d04a6. Summary will update on new commits.

View guided diff Turn on auto-fix

@github-actions

github-actions Bot commented Oct 7, 2026

Copy link
Copy Markdown

API Documentation is up to date.

@cubic-dev-ai cubic-dev-ai Bot left a comment •

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 25 files

Architecture diagram
sequenceDiagram
    participant App as Application
    participant Builder as build_embedder()
    participant Settings as EmbeddingSettings
    participant Cache as EmbeddingCache
    participant FE as FastEmbedEmbedder
    participant ONNX as ONNX Runtime (FastEmbed)
    participant ST as SentenceTransformerEmbedder
    
    Note over App,ST: Dense Embedding Architecture - FastEmbed Default
    
    App->>Builder: build_embedder("model-name")
    alt Model name provided
        Builder->>Settings: Create EmbeddingSettings
        Builder->>FE: FastEmbedEmbedder(settings, tracer)
        Note over FE: NEW: FastEmbed is default dense embedder
        FE-->>App: Embedder instance
    else Embedder instance passed
        Builder-->>App: Return unchanged instance
    end
    
    Note over App,ONNX: Runtime Embedding Flow
    
    App->>FE: embed(texts)
    FE->>Cache: get(text, model, normalize)
    
    alt Cache hit (all texts)
        Cache-->>FE: Cached vectors
        FE-->>App: Return vectors
    else Cache miss (some texts)
        Cache-->>FE: Partial/None
        FE->>FE: _ensure_model_async()
        
        alt Model not loaded yet
            FE->>FE: Acquire model lock
            FE->>ONNX: Load TextEmbedding (lazy, in worker thread)
            Note over ONNX: NEW: Runs ONNX models on CPU<br/>No PyTorch dependency
            ONNX-->>FE: Model instance
            Note over FE: NEW: Lazy load, shared across concurrent calls
        end
        
        FE->>ONNX: embed(texts, batch_size)
        Note over ONNX: Blocking call in worker thread
        ONNX-->>FE: Raw vectors
        FE->>FE: Normalize vectors (if enabled)
        FE->>Cache: set(text, model, normalize, vector)
        FE-->>App: Vectors
    end
    
    Note over App,ONNX: Model Loading Path
    
    FE->>FE: _build_model()
    FE->>ONNX: Import TextEmbedding
    alt Granite default model
        ONNX->>ONNX: Check list_supported_models()
        alt Model not in FastEmbed release
            ONNX->>ONNX: add_custom_model()
            Note over ONNX: NEW: Register granite ONNX export<br/>CLS pooling, raw vectors
        end
    end
    ONNX-->>FE: TextEmbedding instance
    
    Note over App,ST: Alternative Path - Sentence Transformers
    App->>ST: SentenceTransformerEmbedder(settings)
    Note over ST: CHANGED: Requires embed-local extra<br/>Optional alternative to FastEmbed
    ST-->>App: Embedder instance
    
    Note over App,Builder: Key Changes
    Note over Builder: CHANGED: build_embedder creates FastEmbedEmbedder<br/>instead of SentenceTransformerEmbedder
    Note over FE: NEW: FastEmbed is default embedder<br/>- Same interface as before<br/>- Same cache keys<br/>- Same vector compatibility
    Note over ST: CHANGED: Now optional, not default<br/>Use for unsupported models or custom device
Loading

Reply with feedback, questions, or to request a fix.

View guided diff | Turn on auto-fix | Re-trigger cubic

Comment thread docs/docs/api/embedding.md Resolved
Comment thread agrag/embedding/fastembed_dense.py Resolved
Comment thread agrag/embedding/fastembed_dense.py Resolved
Comment thread agrag/embedding/fastembed_dense.py Resolved
@github-actions

github-actions Bot commented Oct 7, 2026

Copy link
Copy Markdown

@codecov

codecov Bot commented Oct 7, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@awinml
awinml force-pushed the fastembed_default_embedder branch from a2909d3 to 31f60db Compare October 7, 2026 18:18
Comment thread agrag/embedding/fastembed_dense.py Resolved
Comment thread agrag/embedding/fastembed_dense.py Resolved
Comment thread agrag/embedding/fastembed_dense.py Resolved
Comment thread agrag/llm/retry.py Resolved

@cubic-dev-ai cubic-dev-ai Bot left a comment •

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 7 files (changes from recent commits).

Reply with feedback, questions, or to request a fix.

View guided diff | Turn on auto-fix | Re-trigger cubic

Comment thread agrag/llm/retry.py Resolved
Comment thread agrag/llm/client_config.py Resolved
@awinml
awinml merged commit fb1eeca into main Oct 7, 2026
33 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant