From 8791883201fcf70d10c57b878b887cc3f9ee1ea5 Mon Sep 17 00:00:00 2001
From: Trayan Azarov <trayan.azarov@amikos.tech>
Date: Sat, 26 Apr 2025 20:47:11 +0300
Subject: [PATCH] fix: Fix Voyage AI EF

---
 chromadb/test/ef/test_voyageai_ef.py          | 19 +++++++++++++++++++
 .../voyageai_embedding_function.py            |  4 +++-
 2 files changed, 22 insertions(+), 1 deletion(-)
 create mode 100644 chromadb/test/ef/test_voyageai_ef.py

diff --git a/chromadb/test/ef/test_voyageai_ef.py b/chromadb/test/ef/test_voyageai_ef.py
new file mode 100644
index 00000000000..bd0976aeee3
--- /dev/null
+++ b/chromadb/test/ef/test_voyageai_ef.py
@@ -0,0 +1,19 @@
+import os
+import pytest
+from chromadb.utils.embedding_functions.voyageai_embedding_function import (
+    VoyageAIEmbeddingFunction,
+)
+
+voyageai = pytest.importorskip("voyageai", reason="voyageai not installed")
+
+
+def test_with_embedding_dimensions() -> None:
+    if os.environ.get("CHROMA_VOYAGE_API_KEY") is None:
+        pytest.skip("CHROMA_VOYAGE_API_KEY not set")
+    ef = VoyageAIEmbeddingFunction(
+        api_key=os.environ["CHROMA_VOYAGE_API_KEY"]
+    )
+    embeddings = ef(["hello world"])
+    assert embeddings is not None
+    assert len(embeddings) == 1
+    assert len(embeddings[0]) == 1536
diff --git a/chromadb/utils/embedding_functions/voyageai_embedding_function.py b/chromadb/utils/embedding_functions/voyageai_embedding_function.py
index a87b83cfffa..8e8d27450dc 100644
--- a/chromadb/utils/embedding_functions/voyageai_embedding_function.py
+++ b/chromadb/utils/embedding_functions/voyageai_embedding_function.py
@@ -55,7 +55,9 @@ def __call__(self, input: Documents) -> Embeddings:
         embeddings = self._client.embed(texts=input, model=self.model_name)
 
         # Convert to numpy arrays
-        return [np.array(embedding, dtype=np.float32) for embedding in embeddings]
+        return [
+            np.array(embedding, dtype=np.float32) for embedding in embeddings.embeddings
+        ]
 
     @staticmethod
     def name() -> str: