Add config option for vertex embedding tasks (#2266)

tests/embeddings/test_vertexai_embeddings.py

@@ -1,5 +1,7 @@
-import pytest
 from unittest.mock import Mock, patch
+
+import pytest
+
 from mem0.embeddings.vertexai import VertexAIEmbedding
 
 
@@ -20,13 +22,23 @@ def mock_os_environ():
 @pytest.fixture
 def mock_config():
     with patch("mem0.configs.embeddings.base.BaseEmbedderConfig") as mock_config:
-        mock_config.vertex_credentials_json = None
+        mock_config.return_value.vertex_credentials_json = "/path/to/credentials.json"
         yield mock_config
+        
+
+@pytest.fixture
+def mock_embedding_types():
+    return ["SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "RETRIEVAL_DOCUMENT", "RETRIEVAL_QUERY", "QUESTION_ANSWERING", "FACT_VERIFICATION", "CODE_RETRIEVAL_QUERY"]
+
+
+@pytest.fixture
+def mock_text_embedding_input():
+    with patch("mem0.embeddings.vertexai.TextEmbeddingInput") as mock_input:
+        yield mock_input
 
 
 @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"
+def test_embed_default_model(mock_text_embedding_model, mock_os_environ, mock_config, mock_text_embedding_input):
     mock_config.return_value.model = "text-embedding-004"
     mock_config.return_value.embedding_dims = 256
 
@@ -37,16 +49,16 @@ def test_embed_default_model(mock_text_embedding_model, mock_os_environ, mock_co
     mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_embedding]
 
     embedder.embed("Hello world")
-
+    mock_text_embedding_input.assert_called_once_with(text="Hello world", task_type="SEMANTIC_SIMILARITY")
     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
+        texts=[mock_text_embedding_input("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"
+def test_embed_custom_model(mock_text_embedding_model, mock_os_environ, mock_config, mock_text_embedding_input):
     mock_config.return_value.model = "custom-embedding-model"
     mock_config.return_value.embedding_dims = 512
 
@@ -58,18 +70,42 @@ def test_embed_custom_model(mock_text_embedding_model, mock_os_environ, mock_con
     mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_embedding]
 
     result = embedder.embed("Test embedding")
-
+    mock_text_embedding_input.assert_called_once_with(text="Test embedding", task_type="SEMANTIC_SIMILARITY")
     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
+        texts=[mock_text_embedding_input("Test embedding")], output_dimensionality=512
     )
 
     assert result == [0.4, 0.5, 0.6]
 
 
+@patch("mem0.embeddings.vertexai.TextEmbeddingModel")  
+def test_embed_with_memory_action(mock_text_embedding_model, mock_os_environ, mock_config, mock_embedding_types, mock_text_embedding_input):
+    mock_config.return_value.model = "text-embedding-004"
+    mock_config.return_value.embedding_dims = 256
+    
+    for embedding_type in mock_embedding_types:
+        
+        mock_config.return_value.memory_add_embedding_type = embedding_type
+        mock_config.return_value.memory_update_embedding_type = embedding_type
+        mock_config.return_value.memory_search_embedding_type = embedding_type
+
+        config = mock_config()
+        embedder = VertexAIEmbedding(config)
+        
+        mock_text_embedding_model.from_pretrained.assert_called_with("text-embedding-004")
+
+        for memory_action in ["add", "update", "search"]:
+            embedder.embed("Hello world", memory_action=memory_action)
+            
+            mock_text_embedding_input.assert_called_with(text="Hello world", task_type=embedding_type)
+            mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_with(
+                texts=[mock_text_embedding_input("Hello world", embedding_type)], output_dimensionality=256
+            )
+            
+
 @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()
     VertexAIEmbedding(config)
@@ -90,7 +126,6 @@ def test_missing_credentials(mock_os, mock_text_embedding_model, mock_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()
@@ -102,3 +137,15 @@ def test_embed_with_different_dimensions(mock_text_embedding_model, mock_os_envi
     result = embedder.embed("Large embedding test")
 
     assert result == [0.1] * 1024
+    
+
+@patch("mem0.embeddings.vertexai.TextEmbeddingModel")    
+def test_invalid_memory_action(mock_text_embedding_model, mock_config):
+    mock_config.return_value.model = "text-embedding-004"
+    mock_config.return_value.embedding_dims = 256
+    
+    config = mock_config()
+    embedder = VertexAIEmbedding(config)
+    
+    with pytest.raises(ValueError):
+        embedder.embed("Hello world", memory_action="invalid_action")

tests/test_main.py

@@ -119,7 +119,7 @@ def test_search(memory_instance, version, enable_graph):
     memory_instance.vector_store.search.assert_called_once_with(
         query=[0.1, 0.2, 0.3], limit=100, filters={"user_id": "test_user"}
     )
-    memory_instance.embedding_model.embed.assert_called_once_with("test query")
+    memory_instance.embedding_model.embed.assert_called_once_with("test query", "search")
 
     if enable_graph:
         memory_instance.graph.search.assert_called_once_with("test query", {"user_id": "test_user"}, 100)