feat(rag): remote TEI embedding and reranker providers (FEAT-024)#93
Conversation
VEKTRA_EMBEDDING_PROVIDER=tei was documented but no provider existed: main.py unconditionally instantiated in-process SentenceTransformers, so every instance duplicates embedding compute and long-window models served by shared TEI instances (bge-m3, 8192-token window vs MiniLM's 128) were unreachable. - TEIEmbeddingProvider (vektra-index): native /embed API, batch chunking to the TEI client batch limit, dimensions from /info with an /embed probe fallback, Bearer auth, health check. - TEIRerankerService (vektra-core): TEI /rerank with the same RerankResult semantics as the in-process service (BUG-015 score propagation, sigmoid normalization when raw scores are detected); shared _build_rerank_result helper and RerankerProtocol for typing. - Dimension plumbing: the Qdrant collection is sized from the active embedding provider instead of the hardcoded 384 default, and ensure_collection fails fast with a clear message on a dimension mismatch with an existing collection. - The cohere rerank option now passes VEKTRA_RERANK_API_KEY through; it never received a key before (dead as wired). Refs: FEAT-024, ADR-0013, ARCH-035, ARCH-036 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
check_embedding_model resolved the provider by the hardcoded 'sentence-transformers' name, so startup failed with a registration error whenever VEKTRA_EMBEDDING_PROVIDER=tei was active (found in the FEAT-024 live smoke). The 'default' alias always points to the active provider. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
bge-m3 via TEI on the eval-full questions (same dual chunks): retrieval hit 82.6% -> 93.5%, MRR 0.7029 -> 0.8478, e2e stable. Configuration reference and .env.example cover the new provider options; the reindex pgvector-hardcoding discovery is recorded in the BUG-021 family note. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fvadicamo
added
enhancement
github-actions
Bot
added
the
documentation
|
Warning Review limit reached
Next review available in: 44 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughAdds remote TEI embedding and reranking providers, configuration and documentation, startup wiring, dynamic Qdrant dimension validation, protocol-based reranking, Cohere API-key forwarding, HTTP dependencies, and mocked TEI tests. ChangesTEI provider integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Startup
participant ProviderRegistry
participant TEIEmbeddingProvider
participant TEIServer
participant Qdrant
Startup->>ProviderRegistry: Register configured embedding provider
Startup->>TEIEmbeddingProvider: Resolve dimensions
TEIEmbeddingProvider->>TEIServer: GET /info or POST /embed
TEIServer-->>TEIEmbeddingProvider: Embedding dimensions
Startup->>Qdrant: Ensure dense collection dimensions
Qdrant-->>Startup: Existing collection validated
sequenceDiagram
participant QueryPipeline
participant TEIRerankerService
participant TEIServer
QueryPipeline->>TEIRerankerService: Rerank query candidates
TEIRerankerService->>TEIServer: POST /rerank
TEIServer-->>TEIRerankerService: Ranked indices and scores
TEIRerankerService-->>QueryPipeline: Reranked results and all scores
Possibly related PRs
Suggested labels: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements FEAT-024, introducing remote embedding and reranking capabilities via HuggingFace Text Embeddings Inference (TEI). It adds TEIEmbeddingProvider and TEIRerankerService, updates configuration options, dynamically sizes Qdrant collections based on active embedding dimensions, and resolves startup warmup issues with alternative providers. The review feedback highlights opportunities to improve the implementation by adopting structlog for logging in compliance with the style guide, processing document embedding batches concurrently using asyncio.gather, robustifying TEI dimension fetching error handling, and ensuring cross-version compatibility with qdrant-client when inspecting collection vector configurations.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
vektra-index/src/vektra_index/startup.py (1)
32-35: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winUse the default embedding alias here
check_provider_registration()still requires("embedding", "sentence-transformers"), butvektra_app/main.pyregisters TEI only as"default"and"tei".ProviderRegistrydoes not alias names, so TEI deployments will fail startup validation. Switch this check to"default"(or register the TEI provider undersentence-transformerstoo).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vektra-index/src/vektra_index/startup.py` around lines 32 - 35, Update the required provider entries in startup validation to use the embedding alias "default" instead of "sentence-transformers", so check_provider_registration() matches the names registered by vektra_app/main.py. Leave the vector_store requirement unchanged.
🧹 Nitpick comments (5)
vektra-index/pyproject.toml (1)
16-16: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnbounded upper version for
httpx.
httpx>=0.27has no upper bound; a futurehttpxmajor release could introduce breaking changes (e.g. toAsyncClient/MockTransportAPIs used throughouttei.pyand its tests) that would only surface at install/runtime.- "httpx>=0.27", + "httpx>=0.27,<1.0",🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vektra-index/pyproject.toml` at line 16, Update the httpx dependency specification in pyproject.toml to include an upper version bound that excludes future breaking major releases while retaining compatibility with versions starting at 0.27.vektra-app/src/vektra_app/main.py (1)
145-154: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMissing
embedding_registeredlog for the sentence-transformers branch.The new TEI branch logs
"embedding_registered"with provider/url, but the existing sentence-transformers branch has no equivalent log call, leaving observability asymmetric between the two providers.♻️ Suggested addition
embedding_provider = SentenceTransformersProvider( model_name=settings.embedding_model ) registry.register("embedding", "default", embedding_provider) registry.register("embedding", "sentence-transformers", embedding_provider) + log.info( + "embedding_registered", + provider="sentence-transformers", + model=settings.embedding_model, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vektra-app/src/vektra_app/main.py` around lines 145 - 154, Add an "embedding_registered" log call to the sentence-transformers branch after both registry.register calls in the embedding provider setup, matching the TEI branch’s logging behavior and including the sentence-transformers provider context.vektra-index/src/vektra_index/providers/tei.py (1)
62-67: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueSequential batch embedding.
embed_documentsawaits each_MAX_BATCHchunk sequentially. For large ingestion batches this serializes network round trips. Consider bounded concurrency (e.g.asyncio.gatherwith a semaphore) if ingestion throughput becomes a bottleneck.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vektra-index/src/vektra_index/providers/tei.py` around lines 62 - 67, Update VektraTEI.embed_documents to process chunked _embed_batch requests with bounded concurrency instead of awaiting each batch sequentially. Use an asyncio semaphore or equivalent limit to cap in-flight requests, gather the batch results, and preserve the original input ordering when flattening them into the returned list.vektra-index/src/vektra_index/providers/qdrant.py (1)
116-132: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winDimension mismatch check silently skipped for malformed/legacy vector config.
If
vectorsisn't a dict (single unnamedVectorParams) or lacks a"dense"key,denseisNone,existing_sizestaysNone, and the mismatch check is skipped entirely — the collection is accepted without validation, exactly the "opaque Qdrant error later" scenario this fix is meant to prevent. Confirmed via Qdrant client behavior thatconfig.params.vectorsis a dict only for named-vector collections; a legacy/unnamed-vector collection would hit this fallback silently.Consider logging a warning (or raising) when
existing_sizecan't be determined, instead of silently returning.♻️ Suggested tweak
dense = vectors.get("dense") if isinstance(vectors, dict) else None existing_size = getattr(dense, "size", None) - if existing_size is not None and existing_size != self._dense_dimensions: + if existing_size is None: + logger.warning( + "Qdrant collection '%s' has no named 'dense' vector config; " + "skipping dimension validation.", + self._collection_name, + ) + elif existing_size != self._dense_dimensions: raise ValueError(...)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vektra-index/src/vektra_index/providers/qdrant.py` around lines 116 - 132, Update the existing-collection validation in the initialization method around self._client.get_collection so configurations where vectors is not a dict or lacks "dense" do not silently return. Raise a clear validation error, or use the established warning logger if available, indicating that the collection’s dense vector size cannot be determined and requires verification before proceeding; preserve the current mismatch error for known sizes.vektra-index/tests/test_tei_provider.py (1)
27-31: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHeader-construction path in
__init__is never actually exercised.Every test (including
test_auth_header_sent_when_api_key_set) constructs the provider with an injected_client, soTEIEmbeddingProvider.__init__'s ownheaders = {"Authorization": ...}computation (tei.py lines 50-53) is dead code from the test's perspective — the test only proves the pre-configured mock client sends the header it was already given, not that the provider builds that header correctly when it constructs its own client. Similarly, the fulldimensions()instance method (which builds its ownhttpx.Clientwithself._api_key) is never tested — only the static_fetch_dimensionshelper is, via a directly-constructed client.Add a test that omits
_clientand lets the provider build its ownAsyncClient/Client, asserting the resulting client sends the expectedAuthorizationheader.Also applies to: 110-127
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vektra-index/tests/test_tei_provider.py` around lines 27 - 31, Add a test alongside the existing TEI provider tests that creates TEIEmbeddingProvider without injecting _client, mocks the provider’s HTTP client construction or request path, and verifies the provider-created client sends the expected Authorization header when an API key is supplied. Exercise the public dimensions() path as needed so both __init__ header construction and the instance method’s self._api_key usage are covered, while leaving injected-client tests unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@vektra-core/src/vektra_core/reranker.py`:
- Line 187: Update the TEI success log in the reranker initialization flow to
pass config.tei_url through the existing _redact_url helper before assigning it
to the url field, matching the redaction pattern used by the registration logs
in main.py.
- Around line 142-153: Add an async aclose() method to TEIRerankerService that
closes its owned httpx.AsyncClient, then invoke this method during application
lifespan shutdown alongside the embedding provider cleanup. Ensure shutdown
awaits the close operation and does not close an externally supplied _client if
ownership is not held.
- Around line 186-188: Wrap the TEI branch in the reranker initialization
function with the same try/except fallback guard used by the local reranker
path. Ensure TEIRerankerService construction failures, including malformed
tei_url errors, are caught and downgraded to returning None rather than aborting
startup, while preserving the existing successful initialization and logging
behavior.
In `@vektra-index/src/vektra_index/providers/tei.py`:
- Around line 73-107: The synchronous dimensions lookup in dimensions() blocks
the event loop when check_embedding_model() validates the TEI provider. Make the
startup dimension retrieval asynchronous, or populate _dimensions before
check_embedding_model() runs, while preserving the existing /info-then-/embed
fallback and cache behavior.
- Around line 119-121: Update the application teardown in main.py to call and
await embedding_provider.aclose() alongside the existing database engine
disposal. Ensure shutdown invokes this cleanup on reload and exit, using the
provider’s aclose method exposed by the TEI provider.
---
Outside diff comments:
In `@vektra-index/src/vektra_index/startup.py`:
- Around line 32-35: Update the required provider entries in startup validation
to use the embedding alias "default" instead of "sentence-transformers", so
check_provider_registration() matches the names registered by
vektra_app/main.py. Leave the vector_store requirement unchanged.
---
Nitpick comments:
In `@vektra-app/src/vektra_app/main.py`:
- Around line 145-154: Add an "embedding_registered" log call to the
sentence-transformers branch after both registry.register calls in the embedding
provider setup, matching the TEI branch’s logging behavior and including the
sentence-transformers provider context.
In `@vektra-index/pyproject.toml`:
- Line 16: Update the httpx dependency specification in pyproject.toml to
include an upper version bound that excludes future breaking major releases
while retaining compatibility with versions starting at 0.27.
In `@vektra-index/src/vektra_index/providers/qdrant.py`:
- Around line 116-132: Update the existing-collection validation in the
initialization method around self._client.get_collection so configurations where
vectors is not a dict or lacks "dense" do not silently return. Raise a clear
validation error, or use the established warning logger if available, indicating
that the collection’s dense vector size cannot be determined and requires
verification before proceeding; preserve the current mismatch error for known
sizes.
In `@vektra-index/src/vektra_index/providers/tei.py`:
- Around line 62-67: Update VektraTEI.embed_documents to process chunked
_embed_batch requests with bounded concurrency instead of awaiting each batch
sequentially. Use an asyncio semaphore or equivalent limit to cap in-flight
requests, gather the batch results, and preserve the original input ordering
when flattening them into the returned list.
In `@vektra-index/tests/test_tei_provider.py`:
- Around line 27-31: Add a test alongside the existing TEI provider tests that
creates TEIEmbeddingProvider without injecting _client, mocks the provider’s
HTTP client construction or request path, and verifies the provider-created
client sends the expected Authorization header when an API key is supplied.
Exercise the public dimensions() path as needed so both __init__ header
construction and the instance method’s self._api_key usage are covered, while
leaving injected-client tests unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 758b46dc-0998-4edb-9ad4-0f9fde26fdcb
⛔ Files ignored due to path filters (3)
.s2s/BACKLOG.mdis excluded by!.s2s/**.s2s/plans/20260712-sprint3-rag-quality.mdis excluded by!.s2s/**uv.lockis excluded by!**/*.lock,!**/*.lock
📒 Files selected for processing (14)
.env.exampleCHANGELOG.mddocs/reference/configuration.mdvektra-app/src/vektra_app/main.pyvektra-core/pyproject.tomlvektra-core/src/vektra_core/advanced_pipeline.pyvektra-core/src/vektra_core/reranker.pyvektra-core/tests/test_reranker.pyvektra-index/pyproject.tomlvektra-index/src/vektra_index/providers/qdrant.pyvektra-index/src/vektra_index/providers/tei.pyvektra-index/src/vektra_index/startup.pyvektra-index/tests/test_tei_provider.pyvektra-shared/src/vektra_shared/config.py
- create_reranker: the tei branch now degrades to None on init/URL errors like the in-process path, and logs the URL redacted (coderabbit #3567707446, #3567707458) - TEIRerankerService gains aclose(); the app lifespan teardown closes remote provider clients best-effort via a new reranker registry entry (coderabbit #3567707442, #3567707464) - _embed_batch warms the dimensions cache so the startup warmup makes the synchronous dimensions() call free (coderabbit #3567707461) - _fetch_dimensions tolerates non-JSON /info responses and falls back to the /embed probe (gemini #3567695402) - ensure_collection reads the dense vector params via getattr fallback for non-dict vector configs (gemini #3567695404) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- wire check_provider_registration into startup (ARCH-057 step 5 was dead code); checks the embedding/default alias - fix the Qdrant healthcheck: the image ships no curl, so the container was permanently unhealthy; bash /dev/tcp probe instead - remove Dockerfile CMD ["server"] so CMD_TARGET=migrate actually migrates (it hung the v0.6.0 rollout script) - stop mocking DocumentChunkOrm in pgvector tests; use the session.add intercept - docs: sync ingest returns new/exists/alias (not "indexed"), markdown is supported, and document_chunks is EMPTY in Qdrant mode (root cause of the BUG-021 reindex no-op) Reviews: Gemini no findings; CodeRabbit no findings Tests: 699 passed; CI green Refs: follow-ups from #93/#95 review + two deployment bugs found during the v0.6.0 rollout
1 participant
What
TEIEmbeddingProvider(vektra-index): remote embedding over a Text Embeddings Inference instance viaVEKTRA_EMBEDDING_PROVIDER=tei(VEKTRA_TEI_URL/VEKTRA_TEI_API_KEY); native/embedAPI, batching to the TEI client limit, dimensions from/infowith an/embedprobe fallback (TEI 1.9.3 does not expose the size in/info)TEIRerankerService(vektra-core):VEKTRA_RERANK_PROVIDER=teiover TEI/rerank, sameRerankResultsemantics as the in-process path (BUG-015 score propagation, sigmoid normalization); shared_build_rerank_resulthelper andRerankerProtocolfor typingsentence-transformersname (startup crashed with any other provider); thecoherererank option never received an API key (nowVEKTRA_RERANK_API_KEY)Why
VEKTRA_EMBEDDING_PROVIDER=teiwas documented but no provider existed: every instance duplicates in-process CPU embedding, and shared host inference (TEI bge-m3) was unreachableDetails
Measurement (eval-full questions, same dual chunks, bge-m3 float32 CPU via TEI)
rescued=3) — full funnel verified with both remote providersrun_reindexstores through hardcoded pgvector, so a reindex in Qdrant mode silently writes nothing (worked around via fresh-namespace ingest)vektra-internal/stack/20260713-feat024-tei-providers.mdChecklist
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation