""" Tests for CohereCrossEncoder. Tests the Cohere cross-encoder implementation, including Azure AI Foundry endpoint support. """ import os from unittest.mock import MagicMock, patch import httpx import pytest from hindsight_api.engine.cross_encoder import CohereCrossEncoder, create_cross_encoder_from_env class TestCohereCrossEncoder: """Test suite for CohereCrossEncoder class.""" @pytest.mark.asyncio async def test_initialization_native_cohere(self): """Test successful initialization with native Cohere API (no base_url).""" encoder = CohereCrossEncoder( api_key="test_key", model="rerank-english-v3.0", ) assert encoder.provider_name == "cohere" assert encoder.api_key == "test_key" assert encoder.model == "rerank-english-v3.0" assert encoder._client is None assert encoder._httpx_client is None # Mock the cohere import mock_cohere = MagicMock() mock_cohere.Client = MagicMock() with patch.dict("sys.modules", {"cohere": mock_cohere}): await encoder.initialize() assert encoder._client is not None assert encoder._httpx_client is None mock_cohere.Client.assert_called_once_with(api_key="test_key", timeout=60.0) @pytest.mark.asyncio async def test_initialization_azure_endpoint(self): """Test initialization with Azure AI Foundry endpoint (uses httpx).""" encoder = CohereCrossEncoder( api_key="test_key", model="cohere-rerank-v3-english", base_url="https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke", ) assert encoder.base_url == "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke" await encoder.initialize() assert encoder._httpx_client is not None assert encoder._client is None assert isinstance(encoder._httpx_client, httpx.Client) @pytest.mark.asyncio async def test_initialization_missing_package(self): """Test initialization fails when cohere package is missing (native API).""" encoder = CohereCrossEncoder( api_key="test_key", model="rerank-english-v3.0", ) with patch.dict("sys.modules", {"cohere": None}): with pytest.raises(ImportError, match="cohere is required"): await encoder.initialize() @pytest.mark.asyncio async def test_initialization_idempotent(self): """Test that calling initialize() multiple times is safe.""" encoder = CohereCrossEncoder( api_key="test_key", model="rerank-english-v3.0", ) mock_cohere = MagicMock() mock_cohere.Client = MagicMock() with patch.dict("sys.modules", {"cohere": mock_cohere}): await encoder.initialize() assert encoder._client is not None # Second call should be no-op await encoder.initialize() # Should only create client once mock_cohere.Client.assert_called_once() @pytest.mark.asyncio async def test_predict_native_cohere_single_query(self): """Test prediction with native Cohere SDK.""" encoder = CohereCrossEncoder( api_key="test_key", model="rerank-english-v3.0", ) # Create mock Cohere response mock_result_1 = MagicMock() mock_result_1.index = 0 mock_result_1.relevance_score = 0.9 mock_result_2 = MagicMock() mock_result_2.index = 1 mock_result_2.relevance_score = 0.7 mock_result_3 = MagicMock() mock_result_3.index = 2 mock_result_3.relevance_score = 0.5 mock_response = MagicMock() mock_response.results = [mock_result_1, mock_result_2, mock_result_3] mock_cohere_client = MagicMock() mock_cohere_client.rerank = MagicMock(return_value=mock_response) mock_cohere = MagicMock() mock_cohere.Client = MagicMock(return_value=mock_cohere_client) with patch.dict("sys.modules", {"cohere": mock_cohere}): 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 rerank was called correctly mock_cohere_client.rerank.assert_called_once() call_args = mock_cohere_client.rerank.call_args assert call_args.kwargs["model"] == "rerank-english-v3.0" assert call_args.kwargs["query"] == "What is Python?" assert len(call_args.kwargs["documents"]) == 3 assert call_args.kwargs["return_documents"] is False @pytest.mark.asyncio async def test_predict_azure_endpoint_single_query(self): """Test prediction with Azure AI Foundry endpoint (httpx direct call).""" encoder = CohereCrossEncoder( api_key="test_key", model="cohere-rerank-v3-english", base_url="https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke", ) await encoder.initialize() # Mock httpx response mock_response = MagicMock() mock_response.json.return_value = { "results": [ {"index": 0, "relevance_score": 0.9}, {"index": 1, "relevance_score": 0.7}, {"index": 2, "relevance_score": 0.5}, ] } encoder._httpx_client.post = MagicMock(return_value=mock_response) 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 httpx.post was called with correct URL and payload encoder._httpx_client.post.assert_called_once() call_args = encoder._httpx_client.post.call_args assert call_args[0][0] == "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke" assert call_args.kwargs["json"]["model"] == "cohere-rerank-v3-english" assert call_args.kwargs["json"]["query"] == "What is Python?" assert len(call_args.kwargs["json"]["documents"]) == 3 assert call_args.kwargs["json"]["return_documents"] is False @pytest.mark.asyncio async def test_predict_multiple_queries(self): """Test prediction with multiple different queries (grouped efficiently).""" encoder = CohereCrossEncoder( api_key="test_key", model="rerank-english-v3.0", ) # First query response mock_result_1_1 = MagicMock() mock_result_1_1.index = 0 mock_result_1_1.relevance_score = 0.9 mock_result_1_2 = MagicMock() mock_result_1_2.index = 1 mock_result_1_2.relevance_score = 0.7 mock_response1 = MagicMock() mock_response1.results = [mock_result_1_1, mock_result_1_2] # Second query response mock_result_2_1 = MagicMock() mock_result_2_1.index = 0 mock_result_2_1.relevance_score = 0.8 mock_response2 = MagicMock() mock_response2.results = [mock_result_2_1] mock_cohere_client = MagicMock() mock_cohere_client.rerank = MagicMock(side_effect=[mock_response1, mock_response2]) mock_cohere = MagicMock() mock_cohere.Client = MagicMock(return_value=mock_cohere_client) with patch.dict("sys.modules", {"cohere": mock_cohere}): 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 rerank was called twice (once per unique query) assert mock_cohere_client.rerank.call_count == 2 @pytest.mark.asyncio async def test_predict_empty_pairs(self): """Test prediction with empty input.""" encoder = CohereCrossEncoder( api_key="test_key", model="rerank-english-v3.0", ) mock_cohere = MagicMock() mock_cohere.Client = MagicMock() with patch.dict("sys.modules", {"cohere": mock_cohere}): 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 = CohereCrossEncoder( api_key="test_key", model="rerank-english-v3.0", ) pairs = [("query", "document")] with pytest.raises(RuntimeError, match="not initialized"): await encoder.predict(pairs) @pytest.mark.asyncio async def test_azure_endpoint_http_error(self): """Test that HTTP errors from Azure endpoint are raised.""" encoder = CohereCrossEncoder( api_key="test_key", model="cohere-rerank-v3-english", base_url="https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke", ) await encoder.initialize() # Mock httpx to raise HTTP error mock_response = MagicMock() mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( "404 Not Found", request=MagicMock(), response=MagicMock(status_code=404), ) encoder._httpx_client.post = MagicMock(return_value=mock_response) pairs = [("What is Python?", "Python is a programming language")] # Should raise the HTTP error with pytest.raises(httpx.HTTPStatusError): await encoder.predict(pairs) class TestFactoryFunction: """Test suite for create_cross_encoder_from_env factory function.""" @pytest.mark.asyncio async def test_create_cohere_from_env(self): """Test creating Cohere cross-encoder from environment variables.""" env_vars = { "HINDSIGHT_API_RERANKER_PROVIDER": "cohere", "HINDSIGHT_API_RERANKER_COHERE_API_KEY": "test_key", "HINDSIGHT_API_RERANKER_COHERE_MODEL": "rerank-english-v3.0", } 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, CohereCrossEncoder) assert encoder.api_key == "test_key" assert encoder.model == "rerank-english-v3.0" assert encoder.base_url is None @pytest.mark.asyncio async def test_create_cohere_with_azure_base_url_from_env(self): """Test creating Cohere cross-encoder with Azure base URL from environment.""" env_vars = { "HINDSIGHT_API_RERANKER_PROVIDER": "cohere", "HINDSIGHT_API_RERANKER_COHERE_API_KEY": "test_key", "HINDSIGHT_API_RERANKER_COHERE_MODEL": "cohere-rerank-v3-english", "HINDSIGHT_API_RERANKER_COHERE_BASE_URL": "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke", } 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, CohereCrossEncoder) assert encoder.api_key == "test_key" assert encoder.model == "cohere-rerank-v3-english" assert encoder.base_url == "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke"