feat(embeddings): add native embed_batch to 5 embedders (LMStudio, Together, HuggingFace, VertexAI, GoogleGenAI) (#5609)

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

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

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

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

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