Merge pull request #5394 from BerriAI/litellm_add_rerank_on_proxy

feat - add rerank on litellm proxy / gateway

litellm/proxy/auth/user_api_key_auth.py

@@ -88,6 +88,8 @@ def _get_bearer_token(
         api_key = api_key.replace("Bearer ", "")  # extract the token
     elif api_key.startswith("Basic "):
         api_key = api_key.replace("Basic ", "")  # handle langfuse input
+    elif api_key.startswith("bearer "):
+        api_key = api_key.replace("bearer ", "")
     else:
         api_key = ""
     return api_key

litellm/proxy/example_config_yaml/otel_test_config.yaml

@@ -4,6 +4,10 @@ model_list:
       model: openai/fake
       api_key: fake-key
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
+  - model_name: rerank-english-v3.0
+    litellm_params:
+      model: cohere/rerank-english-v3.0
+      api_key: os.environ/COHERE_API_KEY
 
 litellm_settings:
   cache: true

litellm/proxy/proxy_config.yaml

@@ -4,10 +4,15 @@ model_list:
       model: openai/fake
       api_key: fake-key
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
+  - model_name: Salesforce/Llama-Rank-V1
+    litellm_params:
+      model: together_ai/Salesforce/Llama-Rank-V1
+      api_key: os.environ/TOGETHERAI_API_KEY
+  - model_name: rerank-english-v3.0
+    litellm_params:
+      model: cohere/rerank-english-v3.0
+      api_key: os.environ/COHERE_API_KEY
 
 # default off mode
 litellm_settings:
-  set_verbose: True
-  cache: True
-  cache_params:
-    mode: default_off
+  set_verbose: True

litellm/router.py

@@ -1641,6 +1641,104 @@ class Router:
                 self.fail_calls[model_name] += 1
             raise e
 
+    async def arerank(self, model: str, **kwargs):
+        try:
+            kwargs["model"] = model
+            kwargs["input"] = input
+            kwargs["original_function"] = self._arerank
+            kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries)
+            timeout = kwargs.get("request_timeout", self.timeout)
+            kwargs.setdefault("metadata", {}).update({"model_group": model})
+
+            response = await self.async_function_with_fallbacks(**kwargs)
+
+            return response
+        except Exception as e:
+            asyncio.create_task(
+                send_llm_exception_alert(
+                    litellm_router_instance=self,
+                    request_kwargs=kwargs,
+                    error_traceback_str=traceback.format_exc(),
+                    original_exception=e,
+                )
+            )
+            raise e
+
+    async def _arerank(self, model: str, **kwargs):
+        model_name = None
+        try:
+            verbose_router_logger.debug(
+                f"Inside _rerank()- model: {model}; kwargs: {kwargs}"
+            )
+            deployment = await self.async_get_available_deployment(
+                model=model,
+                specific_deployment=kwargs.pop("specific_deployment", None),
+            )
+            kwargs.setdefault("metadata", {}).update(
+                {
+                    "deployment": deployment["litellm_params"]["model"],
+                    "model_info": deployment.get("model_info", {}),
+                }
+            )
+            kwargs["model_info"] = deployment.get("model_info", {})
+            data = deployment["litellm_params"].copy()
+            model_name = data["model"]
+            for k, v in self.default_litellm_params.items():
+                if (
+                    k not in kwargs and v is not None
+                ):  # prioritize model-specific params > default router params
+                    kwargs[k] = v
+                elif k == "metadata":
+                    kwargs[k].update(v)
+
+            potential_model_client = self._get_client(
+                deployment=deployment, kwargs=kwargs, client_type="async"
+            )
+            # check if provided keys == client keys #
+            dynamic_api_key = kwargs.get("api_key", None)
+            if (
+                dynamic_api_key is not None
+                and potential_model_client is not None
+                and dynamic_api_key != potential_model_client.api_key
+            ):
+                model_client = None
+            else:
+                model_client = potential_model_client
+            self.total_calls[model_name] += 1
+
+            timeout = (
+                data.get(
+                    "timeout", None
+                )  # timeout set on litellm_params for this deployment
+                or self.timeout  # timeout set on router
+                or kwargs.get(
+                    "timeout", None
+                )  # this uses default_litellm_params when nothing is set
+            )
+
+            response = await litellm.arerank(
+                **{
+                    **data,
+                    "caching": self.cache_responses,
+                    "client": model_client,
+                    "timeout": timeout,
+                    **kwargs,
+                }
+            )
+
+            self.success_calls[model_name] += 1
+            verbose_router_logger.info(
+                f"litellm.arerank(model={model_name})\033[32m 200 OK\033[0m"
+            )
+            return response
+        except Exception as e:
+            verbose_router_logger.info(
+                f"litellm.arerank(model={model_name})\033[31m Exception {str(e)}\033[0m"
+            )
+            if model_name is not None:
+                self.fail_calls[model_name] += 1
+            raise e
+
     def text_completion(
         self,
         model: str,

tests/otel_tests/test_rerank.py

@@ -0,0 +1,65 @@
+import pytest
+import asyncio
+import aiohttp, openai
+from openai import OpenAI, AsyncOpenAI
+from typing import Optional, List, Union
+import uuid
+
+
+async def make_rerank_curl_request(
+    session,
+    key,
+    query,
+    documents,
+    model="rerank-english-v3.0",
+    top_n=3,
+):
+    url = "http://0.0.0.0:4000/rerank"
+    headers = {
+        "Authorization": f"Bearer {key}",
+        "Content-Type": "application/json",
+    }
+
+    data = {
+        "model": model,
+        "query": query,
+        "documents": documents,
+        "top_n": top_n,
+    }
+
+    async with session.post(url, headers=headers, json=data) as response:
+        status = response.status
+        response_text = await response.text()
+
+        if status != 200:
+            raise Exception(response_text)
+
+        return await response.json()
+
+
+@pytest.mark.asyncio
+async def test_basic_rerank_on_proxy():
+    """
+    Test litellm.rerank() on proxy
+
+    This SHOULD NOT call the pass through endpoints :)
+    """
+    async with aiohttp.ClientSession() as session:
+        docs = [
+            "Carson City is the capital city of the American state of Nevada.",
+            "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
+            "Washington, D.C. is the capital of the United States.",
+            "Capital punishment has existed in the United States since before it was a country.",
+        ]
+
+        try:
+            response = await make_rerank_curl_request(
+                session,
+                "sk-1234",
+                query="What is the capital of the United States?",
+                documents=docs,
+            )
+            print("response=", response)
+        except Exception as e:
+            print(e)
+            pytest.fail("Rerank request failed")