Merge pull request #5392 from BerriAI/litellm_add_native_cohere_rerank

[Feat] Add cohere rerank and together ai rerank

docs/my-website/docs/proxy/guardrails/custom_guardrail.md

@@ -55,6 +55,7 @@ class myCustomGuardrail(CustomGuardrail):
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank"
         ],
     ) -> Optional[Union[Exception, str, dict]]:
         """

litellm/__init__.py

@@ -934,6 +934,7 @@ from .proxy.proxy_cli import run_server
 from .router import Router
 from .assistants.main import *
 from .batches.main import *
+from .rerank_api.main import *
 from .fine_tuning.main import *
 from .files.main import *
 from .scheduler import *

litellm/integrations/custom_logger.py

@@ -109,6 +109,7 @@ class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callbac
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ) -> Optional[
         Union[Exception, str, dict]

litellm/llms/cohere/rerank.py

@@ -0,0 +1,78 @@
+"""
+Re rank api
+
+LiteLLM supports the re rank API format, no paramter transformation occurs
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import httpx
+from pydantic import BaseModel
+
+from litellm.llms.base import BaseLLM
+from litellm.llms.custom_httpx.http_handler import (
+    _get_async_httpx_client,
+    _get_httpx_client,
+)
+from litellm.rerank_api.types import RerankRequest, RerankResponse
+
+
+class CohereRerank(BaseLLM):
+    def rerank(
+        self,
+        model: str,
+        api_key: str,
+        query: str,
+        documents: List[Union[str, Dict[str, Any]]],
+        top_n: Optional[int] = None,
+        rank_fields: Optional[List[str]] = None,
+        return_documents: Optional[bool] = True,
+        max_chunks_per_doc: Optional[int] = None,
+        _is_async: Optional[bool] = False,  # New parameter
+    ) -> RerankResponse:
+        request_data = RerankRequest(
+            model=model,
+            query=query,
+            top_n=top_n,
+            documents=documents,
+            rank_fields=rank_fields,
+            return_documents=return_documents,
+            max_chunks_per_doc=max_chunks_per_doc,
+        )
+
+        request_data_dict = request_data.dict(exclude_none=True)
+
+        if _is_async:
+            return self.async_rerank(request_data_dict, api_key)  # type: ignore # Call async method
+
+        client = _get_httpx_client()
+        response = client.post(
+            "https://api.cohere.com/v1/rerank",
+            headers={
+                "accept": "application/json",
+                "content-type": "application/json",
+                "Authorization": f"bearer {api_key}",
+            },
+            json=request_data_dict,
+        )
+
+        return RerankResponse(**response.json())
+
+    async def async_rerank(
+        self,
+        request_data_dict: Dict[str, Any],
+        api_key: str,
+    ) -> RerankResponse:
+        client = _get_async_httpx_client()
+
+        response = await client.post(
+            "https://api.cohere.com/v1/rerank",
+            headers={
+                "accept": "application/json",
+                "content-type": "application/json",
+                "Authorization": f"bearer {api_key}",
+            },
+            json=request_data_dict,
+        )
+
+        return RerankResponse(**response.json())

litellm/llms/togetherai/rerank.py

@@ -0,0 +1,103 @@
+"""
+Re rank api
+
+LiteLLM supports the re rank API format, no paramter transformation occurs
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import httpx
+from pydantic import BaseModel
+
+from litellm.llms.base import BaseLLM
+from litellm.llms.custom_httpx.http_handler import (
+    _get_async_httpx_client,
+    _get_httpx_client,
+)
+from litellm.rerank_api.types import RerankRequest, RerankResponse
+
+
+class TogetherAIRerank(BaseLLM):
+    def rerank(
+        self,
+        model: str,
+        api_key: str,
+        query: str,
+        documents: List[Union[str, Dict[str, Any]]],
+        top_n: Optional[int] = None,
+        rank_fields: Optional[List[str]] = None,
+        return_documents: Optional[bool] = True,
+        max_chunks_per_doc: Optional[int] = None,
+        _is_async: Optional[bool] = False,
+    ) -> RerankResponse:
+        client = _get_httpx_client()
+
+        request_data = RerankRequest(
+            model=model,
+            query=query,
+            top_n=top_n,
+            documents=documents,
+            rank_fields=rank_fields,
+            return_documents=return_documents,
+        )
+
+        # exclude None values from request_data
+        request_data_dict = request_data.dict(exclude_none=True)
+        if max_chunks_per_doc is not None:
+            raise ValueError("TogetherAI does not support max_chunks_per_doc")
+
+        if _is_async:
+            return self.async_rerank(request_data_dict, api_key)  # type: ignore # Call async method
+
+        response = client.post(
+            "https://api.together.xyz/v1/rerank",
+            headers={
+                "accept": "application/json",
+                "content-type": "application/json",
+                "authorization": f"Bearer {api_key}",
+            },
+            json=request_data_dict,
+        )
+
+        if response.status_code != 200:
+            raise Exception(response.text)
+
+        _json_response = response.json()
+
+        response = RerankResponse(
+            id=_json_response.get("id"),
+            results=_json_response.get("results"),
+            meta=_json_response.get("meta") or {},
+        )
+
+        return response
+
+    async def async_rerank(  # New async method
+        self,
+        request_data_dict: Dict[str, Any],
+        api_key: str,
+    ) -> RerankResponse:
+        client = _get_async_httpx_client()  # Use async client
+
+        response = await client.post(
+            "https://api.together.xyz/v1/rerank",
+            headers={
+                "accept": "application/json",
+                "content-type": "application/json",
+                "authorization": f"Bearer {api_key}",
+            },
+            json=request_data_dict,
+        )
+
+        if response.status_code != 200:
+            raise Exception(response.text)
+
+        _json_response = response.json()
+
+        return RerankResponse(
+            id=_json_response.get("id"),
+            results=_json_response.get("results"),
+            meta=_json_response.get("meta") or {},
+        )  # Return response
+
+    pass

litellm/proxy/custom_callbacks1.py

@@ -29,6 +29,7 @@ class MyCustomHandler(
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ):
         return data

litellm/proxy/custom_guardrail.py

@@ -32,6 +32,7 @@ class myCustomGuardrail(CustomGuardrail):
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ) -> Optional[Union[Exception, str, dict]]:
         """

litellm/proxy/example_config_yaml/custom_guardrail.py

@@ -32,6 +32,7 @@ class myCustomGuardrail(CustomGuardrail):
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ) -> Optional[Union[Exception, str, dict]]:
         """

litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py

@@ -32,6 +32,7 @@ class myCustomGuardrail(CustomGuardrail):
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ) -> Optional[Union[Exception, str, dict]]:
         """

litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py

@@ -127,6 +127,7 @@ class lakeraAI_Moderation(CustomGuardrail):
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ):
         if (
@@ -288,6 +289,7 @@ class lakeraAI_Moderation(CustomGuardrail):
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ) -> Optional[Union[Exception, str, Dict]]:
         from litellm.types.guardrails import GuardrailEventHooks

litellm/proxy/hooks/dynamic_rate_limiter.py

@@ -199,6 +199,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ) -> Optional[
         Union[Exception, str, dict]

litellm/proxy/proxy_server.py

@@ -205,6 +205,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
 from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
     router as pass_through_router,
 )
+from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router
 from litellm.proxy.route_llm_request import route_request
 from litellm.proxy.secret_managers.aws_secret_manager import (
     load_aws_kms,
@@ -9881,6 +9882,7 @@ def cleanup_router_config_variables():
 
 
 app.include_router(router)
+app.include_router(rerank_router)
 app.include_router(fine_tuning_router)
 app.include_router(vertex_router)
 app.include_router(gemini_router)

litellm/proxy/rerank_endpoints/endpoints.py

@@ -0,0 +1,124 @@
+#### Rerank Endpoints #####
+from datetime import datetime, timedelta, timezone
+from typing import List, Optional
+
+import fastapi
+import orjson
+from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
+from fastapi.responses import ORJSONResponse
+
+import litellm
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import *
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+router = APIRouter()
+import asyncio
+
+
+@router.post(
+    "/v1/rerank",
+    dependencies=[Depends(user_api_key_auth)],
+    response_class=ORJSONResponse,
+    tags=["rerank"],
+)
+@router.post(
+    "/rerank",
+    dependencies=[Depends(user_api_key_auth)],
+    response_class=ORJSONResponse,
+    tags=["rerank"],
+)
+async def rerank(
+    request: Request,
+    fastapi_response: Response,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    from litellm.proxy.proxy_server import (
+        add_litellm_data_to_request,
+        general_settings,
+        get_custom_headers,
+        llm_router,
+        proxy_config,
+        proxy_logging_obj,
+        route_request,
+        user_model,
+        version,
+    )
+
+    data = {}
+    try:
+        body = await request.body()
+        data = orjson.loads(body)
+
+        # Include original request and headers in the data
+        data = await add_litellm_data_to_request(
+            data=data,
+            request=request,
+            general_settings=general_settings,
+            user_api_key_dict=user_api_key_dict,
+            version=version,
+            proxy_config=proxy_config,
+        )
+
+        ### CALL HOOKS ### - modify incoming data / reject request before calling the model
+        data = await proxy_logging_obj.pre_call_hook(
+            user_api_key_dict=user_api_key_dict, data=data, call_type="rerank"
+        )
+
+        ## ROUTE TO CORRECT ENDPOINT ##
+        llm_call = await route_request(
+            data=data,
+            route_type="arerank",
+            llm_router=llm_router,
+            user_model=user_model,
+        )
+        response = await llm_call
+
+        ### ALERTING ###
+        asyncio.create_task(
+            proxy_logging_obj.update_request_status(
+                litellm_call_id=data.get("litellm_call_id", ""), status="success"
+            )
+        )
+
+        ### RESPONSE HEADERS ###
+        hidden_params = getattr(response, "_hidden_params", {}) or {}
+        model_id = hidden_params.get("model_id", None) or ""
+        cache_key = hidden_params.get("cache_key", None) or ""
+        api_base = hidden_params.get("api_base", None) or ""
+
+        fastapi_response.headers.update(
+            get_custom_headers(
+                user_api_key_dict=user_api_key_dict,
+                model_id=model_id,
+                cache_key=cache_key,
+                api_base=api_base,
+                version=version,
+                model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
+                request_data=data,
+            )
+        )
+
+        return response
+    except Exception as e:
+        await proxy_logging_obj.post_call_failure_hook(
+            user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
+        )
+        verbose_proxy_logger.error(
+            "litellm.proxy.proxy_server.rerank(): Exception occured - {}".format(str(e))
+        )
+        if isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "message", str(e)),
+                type=getattr(e, "type", "None"),
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
+            )
+        else:
+            error_msg = f"{str(e)}"
+            raise ProxyException(
+                message=getattr(e, "message", error_msg),
+                type=getattr(e, "type", "None"),
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", 500),
+            )

litellm/proxy/route_llm_request.py

@@ -33,6 +33,7 @@ ROUTE_ENDPOINT_MAPPING = {
     "aspeech": "/audio/speech",
     "atranscription": "/audio/transcriptions",
     "amoderation": "/moderations",
+    "arerank": "/rerank",
 }
 
 
@@ -48,6 +49,7 @@ async def route_request(
         "aspeech",
         "atranscription",
         "amoderation",
+        "arerank",
     ],
 ):
     """

litellm/proxy/utils.py

@@ -375,6 +375,7 @@ class ProxyLogging:
             "moderation",
             "audio_transcription",
             "pass_through_endpoint",
+            "rerank",
         ],
     ) -> dict:
         """

litellm/rerank_api/main.py

@@ -0,0 +1,166 @@
+import asyncio
+import contextvars
+from functools import partial
+from typing import Any, Coroutine, Dict, List, Literal, Optional, Union
+
+import litellm
+from litellm import get_secret
+from litellm._logging import verbose_logger
+from litellm.llms.cohere.rerank import CohereRerank
+from litellm.llms.togetherai.rerank import TogetherAIRerank
+from litellm.types.router import *
+from litellm.utils import supports_httpx_timeout
+
+from .types import RerankRequest, RerankResponse
+
+####### ENVIRONMENT VARIABLES ###################
+# Initialize any necessary instances or variables here
+cohere_rerank = CohereRerank()
+together_rerank = TogetherAIRerank()
+#################################################
+
+
+async def arerank(
+    model: str,
+    query: str,
+    documents: List[Union[str, Dict[str, Any]]],
+    custom_llm_provider: Optional[Literal["cohere", "together_ai"]] = None,
+    top_n: Optional[int] = None,
+    rank_fields: Optional[List[str]] = None,
+    return_documents: Optional[bool] = True,
+    max_chunks_per_doc: Optional[int] = None,
+    **kwargs,
+) -> Union[RerankResponse, Coroutine[Any, Any, RerankResponse]]:
+    """
+    Async: Reranks a list of documents based on their relevance to the query
+    """
+    try:
+        loop = asyncio.get_event_loop()
+        kwargs["arerank"] = True
+
+        func = partial(
+            rerank,
+            model,
+            query,
+            documents,
+            custom_llm_provider,
+            top_n,
+            rank_fields,
+            return_documents,
+            max_chunks_per_doc,
+            **kwargs,
+        )
+
+        ctx = contextvars.copy_context()
+        func_with_context = partial(ctx.run, func)
+        init_response = await loop.run_in_executor(None, func_with_context)
+
+        if asyncio.iscoroutine(init_response):
+            response = await init_response
+        else:
+            response = init_response
+        return response
+    except Exception as e:
+        raise e
+
+
+def rerank(
+    model: str,
+    query: str,
+    documents: List[Union[str, Dict[str, Any]]],
+    custom_llm_provider: Optional[Literal["cohere", "together_ai"]] = None,
+    top_n: Optional[int] = None,
+    rank_fields: Optional[List[str]] = None,
+    return_documents: Optional[bool] = True,
+    max_chunks_per_doc: Optional[int] = None,
+    **kwargs,
+) -> Union[RerankResponse, Coroutine[Any, Any, RerankResponse]]:
+    """
+    Reranks a list of documents based on their relevance to the query
+    """
+    try:
+        _is_async = kwargs.pop("arerank", False) is True
+        optional_params = GenericLiteLLMParams(**kwargs)
+
+        model, _custom_llm_provider, dynamic_api_key, api_base = (
+            litellm.get_llm_provider(
+                model=model,
+                custom_llm_provider=custom_llm_provider,
+                api_base=optional_params.api_base,
+                api_key=optional_params.api_key,
+            )
+        )
+
+        # Implement rerank logic here based on the custom_llm_provider
+        if _custom_llm_provider == "cohere":
+            # Implement Cohere rerank logic
+            cohere_key = (
+                dynamic_api_key
+                or optional_params.api_key
+                or litellm.cohere_key
+                or get_secret("COHERE_API_KEY")
+                or get_secret("CO_API_KEY")
+                or litellm.api_key
+            )
+
+            if cohere_key is None:
+                raise ValueError(
+                    "Cohere API key is required, please set 'COHERE_API_KEY' in your environment"
+                )
+
+            api_base = (
+                optional_params.api_base
+                or litellm.api_base
+                or get_secret("COHERE_API_BASE")
+                or "https://api.cohere.ai/v1/generate"
+            )
+
+            headers: Dict = litellm.headers or {}
+
+            response = cohere_rerank.rerank(
+                model=model,
+                query=query,
+                documents=documents,
+                top_n=top_n,
+                rank_fields=rank_fields,
+                return_documents=return_documents,
+                max_chunks_per_doc=max_chunks_per_doc,
+                api_key=cohere_key,
+                _is_async=_is_async,
+            )
+            pass
+        elif _custom_llm_provider == "together_ai":
+            # Implement Together AI rerank logic
+            together_key = (
+                dynamic_api_key
+                or optional_params.api_key
+                or litellm.togetherai_api_key
+                or get_secret("TOGETHERAI_API_KEY")
+                or litellm.api_key
+            )
+
+            if together_key is None:
+                raise ValueError(
+                    "TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment"
+                )
+
+            response = together_rerank.rerank(
+                model=model,
+                query=query,
+                documents=documents,
+                top_n=top_n,
+                rank_fields=rank_fields,
+                return_documents=return_documents,
+                max_chunks_per_doc=max_chunks_per_doc,
+                api_key=together_key,
+                _is_async=_is_async,
+            )
+
+        else:
+            raise ValueError(f"Unsupported provider: {_custom_llm_provider}")
+
+        # Placeholder return
+        return response
+    except Exception as e:
+        verbose_logger.error(f"Error in rerank: {str(e)}")
+        raise e

litellm/rerank_api/types.py

@@ -0,0 +1,25 @@
+"""
+LiteLLM Follows the cohere API format for the re rank API
+https://docs.cohere.com/reference/rerank
+
+"""
+
+from typing import List, Optional, Union
+
+from pydantic import BaseModel
+
+
+class RerankRequest(BaseModel):
+    model: str
+    query: str
+    top_n: Optional[int] = None
+    documents: List[Union[str, dict]]
+    rank_fields: Optional[List[str]] = None
+    return_documents: Optional[bool] = None
+    max_chunks_per_doc: Optional[int] = None
+
+
+class RerankResponse(BaseModel):
+    id: str
+    results: List[dict]  # Contains index and relevance_score
+    meta: dict  # Contains api_version and billed_units

litellm/tests/test_rerank.py

@@ -0,0 +1,127 @@
+import json
+import os
+import sys
+import traceback
+
+from dotenv import load_dotenv
+
+load_dotenv()
+import io
+import os
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+
+import os
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+import litellm
+from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+
+
+def assert_response_shape(response, custom_llm_provider):
+    expected_response_shape = {"id": str, "results": list, "meta": dict}
+
+    expected_results_shape = {"index": int, "relevance_score": float}
+
+    expected_meta_shape = {"api_version": dict, "billed_units": dict}
+
+    expected_api_version_shape = {"version": str}
+
+    expected_billed_units_shape = {"search_units": int}
+
+    assert isinstance(response.id, expected_response_shape["id"])
+    assert isinstance(response.results, expected_response_shape["results"])
+    for result in response.results:
+        assert isinstance(result["index"], expected_results_shape["index"])
+        assert isinstance(
+            result["relevance_score"], expected_results_shape["relevance_score"]
+        )
+    assert isinstance(response.meta, expected_response_shape["meta"])
+
+    if custom_llm_provider == "cohere":
+
+        assert isinstance(
+            response.meta["api_version"], expected_meta_shape["api_version"]
+        )
+        assert isinstance(
+            response.meta["api_version"]["version"],
+            expected_api_version_shape["version"],
+        )
+        assert isinstance(
+            response.meta["billed_units"], expected_meta_shape["billed_units"]
+        )
+        assert isinstance(
+            response.meta["billed_units"]["search_units"],
+            expected_billed_units_shape["search_units"],
+        )
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize("sync_mode", [True, False])
+async def test_basic_rerank(sync_mode):
+    if sync_mode is True:
+        response = litellm.rerank(
+            model="cohere/rerank-english-v3.0",
+            query="hello",
+            documents=["hello", "world"],
+            top_n=3,
+        )
+
+        print("re rank response: ", response)
+
+        assert response.id is not None
+        assert response.results is not None
+
+        assert_response_shape(response, custom_llm_provider="cohere")
+    else:
+        response = await litellm.arerank(
+            model="cohere/rerank-english-v3.0",
+            query="hello",
+            documents=["hello", "world"],
+            top_n=3,
+        )
+
+        print("async re rank response: ", response)
+
+        assert response.id is not None
+        assert response.results is not None
+
+        assert_response_shape(response, custom_llm_provider="cohere")
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize("sync_mode", [True, False])
+async def test_basic_rerank_together_ai(sync_mode):
+    if sync_mode is True:
+        response = litellm.rerank(
+            model="together_ai/Salesforce/Llama-Rank-V1",
+            query="hello",
+            documents=["hello", "world"],
+            top_n=3,
+        )
+
+        print("re rank response: ", response)
+
+        assert response.id is not None
+        assert response.results is not None
+
+        assert_response_shape(response, custom_llm_provider="together_ai")
+    else:
+        response = await litellm.arerank(
+            model="together_ai/Salesforce/Llama-Rank-V1",
+            query="hello",
+            documents=["hello", "world"],
+            top_n=3,
+        )
+
+        print("async re rank response: ", response)
+
+        assert response.id is not None
+        assert response.results is not None
+
+        assert_response_shape(response, custom_llm_provider="together_ai")