feat: Add IBM watsonx.ai rerank support

[MateuszOssGit]

Title

Add IBM watsonx.ai rerank support

Relevant issues

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Type

🆕 New Feature
✅ Test

Changes

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Feb 16, 2026 0:14am

[MateuszOssGit]

@greptileai

[greptile-apps]

Greptile Summary

This PR adds IBM watsonx.ai rerank support to LiteLLM by implementing a new IBMWatsonXRerankConfig transformation class that maps between LiteLLM's Cohere-compatible rerank interface and watsonx's /ml/v1/text/rerank endpoint.

• Adds IBMWatsonXRerankConfig in litellm/llms/watsonx/rerank/transformation.py with URL construction, parameter mapping, request transformation, and response parsing
• Registers the new provider in litellm/__init__.py, _lazy_imports_registry.py, and ProviderConfigManager in utils.py
• Adds the WATSONX branch to rerank_api/main.py with credential resolution via IBMWatsonXMixin.get_watsonx_credentials
• Adds the RERANK endpoint to WatsonXAIEndpoint enum
• Includes mock-based unit tests covering URL generation, parameter mapping, request/response transformation, and error handling
• Adds missing __init__.py files to existing watsonx subdirectories (chat/, completion/, embed/)
• Adds documentation with SDK and proxy usage examples

Issues found:

• The test uses the full "watsonx/" prefixed model name in transform_rerank_request, but in production get_llm_provider strips this prefix before it reaches the transformation layer — so the model_id sent to the watsonx API would be incorrect if this behavior is replicated
• "watsonx" is not added to the Literal type hint for custom_llm_provider in the rerank/arerank function signatures
• An inline import in rerank_api/main.py should be moved to module-level per project conventions

Confidence Score: 3/5

• This PR is generally safe to merge but has a test accuracy issue related to model name prefixing that should be addressed to ensure the watsonx API receives correctly formatted model IDs.
• Score of 3 reflects that the core implementation follows established patterns correctly (credential handling, response transformation, config registration), but the test inaccurately validates model_id with the provider prefix included, which masks a potential production issue. The missing Literal type hint and inline import are minor style concerns.
• Pay close attention to litellm/rerank_api/main.py (missing type hint, inline import) and tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py (model name prefix issue in test assertions).

Important Files Changed

Filename	Overview
litellm/llms/watsonx/rerank/transformation.py	New watsonx rerank transformation config. Correctly implements BaseRerankConfig with watsonx-specific parameter mapping and response transformation. Uses type: ignore[override] for validate_environment signature mismatch with IBMWatsonXMixin.
litellm/rerank_api/main.py	Adds watsonx provider branch to rerank handler. Contains an inline import (violates CLAUDE.md style guide) and does not add "watsonx" to the Literal type hint for custom_llm_provider in the function signatures.
litellm/utils.py	Adds watsonx provider mapping in ProviderConfigManager.get_provider_rerank_config. Clean 2-line change following existing patterns.
litellm/init.py	Adds IBMWatsonXRerankConfig import. Follows existing pattern.
litellm/_lazy_imports_registry.py	Registers IBMWatsonXRerankConfig in lazy import registry. Follows existing pattern.
litellm/types/llms/watsonx.py	Adds RERANK endpoint to WatsonXAIEndpoint enum. Simple, correct change.
tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py	Mock-based tests covering URL generation, param mapping, request/response transformation, and error handling. All tests use mocks (no real network calls). Test uses full "watsonx/" prefixed model name which wouldn't match production behavior after get_llm_provider strips the prefix.

Sequence Diagram

sequenceDiagram
    participant User
    participant rerank_api as rerank_api/main.py
    participant ProviderConfigMgr as ProviderConfigManager
    participant IBMWatsonXMixin as IBMWatsonXMixin
    participant IBMWatsonXRerank as IBMWatsonXRerankConfig
    participant HTTPHandler as BaseLLMHTTPHandler
    participant WatsonXAPI as watsonx.ai API

    User->>rerank_api: rerank(model="watsonx/...", query, documents)
    rerank_api->>ProviderConfigMgr: get_provider_rerank_config(provider=WATSONX)
    ProviderConfigMgr-->>rerank_api: IBMWatsonXRerankConfig()
    rerank_api->>IBMWatsonXMixin: get_watsonx_credentials(optional_params, api_key, api_base)
    IBMWatsonXMixin-->>rerank_api: WatsonXCredentials(api_key, api_base, token)
    rerank_api->>HTTPHandler: rerank(model, provider_config, optional_rerank_params, ...)
    HTTPHandler->>IBMWatsonXRerank: validate_environment(headers, model, api_key, optional_params)
    IBMWatsonXRerank-->>HTTPHandler: auth headers
    HTTPHandler->>IBMWatsonXRerank: get_complete_url(api_base, model)
    IBMWatsonXRerank-->>HTTPHandler: /ml/v1/text/rerank?version=...
    HTTPHandler->>IBMWatsonXRerank: transform_rerank_request(model, params, headers)
    IBMWatsonXRerank-->>HTTPHandler: watsonx-formatted request body
    HTTPHandler->>WatsonXAPI: POST /ml/v1/text/rerank
    WatsonXAPI-->>HTTPHandler: JSON response
    HTTPHandler->>IBMWatsonXRerank: transform_rerank_response(model, raw_response, ...)
    IBMWatsonXRerank-->>HTTPHandler: RerankResponse
    HTTPHandler-->>rerank_api: RerankResponse
    rerank_api-->>User: RerankResponse

Loading

_{Last reviewed commit: 1cf1dc6}

[greptile-apps]

_{14 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

litellm/rerank_api/main.py
Missing "watsonx" in type hint
"watsonx" is not included in the Literal type for the custom_llm_provider parameter on both the rerank and arerank function signatures. While this won't cause a runtime error (since litellm.get_llm_provider resolves the provider from the model string), it means type checkers will flag valid watsonx rerank calls. Other providers like "voyage" and "fireworks_ai" are already listed here.

    custom_llm_provider: Optional[
        Literal[
            "cohere",
            "together_ai",
            "azure_ai",
            "infinity",
            "litellm_proxy",
            "hosted_vllm",
            "deepinfra",
            "fireworks_ai",
            "voyage",
            "watsonx",
        ]
    ] = None,

[MateuszOssGit]

@greptileai

[greptile-apps]

Greptile Summary

This PR adds IBM watsonx.ai rerank support to LiteLLM, following the established provider pattern with a transformation config class (IBMWatsonXRerankConfig), credential handling in the rerank API entry point, lazy import registration, and documentation.

• Adds IBMWatsonXRerankConfig extending IBMWatsonXMixin and BaseRerankConfig with watsonx-specific parameter mapping (Cohere format → watsonx format), URL construction, auth handling (IAM tokens, ZenApiKey, Bearer tokens), and response transformation.
• Registers the new provider in litellm/__init__.py, _lazy_imports_registry.py, utils.py (ProviderConfigManager), and rerank_api/main.py.
• Adds RERANK endpoint to WatsonXAIEndpoint enum.
• Adds empty __init__.py files to watsonx/, watsonx/chat/, watsonx/completion/, watsonx/embed/, and the new watsonx/rerank/ directories (converts from implicit namespace packages to explicit packages).
• Includes unit tests with mocks (no real network calls) and new documentation pages.
• The test_transform_rerank_request test passes raw Cohere-format parameters directly instead of the watsonx-mapped format, which means it doesn't accurately reflect the production data flow.

Confidence Score: 4/5

• This PR is safe to merge — it adds a new provider following established patterns with no changes to existing provider behavior.
• The implementation correctly follows the existing rerank provider pattern (similar to Voyage, Fireworks, etc.), properly handles watsonx authentication flows, and confines provider-specific code to the llms/ directory. The main concern is a test that doesn't accurately reflect production data flow (passes unmapped params to transform_rerank_request), which reduces confidence slightly. All other files (registration, documentation, types) are straightforward additions.
• Pay close attention to tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py (test_transform_rerank_request uses unmapped params) and litellm/llms/watsonx/rerank/transformation.py (core transformation logic).

Important Files Changed

Filename	Overview
litellm/llms/watsonx/rerank/transformation.py	Core rerank transformation logic for watsonx.ai; correctly implements BaseRerankConfig and IBMWatsonXMixin, maps Cohere-format params to watsonx format, and transforms responses.
litellm/rerank_api/main.py	Adds watsonx provider branch to the rerank() function with credential handling; import moved to top-level per prior feedback. Structure follows existing provider pattern.
tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py	Tests cover URL generation, param mapping, request/response transformation, and error cases. However, test_transform_rerank_request passes raw Cohere-format params instead of mapped watsonx-format params, not reflecting production behavior.
litellm/init.py	Adds IBMWatsonXRerankConfig to the TYPE_CHECKING lazy imports, consistent with other rerank provider entries.
litellm/_lazy_imports_registry.py	Registers IBMWatsonXRerankConfig in both config names and import map, consistent with existing entries.
litellm/utils.py	Adds watsonx to ProviderConfigManager.get_provider_rerank_config(), returning IBMWatsonXRerankConfig for WATSONX provider.
litellm/types/llms/watsonx.py	Adds RERANK endpoint to WatsonXAIEndpoint enum.

Sequence Diagram

sequenceDiagram
    participant User
    participant RerankAPI as rerank_api/main.py
    participant Mixin as IBMWatsonXMixin
    participant Handler as BaseLLMHTTPHandler
    participant Config as IBMWatsonXRerankConfig
    participant WatsonX as watsonx.ai API

    User->>RerankAPI: rerank(model="watsonx/...", query, documents)
    RerankAPI->>RerankAPI: get_llm_provider() → strips prefix
    RerankAPI->>Config: map_cohere_rerank_params()
    Note over Config: documents → inputs<br/>top_n → parameters.return_options.top_n<br/>return_documents → parameters.return_options.inputs
    RerankAPI->>Mixin: get_watsonx_credentials()
    Mixin-->>RerankAPI: api_key, api_base, token
    RerankAPI->>Handler: rerank(model, params, credentials)
    Handler->>Config: validate_environment(headers, api_key)
    Note over Config: Resolves auth: Bearer token,<br/>ZenApiKey, or IAM token
    Config-->>Handler: headers with Authorization
    Handler->>Config: get_complete_url(api_base, model)
    Config-->>Handler: url + /ml/v1/text/rerank?version=...
    Handler->>Config: transform_rerank_request(model, params)
    Note over Config: _get_api_params() resolves project_id/space_id<br/>_prepare_payload() sets model_id
    Config-->>Handler: watsonx request body
    Handler->>WatsonX: POST /ml/v1/text/rerank
    WatsonX-->>Handler: response JSON
    Handler->>Config: transform_rerank_response()
    Note over Config: Maps results[].score → relevance_score<br/>Maps input → document<br/>Extracts input_token_count
    Config-->>Handler: RerankResponse
    Handler-->>User: RerankResponse

Loading

_{Last reviewed commit: 8b1d937}

[greptile-apps]

_{15 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[MateuszOssGit]

All failed test are not related with this PR

[MateuszOssGit]

@krrishdholakia @Sameerlite Could you please review my PR and merge if possible?

PR Confidence Score: 4/5
All failed test are not related with this PR

[greptile-apps]

Greptile Summary

This PR adds IBM watsonx.ai rerank support to LiteLLM, following the established provider integration pattern. The core transformation logic in IBMWatsonXRerankConfig correctly maps Cohere rerank params to watsonx's /ml/v1/text/rerank API format (documents → inputs, top_n → parameters.return_options.top_n, return_documents → parameters.return_options.inputs). Authentication and credential handling properly reuses the existing IBMWatsonXMixin infrastructure.

• Adds IBMWatsonXRerankConfig class extending IBMWatsonXMixin and BaseRerankConfig with proper MRO handling
• Registers the watsonx provider in the rerank dispatch chain (rerank_api/main.py), config manager (utils.py), and lazy imports registry
• Adds RERANK endpoint to WatsonXAIEndpoint enum
• Includes 8 mock-based unit tests covering URL generation, param mapping, request/response transformation, and error handling
• Adds empty __init__.py files to pre-existing watsonx subdirectories (chat/, completion/, embed/) that were missing them
• Adds documentation with SDK and Proxy usage examples

Confidence Score: 4/5

• This PR is safe to merge — it adds a new provider integration following established patterns with no changes to shared logic.
• The implementation correctly follows existing provider rerank patterns (Voyage, Fireworks, etc.). Provider-specific code is properly isolated in litellm/llms/watsonx/rerank/. Auth/credential handling reuses the existing IBMWatsonXMixin. Tests are mock-based and provide reasonable coverage. Minor deduction for a missing newline at end of docs/my-website/docs/rerank.md and the validate_environment signature override needing a type: ignore comment.
• litellm/llms/watsonx/rerank/transformation.py - core new logic, and litellm/rerank_api/main.py - integration point into the rerank dispatch chain

Important Files Changed

Filename	Overview
litellm/llms/watsonx/rerank/transformation.py	New file: Core transformation logic for watsonx rerank. Properly extends IBMWatsonXMixin and BaseRerankConfig. Maps Cohere params to watsonx format, handles auth, and transforms responses correctly.
litellm/rerank_api/main.py	Adds watsonx provider branch to rerank flow. Follows existing pattern (Voyage, Fireworks). Import moved to top of file per style guide. Credentials flow via IBMWatsonXMixin is correct.
litellm/types/llms/watsonx.py	Adds RERANK endpoint enum value to WatsonXAIEndpoint. Minimal, correct change.
litellm/utils.py	Adds watsonx case to ProviderConfigManager.get_provider_rerank_config(). Follows exact same pattern as other providers.
litellm/init.py	Registers IBMWatsonXRerankConfig for TYPE_CHECKING imports. Standard registration pattern.
litellm/_lazy_imports_registry.py	Adds IBMWatsonXRerankConfig to lazy imports registry. Standard registration pattern.
tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py	New file: 8 test cases covering URL generation, param mapping, request/response transformation, error handling. All mock-based, no network calls. Good coverage of core transformation logic.
docs/my-website/docs/providers/watsonx/rerank.md	New documentation for watsonx rerank with SDK and Proxy usage examples.
docs/my-website/docs/rerank.md	Adds watsonx.ai to the supported providers list and reformats table alignment. No functional changes.

Sequence Diagram

sequenceDiagram
    participant User
    participant rerank as rerank_api/main.py
    participant Mixin as IBMWatsonXMixin
    participant Config as IBMWatsonXRerankConfig
    participant Handler as BaseLLMHTTPHandler
    participant WatsonX as watsonx.ai API

    User->>rerank: rerank(model="watsonx/...", query, documents, project_id)
    rerank->>rerank: get_llm_provider() → strips "watsonx/" prefix
    rerank->>Config: map_cohere_rerank_params()
    Note over Config: documents → inputs<br/>top_n → parameters.return_options.top_n<br/>return_documents → parameters.return_options.inputs
    rerank->>Mixin: get_watsonx_credentials()
    Mixin-->>rerank: api_key, api_base, token
    rerank->>Handler: rerank(model, provider_config, optional_rerank_params)
    Handler->>Config: validate_environment(headers, api_key, optional_params)
    Note over Config: Sets Bearer/ZenApiKey auth header<br/>Pops token from params
    Handler->>Config: get_complete_url(api_base, model)
    Note over Config: {api_base}/ml/v1/text/rerank?version=...
    Handler->>Config: transform_rerank_request(model, params)
    Note over Config: _get_api_params() extracts project_id<br/>_prepare_payload() adds model_id
    Handler->>WatsonX: POST /ml/v1/text/rerank
    WatsonX-->>Handler: JSON response
    Handler->>Config: transform_rerank_response()
    Note over Config: Maps score → relevance_score<br/>input → document<br/>input_token_count → meta.tokens
    Config-->>User: RerankResponse

Loading

_{Last reviewed commit: 8b1d937}

[greptile-apps]

_{15 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}