From 3e2ae734e7fb4c52025f71e8f9471bd00b2a8c0e Mon Sep 17 00:00:00 2001 From: Alok Tripathi Date: Thu, 18 Jun 2026 16:46:31 +0530 Subject: [PATCH] feat(embeddings): add native embed_batch to 5 embedders (LMStudio, Together, HuggingFace, VertexAI, GoogleGenAI) (#5609) --- mem0/embeddings/gemini.py | 17 +++ mem0/embeddings/huggingface.py | 22 ++++ mem0/embeddings/lmstudio.py | 14 ++ mem0/embeddings/together.py | 13 ++ mem0/embeddings/vertexai.py | 21 +++ tests/embeddings/test_gemini_emeddings.py | 59 +++++++++ .../embeddings/test_huggingface_embeddings.py | 68 ++++++++++ tests/embeddings/test_lm_studio_embeddings.py | 53 ++++++++ tests/embeddings/test_together_embeddings.py | 67 ++++++++++ tests/embeddings/test_vertexai_embeddings.py | 121 +++++++++++++++++- 10 files changed, 454 insertions(+), 1 deletion(-) create mode 100644 tests/embeddings/test_together_embeddings.py diff --git a/mem0/embeddings/gemini.py b/mem0/embeddings/gemini.py index 207b2c8aa..9058b90f5 100644 --- a/mem0/embeddings/gemini.py +++ b/mem0/embeddings/gemini.py @@ -37,3 +37,20 @@ class GoogleGenAIEmbedding(EmbeddingBase): response = self.client.models.embed_content(model=self.config.model, contents=text, config=config) return response.embeddings[0].values + + def embed_batch(self, texts, memory_action="add"): + if not texts: + return [] + config = types.EmbedContentConfig(output_dimensionality=self.config.embedding_dims) + MAX_BATCH = 100 + all_embeddings = [] + for i in range(0, len(texts), MAX_BATCH): + chunk = [t.replace("\n", " ") for t in texts[i : i + MAX_BATCH]] + response = self.client.models.embed_content(model=self.config.model, contents=chunk, config=config) + all_embeddings.extend(e.values for e in response.embeddings) + if len(all_embeddings) != len(texts): + raise ValueError( + f"Gemini embed_batch() returned {len(all_embeddings)} embeddings for {len(texts)} texts " + f"using model '{self.config.model}'" + ) + return all_embeddings diff --git a/mem0/embeddings/huggingface.py b/mem0/embeddings/huggingface.py index 705374e86..230fb7e20 100644 --- a/mem0/embeddings/huggingface.py +++ b/mem0/embeddings/huggingface.py @@ -42,3 +42,25 @@ class HuggingFaceEmbedding(EmbeddingBase): ).data[0].embedding else: return self.model.encode(text, convert_to_numpy=True).tolist() + + def embed_batch(self, texts, memory_action="add"): + if not texts: + return [] + if self.config.huggingface_base_url: + response = self.client.embeddings.create(input=texts, model=self.config.model, **self.config.model_kwargs) + sorted_data = sorted(response.data, key=lambda x: x.index) + embeddings = [item.embedding for item in sorted_data] + if len(embeddings) != len(texts): + raise ValueError( + f"HuggingFace embed_batch() returned {len(embeddings)} embeddings for {len(texts)} texts" + f" using model '{self.config.model}'" + ) + return embeddings + else: + result = self.model.encode(texts, convert_to_numpy=True).tolist() + if len(result) != len(texts): + raise ValueError( + f"HuggingFace embed_batch() returned {len(result)} embeddings for {len(texts)} texts" + f" using model '{self.config.model}'" + ) + return result diff --git a/mem0/embeddings/lmstudio.py b/mem0/embeddings/lmstudio.py index 159dce525..f8943e430 100644 --- a/mem0/embeddings/lmstudio.py +++ b/mem0/embeddings/lmstudio.py @@ -27,3 +27,17 @@ class LMStudioEmbedding(EmbeddingBase): """ text = text.replace("\n", " ") return self.client.embeddings.create(input=[text], model=self.config.model).data[0].embedding + + def embed_batch(self, texts, memory_action="add"): + if not texts: + return [] + cleaned = [t.replace("\n", " ") for t in texts] + response = self.client.embeddings.create(input=cleaned, model=self.config.model) + sorted_data = sorted(response.data, key=lambda x: x.index) + embeddings = [item.embedding for item in sorted_data] + if len(embeddings) != len(texts): + raise ValueError( + f"LM Studio embed_batch() returned {len(embeddings)} embeddings for {len(texts)} texts" + f" using model '{self.config.model}'" + ) + return embeddings diff --git a/mem0/embeddings/together.py b/mem0/embeddings/together.py index b3eca0b70..253c40a4c 100644 --- a/mem0/embeddings/together.py +++ b/mem0/embeddings/together.py @@ -29,3 +29,16 @@ class TogetherEmbedding(EmbeddingBase): """ return self.client.embeddings.create(model=self.config.model, input=text).data[0].embedding + + def embed_batch(self, texts, memory_action="add"): + if not texts: + return [] + response = self.client.embeddings.create(model=self.config.model, input=texts) + sorted_data = sorted(response.data, key=lambda x: x.index) + embeddings = [item.embedding for item in sorted_data] + if len(embeddings) != len(texts): + raise ValueError( + f"Together embed_batch() returned {len(embeddings)} embeddings for {len(texts)} texts" + f" using model '{self.config.model}'" + ) + return embeddings diff --git a/mem0/embeddings/vertexai.py b/mem0/embeddings/vertexai.py index 5003b62db..039bfd5f8 100644 --- a/mem0/embeddings/vertexai.py +++ b/mem0/embeddings/vertexai.py @@ -62,3 +62,24 @@ class VertexAIEmbedding(EmbeddingBase): embeddings = self.model.get_embeddings(texts=[text_input], output_dimensionality=self.config.embedding_dims) return embeddings[0].values + + def embed_batch(self, texts, memory_action="add"): + if not texts: + return [] + 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] + all_embeddings = [] + for i in range(0, len(texts), 250): + chunk = texts[i : i + 250] + inputs = [TextEmbeddingInput(text=t, task_type=embedding_type) for t in chunk] + results = self.model.get_embeddings(texts=inputs, output_dimensionality=self.config.embedding_dims) + all_embeddings.extend(r.values for r in results) + if len(all_embeddings) != len(texts): + raise ValueError( + f"Vertex AI embed_batch() returned {len(all_embeddings)} embeddings for {len(texts)} texts" + f" using model '{self.config.model}'" + ) + return all_embeddings diff --git a/tests/embeddings/test_gemini_emeddings.py b/tests/embeddings/test_gemini_emeddings.py index 0792c3c58..78f182612 100644 --- a/tests/embeddings/test_gemini_emeddings.py +++ b/tests/embeddings/test_gemini_emeddings.py @@ -58,3 +58,62 @@ def test_config_initialization(config): assert embedder.config.api_key == "dummy_api_key" assert embedder.config.model == "test_model" assert embedder.config.embedding_dims == 786 + + +def test_embed_batch_single_call(mock_genai, config): + emb0 = type("Embedding", (), {"values": [0.1, 0.2, 0.3]})() + emb1 = type("Embedding", (), {"values": [0.4, 0.5, 0.6]})() + mock_genai.return_value = type("Response", (), {"embeddings": [emb0, emb1]})() + + embedder = GoogleGenAIEmbedding(config) + + texts = ["First text.", "Second text."] + result = embedder.embed_batch(texts) + + mock_genai.assert_called_once_with(model="test_model", contents=texts, config=ANY) + assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + + +def test_embed_batch_empty_list(mock_genai, config): + embedder = GoogleGenAIEmbedding(config) + + result = embedder.embed_batch([]) + + assert result == [] + mock_genai.assert_not_called() + + +def test_embed_batch_count_mismatch_raises(mock_genai, config): + emb0 = type("Embedding", (), {"values": [0.1, 0.2, 0.3]})() + mock_genai.return_value = type("Response", (), {"embeddings": [emb0]})() + + embedder = GoogleGenAIEmbedding(config) + + with pytest.raises(ValueError, match="returned 1 embeddings for 2 texts"): + embedder.embed_batch(["first text", "second text"]) + + +def test_embed_batch_chunks_over_100_texts(mock_genai, config): + def make_chunk_response(**kwargs): + chunk = kwargs["contents"] + emb = type("Embedding", (), {"values": [0.1, 0.2]}) + return type("Response", (), {"embeddings": [emb() for _ in chunk]})() + + mock_genai.side_effect = make_chunk_response + + embedder = GoogleGenAIEmbedding(config) + texts = [f"text {i}" for i in range(150)] + result = embedder.embed_batch(texts) + + assert mock_genai.call_count == 2 + assert len(result) == 150 + + +def test_embed_batch_strips_newlines(mock_genai, config): + emb0 = type("Embedding", (), {"values": [0.1, 0.2, 0.3]})() + mock_genai.return_value = type("Response", (), {"embeddings": [emb0]})() + + embedder = GoogleGenAIEmbedding(config) + embedder.embed_batch(["line one\nline two"]) + + mock_genai.assert_called_once_with(model="test_model", contents=["line one line two"], config=ANY) diff --git a/tests/embeddings/test_huggingface_embeddings.py b/tests/embeddings/test_huggingface_embeddings.py index c7bddd31f..91f3ffe40 100644 --- a/tests/embeddings/test_huggingface_embeddings.py +++ b/tests/embeddings/test_huggingface_embeddings.py @@ -101,3 +101,71 @@ def test_embed_with_huggingface_base_url(): truncate=True, ) assert result == [0.1, 0.2, 0.3] + + +def test_embed_batch_sentence_transformer(mock_sentence_transformer): + config = BaseEmbedderConfig() + embedder = HuggingFaceEmbedding(config) + + mock_sentence_transformer.encode.return_value = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + + texts = ["First text.", "Second text."] + result = embedder.embed_batch(texts) + + mock_sentence_transformer.encode.assert_called_once_with(texts, convert_to_numpy=True) + assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + + +def test_embed_batch_empty_list_sentence_transformer(mock_sentence_transformer): + config = BaseEmbedderConfig() + embedder = HuggingFaceEmbedding(config) + + result = embedder.embed_batch([]) + + assert result == [] + mock_sentence_transformer.encode.assert_not_called() + + +def test_embed_batch_base_url(): + config = BaseEmbedderConfig(huggingface_base_url="http://localhost:8080", model="my-custom-model") + with patch("mem0.embeddings.huggingface.OpenAI") as mock_openai: + mock_client = Mock() + mock_openai.return_value = mock_client + + mock_item0 = Mock(index=0, embedding=[0.1, 0.2, 0.3]) + mock_item1 = Mock(index=1, embedding=[0.4, 0.5, 0.6]) + mock_client.embeddings.create.return_value = Mock(data=[mock_item0, mock_item1]) + + embedder = HuggingFaceEmbedding(config) + texts = ["First text.", "Second text."] + result = embedder.embed_batch(texts) + + mock_client.embeddings.create.assert_called_once_with( + input=texts, model="my-custom-model" + ) + assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + + +def test_embed_batch_count_mismatch_raises_base_url(): + config = BaseEmbedderConfig(huggingface_base_url="http://localhost:8080", model="my-custom-model") + with patch("mem0.embeddings.huggingface.OpenAI") as mock_openai: + mock_client = Mock() + mock_openai.return_value = mock_client + + mock_item0 = Mock(index=0, embedding=[0.1, 0.2, 0.3]) + mock_client.embeddings.create.return_value = Mock(data=[mock_item0]) + + embedder = HuggingFaceEmbedding(config) + + with pytest.raises(ValueError, match="returned 1 embeddings for 2 texts"): + embedder.embed_batch(["first text", "second text"]) + + +def test_embed_batch_count_mismatch_raises_sentence_transformer(mock_sentence_transformer): + config = BaseEmbedderConfig() + embedder = HuggingFaceEmbedding(config) + + mock_sentence_transformer.encode.return_value = np.array([[0.1, 0.2, 0.3]]) + + with pytest.raises(ValueError, match="returned 1 embeddings for 2 texts"): + embedder.embed_batch(["first text", "second text"]) diff --git a/tests/embeddings/test_lm_studio_embeddings.py b/tests/embeddings/test_lm_studio_embeddings.py index e37476c8f..c391fc389 100644 --- a/tests/embeddings/test_lm_studio_embeddings.py +++ b/tests/embeddings/test_lm_studio_embeddings.py @@ -6,6 +6,7 @@ from mem0.configs.embeddings.base import BaseEmbedderConfig from mem0.embeddings.lmstudio import LMStudioEmbedding + @pytest.fixture def mock_lm_studio_client(): with patch("mem0.embeddings.lmstudio.OpenAI") as mock_openai: @@ -27,3 +28,55 @@ def test_embed_text(mock_lm_studio_client): ) assert embedding == [0.1, 0.2, 0.3, 0.4, 0.5] + + +def test_embed_batch_single_call(mock_lm_studio_client): + config = BaseEmbedderConfig(model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf", embedding_dims=512) + embedder = LMStudioEmbedding(config) + + mock_item0 = Mock(index=0, embedding=[0.1, 0.2, 0.3]) + mock_item1 = Mock(index=1, embedding=[0.4, 0.5, 0.6]) + mock_lm_studio_client.embeddings.create.return_value = Mock(data=[mock_item0, mock_item1]) + + texts = ["First text.", "Second text."] + embeddings = embedder.embed_batch(texts) + + mock_lm_studio_client.embeddings.create.assert_called_once_with( + input=texts, model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf" + ) + assert embeddings == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + + +def test_embed_batch_empty_list(mock_lm_studio_client): + config = BaseEmbedderConfig(model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf", embedding_dims=512) + embedder = LMStudioEmbedding(config) + + result = embedder.embed_batch([]) + + assert result == [] + mock_lm_studio_client.embeddings.create.assert_not_called() + + +def test_embed_batch_strips_newlines(mock_lm_studio_client): + config = BaseEmbedderConfig(model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf", embedding_dims=512) + embedder = LMStudioEmbedding(config) + + mock_item0 = Mock(index=0, embedding=[0.1, 0.2, 0.3]) + mock_lm_studio_client.embeddings.create.return_value = Mock(data=[mock_item0]) + + embedder.embed_batch(["line one\nline two"]) + + mock_lm_studio_client.embeddings.create.assert_called_once_with( + input=["line one line two"], model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf" + ) + + +def test_embed_batch_count_mismatch_raises(mock_lm_studio_client): + config = BaseEmbedderConfig(model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf", embedding_dims=512) + embedder = LMStudioEmbedding(config) + + mock_item0 = Mock(index=0, embedding=[0.1, 0.2, 0.3]) + mock_lm_studio_client.embeddings.create.return_value = Mock(data=[mock_item0]) + + with pytest.raises(ValueError, match="returned 1 embeddings for 2 texts"): + embedder.embed_batch(["first text", "second text"]) diff --git a/tests/embeddings/test_together_embeddings.py b/tests/embeddings/test_together_embeddings.py new file mode 100644 index 000000000..86dbe7939 --- /dev/null +++ b/tests/embeddings/test_together_embeddings.py @@ -0,0 +1,67 @@ +from unittest.mock import Mock, patch + +import pytest + +from mem0.configs.embeddings.base import BaseEmbedderConfig +from mem0.embeddings.together import TogetherEmbedding + + +@pytest.fixture +def mock_together_client(): + with patch("mem0.embeddings.together.Together") as mock_together: + mock_client = Mock() + mock_together.return_value = mock_client + yield mock_client + + +def test_embed_text(mock_together_client): + config = BaseEmbedderConfig(model="togethercomputer/m2-bert-80M-8k-retrieval", embedding_dims=768) + embedder = TogetherEmbedding(config) + + mock_together_client.embeddings.create.return_value = Mock(data=[Mock(embedding=[0.1, 0.2, 0.3, 0.4, 0.5])]) + + text = "Sample text to embed." + embedding = embedder.embed(text) + + mock_together_client.embeddings.create.assert_called_once_with( + model="togethercomputer/m2-bert-80M-8k-retrieval", input=text + ) + assert embedding == [0.1, 0.2, 0.3, 0.4, 0.5] + + +def test_embed_batch_single_call(mock_together_client): + config = BaseEmbedderConfig(model="togethercomputer/m2-bert-80M-8k-retrieval", embedding_dims=768) + embedder = TogetherEmbedding(config) + + mock_item0 = Mock(index=0, embedding=[0.1, 0.2, 0.3]) + mock_item1 = Mock(index=1, embedding=[0.4, 0.5, 0.6]) + mock_together_client.embeddings.create.return_value = Mock(data=[mock_item0, mock_item1]) + + texts = ["First text.", "Second text."] + embeddings = embedder.embed_batch(texts) + + mock_together_client.embeddings.create.assert_called_once_with( + model="togethercomputer/m2-bert-80M-8k-retrieval", input=texts + ) + assert embeddings == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + + +def test_embed_batch_empty_list(mock_together_client): + config = BaseEmbedderConfig(model="togethercomputer/m2-bert-80M-8k-retrieval", embedding_dims=768) + embedder = TogetherEmbedding(config) + + result = embedder.embed_batch([]) + + assert result == [] + mock_together_client.embeddings.create.assert_not_called() + + +def test_embed_batch_count_mismatch_raises(mock_together_client): + config = BaseEmbedderConfig(model="togethercomputer/m2-bert-80M-8k-retrieval", embedding_dims=768) + embedder = TogetherEmbedding(config) + + mock_item0 = Mock(index=0, embedding=[0.1, 0.2, 0.3]) + mock_together_client.embeddings.create.return_value = Mock(data=[mock_item0]) + + with pytest.raises(ValueError, match="returned 1 embeddings for 2 texts"): + embedder.embed_batch(["first text", "second text"]) diff --git a/tests/embeddings/test_vertexai_embeddings.py b/tests/embeddings/test_vertexai_embeddings.py index 45644262c..b78bff66e 100644 --- a/tests/embeddings/test_vertexai_embeddings.py +++ b/tests/embeddings/test_vertexai_embeddings.py @@ -1,4 +1,4 @@ -from unittest.mock import Mock, patch +from unittest.mock import ANY, Mock, patch import pytest @@ -159,3 +159,122 @@ def test_invalid_memory_action(mock_text_embedding_model, mock_config): with pytest.raises(ValueError): embedder.embed("Hello world", memory_action="invalid_action") + + +@patch("mem0.embeddings.vertexai.TextEmbeddingModel") +def test_embed_batch_single_call(mock_text_embedding_model, mock_os_environ, mock_config): + mock_config.return_value.model = "gemini-embedding-001" + mock_config.return_value.embedding_dims = 256 + + config = mock_config() + embedder = VertexAIEmbedding(config) + + mock_emb0 = Mock(values=[0.1, 0.2, 0.3]) + mock_emb1 = Mock(values=[0.4, 0.5, 0.6]) + mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_emb0, mock_emb1] + + texts = ["First text.", "Second text."] + result = embedder.embed_batch(texts) + + mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with( + texts=ANY, output_dimensionality=256 + ) + assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + + +@patch("mem0.embeddings.vertexai.TextEmbeddingModel") +def test_embed_batch_empty_list(mock_text_embedding_model, mock_os_environ, mock_config): + mock_config.return_value.model = "gemini-embedding-001" + mock_config.return_value.embedding_dims = 256 + + config = mock_config() + embedder = VertexAIEmbedding(config) + + result = embedder.embed_batch([]) + + assert result == [] + mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_not_called() + + +@patch("mem0.embeddings.vertexai.TextEmbeddingModel") +def test_embed_batch_count_mismatch_raises(mock_text_embedding_model, mock_os_environ, mock_config): + mock_config.return_value.model = "gemini-embedding-001" + mock_config.return_value.embedding_dims = 256 + + config = mock_config() + embedder = VertexAIEmbedding(config) + + mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [Mock(values=[0.1, 0.2, 0.3])] + + with pytest.raises(ValueError, match="returned 1 embeddings for 2 texts"): + embedder.embed_batch(["first text", "second text"]) + + +@patch("mem0.embeddings.vertexai.TextEmbeddingInput") +@patch("mem0.embeddings.vertexai.TextEmbeddingModel") +def test_embed_batch_default_memory_action_uses_add( + mock_text_embedding_model, mock_text_embedding_input, mock_os_environ, mock_config +): + mock_config.return_value.model = "gemini-embedding-001" + mock_config.return_value.embedding_dims = 256 + mock_config.return_value.memory_add_embedding_type = None + + config = mock_config() + embedder = VertexAIEmbedding(config) + + mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [Mock(values=[0.1, 0.2])] + + embedder.embed_batch(["some text"]) # no memory_action — default "add" + + mock_text_embedding_input.assert_called_once_with(text="some text", task_type="RETRIEVAL_DOCUMENT") + + +@patch("mem0.embeddings.vertexai.TextEmbeddingInput") +@patch("mem0.embeddings.vertexai.TextEmbeddingModel") +def test_embed_batch_none_memory_action_uses_default( + mock_text_embedding_model, mock_text_embedding_input, mock_os_environ, mock_config +): + mock_config.return_value.model = "gemini-embedding-001" + mock_config.return_value.embedding_dims = 256 + + config = mock_config() + embedder = VertexAIEmbedding(config) + + mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [Mock(values=[0.1, 0.2])] + + embedder.embed_batch(["some text"], memory_action=None) + + mock_text_embedding_input.assert_called_once_with(text="some text", task_type="SEMANTIC_SIMILARITY") + + +@patch("mem0.embeddings.vertexai.TextEmbeddingModel") +def test_embed_batch_invalid_memory_action_raises(mock_text_embedding_model, mock_os_environ, mock_config): + mock_config.return_value.model = "gemini-embedding-001" + mock_config.return_value.embedding_dims = 256 + + config = mock_config() + embedder = VertexAIEmbedding(config) + + with pytest.raises(ValueError, match="Invalid memory action"): + embedder.embed_batch(["some text"], memory_action="invalid_action") + + +@patch("mem0.embeddings.vertexai.TextEmbeddingModel") +def test_embed_batch_chunking_triggers_two_api_calls(mock_text_embedding_model, mock_os_environ, mock_config): + """300 texts must produce exactly 2 get_embeddings calls (chunks of 250 and 50).""" + mock_config.return_value.model = "gemini-embedding-001" + mock_config.return_value.embedding_dims = 256 + + config = mock_config() + embedder = VertexAIEmbedding(config) + + def make_chunk_response(texts, output_dimensionality): + return [Mock(values=[0.1, 0.2]) for _ in texts] + + mock_text_embedding_model.from_pretrained.return_value.get_embeddings.side_effect = make_chunk_response + + texts = [f"text {i}" for i in range(300)] + result = embedder.embed_batch(texts) + + assert mock_text_embedding_model.from_pretrained.return_value.get_embeddings.call_count == 2 + assert len(result) == 300