Add TEI rerank provider

[joicemjoseph]

Summary

• add a native Hugging Face TEI provider for rerank requests
• send TEI-native /rerank payloads with query and texts instead of OpenAI/vLLM/Cohere-shaped bodies
• register TEI in provider schemas and UI config as a keyless rerank-capable provider

Tests

• go test ./providers/tei
• go test ./... -run '^$'\n- git diff --check\n\nNote: full go test ./... in this clone has unrelated existing failures in MCP stdio fixture tests and a Bedrock rerank assertion.

[CLAassistant]

All committers have signed the CLA.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 98734ea3-2663-470b-ac3f-51bb485dc9bf

📥 Commits

Reviewing files that changed from the base of the PR and between 822b1c4 and 928325c.

📒 Files selected for processing (3)

• core/bifrost.go
• core/schemas/bifrost.go
• core/utils.go

✅ Files skipped from review due to trivial changes (1)

• core/utils.go

🚧 Files skipped from review as they are similar to previous changes (2)

• core/bifrost.go
• core/schemas/bifrost.go

---

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added Hugging Face Text Embeddings Inference (TEI) as a selectable provider across config, standard/provider lists, and UI (icon/label).
  • Enabled TEI embedding and rerank, including /rerank support, configurable request payload building, and optional extra-parameter passthrough.
• Bug Fixes
  • Allowed TEI API keys to be empty when the provider supports it.
• Tests
  • Added TEI test coverage for request/response conversions, upstream HTTP behavior, allowed-request gating, path overrides, and raw payload options.

Walkthrough

Adds TEI as a new Bifrost provider with rerank and embedding support, plus backend wiring, config/schema updates, tests, and UI registration.

Changes

TEI Provider Integration

Layer / File(s)	Summary
TEI provider schema and config registration core/schemas/bifrost.go, core/utils.go, transports/config.schema.json	Registers TEI as a provider, adds it to provider lists, allows empty TEI keys, and extends config schema support for TEI provider entries and base provider type values.
TEI rerank models and conversion helpers core/providers/tei/models.go, core/providers/tei/utils.go	Defines TEI rerank wire models, converts Bifrost rerank requests to TEI format, converts TEI ranks back to Bifrost responses with validation and sorting, and wraps provider errors with request and response bodies.
TEIProvider implementation and request handling core/providers/tei/tei.go, core/bifrost.go	Implements TEI HTTP client setup, rerank execution, embedding delegation, request URL resolution, rerank body building, TEI error parsing, unsupported-operation stubs, and provider factory wiring.
TEI utility and provider tests core/providers/tei/utils_test.go	Covers request/response conversion, rerank and embedding HTTP behavior, request URL overrides, and allowed-request gating.
UI provider registration ui/lib/types/config.ts, ui/lib/constants/logs.ts, ui/lib/constants/config.ts, ui/lib/constants/icons.tsx, ui/app/workspace/providers/fragments/allowedRequestsFields.tsx	Adds TEI to UI provider types, labels, icons, supported requests, allowed-request endpoints, and request selection options.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant Bifrost
  participant TEIProvider
  participant TEI

  Client->>Bifrost: select base provider "tei"
  Bifrost->>TEIProvider: createBaseProvider(config)
  Client->>TEIProvider: Rerank request
  TEIProvider->>TEI: POST /rerank
  TEI-->>TEIProvider: rerank results
  TEIProvider-->>Client: Bifrost rerank response

Loading

Possibly related PRs

• maximhq/bifrost#4425: The embedding path here uses the shared OpenAI-compatible embedding handler, matching the provider auth/header changes in that PR.
• maximhq/bifrost#4601: Extends the same provider-construction path with a new base-provider branch and matching wiring.

Suggested reviewers: akshaydeo, danpiths, roroghost17

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	It has Summary and Tests, but omits most required template sections like Changes, Type, Affected areas, breaking changes, issues, security, and checklist.	Add the missing template sections and briefly fill in type of change, affected areas, testing steps, breaking changes, related issues, security notes, and checklist.
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main change: adding a TEI rerank provider.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🤖 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 `@core/providers/tei/rerank.go`:
- Around line 26-31: In the rerank parameters conversion logic where
bifrostReq.Params is processed, the code currently only forwards ExtraParams and
ReturnDocuments to teiReq but silently ignores TopN, MaxTokensPerDoc, and
Priority from schemas.RerankParameters. Add assignments to also forward these
missing parameters (TopN, MaxTokensPerDoc, and Priority) from bifrostReq.Params
to the corresponding fields in teiReq within the same conditional block,
ensuring all rerank parameters are properly transferred to the TEI request.

In `@transports/config.schema.json`:
- Around line 352-354: The tei provider configuration is using the generic
`#/`$defs/provider schema which requires a keys field with minItems: 1, but TEI
should support keyless configuration. Create a TEI-specific provider schema
definition in the $defs section that removes or relaxes the keys requirement,
then update the tei provider reference to point to this new TEI-specific schema
instead of the generic `#/`$defs/provider reference.

🪄 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 Plus

Run ID: 060956f7-2517-43c7-8a2b-8d0787eff066

📥 Commits

Reviewing files that changed from the base of the PR and between fa9d8f0 and 7c9bd63.

📒 Files selected for processing (16)

• core/bifrost.go
• core/providers/tei/cachedcontents.go
• core/providers/tei/containers.go
• core/providers/tei/models.go
• core/providers/tei/rerank.go
• core/providers/tei/rerank_test.go
• core/providers/tei/tei.go
• core/providers/tei/unsupported.go
• core/schemas/bifrost.go
• core/utils.go
• transports/config.schema.json
• ui/app/workspace/providers/fragments/allowedRequestsFields.tsx
• ui/lib/constants/config.ts
• ui/lib/constants/icons.tsx
• ui/lib/constants/logs.ts
• ui/lib/types/config.ts

[greptile-apps]

Confidence Score: 5/5

Safe to merge; the new TEI provider correctly converts Bifrost rerank requests to TEI format on both the normal and large-payload code paths, handles keyless auth, and honors path overrides.

All issues raised in prior review threads have been resolved in the current revision: first-class rerank parameters are forwarded in ToTEIRerankRequest, buildRequestURL delegates to GetRequestPath for custom path overrides, the large-payload branch explicitly calls ToTEIRerankRequest instead of streaming the raw Bifrost body, and provider_with_tei_config drops the required keys constraint. The only remaining gap is that ProviderResponseHeaders is not populated on the Rerank success path, which is a minor omission relative to the rest of the codebase.

core/providers/tei/tei.go — the Rerank success path does not extract and store provider response headers, unlike every other provider that owns its HTTP call.

Important Files Changed

Filename	Overview
core/providers/tei/tei.go	New TEI provider with Rerank and Embedding operations; correctly handles path overrides, large-payload mode conversion, fasthttp acquire/release pattern, and keyless auth. Missing ProviderResponseHeaders in the Rerank success response.
core/providers/tei/utils.go	ToTEIRerankRequest and ToBifrostRerankResponse converters are pure, handle nil, duplicate-index, out-of-range, sorting, and topN truncation correctly.
core/providers/tei/utils_test.go	Covers converter functions, HTTP round-trips (including large-payload drain and path overrides), and disabled-operation guard; correctly avoids require.* inside handler goroutines.
core/providers/tei/models.go	Clean model definitions; ExtraParams uses json:"-" and GetExtraParams() pattern to support passthrough merging.
transports/config.schema.json	Adds provider_with_tei_config without a required keys field, making keyless TEI deployments schema-valid; TEI added to the provider enum.
core/schemas/bifrost.go	TEI added to ModelProvider constants, SupportedBaseProviders, and StandardProviders consistently.
core/utils.go	TEI added to CanProviderKeyValueBeEmpty, consistent with other keyless providers (VLLM, Ollama, SGL).
ui/lib/types/config.ts	tei added to BaseProvider union type, consistent with Go SupportedBaseProviders.
ui/app/workspace/providers/fragments/allowedRequestsFields.tsx	TEI rerank/embedding path placeholders added; rerank added to the global RequestTypes list.
ui/lib/constants/config.ts	TEI model placeholder, isKeyRequiredByProvider: false, and PROVIDER_SUPPORTED_REQUESTS entry added correctly.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Client
    participant Bifrost
    participant TEIProvider
    participant TEI as TEI Server

    Client->>Bifrost: POST /rerank (BifrostRerankRequest)
    Bifrost->>TEIProvider: Rerank(ctx, key, request)
    
    alt normal payload
        TEIProvider->>TEIProvider: CheckContextAndGetRequestBody → ToTEIRerankRequest
        Note over TEIProvider: {query, texts, top_n, return_text}
    else large payload mode
        TEIProvider->>TEIProvider: ToTEIRerankRequest(request struct)
        Note over TEIProvider: converts from parsed struct, not raw reader
    end

    TEIProvider->>TEI: "POST /rerank {query, texts}"
    TEI-->>TEIProvider: "[{index, score, text?}]"
    
    TEIProvider->>TEIProvider: ToBifrostRerankResponse (sort by score, apply topN)
    
    alt large payload
        TEIProvider->>TEIProvider: DrainLargePayloadRemainder
    end

    TEIProvider-->>Bifrost: BifrostRerankResponse
    Bifrost-->>Client: response

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Client
    participant Bifrost
    participant TEIProvider
    participant TEI as TEI Server

    Client->>Bifrost: POST /rerank (BifrostRerankRequest)
    Bifrost->>TEIProvider: Rerank(ctx, key, request)
    
    alt normal payload
        TEIProvider->>TEIProvider: CheckContextAndGetRequestBody → ToTEIRerankRequest
        Note over TEIProvider: {query, texts, top_n, return_text}
    else large payload mode
        TEIProvider->>TEIProvider: ToTEIRerankRequest(request struct)
        Note over TEIProvider: converts from parsed struct, not raw reader
    end

    TEIProvider->>TEI: "POST /rerank {query, texts}"
    TEI-->>TEIProvider: "[{index, score, text?}]"
    
    TEIProvider->>TEIProvider: ToBifrostRerankResponse (sort by score, apply topN)
    
    alt large payload
        TEIProvider->>TEIProvider: DrainLargePayloadRemainder
    end

    TEIProvider-->>Bifrost: BifrostRerankResponse
    Bifrost-->>Client: response

Loading

_{Reviews (12): Last reviewed commit: "Merge branch 'dev' into feature/tei-rera..." | Re-trigger Greptile}

[rlex]

Since TEI can do both embed and rerank it will be nice to see embedding too

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@core/providers/tei/tei.go`:
- Around line 86-88: The conditional logic in the block with
ApplyLargePayloadRequestBodyWithModelNormalization is causing the TEI converter
output to be bypassed in large-payload mode. Remove the conditional check and
call req.SetBody(jsonData) unconditionally to ensure the properly converted TEI
payload with the correct texts field structure is always set on the request,
regardless of whether large-payload mode is enabled.

🪄 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 Plus

Run ID: 6a4abf2b-2340-4479-a9f5-91f4b596f03b

📥 Commits

Reviewing files that changed from the base of the PR and between 7c9bd63 and e55dd92.

📒 Files selected for processing (16)

• core/bifrost.go
• core/providers/tei/cachedcontents.go
• core/providers/tei/containers.go
• core/providers/tei/models.go
• core/providers/tei/rerank.go
• core/providers/tei/rerank_test.go
• core/providers/tei/tei.go
• core/providers/tei/unsupported.go
• core/schemas/bifrost.go
• core/utils.go
• transports/config.schema.json
• ui/app/workspace/providers/fragments/allowedRequestsFields.tsx
• ui/lib/constants/config.ts
• ui/lib/constants/icons.tsx
• ui/lib/constants/logs.ts
• ui/lib/types/config.ts

🚧 Files skipped from review as they are similar to previous changes (13)

• core/providers/tei/models.go
• ui/app/workspace/providers/fragments/allowedRequestsFields.tsx
• ui/lib/types/config.ts
• core/bifrost.go
• ui/lib/constants/config.ts
• ui/lib/constants/logs.ts
• transports/config.schema.json
• ui/lib/constants/icons.tsx
• core/providers/tei/rerank.go
• core/providers/tei/containers.go
• core/providers/tei/rerank_test.go
• core/providers/tei/cachedcontents.go
• core/providers/tei/unsupported.go

[joicemjoseph]

Since TEI can do both embed and rerank it will be nice to see embedding too

Thank you for the review. as per your suggestion, i have both embed along with rerank.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

core/providers/tei/rerank_test.go (1)

177-215: ⚡ Quick win

Add a negative embedding test for disabled operation config.

Please add one case where TEI embedding is disabled via provider config and assert the call returns an unsupported/blocked error. This will lock in the operation-gating contract and prevent regressions.

As per coding guidelines: "Apply standard Go review practices: ... deterministic tests, and table-driven coverage for behavior changes."

🤖 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 `@core/providers/tei/rerank_test.go` around lines 177 - 215, The
TestTEIProviderEmbedding function only covers the happy path for embedding. Add
a negative test case where the provider configuration disables the embedding
operation, then call provider.Embedding and assert that it returns an
appropriate unsupported or blocked error. This can be implemented by adding a
second test case with modified schemas.ProviderConfig that disables the
embedding operation, or by refactoring into a table-driven test structure to
cover both the enabled and disabled scenarios with NewTEIProvider and the
provider.Embedding method call.

Source: Coding guidelines

🤖 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 `@core/providers/tei/embedding.go`:
- Around line 10-23: The Embedding method in the TEIProvider type currently
dispatches directly to openai.HandleOpenAIEmbeddingRequest without first
validating that embedding operations are allowed in the provider configuration.
Add a check using providerUtils.CheckOperationAllowed before the existing
openai.HandleOpenAIEmbeddingRequest call to enforce the AllowedRequests
configuration, following the same pattern used in the OpenAI provider
implementation. This ensures that embedding requests are gated by the provider's
configured operation permissions before being processed.

---

Nitpick comments:
In `@core/providers/tei/rerank_test.go`:
- Around line 177-215: The TestTEIProviderEmbedding function only covers the
happy path for embedding. Add a negative test case where the provider
configuration disables the embedding operation, then call provider.Embedding and
assert that it returns an appropriate unsupported or blocked error. This can be
implemented by adding a second test case with modified schemas.ProviderConfig
that disables the embedding operation, or by refactoring into a table-driven
test structure to cover both the enabled and disabled scenarios with
NewTEIProvider and the provider.Embedding method call.

🪄 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 Plus

Run ID: 75870620-c3da-4829-aab0-779a45a6b21d

📥 Commits

Reviewing files that changed from the base of the PR and between e55dd92 and be87572.

📒 Files selected for processing (17)

• core/bifrost.go
• core/providers/tei/cachedcontents.go
• core/providers/tei/containers.go
• core/providers/tei/embedding.go
• core/providers/tei/models.go
• core/providers/tei/rerank.go
• core/providers/tei/rerank_test.go
• core/providers/tei/tei.go
• core/providers/tei/unsupported.go
• core/schemas/bifrost.go
• core/utils.go
• transports/config.schema.json
• ui/app/workspace/providers/fragments/allowedRequestsFields.tsx
• ui/lib/constants/config.ts
• ui/lib/constants/icons.tsx
• ui/lib/constants/logs.ts
• ui/lib/types/config.ts

💤 Files with no reviewable changes (6)

• ui/app/workspace/providers/fragments/allowedRequestsFields.tsx
• ui/lib/types/config.ts
• ui/lib/constants/icons.tsx
• ui/lib/constants/config.ts
• transports/config.schema.json
• ui/lib/constants/logs.ts

🚧 Files skipped from review as they are similar to previous changes (8)

• core/utils.go
• core/bifrost.go
• core/providers/tei/cachedcontents.go
• core/providers/tei/models.go
• core/providers/tei/rerank.go
• core/providers/tei/containers.go
• core/schemas/bifrost.go
• core/providers/tei/tei.go

The merge-base changed after approval.

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

The merge-base changed after approval.

[akshaydeo]

Hi @joicemjoseph — thanks for the contribution! This PR is currently blocked because our CLA bot shows the Contributor License Agreement as not yet signed. Could you sign it here so we can move this forward: https://cla-assistant.io/maximhq/bifrost?pullRequest=4469

Let us know if you run into any issues signing.