Merge da60124cfe into 08f2a98250

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:

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)