Add cohere v2/rerank support

docs/my-website/docs/providers/cohere.md

@@ -108,7 +108,7 @@ response = embedding(
 
 ### Usage
 
-
+LiteLLM supports the v1 and v2 clients for Cohere rerank. By default, the `rerank` endpoint uses the v2 client, but you can specify the v1 client by explicitly calling `v1/rerank`
 
 <Tabs>
 <TabItem value="sdk" label="LiteLLM SDK Usage">

docs/my-website/docs/rerank.md

@@ -111,7 +111,7 @@ curl http://0.0.0.0:4000/rerank \
 
 | Provider    | Link to Usage      |
 |-------------|--------------------|
-| Cohere      |   [Usage](#quick-start)                 |
+| Cohere (v1 + v2 clients)      |   [Usage](#quick-start)                 |
 | Together AI|   [Usage](../docs/providers/togetherai)                 |  
 | Azure AI|   [Usage](../docs/providers/azure_ai)                 |  
 | Jina AI|   [Usage](../docs/providers/jina_ai)                 |  

litellm/llms/azure_ai/rerank/transformation.py

@@ -18,6 +18,10 @@ class AzureAIRerankConfig(CohereRerankConfig):
     Azure AI Rerank - Follows the same Spec as Cohere Rerank
     """
 
+    # Azure does not support v2/rerank for cohere yet
+    def __init__(self):
+        super().__init__(True)
+
     def get_complete_url(self, api_base: Optional[str], model: str) -> str:
         if api_base is None:
             raise ValueError(

litellm/llms/base_llm/rerank/transformation.py

@@ -76,6 +76,7 @@ def map_cohere_rerank_params(
         rank_fields: Optional[List[str]] = None,
         return_documents: Optional[bool] = True,
         max_chunks_per_doc: Optional[int] = None,
+        max_tokens_per_doc: Optional[int] = None,
     ) -> OptionalRerankParams:
         pass
 

litellm/llms/cohere/rerank/transformation.py

@@ -18,17 +18,36 @@ class CohereRerankConfig(BaseRerankConfig):
     Reference: https://docs.cohere.com/v2/reference/rerank
     """
 
-    def __init__(self) -> None:
-        pass
-
-    def get_complete_url(self, api_base: Optional[str], model: str) -> str:
+    def __init__(self, api_base: Optional[str], present_version_params: list[str]) -> None:
+        # Default to the v2 client unless the user specifically uses the v1/rerank endpoint or uses v1-specific params
+        uses_v1_params = ("max_chunks_per_doc" in present_version_params) and ('max_tokens_per_doc' not in present_version_params) 
         if api_base:
             # Remove trailing slashes and ensure clean base URL
             api_base = api_base.rstrip("/")
-            if not api_base.endswith("/v1/rerank"):
-                api_base = f"{api_base}/v1/rerank"
+            if api_base.endswith("/v1/rerank") or (uses_v1_params and not api_base.endswith("/v2/rerank")):
+                self.uses_v1_client = True
+                return
+
+        self.uses_v1_client = False
+
+    def get_complete_url(self, api_base: Optional[str], model: str) -> str:
+        if not api_base:
+            return "https://api.cohere.ai/v2/rerank"
+
+        api_base = api_base.rstrip("/")
+
+        # Use /v1/rerank if user intentionally uses the deprecated cohere endpoint
+        if self.uses_v1_client:
+            if api_base.endswith("/v1/rerank"):
+                return api_base
+            else:
+                return f"{api_base}/v1/rerank"
+
+        # By default use the v2 endpoint
+        if api_base.endswith("/v2/rerank"):
             return api_base
-        return "https://api.cohere.ai/v1/rerank"
+
+        return f"{api_base}/v2/rerank"
 
     def get_supported_cohere_rerank_params(self, model: str) -> list:
         return [
@@ -38,6 +57,7 @@ def get_supported_cohere_rerank_params(self, model: str) -> list:
             "max_chunks_per_doc",
             "rank_fields",
             "return_documents",
+            "max_tokens_per_doc"
         ]
 
     def map_cohere_rerank_params(
@@ -52,21 +72,29 @@ def map_cohere_rerank_params(
         rank_fields: Optional[List[str]] = None,
         return_documents: Optional[bool] = True,
         max_chunks_per_doc: Optional[int] = None,
+        max_tokens_per_doc:  Optional[int] = None, 
     ) -> OptionalRerankParams:
         """
         Map Cohere rerank params
 
         No mapping required - returns all supported params
         """
+
+        unique_version_params = (
+            {"max_chunks_per_doc": max_chunks_per_doc} 
+            if self.uses_v1_client 
+            else {"max_tokens_per_doc": max_tokens_per_doc}
+        )
+
         return OptionalRerankParams(
             query=query,
             documents=documents,
             top_n=top_n,
             rank_fields=rank_fields,
             return_documents=return_documents,
-            max_chunks_per_doc=max_chunks_per_doc,
+            **unique_version_params
         )
-
+        
     def validate_environment(
         self,
         headers: dict,
@@ -108,15 +136,23 @@ def transform_rerank_request(
             raise ValueError("query is required for Cohere rerank")
         if "documents" not in optional_rerank_params:
             raise ValueError("documents is required for Cohere rerank")
+
+        unique_version_params = (
+            {"max_chunks_per_doc": optional_rerank_params.get("max_chunks_per_doc", None)} 
+            if self.uses_v1_client 
+            else {"max_tokens_per_doc": optional_rerank_params.get("max_tokens_per_doc", None)}
+        )
+
         rerank_request = RerankRequest(
             model=model,
             query=optional_rerank_params["query"],
             documents=optional_rerank_params["documents"],
             top_n=optional_rerank_params.get("top_n", None),
             rank_fields=optional_rerank_params.get("rank_fields", None),
             return_documents=optional_rerank_params.get("return_documents", None),
-            max_chunks_per_doc=optional_rerank_params.get("max_chunks_per_doc", None),
+            **unique_version_params
         )
+
         return rerank_request.model_dump(exclude_none=True)
 
     def transform_rerank_response(

litellm/llms/custom_httpx/llm_http_handler.py

@@ -708,6 +708,7 @@ def rerank(
         model: str,
         custom_llm_provider: str,
         logging_obj: LiteLLMLoggingObj,
+        provider_config: BaseRerankConfig,
         optional_rerank_params: OptionalRerankParams,
         timeout: Optional[Union[float, httpx.Timeout]],
         model_response: RerankResponse,
@@ -717,10 +718,7 @@ def rerank(
         api_base: Optional[str] = None,
         client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
     ) -> RerankResponse:
-
-        provider_config = ProviderConfigManager.get_provider_rerank_config(
-            model=model, provider=litellm.LlmProviders(custom_llm_provider)
-        )
+
         # get config from model, custom llm provider
         headers = provider_config.validate_environment(
             api_key=api_key,

litellm/llms/infinity/rerank/transformation.py

@@ -20,6 +20,11 @@
 
 
 class InfinityRerankConfig(CohereRerankConfig):
+    # This is set to the v1/rerank endpoint to not break any existing integrations
+    # This should be changed once it is confirmed that infinity supports the v2 endpoint
+    def __init__(self):
+        super().__init__(True)
+
     def get_complete_url(self, api_base: Optional[str], model: str) -> str:
         if api_base is None:
             raise ValueError("api_base is required for Infinity rerank")

litellm/proxy/_types.py

@@ -230,6 +230,7 @@ class LiteLLMRoutes(enum.Enum):
         # rerank
         "/rerank",
         "/v1/rerank",
+        "/v2/rerank"
         # realtime
         "/realtime",
         "/v1/realtime",

litellm/proxy/rerank_endpoints/endpoints.py

@@ -11,7 +11,12 @@
 router = APIRouter()
 import asyncio
 
-
+@router.post(
+    "/v2/rerank",
+    dependencies=[Depends(user_api_key_auth)],
+    response_class=ORJSONResponse,
+    tags=["rerank"],
+)
 @router.post(
     "/v1/rerank",
     dependencies=[Depends(user_api_key_auth)],

litellm/rerank_api/main.py

@@ -83,6 +83,7 @@ def rerank(  # noqa: PLR0915
     rank_fields: Optional[List[str]] = None,
     return_documents: Optional[bool] = True,
     max_chunks_per_doc: Optional[int] = None,
+    max_tokens_per_doc: Optional[int] = None,
     **kwargs,
 ) -> Union[RerankResponse, Coroutine[Any, Any, RerankResponse]]:
     """
@@ -99,6 +100,9 @@ def rerank(  # noqa: PLR0915
     try:
         _is_async = kwargs.pop("arerank", False) is True
         optional_params = GenericLiteLLMParams(**kwargs)
+        # Params that are unique to specific versions of the client for the rerank call
+        unique_version_params = {"max_chunks_per_doc": max_chunks_per_doc, "max_tokens_per_doc": max_tokens_per_doc}
+        present_version_params = [k for k, v in unique_version_params.items() if v is not None]
 
         model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = (
             litellm.get_llm_provider(
@@ -113,6 +117,8 @@ def rerank(  # noqa: PLR0915
             ProviderConfigManager.get_provider_rerank_config(
                 model=model,
                 provider=litellm.LlmProviders(_custom_llm_provider),
+                api_base=optional_params.api_base,
+                present_version_params=present_version_params
             )
         )
 
@@ -127,6 +133,7 @@ def rerank(  # noqa: PLR0915
             rank_fields=rank_fields,
             return_documents=return_documents,
             max_chunks_per_doc=max_chunks_per_doc,
+            max_tokens_per_doc=max_tokens_per_doc,
             non_default_params=kwargs,
         )
 
@@ -173,6 +180,7 @@ def rerank(  # noqa: PLR0915
             response = base_llm_http_handler.rerank(
                 model=model,
                 custom_llm_provider=_custom_llm_provider,
+                provider_config=rerank_provider_config,
                 optional_rerank_params=optional_rerank_params,
                 logging_obj=litellm_logging_obj,
                 timeout=optional_params.timeout,
@@ -194,6 +202,7 @@ def rerank(  # noqa: PLR0915
                 model=model,
                 custom_llm_provider=_custom_llm_provider,
                 optional_rerank_params=optional_rerank_params,
+                provider_config=rerank_provider_config,
                 logging_obj=litellm_logging_obj,
                 timeout=optional_params.timeout,
                 api_key=dynamic_api_key or optional_params.api_key,
@@ -222,6 +231,7 @@ def rerank(  # noqa: PLR0915
             response = base_llm_http_handler.rerank(
                 model=model,
                 custom_llm_provider=_custom_llm_provider,
+                provider_config=rerank_provider_config,
                 optional_rerank_params=optional_rerank_params,
                 logging_obj=litellm_logging_obj,
                 timeout=optional_params.timeout,

litellm/rerank_api/rerank_utils.py

@@ -15,6 +15,7 @@ def get_optional_rerank_params(
     rank_fields: Optional[List[str]] = None,
     return_documents: Optional[bool] = True,
     max_chunks_per_doc: Optional[int] = None,
+    max_tokens_per_doc: Optional[int] = None,
     non_default_params: Optional[dict] = None,
 ) -> OptionalRerankParams:
     return rerank_provider_config.map_cohere_rerank_params(
@@ -27,5 +28,6 @@ def get_optional_rerank_params(
         rank_fields=rank_fields,
         return_documents=return_documents,
         max_chunks_per_doc=max_chunks_per_doc,
+        max_tokens_per_doc=max_tokens_per_doc,
         non_default_params=non_default_params,
     )

litellm/types/rerank.py

@@ -18,6 +18,8 @@ class RerankRequest(BaseModel):
     rank_fields: Optional[List[str]] = None
     return_documents: Optional[bool] = None
     max_chunks_per_doc: Optional[int] = None
+    max_tokens_per_doc: Optional[int] = None
+
 
 
 class OptionalRerankParams(TypedDict, total=False):
@@ -27,6 +29,7 @@ class OptionalRerankParams(TypedDict, total=False):
     rank_fields: Optional[List[str]]
     return_documents: Optional[bool]
     max_chunks_per_doc: Optional[int]
+    max_tokens_per_doc: Optional[int]
 
 
 class RerankBilledUnits(TypedDict, total=False):

litellm/utils.py

@@ -6119,14 +6119,16 @@ def get_provider_embedding_config(
     def get_provider_rerank_config(
         model: str,
         provider: LlmProviders,
+        api_base: Optional[str],
+        present_version_params: list[str],
     ) -> BaseRerankConfig:
         if litellm.LlmProviders.COHERE == provider:
-            return litellm.CohereRerankConfig()
+            return litellm.CohereRerankConfig(api_base, present_version_params)
         elif litellm.LlmProviders.AZURE_AI == provider:
             return litellm.AzureAIRerankConfig()
         elif litellm.LlmProviders.INFINITY == provider:
             return litellm.InfinityRerankConfig()
-        return litellm.CohereRerankConfig()
+        return litellm.CohereRerankConfig(api_base, present_version_params)
 
     @staticmethod
     def get_provider_audio_transcription_config(