feat(gateway): add reranking model support#13890
Merged
Conversation
…rovider support
Add `usage?: { tokens: number }` to the reranking provider spec (V3/V4)
and thread it through the core `rerank()` function and `RerankResult`.
Update Cohere provider to parse `meta.billed_units.input_tokens` and
Together AI provider to return `usage.total_tokens` from responses.
Add `GatewayRerankingModel` class and `rerankingModel()`/`reranking()`
methods to the gateway provider for AI Gateway integration.
Cohere's rerank API bills by search_units, not input_tokens. The response returns `meta.billed_units.search_units` for reranking.
Remove Together AI reranking usage change — keeping this PR focused on Cohere and gateway provider changes only.
Gateway computes billing from request document count directly, so provider-level usage extraction is not needed.
Gateway computes billing from request document count, so usage tracking in the provider spec and core rerank function is not needed.
Align with provider spec which doesn't include usage.
ArshVermaGit
left a comment
There was a problem hiding this comment.
This PR adds reranking model support to the @ai-sdk/gateway package via the new GatewayRerankingModel implementation. The changes are well-scoped and extend the Gateway provider capabilities in a way that is consistent with existing model abstractions in the AI SDK.
Supporting rerankers through the gateway improves search relevance workflows and aligns with common RAG and retrieval optimization patterns.
…eset - gateway-reranking-model.test.ts: headers, o11y, ranking extraction, request body, providerOptions, error handling, URL construction - gateway-provider.test.ts: rerankingModel() and reranking() alias create correct model instances - changeset for @ai-sdk/gateway patch
shaper
approved these changes
Apr 7, 2026
gr2m
pushed a commit
that referenced
this pull request
Apr 7, 2026
Contributor
|
|
Contributor
|
🚀 Published in:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
GatewayRerankingModelclass implementingRerankingModelV4for the@ai-sdk/gatewaypackagererankingModel()andreranking()methods toGatewayProviderinterface and implementationGatewayRerankingModelIdtype (initiallycohere/rerank-v3.5)This enables users to call rerankers through the AI Gateway:
Test plan
GatewayRerankingModelsends requests to/reranking-modelendpoint with correct headersrerankingModel()andreranking()methods return validRerankingModelV4instances