[BUG] fix bug when receive embedding from vertex api

chromadb/test/ef/test_vertex_ef.py

@@ -0,0 +1,25 @@
+from os import getenv
+
+import pytest
+from chromadb.utils.embedding_functions import GoogleVertexEmbeddingFunction
+
+
+def test_api_key() -> None:
+    with pytest.raises(KeyError, match="Request had invalid authentication credentials."):
+        vertex_ef = GoogleVertexEmbeddingFunction(None,
+                                                  'textembedding-gecko-multilingual',
+                                                  getenv('PROJECT_ID'))
+        embeddingsclear = vertex_ef(['Open source is awsome.'])
+
+
+@pytest.mark.skipif(
+    not getenv('GVAI_API_KEY') or not getenv('PROJECT_ID'),
+    reason='API_TOKEN or PROJECT_ID is not set, skipping test.'
+)
+def test_vertex_ef() -> None:
+    vertex_ef = GoogleVertexEmbeddingFunction(getenv('GVAI_API_KEY'),
+                                              'textembedding-gecko-multilingual',
+                                              getenv('PROJECT_ID'))
+    embeddings = vertex_ef(['Open source is awsome.'])
+    assert embeddings is not None
+    assert len(embeddings) > 0

chromadb/utils/embedding_functions.py

@@ -661,8 +661,8 @@ class GoogleVertexEmbeddingFunction(EmbeddingFunction[Documents]):
     def __init__(
         self,
         api_key: str,
+        project_id: str,
         model_name: str = "textembedding-gecko",
-        project_id: str = "cloud-large-language-models",
         region: str = "us-central1",
     ):
         self._api_url = f"https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/goole/models/{model_name}:predict"
@@ -676,8 +676,12 @@ def __call__(self, input: Documents) -> Embeddings:
                 self._api_url, json={"instances": [{"content": text}]}
             ).json()
 
-            if "predictions" in response:
-                embeddings.append(response["predictions"]["embeddings"]["values"])
+            predictions = response.get('predictions')
+            if isinstance(predictions, List) and predictions:
+                for prediction in predictions:
+                    embeddings.append(prediction["embeddings"]["values"])
+            else:
+                raise KeyError(f'Something went wrong -- Response: {response}')
 
         return embeddings