tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-14 14:26:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: cleanup huggingface rerank transformation

Browse Source
This commit is contained in:
Krrish Dholakia
2025-06-06 10:30:44 -07:00
parent b70574017c
commit 1af0b62578
2 changed files with 203 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/huggingface/rerank/transformation.py
+81 -85
View File
@@ -1,10 +1,13 @@
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, TypedDict
import uuid
import os
import uuid
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, TypedDict, Union
import httpx
import litellm
import litellm
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
OptionalRerankParams,
RerankBilledUnits,
@@ -14,16 +17,13 @@ from litellm.types.rerank import (
RerankResponseResult,
RerankTokens,
)
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.utils import token_counter
from ..common_utils import HuggingFaceError
from ..common_utils import HuggingFaceError
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
LoggingClass = LiteLLMLoggingObj
else:
LoggingClass = Any
@@ -31,6 +31,7 @@ else:
class HuggingFaceRerankResponseItem(TypedDict):
"""Type definition for HuggingFace rerank API response items."""
index: int
score: float
text: Optional[str] # Optional, included when return_text=True
@@ -38,6 +39,7 @@ class HuggingFaceRerankResponseItem(TypedDict):
class HuggingFaceRerankResponse(TypedDict):
"""Type definition for HuggingFace rerank API complete response."""
# The response is a list of HuggingFaceRerankResponseItem
pass
@@ -63,12 +65,12 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
"""
# Get base URL from api_base or default
base_url = self.get_api_base(model=model, api_base=api_base)
# Remove trailing slashes and ensure we have the /rerank endpoint
base_url = base_url.rstrip("/")
if not base_url.endswith("/rerank"):
base_url = f"{base_url}/rerank"
return base_url
def get_supported_cohere_rerank_params(self, model: str) -> list:
@@ -76,10 +78,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
"query",
"documents",
"top_n",
"rank_fields",
"return_documents",
"max_chunks_per_doc",
"max_tokens_per_doc",
]
def map_cohere_rerank_params(
@@ -96,15 +95,21 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
return OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
max_tokens_per_doc=max_tokens_per_doc,
)
optional_rerank_params = {}
if non_default_params is not None:
for k, v in non_default_params.items():
if k == "documents" and v is not None:
optional_rerank_params["texts"] = v
elif k == "return_documents" and v is not None and isinstance(v, bool):
optional_rerank_params["return_text"] = v
elif k == "top_n" and v is not None:
optional_rerank_params["top_n"] = v
elif k == "documents" and v is not None:
optional_rerank_params["texts"] = v
elif k == "query" and v is not None:
optional_rerank_params["query"] = v
return OptionalRerankParams(**optional_rerank_params)
def validate_environment(
self,
@@ -115,52 +120,42 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
) -> dict:
# Get API credentials
api_key, api_base = self.get_api_credentials(api_key=api_key, api_base=api_base)
default_headers = {
"accept": "application/json",
"content-type": "application/json",
}
if api_key:
default_headers["Authorization"] = f"Bearer {api_key}"
if "Authorization" in headers:
default_headers["Authorization"] = headers["Authorization"]
return {**default_headers, **headers}
def transform_rerank_request(
self,
model: str,
optional_rerank_params: OptionalRerankParams,
optional_rerank_params: Union[OptionalRerankParams, dict],
headers: dict,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for HuggingFace rerank")
if "documents" not in optional_rerank_params:
raise ValueError("documents 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
return_documents = optional_rerank_params.get("return_documents")
# Default to returning document content unless explicitly set to False
if return_documents is False:
return_text = False
else:
return_text = True
request_body = {
"query": optional_rerank_params["query"],
"texts": optional_rerank_params["documents"],
"raw_scores": False,
"return_text": return_text,
"truncate": False,
"truncation_direction": "Right",
}
if optional_rerank_params.get("top_n") is not None:
request_body["top_n"] = optional_rerank_params["top_n"]
request_body.update(optional_rerank_params)
return request_body
def transform_rerank_response(
@@ -178,20 +173,21 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
raw_response_json: HuggingFaceRerankResponseList = raw_response.json()
except Exception:
raise HuggingFaceError(
message=getattr(raw_response, 'text', str(raw_response)),
status_code=getattr(raw_response, 'status_code', 500)
message=getattr(raw_response, "text", str(raw_response)),
status_code=getattr(raw_response, "status_code", 500),
)
# Use standard litellm token counter for proper token estimation
input_text = request_data.get("query", "")
try:
# Calculate tokens for the raw response JSON string
response_text = str(raw_response_json)
estimated_output_tokens = token_counter(model=model, text=response_text)
# Calculate input tokens from query and documents
query = request_data.get("query", "")
documents = request_data.get("texts", [])
# Convert documents to string if they're not already
documents_text = ""
for doc in documents:
@@ -199,63 +195,65 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
documents_text += doc + " "
elif isinstance(doc, dict) and "text" in doc:
documents_text += doc["text"] + " "
# Calculate input tokens using the same model
input_text = query + " " + documents_text
estimated_input_tokens = token_counter(model=model, text=input_text)
except Exception:
# Fallback to reasonable estimates if token counting fails
estimated_output_tokens = len(raw_response_json) * 10 if raw_response_json else 10
estimated_input_tokens = len(input_text) * 4 if 'input_text' in locals() else 0
estimated_output_tokens = (
len(raw_response_json) * 10 if raw_response_json else 10
)
estimated_input_tokens = (
len(input_text) * 4 if "input_text" in locals() else 0
)
_billed_units = RerankBilledUnits(search_units=1)
_tokens = RerankTokens(
input_tokens=estimated_input_tokens,
output_tokens=estimated_output_tokens
input_tokens=estimated_input_tokens, output_tokens=estimated_output_tokens
)
rerank_meta = RerankResponseMeta(
api_version={"version": "1.0"},
billed_units=_billed_units,
tokens=_tokens
api_version={"version": "1.0"}, billed_units=_billed_units, tokens=_tokens
)
# Check if documents should be returned based on request parameters
should_return_documents = request_data.get("return_text", False) or request_data.get("return_documents", False)
should_return_documents = request_data.get(
"return_text", False
) or request_data.get("return_documents", False)
original_documents = request_data.get("texts", [])
results = []
for item in raw_response_json:
# Extract required fields with defaults to handle None values
index = item.get("index")
score = item.get("score")
# Skip items that don't have required fields
if index is None or score is None:
continue
# Create RerankResponseResult with required fields
result = RerankResponseResult(
index=index,
relevance_score=score
)
result = RerankResponseResult(index=index, relevance_score=score)
# Add optional document field if needed
if should_return_documents:
text_content = item.get("text", "")
# 1. First try to use text returned directly from API if available
if text_content:
result["document"] = RerankResponseDocument(text=text_content)
# 2. If no text in API response but original documents are available, use those
elif original_documents and 0 <= item.get("index", -1) < len(original_documents):
elif original_documents and 0 <= item.get("index", -1) < len(
original_documents
):
doc = original_documents[item.get("index")]
if isinstance(doc, str):
result["document"] = RerankResponseDocument(text=doc)
elif isinstance(doc, dict) and "text" in doc:
result["document"] = RerankResponseDocument(text=doc["text"])
results.append(result)
return RerankResponse(
id=str(uuid.uuid4()),
results=results,
@@ -271,28 +269,26 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
self,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> tuple[Optional[str], Optional[str]]:
) -> Tuple[Optional[str], Optional[str]]:
"""
Get API key and base URL from multiple sources.
Returns tuple of (api_key, api_base).
Parameters:
api_key: API key provided directly to this function, takes precedence over all other sources
api_base: API base provided directly to this function, takes precedence over all other sources
"""
# Get API key from multiple sources
final_api_key = (
api_key or
litellm.huggingface_key or
get_secret_str("HUGGINGFACE_API_KEY")
api_key or litellm.huggingface_key or get_secret_str("HUGGINGFACE_API_KEY")
)
# Get API base from multiple sources
# Get API base from multiple sources
final_api_base = (
api_base or
litellm.api_base or
get_secret_str("HF_API_BASE") or
get_secret_str("HUGGINGFACE_API_BASE")
api_base
or litellm.api_base
or get_secret_str("HF_API_BASE")
or get_secret_str("HUGGINGFACE_API_BASE")
)
return final_api_key, final_api_base
return final_api_key, final_api_base
tests/test_litellm/llms/huggingface/test_rerank.py → tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py
+122 -132
View File
@@ -4,19 +4,20 @@ Based on the test patterns from other rerank providers and the current HuggingFa
"""
import asyncio
import json
from unittest.mock import patch, MagicMock, AsyncMock
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
def assert_response_shape(response, custom_llm_provider):
"""Helper function to validate response structure"""
assert hasattr(response, 'id')
assert hasattr(response, 'results')
assert hasattr(response, 'meta')
assert hasattr(response, "id")
assert hasattr(response, "results")
assert hasattr(response, "meta")
assert isinstance(response.results, list)
for result in response.results:
assert "index" in result
assert "relevance_score" in result
@@ -25,21 +26,18 @@ def assert_response_shape(response, custom_llm_provider):
@pytest.mark.parametrize("sync_mode", [True, False])
@patch('litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post')
@patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post')
@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_basic_rerank_huggingface(mock_sync_post, mock_async_post, sync_mode):
"""Test basic HuggingFace rerank functionality."""
# Mock response data that matches HuggingFace rerank API format
mock_response_data = [
{"index": 0, "score": 0.9},
{"index": 1, "score": 0.1}
]
mock_response_data = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}]
def return_val():
return mock_response_data
api_key = "test_hf_api_key"
if sync_mode:
# Create mock response object for sync
mock_response = MagicMock()
@@ -47,7 +45,7 @@ def test_basic_rerank_huggingface(mock_sync_post, mock_async_post, sync_mode):
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_sync_post.return_value = mock_response
response = litellm.rerank(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
@@ -59,52 +57,51 @@ def test_basic_rerank_huggingface(mock_sync_post, mock_async_post, sync_mode):
else:
# Create mock response object for async
mock_response = AsyncMock()
def return_val():
return mock_response_data
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_async_post.return_value = mock_response
response = asyncio.run(litellm.arerank(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
documents=["hello", "world"],
top_n=2,
api_key=api_key,
))
response = asyncio.run(
litellm.arerank(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
documents=["hello", "world"],
top_n=2,
api_key=api_key,
)
)
mock_async_post.assert_called_once()
assert response.results is not None
assert len(response.results) == 2
assert response.results[0]["index"] == 0
assert response.results[0]["relevance_score"] == 0.9
assert_response_shape(response, custom_llm_provider="huggingface")
@pytest.mark.parametrize("sync_mode", [True, False])
@patch('litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post')
@patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post')
@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_huggingface_rerank_custom_api_base(mock_sync_post, mock_async_post, sync_mode):
"""Test HuggingFace rerank with custom API base."""
mock_response_data = [
{"index": 0, "score": 0.9},
{"index": 1, "score": 0.1}
]
mock_response_data = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}]
def return_val():
return mock_response_data
if sync_mode:
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_sync_post.return_value = mock_response
response = litellm.rerank(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
@@ -113,7 +110,7 @@ def test_huggingface_rerank_custom_api_base(mock_sync_post, mock_async_post, syn
api_base="https://my-custom-hf-endpoint.com",
api_key="test_api_key",
)
mock_sync_post.assert_called_once()
call_url = mock_sync_post.call_args.kwargs["url"]
assert "my-custom-hf-endpoint.com" in call_url
@@ -121,24 +118,26 @@ def test_huggingface_rerank_custom_api_base(mock_sync_post, mock_async_post, syn
assert len(response.results) == 2
else:
mock_response = AsyncMock()
def return_val():
return mock_response_data
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_async_post.return_value = mock_response
response = asyncio.run(litellm.arerank(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
documents=["hello", "world"],
top_n=2,
api_base="https://my-custom-hf-endpoint.com",
api_key="test_api_key",
))
response = asyncio.run(
litellm.arerank(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
documents=["hello", "world"],
top_n=2,
api_base="https://my-custom-hf-endpoint.com",
api_key="test_api_key",
)
)
mock_async_post.assert_called_once()
call_url = mock_async_post.call_args.kwargs["url"]
assert "my-custom-hf-endpoint.com" in call_url
@@ -146,61 +145,58 @@ def test_huggingface_rerank_custom_api_base(mock_sync_post, mock_async_post, syn
assert len(response.results) == 2
@patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post')
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_huggingface_rerank_with_env_vars(mock_post, monkeypatch):
"""Test HuggingFace rerank with environment variable configuration."""
monkeypatch.setenv("HUGGINGFACE_API_KEY", "env_test_key")
monkeypatch.setenv("HUGGINGFACE_API_BASE", "https://env-hf-endpoint.com")
mock_response_data = [
{"index": 0, "score": 0.9},
{"index": 1, "score": 0.1}
]
mock_response_data = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}]
def return_val():
return mock_response_data
mock_response = 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(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
documents=["hello", "world"],
top_n=2,
)
mock_post.assert_called_once()
call_url = mock_post.call_args.kwargs["url"]
assert "env-hf-endpoint.com" in call_url
headers = mock_post.call_args.kwargs.get("headers", {})
assert "env_test_key" in str(headers.get("Authorization", ""))
assert response.results is not None
assert len(response.results) == 2
@patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post')
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_huggingface_rerank_return_documents(mock_post):
"""Test HuggingFace rerank with return_documents=True."""
mock_response_data = [
{"index": 0, "score": 0.9, "text": "hello"},
{"index": 1, "score": 0.1, "text": "world"}
{"index": 1, "score": 0.1, "text": "world"},
]
def return_val():
return mock_response_data
mock_response = 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(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
@@ -209,11 +205,11 @@ def test_huggingface_rerank_return_documents(mock_post):
return_documents=True,
api_key="test_api_key",
)
mock_post.assert_called_once()
request_data = json.loads(mock_post.call_args.kwargs["data"])
assert request_data.get("return_text") is True
assert response.results is not None
assert len(response.results) == 2
# Check that documents are included in response
@@ -222,18 +218,19 @@ def test_huggingface_rerank_return_documents(mock_post):
assert "text" in result["document"]
@patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post')
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_huggingface_rerank_error_handling(mock_post):
"""Test HuggingFace rerank error handling."""
def return_val():
return {"error": "Unauthorized"}
mock_response = MagicMock()
mock_response.status_code = 401
mock_response.json = return_val
mock_response.text = "Unauthorized"
mock_post.return_value = mock_response
with pytest.raises(Exception):
litellm.rerank(
model="huggingface/BAAI/bge-reranker-base",
@@ -247,59 +244,61 @@ def test_huggingface_rerank_error_handling(mock_post):
def test_huggingface_rerank_config():
"""Test HuggingFaceRerankConfig class functionality."""
from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig
config = 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://api-inference.huggingface.co/rerank"
)
# Test custom API base
custom_url = config.get_complete_url("https://custom.huggingface.co", "test")
assert custom_url == "https://custom.huggingface.co/rerank"
# Test supported parameters
supported_params = 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
# Test parameter mapping
params = config.map_cohere_rerank_params(
non_default_params={},
non_default_params={
"query": "hello",
"documents": ["hello", "world"],
"top_n": 2,
"return_documents": True,
},
model="test",
drop_params=False,
query="hello",
documents=["hello", "world"],
top_n=2,
return_documents=True,
)
print(f"params: {params}")
assert params["query"] == "hello"
assert params["documents"] == ["hello", "world"]
assert params["texts"] == ["hello", "world"]
assert params["top_n"] == 2
assert params["return_documents"] is True
assert params["return_text"] is True
def test_request_transformation():
"""Test request transformation logic."""
from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig
from litellm.types.rerank import OptionalRerankParams
config = HuggingFaceRerankConfig()
optional_params = OptionalRerankParams(
query="hello",
documents=["hello", "world"],
top_n=2,
return_documents=True
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={}
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
@@ -313,39 +312,39 @@ def test_response_transformation():
"""Test response transformation logic."""
from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig
from litellm.types.rerank import RerankResponse
config = HuggingFaceRerankConfig()
# Mock HuggingFace response
hf_response_data = [
{"index": 0, "score": 0.9, "text": "hello"},
{"index": 1, "score": 0.1, "text": "world"}
{"index": 1, "score": 0.1, "text": "world"},
]
def return_val():
return hf_response_data
# Create mock httpx response
mock_response = MagicMock()
mock_response.json = return_val
model_response = RerankResponse()
transformed_response = config.transform_rerank_response(
model="test",
raw_response=mock_response,
model_response=model_response,
logging_obj=None,
request_data={"return_text": True}
request_data={"return_text": True},
)
assert transformed_response.results is not None
assert len(transformed_response.results) == 2
assert transformed_response.results[0]["index"] == 0
assert transformed_response.results[0]["relevance_score"] == 0.9
assert transformed_response.results[1]["index"] == 1
assert transformed_response.results[1]["relevance_score"] == 0.1
# Check documents are included when return_text is True
for result in transformed_response.results:
if "document" in result:
@@ -355,70 +354,61 @@ def test_response_transformation():
def test_validate_environment():
"""Test environment validation logic."""
from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig
config = HuggingFaceRerankConfig()
# Test with API key
headers = config.validate_environment(
headers={},
model="test",
api_key="test_key"
)
headers = config.validate_environment(headers={}, model="test", api_key="test_key")
assert "Authorization" in headers
assert "Bearer test_key" in headers["Authorization"]
assert headers["accept"] == "application/json"
assert headers["content-type"] == "application/json"
# Test headers override
custom_headers = {"custom": "header"}
headers = config.validate_environment(
headers=custom_headers,
model="test",
api_key="test_key"
headers=custom_headers, model="test", api_key="test_key"
)
assert "custom" in headers
assert headers["custom"] == "header"
@patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post')
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_huggingface_rerank_request_payload(mock_post):
"""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 = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}]
def return_val():
return mock_response_data
mock_response = 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(
model="huggingface/BAAI/bge-reranker-base",
query="hello",
documents=["hello", "world"],
api_key="test_api_key",
top_n=2,
return_documents=True,
api_key="test_api_key",
)
mock_post.assert_called_once()
# Verify URL
call_url = mock_post.call_args.kwargs["url"]
assert call_url == "https://api-inference.huggingface.co/rerank"
# Verify headers
headers = 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 = {
@@ -428,11 +418,11 @@ def test_huggingface_rerank_request_payload(mock_post):
"return_text": True,
"truncate": False,
"truncation_direction": "Right",
"top_n": 2
"top_n": 2,
}
for key, value in expected_request.items():
assert request_data[key] == value
assert response.results is not None
assert len(response.results) == 2