Docs Update (#2591)

evaluation/src/memzero/add.py

@@ -0,0 +1,141 @@
+from mem0 import MemoryClient
+import json
+import time
+import os
+import threading
+from tqdm import tqdm
+from concurrent.futures import ThreadPoolExecutor
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+# Update custom instructions
+custom_instructions ="""
+Generate personal memories that follow these guidelines:
+
+1. Each memory should be self-contained with complete context, including:
+   - The person's name, do not use "user" while creating memories
+   - Personal details (career aspirations, hobbies, life circumstances)
+   - Emotional states and reactions
+   - Ongoing journeys or future plans
+   - Specific dates when events occurred
+
+2. Include meaningful personal narratives focusing on:
+   - Identity and self-acceptance journeys
+   - Family planning and parenting
+   - Creative outlets and hobbies
+   - Mental health and self-care activities
+   - Career aspirations and education goals
+   - Important life events and milestones
+
+3. Make each memory rich with specific details rather than general statements
+   - Include timeframes (exact dates when possible)
+   - Name specific activities (e.g., "charity race for mental health" rather than just "exercise")
+   - Include emotional context and personal growth elements
+
+4. Extract memories only from user messages, not incorporating assistant responses
+
+5. Format each memory as a paragraph with a clear narrative structure that captures the person's experience, challenges, and aspirations
+"""
+
+
+class MemoryADD:
+    def __init__(self, data_path=None, batch_size=2, is_graph=False):
+        self.mem0_client = MemoryClient(
+            api_key=os.getenv("MEM0_API_KEY"),
+            org_id=os.getenv("MEM0_ORGANIZATION_ID"),
+            project_id=os.getenv("MEM0_PROJECT_ID")
+        )
+
+        self.mem0_client.update_project(custom_instructions=custom_instructions)
+        self.batch_size = batch_size
+        self.data_path = data_path
+        self.data = None
+        self.is_graph = is_graph
+        if data_path:
+            self.load_data()
+
+    def load_data(self):
+        with open(self.data_path, 'r') as f:
+            self.data = json.load(f)
+        return self.data
+
+    def add_memory(self, user_id, message, metadata, retries=3):
+        for attempt in range(retries):
+            try:
+                _ = self.mem0_client.add(message, user_id=user_id, version="v2",
+                                         metadata=metadata, enable_graph=self.is_graph)
+                return
+            except Exception as e:
+                if attempt < retries - 1:
+                    time.sleep(1)  # Wait before retrying
+                    continue
+                else:
+                    raise e
+
+    def add_memories_for_speaker(self, speaker, messages, timestamp, desc):
+        for i in tqdm(range(0, len(messages), self.batch_size), desc=desc):
+            batch_messages = messages[i:i+self.batch_size]
+            self.add_memory(speaker, batch_messages, metadata={"timestamp": timestamp})
+
+    def process_conversation(self, item, idx):
+        conversation = item['conversation']
+        speaker_a = conversation['speaker_a']
+        speaker_b = conversation['speaker_b']
+
+        speaker_a_user_id = f"{speaker_a}_{idx}"
+        speaker_b_user_id = f"{speaker_b}_{idx}"
+
+        # delete all memories for the two users
+        self.mem0_client.delete_all(user_id=speaker_a_user_id)
+        self.mem0_client.delete_all(user_id=speaker_b_user_id)
+
+        for key in conversation.keys():
+            if key in ['speaker_a', 'speaker_b'] or "date" in key or "timestamp" in key:
+                continue
+
+            date_time_key = key + "_date_time"
+            timestamp = conversation[date_time_key]
+            chats = conversation[key]
+
+            messages = []
+            messages_reverse = []
+            for chat in chats:
+                if chat['speaker'] == speaker_a:
+                    messages.append({"role": "user", "content": f"{speaker_a}: {chat['text']}"})
+                    messages_reverse.append({"role": "assistant", "content": f"{speaker_a}: {chat['text']}"})
+                elif chat['speaker'] == speaker_b:
+                    messages.append({"role": "assistant", "content": f"{speaker_b}: {chat['text']}"})
+                    messages_reverse.append({"role": "user", "content": f"{speaker_b}: {chat['text']}"})
+                else:
+                    raise ValueError(f"Unknown speaker: {chat['speaker']}")
+
+            # add memories for the two users on different threads
+            thread_a = threading.Thread(
+                target=self.add_memories_for_speaker,
+                args=(speaker_a_user_id, messages, timestamp, "Adding Memories for Speaker A")
+            )
+            thread_b = threading.Thread(
+                target=self.add_memories_for_speaker,
+                args=(speaker_b_user_id, messages_reverse, timestamp, "Adding Memories for Speaker B")
+            )
+
+            thread_a.start()
+            thread_b.start()
+            thread_a.join()
+            thread_b.join()
+
+        print("Messages added successfully")
+
+    def process_all_conversations(self, max_workers=10):
+        if not self.data:
+            raise ValueError("No data loaded. Please set data_path and call load_data() first.")
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = [
+                executor.submit(self.process_conversation, item, idx)
+                for idx, item in enumerate(self.data)
+            ]
+
+            for future in futures:
+                future.result()