Fix: Migrate Gemini Embeddings (#3002)

Co-authored-by: Dev-Khant <devkhant24@gmail.com>

mem0/embeddings/gemini.py

@@ -1,7 +1,7 @@
 import os
 from typing import Literal, Optional
 
-import google.generativeai as genai
+import google.genai as genai
 
 from mem0.configs.embeddings.base import BaseEmbedderConfig
 from mem0.embeddings.base import EmbeddingBase
@@ -12,23 +12,28 @@ class GoogleGenAIEmbedding(EmbeddingBase):
         super().__init__(config)
 
         self.config.model = self.config.model or "models/text-embedding-004"
-        self.config.embedding_dims = self.config.embedding_dims or 768
+        self.config.embedding_dims = self.config.embedding_dims or self.config.output_dimensionality or 768
 
         api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
 
-        genai.configure(api_key=api_key)
+        if api_key:
+            self.client = genai.Client(api_key="api_key")
+        else:
+            self.client = genai.Client()
 
     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.
+            memory_action (optional): The type of embedding to use. (Currently not used by Gemini for task_type)
         Returns:
             list: The embedding vector.
         """
         text = text.replace("\n", " ")
-        response = genai.embed_content(
+
+        response = self.client.models.embed_content(
             model=self.config.model, content=text, output_dimensionality=self.config.embedding_dims
         )
+
         return response["embedding"]

mem0/llms/gemini.py

@@ -4,7 +4,7 @@ from typing import Dict, List, Optional
 try:
     from google import genai
     from google.genai import types
-    
+
 except ImportError:
     raise ImportError(
         "The 'google-generativeai' library is required. Please install it using 'pip install google-generativeai'."
@@ -49,16 +49,17 @@ class GeminiLLM(LLMBase):
             for part in candidate.content.parts:
                 fn = getattr(part, "function_call", None)
                 if fn:
-                    processed_response["tool_calls"].append({
-                        "name": fn.name,
-                        "arguments": fn.args,
-                    })
+                    processed_response["tool_calls"].append(
+                        {
+                            "name": fn.name,
+                            "arguments": fn.args,
+                        }
+                    )
 
             return processed_response
 
         return content
 
-
     def _reformat_messages(self, messages: List[Dict[str, str]]) -> List[types.Content]:
         """
         Reformat messages for Gemini using google.genai.types.
@@ -78,15 +79,11 @@ class GeminiLLM(LLMBase):
                 content = message["content"]
 
             new_messages.append(
-                types.Content(
-                    role="model" if message["role"] == "model" else "user",
-                    parts=[types.Part(text=content)]
-                )
+                types.Content(role="model" if message["role"] == "model" else "user", parts=[types.Part(text=content)])
             )
 
         return new_messages
 
-
     def _reformat_tools(self, tools: Optional[List[Dict]]):
         """
         Reformat tools for Gemini.
@@ -131,7 +128,6 @@ class GeminiLLM(LLMBase):
         tools: Optional[List[Dict]] = None,
         tool_choice: str = "auto",
     ):
-        
         """
         Generate a response based on the given messages using Gemini.
 
@@ -161,31 +157,22 @@ class GeminiLLM(LLMBase):
             tool_config = types.ToolConfig(
                 function_calling_config=types.FunctionCallingConfig(
                     mode=tool_choice.upper(),  # Assuming 'any' should become 'ANY', etc.
-                    allowed_function_names=[
-                        tool["function"]["name"] for tool in tools
-                    ] if tool_choice == "any" else None
+                    allowed_function_names=[tool["function"]["name"] for tool in tools]
+                    if tool_choice == "any"
+                    else None,
                 )
             )
 
-        print(f"Tool config: {tool_config}")
-        print(f"Params: {params}" )
-        print(f"Messages: {messages}")
-        print(f"Tools: {tools}")
-        print(f"Reformatted messages: {self._reformat_messages(messages)}")
-        print(f"Reformatted tools: {self._reformat_tools(tools)}")
-
         response = self.client_gemini.models.generate_content(
-                model=self.config.model,
-                contents=self._reformat_messages(messages),
-                config=types.GenerateContentConfig(
-                    temperature= self.config.temperature,
-                    max_output_tokens= self.config.max_tokens,
-                    top_p= self.config.top_p,
-                    tools=self._reformat_tools(tools),
-                    tool_config=tool_config,
-
-                ),
-            )
-        print(f"Response test: {response}")
+            model=self.config.model,
+            contents=self._reformat_messages(messages),
+            config=types.GenerateContentConfig(
+                temperature=self.config.temperature,
+                max_output_tokens=self.config.max_tokens,
+                top_p=self.config.top_p,
+                tools=self._reformat_tools(tools),
+                tool_config=tool_config,
+            ),
+        )
 
         return self._parse_response(response, tools)