feat: VoyageAI encoder

pyproject.toml

@@ -38,6 +38,7 @@ matplotlib = { version = "^3.8.3", optional = true}
 qdrant-client = {version = "^1.8.0", optional = true}
 google-cloud-aiplatform = {version = "^1.45.0", optional = true}
 requests-mock = "^1.12.1"
+voyageai = {version = "^0.2.2", optional = true}
 
 [tool.poetry.extras]
 hybrid = ["pinecone-text"]
@@ -49,6 +50,7 @@ processing = ["matplotlib"]
 mistralai = ["mistralai"]
 qdrant = ["qdrant-client"]
 google = ["google-cloud-aiplatform"]
+voyageai = ["voyageai"]
 
 [tool.poetry.group.dev.dependencies]
 ipykernel = "^6.25.0"

semantic_router/encoders/__init__.py

@@ -11,6 +11,7 @@
 from semantic_router.encoders.tfidf import TfidfEncoder
 from semantic_router.encoders.vit import VitEncoder
 from semantic_router.encoders.zure import AzureOpenAIEncoder
+from semantic_router.encoders.voyageai import VoyageAIEncoder
 
 __all__ = [
     "BaseEncoder",
@@ -26,4 +27,5 @@
     "VitEncoder",
     "CLIPEncoder",
     "GoogleEncoder",
+    "VoyageAIEncoder",
 ]

semantic_router/encoders/voyageai.py

@@ -0,0 +1,73 @@
+import os
+from time import sleep
+from typing import Any, List, Optional
+
+from pydantic.v1 import PrivateAttr
+
+from semantic_router.encoders import BaseEncoder
+from semantic_router.utils.defaults import EncoderDefault
+from semantic_router.utils.logger import logger
+
+
+class VoyageAIEncoder(BaseEncoder):
+    _client: Any = PrivateAttr()
+    type: str = "voyageai"
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        voyage_api_key: Optional[str] = None,
+        score_threshold: float = 0.82,
+    ):
+        if name is None:
+            name = EncoderDefault.VOYAGE.value["embedding_model"]
+        super().__init__(name=name, score_threshold=score_threshold)
+        self._client = self._initialize_client(api_key=voyage_api_key)
+
+    def _initialize_client(self, api_key: Optional[str] = None):
+        try:
+            import voyageai
+        except ImportError:
+            raise ImportError(
+                "Please install VoyageAI to use VoyageAIEncoder. "
+                "You can install it with: "
+                "`pip install 'semantic-router[voyageai]'`"
+            )
+
+        api_key = api_key or os.getenv("VOYAGEAI_API_KEY")
+        if api_key is None:
+            raise ValueError("VoyageAI API key not provided")
+        try:
+            client = voyageai.Client(api_key=api_key)
+        except Exception as e:
+            raise ValueError(f"Unable to connect to VoyageAI {e.args}: {e}") from e
+        return client
+
+    def __call__(self, docs: List[str]) -> List[List[float]]:
+        if self._client == PrivateAttr():
+            raise ValueError("VoyageAI client is not initialized.")
+        embeds = None
+        error_message = ""
+
+        # Exponential backoff
+        for j in range(1, 7):
+            try:
+                embeds = self._client.embed(
+                    texts=docs,
+                    model=self.name,
+                    input_type="query",  # query or document
+                )
+                if embeds.embeddings:
+                    break
+                else:
+                    sleep(2**j)
+                    logger.warning(f"Retrying in {2**j} seconds...")
+
+            except Exception as e:
+                logger.error(f"VoyageAI API call failed. Error: {error_message}")
+                raise ValueError(f"VoyageAI API call failed. Error: {e}") from e
+
+        if not embeds or not embeds.embeddings:
+            raise ValueError("VoyageAI API call failed. Error: No embeddings found.")
+
+        return embeds.embeddings

semantic_router/utils/defaults.py

@@ -3,6 +3,7 @@
 
 
 class EncoderDefault(Enum):
+    VOYAGE = {"embedding_model": "voyage-lite-02-instruct"}
     FASTEMBED = {
         "embedding_model": "BAAI/bge-small-en-v1.5",
         "language_model": "BAAI/bge-small-en-v1.5",

tests/unit/encoders/test_voyageai.py

@@ -0,0 +1,100 @@
+from unittest.mock import patch
+
+import pytest
+
+from pydantic.v1 import PrivateAttr
+
+from semantic_router.encoders import VoyageAIEncoder
+
+
+@pytest.fixture
+def voyageai_encoder(mocker):
+    mocker.patch("voyageai.Client")
+    return VoyageAIEncoder(voyage_api_key="test_api_key")
+
+
+class TestVoyageAIEncoder:
+    def test_voyageai_encoder_import_error(self):
+        with patch.dict("sys.modules", {"voyageai": None}):
+            with pytest.raises(ImportError) as error:
+                VoyageAIEncoder()
+
+        assert "pip install 'semantic-router[voyageai]'" in str(error.value)
+
+    def test_voyageai_encoder_init_success(self, mocker):
+        side_effect = ["fake-model-name", "fake-api-key"]
+        mocker.patch("os.getenv", side_effect=side_effect)
+        encoder = VoyageAIEncoder()
+        assert encoder._client is not PrivateAttr()
+
+    def test_voyageai_encoder_init_no_api_key(self, mocker):
+        mocker.patch("os.getenv", return_value=None)
+        with pytest.raises(ValueError) as _:
+            VoyageAIEncoder()
+
+    def test_voyageai_encoder_call_uninitialized_client(self, voyageai_encoder):
+        voyageai_encoder._client = PrivateAttr()
+        with pytest.raises(ValueError) as e:
+            voyageai_encoder(["test document"])
+        assert "VoyageAI client is not initialized." in str(e.value)
+
+    def test_voyageai_encoder_init_exception(self, mocker):
+        mocker.patch("os.getenv", return_value="fake-api-key")
+        mocker.patch("voyageai.Client", side_effect=Exception("Initialization error"))
+        with pytest.raises(ValueError) as e:
+            VoyageAIEncoder()
+        assert (
+            "VOYAGE API client failed to initialize. Error: Initialization error"
+            in str(e.value)
+        )
+
+    def test_voyageai_encoder_call_success(self, voyageai_encoder, mocker):
+        mock_response = mocker.Mock()
+        mock_response.embeddings = [[0.1, 0.2]]
+
+        mocker.patch("os.getenv", return_value="fake-api-key", autospec=True)
+        mocker.patch("time.sleep", return_value=None)
+
+        mocker.patch.object(
+            voyageai_encoder._client, "embed", return_value=mock_response
+        )
+        embeddings = voyageai_encoder(["test document"])
+        assert embeddings == [[0.1, 0.2]]
+
+    def test_voyageai_encoder_call_with_retries(self, voyageai_encoder, mocker):
+        error = Exception("Network error")
+        mocker.patch("os.getenv", return_value="fake-api-key")
+        mocker.patch("time.sleep", return_value=None)
+        mocker.patch.object(
+            voyageai_encoder._client,
+            "embed",
+            side_effect=[error, error, mocker.Mock(embeddings=[[0.1, 0.2]])],
+        )
+        embeddings = voyageai_encoder(["test document"])
+        assert embeddings == [[0.1, 0.2]]
+
+    def test_voyageai_encoder_call_failure_non_voyage_error(
+        self, voyageai_encoder, mocker
+    ):
+        mocker.patch("os.getenv", return_value="fake-api-key")
+        mocker.patch("time.sleep", return_value=None)
+        mocker.patch.object(
+            voyageai_encoder._client,
+            "embed",
+            side_effect=Exception("General error"),
+        )
+        with pytest.raises(ValueError) as e:
+            voyageai_encoder(["test document"])
+        assert "VoyageAI API call failed. Error: General error" in str(e.value)
+
+    def test_voyageai_encoder_call_successful_retry(self, voyageai_encoder, mocker):
+        mock_response = mocker.Mock()
+        mock_response.embeddings = [[0.1, 0.2]]
+
+        mocker.patch("os.getenv", return_value="fake-api-key")
+        mocker.patch("time.sleep", return_value=None)
+
+        responses = [Exception("Temporary error"), mock_response]
+        mocker.patch.object(voyageai_encoder._client, "embed", side_effect=responses)
+        embeddings = voyageai_encoder(["test document"])
+        assert embeddings == [[0.1, 0.2]]