diff --git a/hindsight-api-slim/hindsight_api/engine/cross_encoder.py b/hindsight-api-slim/hindsight_api/engine/cross_encoder.py index 6081691f..fe99928c 100644 --- a/hindsight-api-slim/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api-slim/hindsight_api/engine/cross_encoder.py @@ -544,6 +544,7 @@ class CohereCrossEncoder(CrossEncoderModel): self.base_url = base_url self.timeout = timeout self._client = None + self._httpx_client: httpx.Client | None = None @property def provider_name(self) -> str: @@ -551,23 +552,32 @@ class CohereCrossEncoder(CrossEncoderModel): async def initialize(self) -> None: """Initialize the Cohere client.""" - if self._client is not None: + if self._client is not None or self._httpx_client is not None: return - try: - import cohere - except ImportError: - raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere") - base_url_msg = f" at {self.base_url}" if self.base_url else "" logger.info(f"Reranker: initializing Cohere provider with model {self.model}{base_url_msg}") - # Build client kwargs, only including base_url if set (for Azure or custom endpoints) - client_kwargs = {"api_key": self.api_key, "timeout": self.timeout} if self.base_url: - client_kwargs["base_url"] = self.base_url - self._client = cohere.Client(**client_kwargs) - logger.info("Reranker: Cohere provider initialized") + # For custom endpoints (Azure AI Foundry), use httpx directly to avoid SDK path appending + # Azure endpoints already include the full path (e.g., /models/.../invoke) + self._httpx_client = httpx.Client( + timeout=self.timeout, + headers={ + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + ) + logger.info("Reranker: Cohere provider initialized (using httpx for custom endpoint)") + else: + # For native Cohere API, use the official SDK + try: + import cohere + except ImportError: + raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere") + + self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout) + logger.info("Reranker: Cohere provider initialized") async def predict(self, pairs: list[tuple[str, str]]) -> list[float]: """ @@ -579,7 +589,7 @@ class CohereCrossEncoder(CrossEncoderModel): Returns: List of relevance scores """ - if self._client is None: + if self._client is None and self._httpx_client is None: raise RuntimeError("Reranker not initialized. Call initialize() first.") if not pairs: @@ -605,18 +615,40 @@ class CohereCrossEncoder(CrossEncoderModel): texts = [text for _, text in indexed_texts] indices = [idx for idx, _ in indexed_texts] - response = self._client.rerank( - query=query, - documents=texts, - model=self.model, - return_documents=False, - ) + if self._httpx_client: + # Direct HTTP request for custom endpoints (Azure AI Foundry) + response = self._httpx_client.post( + self.base_url, + json={ + "model": self.model, + "query": query, + "documents": texts, + "return_documents": False, + }, + ) + response.raise_for_status() + result = response.json() - # Map scores back to original positions - for result in response.results: - original_idx = result.index - score = result.relevance_score - all_scores[indices[original_idx]] = score + # Map scores back to original positions + # Azure Cohere response format: {"results": [{"index": 0, "relevance_score": 0.9}, ...]} + for item in result.get("results", []): + original_idx = item["index"] + score = item["relevance_score"] + all_scores[indices[original_idx]] = score + else: + # Native Cohere SDK for standard API + response = self._client.rerank( + query=query, + documents=texts, + model=self.model, + return_documents=False, + ) + + # Map scores back to original positions + for result in response.results: + original_idx = result.index + score = result.relevance_score + all_scores[indices[original_idx]] = score return all_scores diff --git a/hindsight-api-slim/tests/test_cohere_cross_encoder.py b/hindsight-api-slim/tests/test_cohere_cross_encoder.py new file mode 100644 index 00000000..d0a14bdc --- /dev/null +++ b/hindsight-api-slim/tests/test_cohere_cross_encoder.py @@ -0,0 +1,339 @@ +""" +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"