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..1d28ecdc92 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 @@ -22,10 +24,15 @@ def mock_config(): with patch("mem0.configs.embeddings.base.BaseEmbedderConfig") as mock_config: mock_config.vertex_credentials_json = None yield mock_config + +@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): +def test_embed_default_model(mock_text_embedding_model, mock_os_environ, mock_config, mock_text_embedding_input): mock_config.vertex_credentials_json = "/path/to/credentials.json" mock_config.return_value.model = "text-embedding-004" mock_config.return_value.embedding_dims = 256 @@ -37,15 +44,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): +def test_embed_custom_model(mock_text_embedding_model, mock_os_environ, mock_config, mock_text_embedding_input): mock_config.vertex_credentials_json = "/path/to/credentials.json" mock_config.return_value.model = "custom-embedding-model" mock_config.return_value.embedding_dims = 512 @@ -58,10 +66,10 @@ 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] @@ -69,7 +77,7 @@ def test_embed_custom_model(mock_text_embedding_model, mock_os_environ, mock_con @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_os.getenv.return_value = "/path/to/credentials.json" mock_config.vertex_credentials_json = None config = mock_config() VertexAIEmbedding(config) 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)