Repository: mem0ai/mem0
Pull request: Add config option for vertex embedding tasks
## Description
__Feature__  
The feature provides users with multiple options for embedding tasks in Vertex AI, and modifies `EmbeddingBase.embed` to include the memory action that triggered the function call.
    
__Motivation__  
Vertex AI offers multiple options for embedding tasks, allowing the extraction of text embeddings for specific use cases, such as distinguishing between a "retrieval document" and a "retrieval query."  
If this feature is added, it would give developers greater flexibility in using mem0 based on their needs.  
  
__Note__  
- To use the feature, need to specify the embedding task type in`mem0.configs.embeddings.base.BaseEmbedderConfig`. 
- I modified `mem0.embeddings.base.EmbeddingBase` structure and codes related to it

Fixes #2259 

## Type of change

Please delete options that are not relevant.
- [x] New feature (non-breaking change which adds functionality)

## How Has This Been Tested?
Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

Please delete options that are not relevant.

- [x] Unit Test
- [x] Test Script (please provide)

## Checklist:

- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] Any dependent changes have been merged and published in downstream modules
- [x] I have checked my code and corrected any misspellings


Problem statement:
Add multiple options for vertex ai embedding tasks
### 🚀 The feature

The feature will provide users with multiple options for embedding tasks of Vertex AI. 

### Motivation, pitch

Vertex AI offers multiple options for embedding tasks, allowing the extraction of text embeddings for specific use cases, such as distinguishing between a "retrieval document" and a "retrieval query."
If this feature is added, it would give developers greater flexibility in using mem0 based on their needs.

Hints from issue discussion:
@Dev-Khant I would like to contribute this feature. Could you please assign this to me?
Hey @rst0070 Great. Feel free to work on it.

Commit message: Add test

Patch to review:
diff --git a/docs/components/embedders/config.mdx b/docs/components/embedders/config.mdx
index a5ad8fa049..82756e2034 100644
--- a/docs/components/embedders/config.mdx
+++ b/docs/components/embedders/config.mdx
@@ -58,6 +58,9 @@ Here's a comprehensive list of all parameters that can be used across different
 | `azure_kwargs` | Key-Value arguments for the AzureOpenAI embedding model |
 | `openai_base_url`    | Base URL for OpenAI API                       | OpenAI            |
 | `vertex_credentials_json` | Path to the Google Cloud credentials JSON file for VertexAI |
+| `memory_add_embedding_type` | The type of embedding to use for the add memory action |
+| `memory_update_embedding_type` | The type of embedding to use for the update memory action |
+| `memory_search_embedding_type` | The type of embedding to use for the search memory action |
 
 
 ## Supported Embedding Models
diff --git a/docs/components/embedders/models/vertexai.mdx b/docs/components/embedders/models/vertexai.mdx
index b30c12911c..f2a0a896fc 100644
--- a/docs/components/embedders/models/vertexai.mdx
+++ b/docs/components/embedders/models/vertexai.mdx
@@ -16,7 +16,10 @@ config = {
     "embedder": {
         "provider": "vertexai",
         "config": {
-            "model": "text-embedding-004"
+            "model": "text-embedding-004",
+            "memory_add_embedding_type": "RETRIEVAL_DOCUMENT",
+            "memory_update_embedding_type": "RETRIEVAL_DOCUMENT",
+            "memory_search_embedding_type": "RETRIEVAL_QUERY"
         }
     }
 }
@@ -24,7 +27,14 @@ config = {
 m = Memory.from_config(config)
 m.add("I'm visiting Paris", user_id="john")
 ```
-
+The embedding types can be one of the following:
+- SEMANTIC_SIMILARITY
+- CLASSIFICATION
+- CLUSTERING
+- RETRIEVAL_DOCUMENT, RETRIEVAL_QUERY, QUESTION_ANSWERING, FACT_VERIFICATION
+- CODE_RETRIEVAL_QUERY  
+Check out the [Vertex AI documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types#supported_task_types) for more information.  
+  
 ### Config
 
 Here are the parameters available for configuring the Vertex AI embedder:
@@ -34,3 +44,6 @@ Here are the parameters available for configuring the Vertex AI embedder:
 | `model`                   | The name of the Vertex AI embedding model to use | `text-embedding-004` |
 | `vertex_credentials_json` | Path to the Google Cloud credentials JSON file   | `None`               |
 | `embedding_dims`          | Dimensions of the embedding model                | `256`                |
+| `memory_add_embedding_type` | The type of embedding to use for the add memory action | `RETRIEVAL_DOCUMENT` |
+| `memory_update_embedding_type` | The type of embedding to use for the update memory action | `RETRIEVAL_DOCUMENT` |
+| `memory_search_embedding_type` | The type of embedding to use for the search memory action | `RETRIEVAL_QUERY` |
diff --git a/mem0/configs/embeddings/base.py b/mem0/configs/embeddings/base.py
index a3d989ee1c..a234e3895d 100644
--- a/mem0/configs/embeddings/base.py
+++ b/mem0/configs/embeddings/base.py
@@ -27,6 +27,9 @@ def __init__(
         http_client_proxies: Optional[Union[Dict, str]] = None,
         # VertexAI specific
         vertex_credentials_json: Optional[str] = None,
+        memory_add_embedding_type: Optional[str] = None,
+        memory_update_embedding_type: Optional[str] = None,
+        memory_search_embedding_type: Optional[str] = None,
     ):
         """
         Initializes a configuration class instance for the Embeddings.
@@ -47,6 +50,14 @@ def __init__(
         :type azure_kwargs: Optional[Dict[str, Any]], defaults a dict inside init
         :param http_client_proxies: The proxy server settings used to create self.http_client, defaults to None
         :type http_client_proxies: Optional[Dict | str], optional
+        :param vertex_credentials_json: The path to the Vertex AI credentials JSON file, defaults to None
+        :type vertex_credentials_json: Optional[str], optional
+        :param memory_add_embedding_type: The type of embedding to use for the add memory action, defaults to None
+        :type memory_add_embedding_type: Optional[str], optional
+        :param memory_update_embedding_type: The type of embedding to use for the update memory action, defaults to None
+        :type memory_update_embedding_type: Optional[str], optional
+        :param memory_search_embedding_type: The type of embedding to use for the search memory action, defaults to None
+        :type memory_search_embedding_type: Optional[str], optional
         """
 
         self.model = model
@@ -68,3 +79,6 @@ def __init__(
 
         # VertexAI specific
         self.vertex_credentials_json = vertex_credentials_json
+        self.memory_add_embedding_type = memory_add_embedding_type
+        self.memory_update_embedding_type = memory_update_embedding_type
+        self.memory_search_embedding_type = memory_search_embedding_type
diff --git a/mem0/embeddings/azure_openai.py b/mem0/embeddings/azure_openai.py
index 4b33b1a2f5..8b0e89f04d 100644
--- a/mem0/embeddings/azure_openai.py
+++ b/mem0/embeddings/azure_openai.py
@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Literal, Optional
 
 from openai import AzureOpenAI
 
@@ -26,13 +26,13 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
             default_headers=default_headers,
         )
 
-    def embed(self, text):
+    def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None):
         """
         Get the embedding for the given text using OpenAI.
 
         Args:
             text (str): The text to embed.
-
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
         Returns:
             list: The embedding vector.
         """
diff --git a/mem0/embeddings/base.py b/mem0/embeddings/base.py
index c5d12b6e23..bb924c4602 100644
--- a/mem0/embeddings/base.py
+++ b/mem0/embeddings/base.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Literal, Optional
 
 from mem0.configs.embeddings.base import BaseEmbedderConfig
 
@@ -18,13 +18,13 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
             self.config = config
 
     @abstractmethod
-    def embed(self, text):
+    def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]]):
         """
         Get the embedding for the given text.
 
         Args:
             text (str): The text to embed.
-
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
         Returns:
             list: The embedding vector.
         """
diff --git a/mem0/embeddings/gemini.py b/mem0/embeddings/gemini.py
index 210848e39a..7af8606154 100644
--- a/mem0/embeddings/gemini.py
+++ b/mem0/embeddings/gemini.py
@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Literal, Optional
 
 import google.generativeai as genai
 
@@ -18,11 +18,12 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
 
         genai.configure(api_key=api_key)
 
-    def embed(self, text):
+    def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None):
         """
         Get the embedding for the given text using Google Generative AI.
         Args:
             text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
         Returns:
             list: The embedding vector.
         """
diff --git a/mem0/embeddings/huggingface.py b/mem0/embeddings/huggingface.py
index b8641cacab..88b4579c62 100644
--- a/mem0/embeddings/huggingface.py
+++ b/mem0/embeddings/huggingface.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Literal, Optional
 
 from sentence_transformers import SentenceTransformer
 
@@ -16,13 +16,13 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
 
         self.config.embedding_dims = self.config.embedding_dims or self.model.get_sentence_embedding_dimension()
 
-    def embed(self, text):
+    def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None):
         """
         Get the embedding for the given text using Hugging Face.
 
         Args:
             text (str): The text to embed.
-
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
         Returns:
             list: The embedding vector.
         """
diff --git a/mem0/embeddings/ollama.py b/mem0/embeddings/ollama.py
index 30034365c4..2f824928d0 100644
--- a/mem0/embeddings/ollama.py
+++ b/mem0/embeddings/ollama.py
@@ -1,6 +1,6 @@
 import subprocess
 import sys
-from typing import Optional
+from typing import Literal, Optional
 
 from mem0.configs.embeddings.base import BaseEmbedderConfig
 from mem0.embeddings.base import EmbeddingBase
@@ -39,13 +39,13 @@ def _ensure_model_exists(self):
         if not any(model.get("name") == self.config.model for model in local_models):
             self.client.pull(self.config.model)
 
-    def embed(self, text):
+    def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None):
         """
         Get the embedding for the given text using Ollama.
 
         Args:
             text (str): The text to embed.
-
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
         Returns:
             list: The embedding vector.
         """
diff --git a/mem0/embeddings/openai.py b/mem0/embeddings/openai.py
index b68b8ffc09..def289fa69 100644
--- a/mem0/embeddings/openai.py
+++ b/mem0/embeddings/openai.py
@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Literal, Optional
 
 from openai import OpenAI
 
@@ -18,13 +18,13 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         base_url = self.config.openai_base_url or os.getenv("OPENAI_API_BASE")
         self.client = OpenAI(api_key=api_key, base_url=base_url)
 
-    def embed(self, text):
+    def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None):
         """
         Get the embedding for the given text using OpenAI.
 
         Args:
             text (str): The text to embed.
-
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
         Returns:
             list: The embedding vector.
         """
diff --git a/mem0/embeddings/together.py b/mem0/embeddings/together.py
index 76b9124bd7..86be9050ae 100644
--- a/mem0/embeddings/together.py
+++ b/mem0/embeddings/together.py
@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Literal, Optional
 
 from together import Together
 
@@ -17,13 +17,13 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         self.config.embedding_dims = self.config.embedding_dims or 768
         self.client = Together(api_key=api_key)
 
-    def embed(self, text):
+    def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None):
         """
         Get the embedding for the given text using OpenAI.
 
         Args:
             text (str): The text to embed.
-
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
         Returns:
             list: The embedding vector.
         """
diff --git a/mem0/embeddings/vertexai.py b/mem0/embeddings/vertexai.py
index 740e630301..e309decae1 100644
--- a/mem0/embeddings/vertexai.py
+++ b/mem0/embeddings/vertexai.py
@@ -1,7 +1,7 @@
 import os
-from typing import Optional
+from typing import Literal, Optional
 
-from vertexai.language_models import TextEmbeddingModel
+from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
 
 from mem0.configs.embeddings.base import BaseEmbedderConfig
 from mem0.embeddings.base import EmbeddingBase
@@ -13,7 +13,13 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
 
         self.config.model = self.config.model or "text-embedding-004"
         self.config.embedding_dims = self.config.embedding_dims or 256
-
+        
+        self.embedding_types = {
+            "add": self.config.memory_add_embedding_type or "RETRIEVAL_DOCUMENT",
+            "update": self.config.memory_update_embedding_type or "RETRIEVAL_DOCUMENT",
+            "search": self.config.memory_search_embedding_type or "RETRIEVAL_QUERY"
+        }
+        
         credentials_path = self.config.vertex_credentials_json
 
         if credentials_path:
@@ -25,16 +31,24 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
 
         self.model = TextEmbeddingModel.from_pretrained(self.config.model)
 
-    def embed(self, text):
+    def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None):
         """
         Get the embedding for the given text using Vertex AI.
 
         Args:
             text (str): The text to embed.
-
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
         Returns:
             list: The embedding vector.
         """
-        embeddings = self.model.get_embeddings(texts=[text], output_dimensionality=self.config.embedding_dims)
+        embedding_type = "SEMANTIC_SIMILARITY"
+        if memory_action is not None:
+            if memory_action not in self.embedding_types:
+                raise ValueError(f"Invalid memory action: {memory_action}")
+            
+            embedding_type = self.embedding_types[memory_action]
+            
+        text_input = TextEmbeddingInput(text=text, task_type=embedding_type)
+        embeddings = self.model.get_embeddings(texts=[text_input], output_dimensionality=self.config.embedding_dims)
 
         return embeddings[0].values
diff --git a/mem0/memory/main.py b/mem0/memory/main.py
index c6a0f2b668..1c3a4b03d5 100644
--- a/mem0/memory/main.py
+++ b/mem0/memory/main.py
@@ -9,7 +9,7 @@
 
 import pytz
 from pydantic import ValidationError
-from mem0.memory.utils import parse_vision_messages
+
 from mem0.configs.base import MemoryConfig, MemoryItem
 from mem0.configs.prompts import get_update_memory_messages
 from mem0.memory.base import MemoryBase
@@ -19,6 +19,7 @@
 from mem0.memory.utils import (
     get_fact_retrieval_messages,
     parse_messages,
+    parse_vision_messages,
     remove_code_blocks,
 )
 from mem0.utils.factory import EmbedderFactory, LlmFactory, VectorStoreFactory
@@ -167,7 +168,7 @@ def _add_to_vector_store(self, messages, metadata, filters):
         retrieved_old_memory = []
         new_message_embeddings = {}
         for new_mem in new_retrieved_facts:
-            messages_embeddings = self.embedding_model.embed(new_mem)
+            messages_embeddings = self.embedding_model.embed(new_mem, "add")
             new_message_embeddings[new_mem] = messages_embeddings
             existing_memories = self.vector_store.search(
                 query=messages_embeddings,
@@ -446,7 +447,7 @@ def search(self, query, user_id=None, agent_id=None, run_id=None, limit=100, fil
             return original_memories
 
     def _search_vector_store(self, query, filters, limit):
-        embeddings = self.embedding_model.embed(query)
+        embeddings = self.embedding_model.embed(query, "search")
         memories = self.vector_store.search(query=embeddings, limit=limit, filters=filters)
 
         excluded_keys = {
@@ -494,7 +495,7 @@ def update(self, memory_id, data):
         """
         capture_event("mem0.update", self, {"memory_id": memory_id})
 
-        existing_embeddings = {data: self.embedding_model.embed(data)}
+        existing_embeddings = {data: self.embedding_model.embed(data, "update")}
 
         self._update_memory(memory_id, data, existing_embeddings)
         return {"message": "Memory updated successfully!"}
@@ -562,7 +563,7 @@ def _create_memory(self, data, existing_embeddings, metadata=None):
         if data in existing_embeddings:
             embeddings = existing_embeddings[data]
         else:
-            embeddings = self.embedding_model.embed(data)
+            embeddings = self.embedding_model.embed(data, "add")
         memory_id = str(uuid.uuid4())
         metadata = metadata or {}
         metadata["data"] = data
@@ -603,7 +604,7 @@ def _update_memory(self, memory_id, data, existing_embeddings, metadata=None):
         if data in existing_embeddings:
             embeddings = existing_embeddings[data]
         else:
-            embeddings = self.embedding_model.embed(data)
+            embeddings = self.embedding_model.embed(data, "update")
         self.vector_store.update(
             vector_id=memory_id,
             vector=embeddings,
diff --git a/tests/embeddings/test_vertexai_embeddings.py b/tests/embeddings/test_vertexai_embeddings.py
index 861c84bd15..3353d673db 100644
--- a/tests/embeddings/test_vertexai_embeddings.py
+++ b/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")
\ No newline at end of file
diff --git a/tests/test_main.py b/tests/test_main.py
index a311f854b2..29e662001d 100644
--- a/tests/test_main.py
+++ b/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)


Code review comments:
diff_hunk: @@ -58,6 +58,9 @@ Here's a comprehensive list of all parameters that can be used across different
 | `azure_kwargs` | Key-Value arguments for the AzureOpenAI embedding model |
 | `openai_base_url`    | Base URL for OpenAI API                       | OpenAI            |
 | `vertex_credentials_json` | Path to the Google Cloud credentials JSON file for VertexAI |
+| `memory_add_embedding_type` | The type of embedding to use for the add memory action |
path: docs/components/embedders/config.mdx
text: @user1:
Here can you please a add seperate column at the end to show for which `provider` does it supports? 

@author:
Thank you for the review! I just added the commit