Adding Gemini (#1862)

mem0/embeddings/configs.py

@@ -13,7 +13,7 @@ class EmbedderConfig(BaseModel):
     @field_validator("config")
     def validate_config(cls, v, values):
         provider = values.data.get("provider")
-        if provider in ["openai", "ollama", "huggingface", "azure_openai", "vertexai"]:
+        if provider in ["openai", "ollama", "huggingface", "azure_openai", "gemini", "vertexai"]:
             return v
         else:
             raise ValueError(f"Unsupported embedding provider: {provider}")

mem0/embeddings/gemini.py

@@ -0,0 +1,27 @@
+import os
+from typing import Optional
+import google.generativeai as genai
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+class GoogleGenAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+        if self.config.model is None:
+            self.config.model = "models/text-embedding-004" # embedding-dim = 768
+
+        genai.configure(api_key=self.config.api_key or os.getenv("GOOGLE_API_KEY"))
+
+    def embed(self, text):
+        """
+        Get the embedding for the given text using Google Generative AI.
+        Args:
+            text (str): The text to embed.
+        Returns:
+            list: The embedding vector.
+        """
+        text = text.replace("\n", " ")
+        response = genai.embed_content(model=self.config.model, content=text)
+        return response['embedding']

mem0/utils/factory.py

@@ -41,6 +41,7 @@ class EmbedderFactory:
         "ollama": "mem0.embeddings.ollama.OllamaEmbedding",
         "huggingface": "mem0.embeddings.huggingface.HuggingFaceEmbedding",
         "azure_openai": "mem0.embeddings.azure_openai.AzureOpenAIEmbedding",
+        "gemini": "mem0.embeddings.gemini.GoogleGenAIEmbedding",
     }
 
     @classmethod