From 5f2f7aa7547a018b5ae2db4fc444ba30ca3cfe5a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 27 Aug 2024 17:36:40 -0700 Subject: [PATCH 1/2] feat - add rerank on proxy --- litellm/proxy/proxy_config.yaml | 13 +++-- litellm/router.py | 98 +++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index ef11d798e5d..536f6e2e57f 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/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 \ No newline at end of file diff --git a/litellm/router.py b/litellm/router.py index 2d180f07637..bc23972b63d 100644 --- a/litellm/router.py +++ b/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, From c27640e6e42063fcd23b020c18b6e2f668ba52f6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 27 Aug 2024 17:50:37 -0700 Subject: [PATCH 2/2] add /rerank test --- litellm/proxy/auth/user_api_key_auth.py | 2 + .../example_config_yaml/otel_test_config.yaml | 4 ++ tests/otel_tests/test_rerank.py | 65 +++++++++++++++++++ 3 files changed, 71 insertions(+) create mode 100644 tests/otel_tests/test_rerank.py diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index fde5d5ca510..9623267ee34 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/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 diff --git a/litellm/proxy/example_config_yaml/otel_test_config.yaml b/litellm/proxy/example_config_yaml/otel_test_config.yaml index a041a2bd0ce..7d8f6d4fe95 100644 --- a/litellm/proxy/example_config_yaml/otel_test_config.yaml +++ b/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 diff --git a/tests/otel_tests/test_rerank.py b/tests/otel_tests/test_rerank.py new file mode 100644 index 00000000000..47fe109e7f6 --- /dev/null +++ b/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")