From 61dd5a5ea4fb3629ae0b00879a379abc26f6f9f0 Mon Sep 17 00:00:00 2001
From: dbcontributions <navyaalapati13@gmail.com>
Date: Thu, 26 Sep 2024 11:03:55 -0500
Subject: [PATCH] Add vertexai test cases (#1907)

---
 Makefile                                     |   2 +-
 tests/embeddings/test_vertexai_embeddings.py | 119 +++++++++++++++++++
 2 files changed, 120 insertions(+), 1 deletion(-)
 create mode 100644 tests/embeddings/test_vertexai_embeddings.py

diff --git a/Makefile b/Makefile
index 965a719301..c3fb713da3 100644
--- a/Makefile
+++ b/Makefile
@@ -12,7 +12,7 @@ install:
 
 install_all:
 	poetry install
-	poetry run pip install groq together boto3 litellm ollama chromadb sentence_transformers
+	poetry run pip install groq together boto3 litellm ollama chromadb sentence_transformers vertexai
 
 # Format code with ruff
 format:
diff --git a/tests/embeddings/test_vertexai_embeddings.py b/tests/embeddings/test_vertexai_embeddings.py
new file mode 100644
index 0000000000..4b85e077e3
--- /dev/null
+++ b/tests/embeddings/test_vertexai_embeddings.py
@@ -0,0 +1,119 @@
+import pytest
+from unittest.mock import Mock, patch
+from mem0.embeddings.vertexai import VertexAI
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+
+
+@pytest.fixture
+def mock_text_embedding_model():
+    with patch("mem0.embeddings.vertexai.TextEmbeddingModel") as mock_model:
+        mock_instance = Mock()
+        mock_model.from_pretrained.return_value = mock_instance
+        yield mock_instance
+
+
+@pytest.fixture
+def mock_os_environ():
+    with patch("mem0.embeddings.vertexai.os.environ", {}) as mock_environ:
+        yield mock_environ
+
+
+@pytest.fixture
+def mock_config():
+    with patch("mem0.configs.embeddings.base.BaseEmbedderConfig") as mock_config:
+        mock_config.vertex_credentials_json = None
+        yield mock_config
+
+
+@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
+def test_embed_default_model(mock_text_embedding_model, mock_os_environ, mock_config):
+    mock_config.vertex_credentials_json = "/path/to/credentials.json"
+    mock_config.return_value.model = "text-embedding-004"
+    mock_config.return_value.embedding_dims = 256
+
+    config = mock_config()
+    embedder = VertexAI(config)
+
+    mock_embedding = Mock(values=[0.1, 0.2, 0.3])
+    mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [
+        mock_embedding
+    ]
+
+    result = embedder.embed("Hello world")
+
+    mock_text_embedding_model.from_pretrained.assert_called_once_with(
+        "text-embedding-004"
+    )
+    mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with(
+        texts=["Hello world"], output_dimensionality=256
+    )
+
+
+@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
+def test_embed_custom_model(mock_text_embedding_model, mock_os_environ, mock_config):
+    mock_config.vertex_credentials_json = "/path/to/credentials.json"
+    mock_config.return_value.model = "custom-embedding-model"
+    mock_config.return_value.embedding_dims = 512
+
+    config = mock_config()
+
+    embedder = VertexAI(config)
+
+    mock_embedding = Mock(values=[0.4, 0.5, 0.6])
+    mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [
+        mock_embedding
+    ]
+
+    result = embedder.embed("Test embedding")
+
+    mock_text_embedding_model.from_pretrained.assert_called_with(
+        "custom-embedding-model"
+    )
+    mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with(
+        texts=["Test embedding"], output_dimensionality=512
+    )
+
+    assert result == [0.4, 0.5, 0.6]
+
+
+@patch("mem0.embeddings.vertexai.os")
+def test_credentials_from_environment(mock_os, mock_text_embedding_model, mock_config):
+    mock_os.getenv.return_value = "/path/to/env/credentials.json"
+    mock_config.vertex_credentials_json = None
+    config = mock_config()
+    VertexAI(config)
+
+    mock_os.environ.setitem.assert_not_called()
+
+
+@patch("mem0.embeddings.vertexai.os")
+def test_missing_credentials(mock_os, mock_text_embedding_model, mock_config):
+    mock_os.getenv.return_value = None
+    mock_config.return_value.vertex_credentials_json = None
+
+    config = mock_config()
+
+    with pytest.raises(
+        ValueError, match="Google application credentials JSON is not provided"
+    ):
+        VertexAI(config)
+
+
+@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
+def test_embed_with_different_dimensions(
+    mock_text_embedding_model, mock_os_environ, mock_config
+):
+    mock_config.vertex_credentials_json = "/path/to/credentials.json"
+    mock_config.return_value.embedding_dims = 1024
+
+    config = mock_config()
+    embedder = VertexAI(config)
+
+    mock_embedding = Mock(values=[0.1] * 1024)
+    mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [
+        mock_embedding
+    ]
+
+    result = embedder.embed("Large embedding test")
+
+    assert result == [0.1] * 1024