""" Tests for LiteLLMSDKCrossEncoder. Tests the LiteLLM SDK-based cross-encoder implementation for reranking. """ import os from unittest.mock import AsyncMock, MagicMock, patch import pytest from hindsight_api.engine.cross_encoder import LiteLLMSDKCrossEncoder, create_cross_encoder_from_env class TestLiteLLMSDKCrossEncoder: """Test suite for LiteLLMSDKCrossEncoder class.""" @pytest.mark.asyncio async def test_initialization_success(self): """Test successful initialization with valid config.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="deepinfra/Qwen3-reranker-8B", ) assert encoder.provider_name == "litellm-sdk" assert encoder.api_key == "test_key" assert encoder.model == "deepinfra/Qwen3-reranker-8B" assert encoder._initialized is False # Mock the litellm import mock_litellm = MagicMock() with patch.dict("sys.modules", {"litellm": mock_litellm}): await encoder.initialize() assert encoder._initialized is True @pytest.mark.asyncio async def test_initialization_missing_package(self): """Test initialization fails when litellm package is missing.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="cohere/rerank-english-v3.0", ) with patch.dict("sys.modules", {"litellm": None}): with pytest.raises(ImportError, match="litellm is required"): await encoder.initialize() @pytest.mark.asyncio async def test_initialization_idempotent(self): """Test that calling initialize() multiple times is safe.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="cohere/rerank-english-v3.0", ) mock_litellm = MagicMock() with patch.dict("sys.modules", {"litellm": mock_litellm}): await encoder.initialize() assert encoder._initialized is True # Second call should be no-op await encoder.initialize() assert encoder._initialized is True @pytest.mark.asyncio async def test_predict_single_query(self): """Test prediction with a single query and multiple documents.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="deepinfra/Qwen3-reranker-8B", ) # Create mock response with results as TypedDicts mock_response = MagicMock() mock_response.results = [ {"index": 0, "relevance_score": 0.9}, {"index": 1, "relevance_score": 0.7}, {"index": 2, "relevance_score": 0.5}, ] mock_litellm = MagicMock() mock_litellm.arerank = AsyncMock(return_value=mock_response) with patch.dict("sys.modules", {"litellm": mock_litellm}): await encoder.initialize() pairs = [ ("What is Python?", "Python is a programming language"), ("What is Python?", "Python is a snake"), ("What is Python?", "Python is a British comedy group"), ] scores = await encoder.predict(pairs) assert len(scores) == 3 assert scores == [0.9, 0.7, 0.5] # Verify arerank was called correctly mock_litellm.arerank.assert_called_once() call_args = mock_litellm.arerank.call_args assert call_args.kwargs["model"] == "deepinfra/Qwen3-reranker-8B" assert call_args.kwargs["query"] == "What is Python?" assert len(call_args.kwargs["documents"]) == 3 assert call_args.kwargs["api_key"] == "test_key" @pytest.mark.asyncio async def test_predict_multiple_queries(self): """Test prediction with multiple different queries (grouped efficiently).""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="cohere/rerank-english-v3.0", ) # First query response mock_response1 = MagicMock() mock_response1.results = [ {"index": 0, "relevance_score": 0.9}, {"index": 1, "relevance_score": 0.7}, ] # Second query response mock_response2 = MagicMock() mock_response2.results = [ {"index": 0, "relevance_score": 0.8}, ] mock_litellm = MagicMock() mock_litellm.arerank = AsyncMock(side_effect=[mock_response1, mock_response2]) with patch.dict("sys.modules", {"litellm": mock_litellm}): await encoder.initialize() pairs = [ ("What is Python?", "Python is a programming language"), ("What is Python?", "Python is a snake"), ("What is Java?", "Java is a programming language"), ] scores = await encoder.predict(pairs) assert len(scores) == 3 assert scores[0] == 0.9 # First query, first doc assert scores[1] == 0.7 # First query, second doc assert scores[2] == 0.8 # Second query, first doc # Verify arerank was called twice (once per unique query) assert mock_litellm.arerank.call_count == 2 @pytest.mark.asyncio async def test_predict_empty_pairs(self): """Test prediction with empty input.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="cohere/rerank-english-v3.0", ) mock_litellm = MagicMock() with patch.dict("sys.modules", {"litellm": mock_litellm}): await encoder.initialize() scores = await encoder.predict([]) assert scores == [] @pytest.mark.asyncio async def test_predict_not_initialized(self): """Test that predict fails if encoder not initialized.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="cohere/rerank-english-v3.0", ) pairs = [("query", "document")] with pytest.raises(RuntimeError, match="not initialized"): await encoder.predict(pairs) @pytest.mark.asyncio async def test_predict_error_handling(self): """Test that errors during prediction are raised.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="cohere/rerank-english-v3.0", ) # Mock litellm to raise an error mock_litellm = MagicMock() mock_litellm.arerank = AsyncMock(side_effect=Exception("API Error")) with patch.dict("sys.modules", {"litellm": mock_litellm}): await encoder.initialize() pairs = [ ("What is Python?", "Python is a programming language"), ] # Should raise the exception with pytest.raises(Exception, match="API Error"): await encoder.predict(pairs) @pytest.mark.asyncio async def test_custom_api_base(self): """Test that custom API base URL is passed to rerank calls.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="cohere/rerank-english-v3.0", api_base="https://custom.api.example.com", ) mock_response = MagicMock() mock_response.results = [ {"index": 0, "relevance_score": 0.9}, ] mock_litellm = MagicMock() mock_litellm.arerank = AsyncMock(return_value=mock_response) with patch.dict("sys.modules", {"litellm": mock_litellm}): await encoder.initialize() # Test that api_base is passed to arerank pairs = [("query", "document")] scores = await encoder.predict(pairs) assert scores == [0.9] mock_litellm.arerank.assert_called_once() call_args = mock_litellm.arerank.call_args assert call_args.kwargs["api_base"] == "https://custom.api.example.com" @pytest.mark.asyncio async def test_response_with_direct_score_list(self): """Test handling of response format with direct score list.""" encoder = LiteLLMSDKCrossEncoder( api_key="test_key", model="some-provider/model", ) # Mock litellm to return direct list of scores mock_litellm = MagicMock() mock_litellm.arerank = AsyncMock(return_value=[0.9, 0.7, 0.5]) with patch.dict("sys.modules", {"litellm": mock_litellm}): await encoder.initialize() pairs = [ ("query", "doc1"), ("query", "doc2"), ("query", "doc3"), ] scores = await encoder.predict(pairs) assert scores == [0.9, 0.7, 0.5] class TestFactoryFunction: """Test suite for create_cross_encoder_from_env factory function.""" @pytest.mark.asyncio async def test_create_litellm_sdk_from_env(self): """Test creating LiteLLM SDK cross-encoder from environment variables.""" env_vars = { "HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk", "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY": "test_key", "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "deepinfra/Qwen3-reranker-8B", } with patch.dict(os.environ, env_vars, clear=False): # Need to reload config to pick up env vars from hindsight_api.config import HindsightConfig config = HindsightConfig.from_env() with patch("hindsight_api.config.get_config", return_value=config): encoder = create_cross_encoder_from_env() assert isinstance(encoder, LiteLLMSDKCrossEncoder) assert encoder.api_key == "test_key" assert encoder.model == "deepinfra/Qwen3-reranker-8B" @pytest.mark.asyncio async def test_create_litellm_sdk_missing_api_key(self): """Test that factory raises error when API key is missing.""" env_vars = { "HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk", "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "deepinfra/Qwen3-reranker-8B", } with patch.dict(os.environ, env_vars, clear=False): # Remove API key if set if "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY" in os.environ: del os.environ["HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"] from hindsight_api.config import HindsightConfig config = HindsightConfig.from_env() with patch("hindsight_api.config.get_config", return_value=config): with pytest.raises(ValueError, match="HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY is required"): create_cross_encoder_from_env() @pytest.mark.asyncio async def test_create_litellm_sdk_with_custom_api_base(self): """Test creating LiteLLM SDK cross-encoder with custom API base.""" env_vars = { "HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk", "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY": "test_key", "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "cohere/rerank-english-v3.0", "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE": "https://custom.api.example.com", } with patch.dict(os.environ, env_vars, clear=False): from hindsight_api.config import HindsightConfig config = HindsightConfig.from_env() with patch("hindsight_api.config.get_config", return_value=config): encoder = create_cross_encoder_from_env() assert isinstance(encoder, LiteLLMSDKCrossEncoder) assert encoder.api_base == "https://custom.api.example.com" class TestLiteLLMSDKCohereCrossEncoder: """Tests for LiteLLM SDK calling Cohere (runs in CI with COHERE_API_KEY).""" @pytest.fixture async def litellm_cohere_cross_encoder(self): """Create LiteLLM SDK cross-encoder instance for Cohere.""" if not os.environ.get("COHERE_API_KEY"): pytest.skip("Cohere API key not available (set COHERE_API_KEY)") encoder = LiteLLMSDKCrossEncoder( api_key=os.environ["COHERE_API_KEY"], model="cohere/rerank-english-v3.0", ) await encoder.initialize() return encoder @pytest.mark.asyncio async def test_litellm_sdk_cohere_initialization(self, litellm_cohere_cross_encoder): """Test that LiteLLM SDK Cohere cross-encoder initializes correctly.""" assert litellm_cohere_cross_encoder.provider_name == "litellm-sdk" assert litellm_cohere_cross_encoder.model == "cohere/rerank-english-v3.0" @pytest.mark.asyncio async def test_litellm_sdk_cohere_predict(self, litellm_cohere_cross_encoder): """Test that LiteLLM SDK can call Cohere rerank API.""" pairs = [ ("What is the capital of France?", "Paris is the capital of France."), ("What is the capital of France?", "The Eiffel Tower is in Paris."), ("What is the capital of France?", "Python is a programming language."), ] scores = await litellm_cohere_cross_encoder.predict(pairs) assert len(scores) == 3 assert all(isinstance(s, float) for s in scores) # The first result should be most relevant assert scores[0] > scores[2], "Direct answer should score higher than unrelated text" # All scores should be in valid range assert all(0.0 <= score <= 1.0 for score in scores) class TestIntegration: """Integration tests with real API (optional - requires API keys).""" @pytest.mark.skipif( not os.environ.get("DEEPINFRA_API_KEY"), reason="DEEPINFRA_API_KEY not set - skipping integration test", ) @pytest.mark.asyncio async def test_real_deepinfra_api(self): """Test with real DeepInfra API (requires DEEPINFRA_API_KEY env var).""" encoder = LiteLLMSDKCrossEncoder( api_key=os.environ["DEEPINFRA_API_KEY"], model="deepinfra/Qwen3-reranker-8B", ) await encoder.initialize() pairs = [ ("What is Python?", "Python is a high-level programming language"), ("What is Python?", "Python is a species of snake"), ("What is Python?", "Python is unrelated text about cars"), ] scores = await encoder.predict(pairs) # First doc should have highest score (most relevant) assert len(scores) == 3 assert scores[0] > scores[1] assert scores[1] > scores[2] assert all(0.0 <= score <= 1.0 for score in scores)