[WIP] feat: add Voyage embeddings

libs/kotaemon/kotaemon/embeddings/__init__.py

@@ -9,6 +9,7 @@
 )
 from .openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
 from .tei_endpoint_embed import TeiEndpointEmbeddings
+from .voyageai import VoyageAIEmbeddings
 
 __all__ = [
     "BaseEmbeddings",
@@ -21,4 +22,5 @@
     "OpenAIEmbeddings",
     "AzureOpenAIEmbeddings",
     "FastEmbedEmbeddings",
+    "VoyageAIEmbeddings",
 ]

libs/kotaemon/kotaemon/embeddings/voyageai.py

@@ -0,0 +1,65 @@
+"""Implements embeddings from [Voyage AI](https://voyageai.com).
+"""
+
+import importlib
+
+from kotaemon.base import Document, DocumentWithEmbedding, Param
+
+from .base import BaseEmbeddings
+
+
+vo = None
+
+
+def _import_voyageai():
+    global vo
+    if not vo:
+        vo = importlib.import_module("voyageai")
+    return vo
+
+
+def _format_output(texts: list[str], embeddings: list[list]):
+    """Formats the output of all `.embed` calls.
+
+    Args:
+        texts: List of original documents
+        embeddings: Embeddings corresponding to each document
+    """
+    return [
+        DocumentWithEmbedding(content=text, embedding=embedding)
+        for text, embedding in zip(texts, embeddings)
+    ]
+
+
+class VoyageAIEmbeddings(BaseEmbeddings):
+    """Voyage AI provides best-in-class embedding models and rerankers."""
+
+    api_key: str = Param(None, help="Voyage API key", required=False)
+    model_name: str = Param(
+        "voyage-3",
+        help=(
+            "Model name to use. The Voyage "
+            "[documentation](https://docs.voyageai.com/docs/embeddings) "
+            "provides a list of all available embedding models."
+        ),
+        required=True,
+    )
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._client = _import_voyageai().Client(api_key=self.api_key)
+        self._aclient = _import_voyageai().AsyncClient(api_key=self.api_key)
+
+    def invoke(
+        self, text: str | list[str] | Document | list[Document], *args, **kwargs
+    ) -> list[DocumentWithEmbedding]:
+        texts = [t.content for t in self.prepare_input(text)]
+        embeddings = self._client.embed(texts, model=self.model_name).embeddings
+        return _format_output(texts, embeddings)
+
+    async def ainvoke(
+        self, text: str | list[str] | Document | list[Document], *args, **kwargs
+    ) -> list[DocumentWithEmbedding]:
+        texts = [t.content for t in self.prepare_input(text)]
+        embeddings = await self._aclient.embed(texts, model=self.model_name).embeddings
+        return _format_output(texts, embeddings)

libs/kotaemon/pyproject.toml

@@ -80,6 +80,7 @@ adv = [
     "tabulate",
     "unstructured>=0.15.8,<0.16",
     "wikipedia>=1.4.0,<1.5",
+    "voyageai>=0.2.0",
 ]
 dev = [
     "black",

libs/kotaemon/tests/conftest.py

@@ -70,6 +70,15 @@ def if_llama_cpp_not_installed():
         return False
 
 
+def if_voyageai_not_installed():
+    try:
+        import voyageai  # noqa: F401
+    except ImportError:
+        return True
+    else:
+        return False
+
+
 skip_when_haystack_not_installed = pytest.mark.skipif(
     if_haystack_not_installed(), reason="Haystack is not installed"
 )
@@ -97,3 +106,7 @@ def if_llama_cpp_not_installed():
 skip_llama_cpp_not_installed = pytest.mark.skipif(
     if_llama_cpp_not_installed(), reason="llama_cpp is not installed"
 )
+
+skip_when_voyageai_not_installed = pytest.mark.skipif(
+    if_voyageai_not_installed(), reason="voyageai is not installed"
+)

libs/kotaemon/tests/test_embedding_models.py

@@ -11,12 +11,14 @@
     LCCohereEmbeddings,
     LCHuggingFaceEmbeddings,
     OpenAIEmbeddings,
+    VoyageAIEmbeddings,
 )
 
 from .conftest import (
     skip_when_cohere_not_installed,
     skip_when_fastembed_not_installed,
     skip_when_sentence_bert_not_installed,
+    skip_when_voyageai_not_installed,
 )
 
 with open(Path(__file__).parent / "resources" / "embedding_openai_batch.json") as f:
@@ -155,3 +157,18 @@ def test_fastembed_embeddings():
     model = FastEmbedEmbeddings()
     output = model("Hello World")
     assert_embedding_result(output)
+
+
+@skip_when_voyageai_not_installed
+@patch(
+    "voyageai.Client.embed",
+    side_effect=lambda *args, **kwargs: [[1.0, 2.1, 3.2]],
+)
+@patch(
+    "voyageai.AsyncClient.embed",
+    side_effect=lambda *args, **kwargs: [[1.0, 2.1, 3.2]],
+)
+def test_voyageai_embeddings():
+    model = VoyageAIEmbeddings()
+    output = model("Hello, world!")
+    assert_embedding_result(output)