mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
262 lines
9.1 KiB
Python
262 lines
9.1 KiB
Python
"""
|
|
Tests for IBM watsonx.ai rerank transformation functionality.
|
|
"""
|
|
|
|
import json
|
|
import re
|
|
import uuid
|
|
from unittest.mock import MagicMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from litellm.llms.watsonx.common_utils import (
|
|
WatsonXAIError,
|
|
)
|
|
from litellm.llms.watsonx.rerank.transformation import IBMWatsonXRerankConfig
|
|
from litellm.types.rerank import RerankResponse
|
|
|
|
|
|
class TestIBMWatsonXRerankTransform:
|
|
def setup_method(self):
|
|
self.config = IBMWatsonXRerankConfig()
|
|
self.model = "watsonx/cross-encoder/ms-marco-minilm-l-12-v2"
|
|
|
|
def test_get_complete_url(self):
|
|
"""Test URL generation for IBM watsonx.ai rerank API."""
|
|
|
|
api_base = "https://us-south.ml.cloud.ibm.com"
|
|
model = "watsonx/cross-encoder/ms-marco-minilm-l-12-v2"
|
|
url = self.config.get_complete_url(api_base, model)
|
|
assert (
|
|
url
|
|
== "https://us-south.ml.cloud.ibm.com/ml/v1/text/rerank?version=2024-03-13"
|
|
)
|
|
|
|
def test_map_cohere_rerank_params_basic(self):
|
|
"""Test basic parameter mapping for IBM watsonx.ai rerank."""
|
|
params = self.config.map_cohere_rerank_params(
|
|
non_default_params={
|
|
"query": "hello",
|
|
"documents": ["hello", "world"],
|
|
"top_n": 2,
|
|
"return_documents": True,
|
|
"max_tokens_per_doc": 100,
|
|
},
|
|
model="test",
|
|
drop_params=False,
|
|
query="hello",
|
|
documents=["hello", "world"],
|
|
)
|
|
assert params["query"] == "hello"
|
|
assert params["inputs"] == [{"text": "hello"}, {"text": "world"}]
|
|
assert params["parameters"]["return_options"]["top_n"] == 2
|
|
assert params["parameters"]["return_options"]["inputs"] is True
|
|
assert params["parameters"]["truncate_input_tokens"] == 100
|
|
|
|
def test_transform_rerank_request(self):
|
|
"""Test request transformation for IBM watsonx.ai format."""
|
|
optional_params = {
|
|
"query": "What is the capital of France?",
|
|
"documents": [
|
|
"Paris is the capital of France.",
|
|
"France is a country in Europe.",
|
|
],
|
|
"top_n": 2,
|
|
"return_documents": True,
|
|
"project_id": uuid.uuid4(),
|
|
}
|
|
|
|
request_body = self.config.transform_rerank_request(
|
|
model="cross-encoder/ms-marco-minilm-l-12-v2",
|
|
optional_rerank_params=optional_params,
|
|
headers={},
|
|
)
|
|
|
|
assert request_body["model_id"] == "cross-encoder/ms-marco-minilm-l-12-v2"
|
|
assert request_body["project_id"] is not None
|
|
assert request_body["query"] == "What is the capital of France?"
|
|
assert request_body["documents"] == optional_params["documents"]
|
|
assert request_body["top_n"] == 2
|
|
assert request_body["return_documents"] is True
|
|
|
|
def test_transform_rerank_response_success(self):
|
|
"""Test successful response transformation."""
|
|
# Mock IBM watsonx.ai response format
|
|
response_data = {
|
|
"model_id": self.model,
|
|
"results": [
|
|
{
|
|
"index": 0,
|
|
"score": 6.53515625,
|
|
"input": {
|
|
"text": "Python is great for beginners due to simple syntax."
|
|
},
|
|
},
|
|
{
|
|
"index": 1,
|
|
"score": -7.1875,
|
|
"input": {"text": "JavaScript runs in browsers and is versatile."},
|
|
},
|
|
],
|
|
"input_token_count": 62,
|
|
}
|
|
|
|
# Create mock httpx response
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.json.return_value = response_data
|
|
mock_response.status_code = 200
|
|
mock_response.headers = {}
|
|
|
|
# Create mock logging object
|
|
mock_logging = MagicMock()
|
|
|
|
model_response = RerankResponse()
|
|
|
|
result = self.config.transform_rerank_response(
|
|
model=self.model,
|
|
raw_response=mock_response,
|
|
model_response=model_response,
|
|
logging_obj=mock_logging,
|
|
)
|
|
|
|
assert uuid.UUID(result.id).version == 4
|
|
assert len(result.results) == 2
|
|
assert result.results[0]["index"] == 0
|
|
assert result.results[0]["relevance_score"] == 6.53515625
|
|
assert (
|
|
result.results[0]["document"]["text"]
|
|
== "Python is great for beginners due to simple syntax."
|
|
)
|
|
assert result.results[1]["index"] == 1
|
|
assert result.results[1]["relevance_score"] == -7.1875
|
|
assert (
|
|
result.results[1]["document"]["text"]
|
|
== "JavaScript runs in browsers and is versatile."
|
|
)
|
|
|
|
# # Verify metadata
|
|
assert result.meta["tokens"]["input_tokens"] == 62
|
|
|
|
def test_transform_rerank_response_without_documents(self):
|
|
"""Test response transformation when return_documents is False."""
|
|
response_data = {
|
|
"model_id": self.model,
|
|
"results": [
|
|
{
|
|
"index": 0,
|
|
"score": 6.53515625,
|
|
},
|
|
{
|
|
"index": 1,
|
|
"score": -7.1875,
|
|
},
|
|
],
|
|
"input_token_count": 62,
|
|
}
|
|
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.json.return_value = response_data
|
|
mock_response.status_code = 200
|
|
mock_response.headers = {}
|
|
|
|
mock_logging = MagicMock()
|
|
model_response = RerankResponse()
|
|
|
|
result = self.config.transform_rerank_response(
|
|
model=self.model,
|
|
raw_response=mock_response,
|
|
model_response=model_response,
|
|
logging_obj=mock_logging,
|
|
)
|
|
|
|
assert uuid.UUID(result.id).version == 4
|
|
assert len(result.results) == 2
|
|
|
|
assert result.results[0]["index"] == 0
|
|
assert result.results[0]["relevance_score"] == 6.53515625
|
|
assert "document" not in result.results[0]
|
|
|
|
assert result.results[1]["index"] == 1
|
|
assert result.results[1]["relevance_score"] == -7.1875
|
|
assert "document" not in result.results[1]
|
|
|
|
def test_transform_rerank_response_missing_results(self):
|
|
"""Test that missing results raises ValueError."""
|
|
response_data = {
|
|
"model": self.model,
|
|
"usage": {"total_tokens": 10},
|
|
}
|
|
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.json.return_value = response_data
|
|
mock_response.status_code = 200
|
|
mock_response.headers = {}
|
|
|
|
mock_logging = MagicMock()
|
|
model_response = RerankResponse()
|
|
|
|
expected_error_msg = re.escape("No results found")
|
|
|
|
with pytest.raises(ValueError, match=expected_error_msg):
|
|
self.config.transform_rerank_response(
|
|
model=self.model,
|
|
raw_response=mock_response,
|
|
model_response=model_response,
|
|
logging_obj=mock_logging,
|
|
)
|
|
|
|
def test_transform_rerank_response_invalid_json(self):
|
|
"""Test error handling for invalid JSON response."""
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "doc", 0)
|
|
mock_response.text = "Invalid JSON response"
|
|
mock_response.status_code = 500
|
|
mock_response.headers = {}
|
|
|
|
mock_logging = MagicMock()
|
|
model_response = RerankResponse()
|
|
|
|
expected_error_msg = re.escape("Failed to parse response")
|
|
|
|
with pytest.raises(Exception, match=expected_error_msg):
|
|
self.config.transform_rerank_response(
|
|
model=self.model,
|
|
raw_response=mock_response,
|
|
model_response=model_response,
|
|
logging_obj=mock_logging,
|
|
)
|
|
|
|
def test_transform_rerank_response_without_id_stamps_a_fresh_id_per_call(self):
|
|
response_data = {
|
|
"model_id": self.model,
|
|
"results": [{"index": 0, "score": 1.5}],
|
|
"input_token_count": 12,
|
|
}
|
|
|
|
def transform() -> str:
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.json.return_value = response_data
|
|
mock_response.status_code = 200
|
|
mock_response.headers = {}
|
|
return self.config.transform_rerank_response(
|
|
model=self.model,
|
|
raw_response=mock_response,
|
|
model_response=RerankResponse(),
|
|
logging_obj=MagicMock(),
|
|
).id
|
|
|
|
first, second = transform(), transform()
|
|
|
|
assert first != second
|
|
assert self.model not in (first, second)
|
|
|
|
def test_get_supported_cohere_rerank_params(self):
|
|
"""Test getting supported parameters for IBM watsonx.ai rerank."""
|
|
supported_params = self.config.get_supported_cohere_rerank_params(self.model)
|
|
assert "query" in supported_params
|
|
assert "documents" in supported_params
|
|
assert "top_n" in supported_params
|
|
assert "return_documents" in supported_params
|
|
assert "max_tokens_per_doc" in supported_params
|
|
assert len(supported_params) == 5
|