fix: enable rerank for custom OpenAI-compatible providers

[eyeveil]

Problem

Custom providers using base_provider_type: openai cannot serve /v1/rerank. Requests fail before reaching the upstream provider with:

{"code":"unsupported_operation","message":"rerank is not supported by hawk provider"}

This blocks OpenAI-compatible local/upstream runtimes (omlx, llama.cpp, vLLM, ...) that implement /v1/rerank even though OpenAI itself does not.

Root cause

Custom OpenAI-compatible providers are instantiated as OpenAIProvider, but OpenAIProvider.Rerank unconditionally returned unsupported_operation. As a result custom_provider_config.allowed_requests.rerank had no effect and the request never reached upstream.

Fix

• Native OpenAI rerank stays unsupported (no custom_provider_config ⇒ unchanged behavior).
• For custom OpenAI-compatible providers:
  • honor custom_provider_config.allowed_requests via the existing CheckOperationAllowed helper (no gating configured ⇒ allowed, matching other operations);
  • route rerank to /v1/rerank by default, honoring request_path_overrides through the existing OpenAI request URL builder;
  • parse Cohere-style rerank responses (results[].index, relevance_score, optional document, meta.tokens/billed_units, top-level usage) into Bifrost rerank responses;
  • preserve extra params passthrough, raw request/response options, latency, usage, and provider response headers.

Testing

Added offline unit tests (fake upstream via local HTTP): custom-provider rerank routing to /v1/rerank, auth header forwarding, document-object forwarding, extra-params passthrough, local return_documents handling, result ordering, usage extraction, response headers, native OpenAI still unsupported, and allowed_requests gating.

• go build ./... (core) — pass
• go vet ./providers/openai/ — pass
• go test ./providers/openai/ — pass

Fixes #4834

[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: 1a87a086-7d40-45dd-a570-541c1d202e76

📥 Commits

Reviewing files that changed from the base of the PR and between 6b69501 and 61f3a7b.

📒 Files selected for processing (5)

• core/providers/openai/openai.go
• core/providers/openai/rerank.go
• core/providers/openai/types.go
• docs/providers/custom-providers.mdx
• docs/quickstart/gateway/reranking.mdx

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

• docs/providers/custom-providers.mdx
• docs/quickstart/gateway/reranking.mdx

---

📝 Walkthrough

Summary by CodeRabbit

• New Features

  • Added support for reranking requests in custom OpenAI-compatible setups (when enabled by allowed requests).
  • Supports configurable rerank endpoint paths and richer rerank responses (including optional document details).

• Bug Fixes

  • Improves rerank result ordering by relevance score and index.
  • Preserves upstream-provided document payloads and backfills missing documents when needed.
  • Enhances usage reporting from alternative billing/token fields when direct counts are unavailable.
  • Improves reliability for large payloads and smaller “large response” scenarios.

• Documentation

  • Documented custom OpenAI-compatible provider reranking configuration and allowed request types.

• Tests

  • Expanded reranking coverage for correct request/response mapping, headers, and streaming behavior.

Walkthrough

This PR adds rerank support for custom OpenAI-compatible providers, including request/response conversion, upstream HTTP handling, tests, and docs updates. OpenAIProvider.Rerank now routes permitted custom-provider requests to a rerank handler instead of always returning unsupported.

Changes

Custom OpenAI provider rerank support

Layer / File(s)	Summary
Rerank request/response conversion core/providers/openai/types.go, core/providers/openai/rerank.go	Adds OpenAI rerank request/response types, request conversion, result mapping, document parsing, stable sorting, document backfill, and usage extraction.
Rerank entry point and HTTP handling core/providers/openai/openai.go	Rerank now permits custom-provider rerank requests when allowed and delegates to HandleOpenAIRerankRequest, which builds the JSON POST, performs the upstream call, finalizes the response, and fills rerank metadata.
Tests and docs for custom rerank core/providers/openai/rerank_test.go, docs/providers/custom-providers.mdx, docs/quickstart/gateway/reranking.mdx	Adds rerank coverage for upstream forwarding, response parsing, usage mapping, large payload streaming, unsupported cases, and updates custom-provider docs to describe rerank configuration.

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

Suggested reviewers: danpiths, akshaydeo, roroghost17

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: enabling rerank for custom OpenAI-compatible providers.
Description check	✅ Passed	It covers the problem, root cause, fix, testing, and related issue, though it doesn't follow every template section.
Linked Issues check	✅ Passed	The code now forwards rerank for custom OpenAI-based providers and keeps native OpenAI rerank unsupported, matching #4834.
Out of Scope Changes check	✅ Passed	The added docs and tests support the rerank feature, and no unrelated code changes stand out.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ 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.}

[greptile-apps]

Confidence Score: 5/5

The change is safe to merge; it adds a new code path only when customProviderConfig is present and never touches the native OpenAI path.

Both previously flagged bugs (large-response streaming dropping results, and upstream document backfill overwrite) were fixed before this review. The new rerank handler correctly uses the unary fasthttp client, skips PrepareResponseStreaming so the response body is always buffered for in-process parsing, follows the existing acquire/release lifecycle pattern, and handles all three usage shapes (top-level usage, meta.tokens, meta.billed_units.search_units). Test coverage is thorough and directly exercises the corner cases that caused the prior issues.

No files require special attention.

Important Files Changed

Filename	Overview
core/providers/openai/openai.go	Extends Rerank() to gate on customProviderConfig and delegates to HandleOpenAIRerankRequest; native OpenAI still returns unsupported. The fasthttp acquire/release lifecycle and respOwned flag are handled correctly.
core/providers/openai/rerank.go	New file implementing the rerank request/response pipeline. PrepareResponseStreaming is intentionally skipped so the unary client always buffers the response body, preventing the lpResult-nil body issue. Cohere-style usage (tokens + billed_units.search_units) and document backfill logic are correct.
core/providers/openai/rerank_test.go	Comprehensive offline unit tests covering endpoint routing, auth, document forwarding, extra-params passthrough, return_documents handling, result ordering, usage extraction, response headers, large-response threshold, native OpenAI block, allowed_requests gating, upstream document preservation, and large payload request streaming.
core/providers/openai/types.go	Adds OpenAIRerankRequest/Response types. ExtraParams tagged json:"-" with GetExtraParams() satisfying RequestBodyWithExtraParams for passthrough merge; ReturnDocuments intentionally omitted since Bifrost handles it locally.
docs/providers/custom-providers.mdx	Adds rerank to the allowed operations list with an accurate "(OpenAI-compatible base only)" qualifier.
docs/quickstart/gateway/reranking.mdx	New section accurately describes allowed_requests semantics, default /v1/rerank path, and request_path_overrides override mechanism for custom OpenAI-compatible providers.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Client
    participant BifrostHTTP
    participant OpenAIProvider
    participant HandleOpenAIRerankRequest
    participant UpstreamAPI

    Client->>BifrostHTTP: POST /v1/rerank
    BifrostHTTP->>OpenAIProvider: Rerank(ctx, key, request)

    alt "customProviderConfig == nil (native OpenAI)"
        OpenAIProvider-->>BifrostHTTP: unsupported_operation error
    else customProviderConfig present
        OpenAIProvider->>OpenAIProvider: CheckOperationAllowed()
        alt "AllowedRequests.Rerank == false"
            OpenAIProvider-->>BifrostHTTP: unsupported_operation error
        else allowed
            OpenAIProvider->>HandleOpenAIRerankRequest: buildRequestURL("/v1/rerank")
            HandleOpenAIRerankRequest->>HandleOpenAIRerankRequest: ToOpenAIRerankRequest() [excl. ReturnDocuments]
            HandleOpenAIRerankRequest->>HandleOpenAIRerankRequest: MergeExtraParamsIntoJSON (if passthrough)
            HandleOpenAIRerankRequest->>UpstreamAPI: POST /v1/rerank (Bearer auth)
            UpstreamAPI-->>HandleOpenAIRerankRequest: Cohere-style JSON response
            HandleOpenAIRerankRequest->>HandleOpenAIRerankRequest: finalizeOpenAIResponse (unary, buffered)
            HandleOpenAIRerankRequest->>HandleOpenAIRerankRequest: ToBifrostRerankResponse()
            note right of HandleOpenAIRerankRequest: sort by relevance_score desc,<br/>backfill docs from request if returnDocuments=true<br/>and upstream omitted them
            HandleOpenAIRerankRequest-->>OpenAIProvider: BifrostRerankResponse
            OpenAIProvider-->>BifrostHTTP: BifrostRerankResponse
        end
    end

    BifrostHTTP-->>Client: rerank 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 BifrostHTTP
    participant OpenAIProvider
    participant HandleOpenAIRerankRequest
    participant UpstreamAPI

    Client->>BifrostHTTP: POST /v1/rerank
    BifrostHTTP->>OpenAIProvider: Rerank(ctx, key, request)

    alt "customProviderConfig == nil (native OpenAI)"
        OpenAIProvider-->>BifrostHTTP: unsupported_operation error
    else customProviderConfig present
        OpenAIProvider->>OpenAIProvider: CheckOperationAllowed()
        alt "AllowedRequests.Rerank == false"
            OpenAIProvider-->>BifrostHTTP: unsupported_operation error
        else allowed
            OpenAIProvider->>HandleOpenAIRerankRequest: buildRequestURL("/v1/rerank")
            HandleOpenAIRerankRequest->>HandleOpenAIRerankRequest: ToOpenAIRerankRequest() [excl. ReturnDocuments]
            HandleOpenAIRerankRequest->>HandleOpenAIRerankRequest: MergeExtraParamsIntoJSON (if passthrough)
            HandleOpenAIRerankRequest->>UpstreamAPI: POST /v1/rerank (Bearer auth)
            UpstreamAPI-->>HandleOpenAIRerankRequest: Cohere-style JSON response
            HandleOpenAIRerankRequest->>HandleOpenAIRerankRequest: finalizeOpenAIResponse (unary, buffered)
            HandleOpenAIRerankRequest->>HandleOpenAIRerankRequest: ToBifrostRerankResponse()
            note right of HandleOpenAIRerankRequest: sort by relevance_score desc,<br/>backfill docs from request if returnDocuments=true<br/>and upstream omitted them
            HandleOpenAIRerankRequest-->>OpenAIProvider: BifrostRerankResponse
            OpenAIProvider-->>BifrostHTTP: BifrostRerankResponse
        end
    end

    BifrostHTTP-->>Client: rerank response

Loading

_{Reviews (5): Last reviewed commit: "[refactor]: OpenAI provider - move reran..." | Re-trigger Greptile}

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

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

92-99: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Missing client.Shutdown() call for created providers.

None of the three tests call provider.Shutdown() after use. Based on learnings, provider tests in core/providers/<provider>/*_test.go should call client.Shutdown() at the end of each test function (not via defer).

🧹 Proposed fix

 	if response.ExtraFields.ProviderResponseHeaders["X-Test-Header"] != "present" {
 		t.Fatalf("expected provider response headers, got %#v", response.ExtraFields.ProviderResponseHeaders)
 	}
+	provider.Shutdown()
 }

Apply similarly at the end of TestOpenAIRerankUnsupportedForNativeProvider and TestCustomOpenAIRerankHonorsAllowedRequests.

Also applies to: 155-157, 163-170

🤖 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/openai/rerank_test.go` around lines 92 - 99, The OpenAI
provider tests create a provider instance but never clean it up, so add an
explicit provider.Shutdown() call at the end of each test function that uses
NewOpenAIProvider, including TestOpenAIRerankUnsupportedForNativeProvider and
TestCustomOpenAIRerankHonorsAllowedRequests. Keep the shutdown call as a normal
अंतिम statement in each test rather than using defer, and apply the same pattern
anywhere in core/providers/openai/rerank_test.go where a provider is constructed
and used.

Source: Learnings

🤖 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/openai/rerank.go`:
- Around line 61-69: The rerank usage parser in
openAIRerankMeta/openAIRerankTokenUsage only maps input_tokens and
output_tokens, so Cohere-shaped rerank responses lose billing data. Update the
rerank metadata mapping to handle billed_units.search_units specifically, either
by adding a search_units field to openAIRerankTokenUsage or by introducing a
billed-units-specific mapper used by the rerank path. Keep the change localized
to the rerank parsing logic in rerank.go so Cohere billing is preserved when
parsing meta.billed_units.
- Around line 219-229: The rerank request path in CheckContextAndGetRequestBody
handling still always calls req.SetBody(jsonData), so large-payload mode never
switches to streaming and sends an empty upstream body. Update the rerank flow
around toOpenAIRerankRequest and the req setup to detect when
CheckContextAndGetRequestBody returns nil for passthrough, then use the existing
body-stream path instead of SetBody, while keeping the normal JSON body path
unchanged.
- Around line 245-256: The large-response branch in finalizeOpenAIResponse
handling inside rerank.go is dropping the rerank ranking list by returning a
BifrostRerankResponse with only Model, Usage, and ExtraFields. Update this
lpResult != nil path so it either preserves and returns the parsed Results from
the rerank flow or explicitly rejects large-response mode in the same way the
responses lifecycle does, ensuring rerank callers never receive an incomplete
payload.

---

Nitpick comments:
In `@core/providers/openai/rerank_test.go`:
- Around line 92-99: The OpenAI provider tests create a provider instance but
never clean it up, so add an explicit provider.Shutdown() call at the end of
each test function that uses NewOpenAIProvider, including
TestOpenAIRerankUnsupportedForNativeProvider and
TestCustomOpenAIRerankHonorsAllowedRequests. Keep the shutdown call as a normal
अंतिम statement in each test rather than using defer, and apply the same pattern
anywhere in core/providers/openai/rerank_test.go where a provider is constructed
and used.

🪄 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: 7f71eb5c-0e0a-4a6a-a801-ed9fab924641

📥 Commits

Reviewing files that changed from the base of the PR and between 6484317 and 5a642b7.

📒 Files selected for processing (3)

• core/providers/openai/openai.go
• core/providers/openai/rerank.go
• core/providers/openai/rerank_test.go

[eyeveil]

@coderabbitai review — all three actionable items predate the current head: search_units billing mapping and the large-payload passthrough body landed in 8bacf4f, and the large-response branch was removed in 9508a58 (rerank now always parses the JSON body in-process, so Results are never dropped); each has a dedicated test. On the Shutdown() nitpick: no provider type in this repo exposes Shutdown() — the client.Shutdown() convention applies to bifrost.Bifrost instances, which these tests do not construct; the httptest servers are closed via defer. Skipping rather than adding new exported API to the core provider in a scoped bugfix.

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
git log --oneline -15
echo "---"
git show 8bacf4f1e --stat 2>/dev/null | head -50
echo "---"
git show 9508a58c5 --stat 2>/dev/null | head -50

Length of output: 3122

---

Confirmed — both commits are present at the current head and match the described fixes (search_units billing + passthrough body in 8bacf4f; large-response branch removal in 9508a58). Thanks for the detailed breakdown, especially the clarification on Shutdown() scope — that makes sense given no provider type exposes it here.

[resolve]

[Pratham-Mishra04]

Hey @nnNyx thanks for PR, I have added some comments to make it more compatible with the code conventions + would also require relevant doc updates

[eyeveil]

Oh sorry, updated now, hope all is well