""" Tests for DeepInfra rerank transformation functionality. Based on the test patterns from other rerank providers and the current DeepInfra implementation. """ import json from unittest.mock import MagicMock import httpx import pytest from litellm.llms.deepinfra.rerank.transformation import DeepinfraRerankConfig from litellm.types.rerank import ( OptionalRerankParams, RerankResponse, ) class TestDeepinfraRerankTransform: def setup_method(self): self.config = DeepinfraRerankConfig() self.model = "deepinfra/Qwen/Qwen3-Reranker-0.6B" def test_get_complete_url(self): """Test URL generation for DeepInfra rerank API.""" # Test basic URL generation api_base = "https://api.deepinfra.com" model = "Qwen/Qwen3-Reranker-0.6B" url = self.config.get_complete_url(api_base, model) assert url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B" # Test URL with slash at the end api_base_with_slash = "https://api.deepinfra.com/" url = self.config.get_complete_url(api_base_with_slash, model) assert url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B" # Test URL with openai replacement api_base_openai = "https://api.deepinfra.com/openai" url = self.config.get_complete_url(api_base_openai, model) assert url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B" # Test error when api_base is None with pytest.raises(ValueError, match="Deepinfra API Base is required"): self.config.get_complete_url(None, model) def test_map_cohere_rerank_params_basic(self): """Test basic parameter mapping for DeepInfra rerank.""" params = self.config.map_cohere_rerank_params( non_default_params={"documents": ["doc1", "doc2"]}, model=self.model, drop_params=False, query="test query", documents=["doc1", "doc2"], ) assert params["queries"] == [ "test query", "test query", ] # DeepInfra requires queries to match documents length assert params["documents"] == ["doc1", "doc2"] def test_map_cohere_rerank_params_with_non_default(self): """Test parameter mapping with DeepInfra-specific parameters.""" non_default_params = { "queries": ["custom query"], "documents": ["doc1", "doc2", "doc3"], "service_tier": "premium", "instruction": "custom instruction", "webhook": "https://webhook.example.com", } params = self.config.map_cohere_rerank_params( non_default_params=non_default_params, model=self.model, drop_params=False, query="test query", documents=["doc1", "doc2"], ) # queries should override the query parameter (custom queries take precedence) assert params["queries"] == ["custom query"] assert params["documents"] == ["doc1", "doc2", "doc3"] assert params["service_tier"] == "premium" assert params["instruction"] == "custom instruction" assert params["webhook"] == "https://webhook.example.com" def test_transform_rerank_request(self): """Test request transformation for DeepInfra format.""" optional_params = OptionalRerankParams( queries=["test query"], documents=["doc1", "doc2"], service_tier="default", ) request_body = self.config.transform_rerank_request( model=self.model, optional_rerank_params=optional_params, headers={} ) assert request_body["queries"] == ["test query"] assert request_body["documents"] == ["doc1", "doc2"] assert request_body["service_tier"] == "default" def test_transform_rerank_request_missing_documents(self): """Test that transform_rerank_request handles missing documents gracefully.""" optional_params = OptionalRerankParams(queries=["test query"]) # The current implementation doesn't validate documents, it just returns the params result = self.config.transform_rerank_request( model=self.model, optional_rerank_params=optional_params, headers={} ) assert result == optional_params def test_transform_rerank_response_success(self): """Test successful response transformation.""" # Mock DeepInfra response format response_data = { "scores": [0.9, 0.7, 0.3], "input_tokens": 42, "request_id": "test-request-123", "inference_status": { "status": "success", "runtime_ms": 150, "cost": 0.0001, "tokens_generated": 0, "tokens_input": 42, }, } # Create mock httpx response mock_response = MagicMock(spec=httpx.Response) mock_response.json.return_value = response_data mock_response.text = json.dumps(response_data) # 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, ) # Verify response structure assert result.id == "test-request-123" assert len(result.results) == 3 assert result.results[0]["index"] == 0 assert result.results[0]["relevance_score"] == 0.9 assert result.results[1]["index"] == 1 assert result.results[1]["relevance_score"] == 0.7 assert result.results[2]["index"] == 2 assert result.results[2]["relevance_score"] == 0.3 # Verify metadata assert result.meta["tokens"]["input_tokens"] == 42 assert result.meta["tokens"]["output_tokens"] == 0 assert result.meta["billed_units"]["total_tokens"] == 42 # Verify hidden params assert result._hidden_params["status"] == "success" assert result._hidden_params["runtime_ms"] == 150 assert result._hidden_params["cost"] == 0.0001 assert result._hidden_params["tokens_generated"] == 0 assert result._hidden_params["tokens_input"] == 42 assert result._hidden_params["model"] == self.model # Verify logging was called mock_logging.post_call.assert_called_once_with( original_response=mock_response.text ) def test_transform_rerank_response_minimal(self): """Test response transformation with minimal data.""" response_data = { "scores": [0.8, 0.2], "input_tokens": 20, } mock_response = MagicMock(spec=httpx.Response) mock_response.json.return_value = response_data mock_response.text = json.dumps(response_data) 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, ) # Should generate UUID when request_id is missing assert result.id is not None assert len(result.id) > 0 # Should handle missing inference_status gracefully assert result._hidden_params["status"] == "unknown" assert result._hidden_params["runtime_ms"] == 0 assert result._hidden_params["cost"] == 0.0 def test_transform_rerank_response_error_fallback(self): """Test error handling and fallback in response transformation.""" # Create a response that will cause JSON parsing to fail mock_response = MagicMock(spec=httpx.Response) mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "doc", 0) mock_response.text = "Invalid JSON response" mock_logging = MagicMock() model_response = RerankResponse() # The current implementation should handle JSON parsing errors gracefully # by falling back to the parent implementation result = self.config.transform_rerank_response( model=self.model, raw_response=mock_response, model_response=model_response, logging_obj=mock_logging, ) # Should return the original model_response when fallback occurs assert result == model_response def test_get_supported_cohere_rerank_params(self): """Test getting supported parameters for DeepInfra rerank.""" supported_params = self.config.get_supported_cohere_rerank_params(self.model) assert "query" in supported_params assert "documents" in supported_params assert len(supported_params) == 2 def test_query_replication_for_deepinfra_requirement(self): """Test that queries are replicated to match documents length as required by DeepInfra.""" # Test with different document lengths test_cases = [ (["doc1"], ["query1"]), (["doc1", "doc2"], ["query1", "query1"]), (["doc1", "doc2", "doc3"], ["query1", "query1", "query1"]), ] for documents, expected_queries in test_cases: params = self.config.map_cohere_rerank_params( non_default_params={}, model=self.model, drop_params=False, query="query1", documents=documents, ) assert ( params["queries"] == expected_queries ), f"Failed for {len(documents)} documents" assert len(params["queries"]) == len( documents ), "Queries length must match documents length" def test_get_error_class_basic(self): """Test error class generation for basic error.""" error_message = "Authentication failed" status_code = 401 headers = {"content-type": "application/json"} with pytest.raises(Exception) as exc_info: self.config.get_error_class(error_message, status_code, headers) # The method should raise a BaseLLMException assert exc_info.value.args[0] == error_message def test_get_error_class_with_detail(self): """Test error class generation with DeepInfra error format.""" error_data = {"detail": {"error": "Model not found"}} error_message = json.dumps(error_data) status_code = 404 headers = {"content-type": "application/json"} with pytest.raises(Exception) as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should extract the nested error message assert "Model not found" in str(exc_info.value) def test_get_error_class_with_string_detail(self): """Test error class generation with string detail.""" error_data = {"detail": "Service unavailable"} error_message = json.dumps(error_data) status_code = 503 headers = {"content-type": "application/json"} with pytest.raises(Exception) as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should extract the string detail assert "Service unavailable" in str(exc_info.value) def test_get_error_class_invalid_json(self): """Test error class generation with invalid JSON.""" error_message = "Invalid JSON error message" status_code = 500 headers = {"content-type": "application/json"} with pytest.raises(Exception) as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should use the original error message when JSON parsing fails assert "Invalid JSON error message" in str(exc_info.value)