diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 476e22fc088..fa67bfad732 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -32,6 +32,9 @@ else: LoggingClass = Any +ROUTER_BASE_URL: Final = "https://router.huggingface.co" + + class HuggingFaceRerankResponseItem(TypedDict): """Type definition for HuggingFace rerank API response items.""" @@ -59,7 +62,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): elif os.getenv("HUGGINGFACE_API_BASE") is not None: return os.getenv("HUGGINGFACE_API_BASE", "") else: - return "https://api-inference.huggingface.co" + return ROUTER_BASE_URL def get_complete_url( self, @@ -70,7 +73,17 @@ class HuggingFaceRerankConfig(BaseRerankConfig): """ Get the complete URL for the API call, including the /rerank suffix if necessary. """ - return ensure_api_route(self.get_api_base(model=model, api_base=api_base), "rerank") + if model.startswith(("http://", "https://")): + return model + + resolved_api_base: Final = self.get_api_base(model=model, api_base=api_base) + clean_model: Final = model.removeprefix("huggingface/") + base: Final = resolved_api_base.rstrip("/") + if base == ROUTER_BASE_URL: + return f"{ROUTER_BASE_URL}/hf-inference/models/{clean_model}" + if "/models/" in base or base.endswith("/rerank"): + return resolved_api_base + return ensure_api_route(resolved_api_base, "rerank") def get_supported_cohere_rerank_params(self, model: str) -> list: return [ @@ -153,9 +166,15 @@ class HuggingFaceRerankConfig(BaseRerankConfig): raise ValueError("query is required for HuggingFace rerank") if "texts" not in optional_rerank_params: raise ValueError("Cohere 'documents' param is required for HuggingFace rerank") - # Ensure return_text is a boolean value - # HuggingFace API expects return_text parameter, corresponding to our return_documents parameter + + query: Final = optional_rerank_params["query"] + texts: Final = optional_rerank_params["texts"] + request_body: Final = { + "inputs": { + "query": query, + "texts": texts, + }, "raw_scores": False, "truncation_direction": "Right", } @@ -183,6 +202,12 @@ class HuggingFaceRerankConfig(BaseRerankConfig): status_code=getattr(raw_response, "status_code", 500), ) + if isinstance(raw_response_json, dict) and "error" in raw_response_json: + raise HuggingFaceError( + message=str(raw_response_json["error"]), + status_code=getattr(raw_response, "status_code", 400), + ) + # Use standard litellm token counter for proper token estimation input_text = request_data.get("query", "") try: diff --git a/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py index 00a3cbfdb7f..f518274fd04 100644 --- a/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py +++ b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py @@ -229,7 +229,15 @@ def test_huggingface_rerank_config() -> None: config: Final = HuggingFaceRerankConfig() # Test complete URL generation - assert config.get_complete_url(None, "test") == "https://api-inference.huggingface.co/rerank" + assert config.get_complete_url(None, "test") == "https://router.huggingface.co/hf-inference/models/test" + assert ( + config.get_complete_url(None, "huggingface/BAAI/bge-reranker-base") + == "https://router.huggingface.co/hf-inference/models/BAAI/bge-reranker-base" + ) + assert ( + config.get_complete_url("https://router.huggingface.co", "BAAI/bge-reranker-base") + == "https://router.huggingface.co/hf-inference/models/BAAI/bge-reranker-base" + ) # Test custom API base custom_url = config.get_complete_url("https://custom.huggingface.co", "test") @@ -257,7 +265,6 @@ def test_huggingface_rerank_config() -> None: query="hello", documents=["hello", "world"], ) - print(f"params: {params}") assert params["query"] == "hello" assert params["texts"] == ["hello", "world"] assert params["top_n"] == 2 @@ -278,6 +285,10 @@ def test_request_transformation() -> None: ) expected_request: Final = { + "inputs": { + "query": "hello", + "texts": ["hello", "world"], + }, "query": "hello", "texts": ["hello", "world"], "top_n": 2, @@ -380,7 +391,7 @@ def test_huggingface_rerank_request_payload(mock_post: MagicMock) -> None: # Verify URL call_url: Final = mock_post.call_args.kwargs["url"] - assert call_url == "https://api-inference.huggingface.co/rerank" + assert call_url == "https://router.huggingface.co/hf-inference/models/BAAI/bge-reranker-base" # Verify headers headers: Final = mock_post.call_args.kwargs["headers"] @@ -390,6 +401,10 @@ def test_huggingface_rerank_request_payload(mock_post: MagicMock) -> None: # Verify request body request_data: Final = json.loads(mock_post.call_args.kwargs["data"]) expected_request: Final = { + "inputs": { + "query": "hello", + "texts": ["hello", "world"], + }, "query": "hello", "texts": ["hello", "world"], "raw_scores": False, @@ -406,6 +421,10 @@ def test_huggingface_rerank_request_payload(mock_post: MagicMock) -> None: def test_huggingface_rerank_forwards_truncate_options() -> None: expected_request: Final = { + "inputs": { + "query": "hello", + "texts": ["hello", "world"], + }, "query": "hello", "texts": ["hello", "world"], "raw_scores": False, @@ -415,7 +434,7 @@ def test_huggingface_rerank_forwards_truncate_options() -> None: } def respond(request: httpx.Request) -> httpx.Response: - assert request.url == "https://api-inference.huggingface.co/rerank" + assert str(request.url) == "https://router.huggingface.co/hf-inference/models/BAAI/bge-reranker-base" assert request.headers["Authorization"] == "Bearer test_api_key" assert json.loads(request.content) == expected_request return httpx.Response(200, json=[{"index": 0, "score": 0.9}]) @@ -434,3 +453,24 @@ def test_huggingface_rerank_forwards_truncate_options() -> None: assert response.results is not None assert response.results[0]["relevance_score"] == 0.9 + + +def test_huggingface_rerank_error_response_raises_huggingface_error() -> None: + from litellm.llms.huggingface.common_utils import HuggingFaceError + from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig + from litellm.types.rerank import RerankResponse + + config: Final = HuggingFaceRerankConfig() + mock_response = MagicMock() + mock_response.json = lambda: {"error": "Model BAAI/bge-reranker-base is currently loading"} + mock_response.status_code = 503 + + with pytest.raises(HuggingFaceError) as exc_info: + config.transform_rerank_response( + model="huggingface/BAAI/bge-reranker-base", + raw_response=mock_response, + model_response=RerankResponse(), + logging_obj=None, + ) + + assert "Model BAAI/bge-reranker-base is currently loading" in str(exc_info.value)