Fix: Add Google Genai library support (mem0ai#2941)

Author: akshat1423

docs/components/llms/models/gemini.mdx

@@ -4,36 +4,44 @@ title: Gemini
 
 <Snippet file="paper-release.mdx" />
 
-To use Gemini model, you have to set the `GEMINI_API_KEY` environment variable. You can obtain the Gemini API key from the [Google AI Studio](https://aistudio.google.com/app/apikey)
+To use the Gemini model, set the `GEMINI_API_KEY` environment variable. You can obtain the Gemini API key from [Google AI Studio](https://aistudio.google.com/app/apikey).
+
+> **Note:** As of the latest release, Mem0 uses the new `google.genai` SDK instead of the deprecated `google.generativeai`. All message formatting and model interaction now use the updated `types` module from `google.genai`.
+
+> **Note:** Some Gemini models are being deprecated and will retire soon. It is recommended to migrate to the latest stable models like `"gemini-2.0-flash-001"` or `"gemini-2.0-flash-lite-001"` to ensure ongoing support and improvements.
 
 ## Usage
 
 ```python
 import os
 from mem0 import Memory
 
-os.environ["OPENAI_API_KEY"] = "your-api-key" # used for embedding model
-os.environ["GEMINI_API_KEY"] = "your-api-key"
+os.environ["OPENAI_API_KEY"] = "your-openai-api-key"  # Used for embedding model
+os.environ["GEMINI_API_KEY"] = "your-gemini-api-key"
 
 config = {
     "llm": {
         "provider": "gemini",
         "config": {
-            "model": "gemini-1.5-flash-latest",
+            "model": "gemini-2.0-flash-001",
             "temperature": 0.2,
             "max_tokens": 2000,
+            "top_p": 1.0
         }
     }
 }
 
 m = Memory.from_config(config)
+
 messages = [
     {"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
-    {"role": "assistant", "content": "How about a thriller movies? They can be quite engaging."},
-    {"role": "user", "content": "I’m not a big fan of thriller movies but I love sci-fi movies."},
-    {"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
+    {"role": "assistant", "content": "How about thriller movies? They can be quite engaging."},
+    {"role": "user", "content": "I’m not a big fan of thrillers, but I love sci-fi movies."},
+    {"role": "assistant", "content": "Got it! I'll avoid thrillers and suggest sci-fi movies instead."}
 ]
+
 m.add(messages, user_id="alice", metadata={"category": "movies"})
+
 ```
 
 ## Config

docs/open-source/graph_memory/overview.mdx

@@ -238,16 +238,24 @@ The Mem0's graph supports the following operations:
 ### Add Memories
 
 <Note>
-If you are using Mem0 with Graph Memory, it is recommended to pass `user_id`. Use `userId` in NodeSDK.
+Mem0 with Graph Memory supports both `user_id` and `agent_id` parameters. You can use either or both to organize your memories. Use `userId` and `agentId` in NodeSDK.
 </Note>
 
 <CodeGroup>
 ```python Python
+# Using only user_id
 m.add("I like pizza", user_id="alice")
+
+# Using both user_id and agent_id
+m.add("I like pizza", user_id="alice", agent_id="food-assistant")
 ```
 
 ```typescript TypeScript
+// Using only userId
 memory.add("I like pizza", { userId: "alice" });
+
+// Using both userId and agentId
+memory.add("I like pizza", { userId: "alice", agentId: "food-assistant" });
 ```
 
 ```json Output
@@ -260,11 +268,19 @@ memory.add("I like pizza", { userId: "alice" });
 
 <CodeGroup>
 ```python Python
+# Get all memories for a user
 m.get_all(user_id="alice")
+
+# Get all memories for a specific agent belonging to a user
+m.get_all(user_id="alice", agent_id="food-assistant")
 ```
 
 ```typescript TypeScript
+// Get all memories for a user
 memory.getAll({ userId: "alice" });
+
+// Get all memories for a specific agent belonging to a user
+memory.getAll({ userId: "alice", agentId: "food-assistant" });
 ```
 
 ```json Output
@@ -277,7 +293,8 @@ memory.getAll({ userId: "alice" });
             'metadata': None,
             'created_at': '2024-08-20T14:09:27.588719-07:00',
             'updated_at': None,
-            'user_id': 'alice'
+            'user_id': 'alice',
+            'agent_id': 'food-assistant'
         }
     ],
     'entities': [
@@ -295,11 +312,19 @@ memory.getAll({ userId: "alice" });
 
 <CodeGroup>
 ```python Python
+# Search memories for a user
 m.search("tell me my name.", user_id="alice")
+
+# Search memories for a specific agent belonging to a user
+m.search("tell me my name.", user_id="alice", agent_id="food-assistant")
 ```
 
 ```typescript TypeScript
+// Search memories for a user
 memory.search("tell me my name.", { userId: "alice" });
+
+// Search memories for a specific agent belonging to a user
+memory.search("tell me my name.", { userId: "alice", agentId: "food-assistant" });
 ```
 
 ```json Output
@@ -312,7 +337,8 @@ memory.search("tell me my name.", { userId: "alice" });
             'metadata': None,
             'created_at': '2024-08-20T14:09:27.588719-07:00',
             'updated_at': None,
-            'user_id': 'alice'
+            'user_id': 'alice',
+            'agent_id': 'food-assistant'
         }
     ],
     'entities': [
@@ -331,11 +357,19 @@ memory.search("tell me my name.", { userId: "alice" });
 
 <CodeGroup>
 ```python Python
+# Delete all memories for a user
 m.delete_all(user_id="alice")
+
+# Delete all memories for a specific agent belonging to a user
+m.delete_all(user_id="alice", agent_id="food-assistant")
 ```
 
 ```typescript TypeScript
+// Delete all memories for a user
 memory.deleteAll({ userId: "alice" });
+
+// Delete all memories for a specific agent belonging to a user
+memory.deleteAll({ userId: "alice", agentId: "food-assistant" });
 ```
 </CodeGroup>
 
@@ -516,6 +550,42 @@ memory.search("Who is spiderman?", { userId: "alice123" });
 
 > **Note:** The Graph Memory implementation is not standalone. You will be adding/retrieving memories to the vector store and the graph store simultaneously.
 
+## Using Multiple Agents with Graph Memory
+
+When working with multiple agents, you can use the `agent_id` parameter to organize memories by both user and agent. This allows you to:
+
+1. Create agent-specific knowledge graphs
+2. Share common knowledge between agents
+3. Isolate sensitive or specialized information to specific agents
+
+### Example: Multi-Agent Setup
+
+<CodeGroup>
+```python Python
+# Add memories for different agents
+m.add("I prefer Italian cuisine", user_id="bob", agent_id="food-assistant")
+m.add("I'm allergic to peanuts", user_id="bob", agent_id="health-assistant")
+m.add("I live in Seattle", user_id="bob")  # Shared across all agents
+
+# Search within specific agent context
+food_preferences = m.search("What food do I like?", user_id="bob", agent_id="food-assistant")
+health_info = m.search("What are my allergies?", user_id="bob", agent_id="health-assistant")
+location = m.search("Where do I live?", user_id="bob")  # Searches across all agents
+```
+
+```typescript TypeScript
+// Add memories for different agents
+memory.add("I prefer Italian cuisine", { userId: "bob", agentId: "food-assistant" });
+memory.add("I'm allergic to peanuts", { userId: "bob", agentId: "health-assistant" });
+memory.add("I live in Seattle", { userId: "bob" });  // Shared across all agents
+
+// Search within specific agent context
+const foodPreferences = memory.search("What food do I like?", { userId: "bob", agentId: "food-assistant" });
+const healthInfo = memory.search("What are my allergies?", { userId: "bob", agentId: "health-assistant" });
+const location = memory.search("Where do I live?", { userId: "bob" });  // Searches across all agents
+```
+</CodeGroup>
+
 If you want to use a managed version of Mem0, please check out [Mem0](https://mem0.dev/pd). If you have any questions, please feel free to reach out to us using one of the following methods:
 
-<Snippet file="get-help.mdx" />
+<Snippet file="get-help.mdx" />

embedchain/poetry.lock

@@ -2,9 +2,9 @@
 from typing import Dict, List, Optional
 
 try:
-    import google.generativeai as genai
-    from google.generativeai import GenerativeModel, protos
-    from google.generativeai.types import content_types
+    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'."
@@ -22,66 +22,71 @@ def __init__(self, config: Optional[BaseLlmConfig] = None):
             self.config.model = "gemini-1.5-flash-latest"
 
         api_key = self.config.api_key or os.getenv("GEMINI_API_KEY")
-        genai.configure(api_key=api_key)
-        self.client = GenerativeModel(model_name=self.config.model)
+        self.client_gemini = genai.Client(
+            api_key=api_key,
+        )
 
     def _parse_response(self, response, tools):
         """
         Process the response based on whether tools are used or not.
 
         Args:
-            response: The raw response from API.
+            response: The raw response from the API.
             tools: The list of tools provided in the request.
 
         Returns:
             str or dict: The processed response.
         """
+        candidate = response.candidates[0]
+        content = candidate.content.parts[0].text if candidate.content.parts else None
+
         if tools:
             processed_response = {
-                "content": (content if (content := response.candidates[0].content.parts[0].text) else None),
+                "content": content,
                 "tool_calls": [],
             }
 
-            for part in response.candidates[0].content.parts:
-                if fn := part.function_call:
-                    if isinstance(fn, protos.FunctionCall):
-                        fn_call = type(fn).to_dict(fn)
-                        processed_response["tool_calls"].append({"name": fn_call["name"], "arguments": fn_call["args"]})
-                        continue
-                    processed_response["tool_calls"].append({"name": fn.name, "arguments": fn.args})
+            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,
+                    })
 
             return processed_response
-        else:
-            return response.candidates[0].content.parts[0].text
 
-    def _reformat_messages(self, messages: List[Dict[str, str]]):
+        return content
+
+
+    def _reformat_messages(self, messages: List[Dict[str, str]]) -> List[types.Content]:
         """
-        Reformat messages for Gemini.
+        Reformat messages for Gemini using google.genai.types.
 
         Args:
             messages: The list of messages provided in the request.
 
         Returns:
-            list: The list of messages in the required format.
+            list: A list of types.Content objects with proper role and parts.
         """
         new_messages = []
 
         for message in messages:
             if message["role"] == "system":
                 content = "THIS IS A SYSTEM PROMPT. YOU MUST OBEY THIS: " + message["content"]
-
             else:
                 content = message["content"]
 
             new_messages.append(
-                {
-                    "parts": content,
-                    "role": "model" if message["role"] == "model" else "user",
-                }
+                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.
@@ -126,6 +131,7 @@ def generate_response(
         tools: Optional[List[Dict]] = None,
         tool_choice: str = "auto",
     ):
+        
         """
         Generate a response based on the given messages using Gemini.
 
@@ -149,23 +155,37 @@ def generate_response(
             params["response_mime_type"] = "application/json"
             if "schema" in response_format:
                 params["response_schema"] = response_format["schema"]
+
+        tool_config = None
         if tool_choice:
-            tool_config = content_types.to_tool_config(
-                {
-                    "function_calling_config": {
-                        "mode": tool_choice,
-                        "allowed_function_names": (
-                            [tool["function"]["name"] for tool in tools] if tool_choice == "any" else None
-                        ),
-                    }
-                }
+            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
+                )
             )
 
-        response = self.client.generate_content(
-            contents=self._reformat_messages(messages),
-            tools=self._reformat_tools(tools),
-            generation_config=genai.GenerationConfig(**params),
-            tool_config=tool_config,
-        )
+        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}")
 
         return self._parse_response(response, tools)