diff --git a/tests/integration/providers/test_huggingface_rerank_truncate_wire.py b/tests/integration/providers/test_huggingface_rerank_truncate_wire.py new file mode 100644 index 00000000000..0d29dbf0ded --- /dev/null +++ b/tests/integration/providers/test_huggingface_rerank_truncate_wire.py @@ -0,0 +1,74 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway +from integration._support.wire import Reply, Request, wire_server + +MODEL: Final = "huggingface/BAAI/bge-reranker-base" +QUERY: Final = "what is a fox" +TEXT: Final = "a short doc about foxes" +RESPONSE: Final = json.dumps([{"index": 0, "score": 0.82}]).encode() +TRUNCATION_CASES: Final = ( + pytest.param( + MappingProxyType({}), + MappingProxyType({"truncation_direction": "Right"}), + id="omitted-truncate-uses-server-default", + ), + pytest.param( + MappingProxyType({"truncate": True, "truncation_direction": "Left"}), + MappingProxyType({"truncate": True, "truncation_direction": "Left"}), + id="explicit-true-and-left-are-forwarded", + ), + pytest.param( + MappingProxyType({"truncate": False}), + MappingProxyType({"truncate": False, "truncation_direction": "Right"}), + id="explicit-false-is-forwarded", + ), +) + + +def rerank_peer(request: Request, expected_body: Mapping[str, object]) -> Reply: + assert request.method == "POST" and request.target == "/rerank", request.target + assert request.headers["authorization"] == "Bearer integration-provider-key" + body: Final = JSON_OBJECT.validate_json(request.body) + assert body == expected_body, body + return Reply(body=RESPONSE) + + +@pytest.mark.parametrize(("request_params", "expected_truncation_params"), TRUNCATION_CASES) +def test_huggingface_rerank_forwards_truncation_options( + gateway: Gateway, + request_params: Mapping[str, bool | str], + expected_truncation_params: Mapping[str, bool | str], +) -> None: + expected_body: Final = { + "raw_scores": False, + "truncation_direction": "Right", + "query": QUERY, + "texts": [TEXT], + "return_text": False, + **expected_truncation_params, + } + with wire_server(lambda request: rerank_peer(request, expected_body)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=MODEL, + api_base=wire.url, + model_info={"mode": "rerank"}, + ) + response: Final = gateway.request( + "POST", + "/v1/rerank", + { + "model": model, + "query": QUERY, + "documents": [TEXT], + "return_documents": False, + **request_params, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["results"] == [{"index": 0, "relevance_score": 0.82}], response.text + assert len(wire.drain()) == 1, "Expected exactly one provider rerank call" 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 6fd2b006fef..2361b113433 100644 --- a/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py +++ b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py @@ -5,6 +5,7 @@ Based on the test patterns from other rerank providers and the current HuggingFa import asyncio import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -219,11 +220,11 @@ def test_huggingface_rerank_return_documents(mock_post): assert "text" in result["document"] -def test_huggingface_rerank_config(): +def test_huggingface_rerank_config() -> None: """Test HuggingFaceRerankConfig class functionality.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig - config = HuggingFaceRerankConfig() + config: Final = HuggingFaceRerankConfig() # Test complete URL generation assert config.get_complete_url(None, "test") == "https://api-inference.huggingface.co/rerank" @@ -233,14 +234,16 @@ def test_huggingface_rerank_config(): assert custom_url == "https://custom.huggingface.co/rerank" # Test supported parameters - supported_params = config.get_supported_cohere_rerank_params("test") + supported_params: Final = config.get_supported_cohere_rerank_params("test") assert "query" in supported_params assert "documents" in supported_params assert "top_n" in supported_params assert "return_documents" in supported_params + assert "truncate" in supported_params + assert "truncation_direction" in supported_params # Test parameter mapping - params = config.map_cohere_rerank_params( + params: Final = config.map_cohere_rerank_params( non_default_params={ "query": "hello", "documents": ["hello", "world"], @@ -259,24 +262,28 @@ def test_huggingface_rerank_config(): assert params["return_text"] is True -def test_request_transformation(): +def test_request_transformation() -> None: """Test request transformation logic.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig from litellm.types.rerank import OptionalRerankParams - config = HuggingFaceRerankConfig() + config: Final = HuggingFaceRerankConfig() - optional_params = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) + optional_params: Final = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) - request_body = config.transform_rerank_request(model="test", optional_rerank_params=optional_params, headers={}) + request_body: Final = config.transform_rerank_request( + model="test", optional_rerank_params=optional_params, headers={} + ) - assert request_body["query"] == "hello" - assert request_body["texts"] == ["hello", "world"] - assert request_body["top_n"] == 2 - assert request_body["return_text"] is True - assert request_body["raw_scores"] is False - assert request_body["truncate"] is False - assert request_body["truncation_direction"] == "Right" + expected_request: Final = { + "query": "hello", + "texts": ["hello", "world"], + "top_n": 2, + "return_text": True, + "raw_scores": False, + "truncation_direction": "Right", + } + assert request_body == expected_request def test_response_transformation(): @@ -345,20 +352,20 @@ def test_validate_environment(): @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_huggingface_rerank_request_payload(mock_post): +def test_huggingface_rerank_request_payload(mock_post: MagicMock) -> None: """Test that the request payload is correctly formatted for HuggingFace API.""" - mock_response_data = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}] + mock_response_data: Final = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}] def return_val(): return mock_response_data - mock_response = MagicMock() + mock_response: Final = MagicMock() mock_response.json = return_val mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} mock_post.return_value = mock_response - response = litellm.rerank( + response: Final = litellm.rerank( model="huggingface/BAAI/bge-reranker-base", query="hello", documents=["hello", "world"], @@ -370,28 +377,60 @@ def test_huggingface_rerank_request_payload(mock_post): mock_post.assert_called_once() # Verify URL - call_url = mock_post.call_args.kwargs["url"] + call_url: Final = mock_post.call_args.kwargs["url"] assert call_url == "https://api-inference.huggingface.co/rerank" # Verify headers - headers = mock_post.call_args.kwargs["headers"] + headers: Final = mock_post.call_args.kwargs["headers"] assert "Bearer test_api_key" in headers["Authorization"] assert headers["content-type"] == "application/json" # Verify request body - request_data = json.loads(mock_post.call_args.kwargs["data"]) - expected_request = { + request_data: Final = json.loads(mock_post.call_args.kwargs["data"]) + expected_request: Final = { "query": "hello", "texts": ["hello", "world"], "raw_scores": False, "return_text": True, - "truncate": False, "truncation_direction": "Right", "top_n": 2, } - for key, value in expected_request.items(): - assert request_data[key] == value + assert request_data == expected_request assert response.results is not None assert len(response.results) == 2 + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") +def test_huggingface_rerank_forwards_truncate_options(mock_post: MagicMock) -> None: + mock_response: Final = MagicMock() + mock_response.json.return_value = [{"index": 0, "score": 0.9}] + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_post.return_value = mock_response + + response: Final = litellm.rerank( + model="huggingface/BAAI/bge-reranker-base", + query="hello", + documents=["hello", "world"], + api_key="test_api_key", + return_documents=False, + truncate=True, + truncation_direction="Left", + ) + + mock_post.assert_called_once() + + request_data: Final = json.loads(mock_post.call_args.kwargs["data"]) + expected_request: Final = { + "query": "hello", + "texts": ["hello", "world"], + "raw_scores": False, + "return_text": False, + "truncation_direction": "Left", + "truncate": True, + } + assert request_data == expected_request + assert response.results is not None + assert response.results[0]["relevance_score"] == 0.9