Docs Update (#2591)

2025-04-29 08:15:25 -07:00
parent 6d13e83001
commit 393a4fd5a6
111 changed files with 2296 additions and 99 deletions
@@ -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()
@@ -0,0 +1,189 @@
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+from mem0 import MemoryClient
+import json
+import time
+from jinja2 import Template
+from openai import OpenAI
+from prompts import ANSWER_PROMPT_GRAPH, ANSWER_PROMPT
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+class MemorySearch:
+
+    def __init__(self, output_path='results.json', top_k=10, filter_memories=False, 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.top_k = top_k
+        self.openai_client = OpenAI()
+        self.results = defaultdict(list)
+        self.output_path = output_path
+        self.filter_memories = filter_memories
+        self.is_graph = is_graph
+
+        if self.is_graph:
+            self.ANSWER_PROMPT = ANSWER_PROMPT_GRAPH
+        else:
+            self.ANSWER_PROMPT = ANSWER_PROMPT
+
+    def search_memory(self, user_id, query, max_retries=3, retry_delay=1):
+        start_time = time.time()
+        retries = 0
+        while retries < max_retries:
+            try:
+                if self.is_graph:
+                    print("Searching with graph")
+                    memories = self.mem0_client.search(query, user_id=user_id, top_k=self.top_k,
+                                                filter_memories=self.filter_memories, enable_graph=True, output_format='v1.1')
+                else:
+                    memories = self.mem0_client.search(query, user_id=user_id, top_k=self.top_k,
+                                                filter_memories=self.filter_memories)
+                break
+            except Exception as e:
+                print("Retrying...")
+                retries += 1
+                if retries >= max_retries:
+                    raise e
+                time.sleep(retry_delay)
+
+        end_time = time.time()
+        if not self.is_graph:
+            semantic_memories = [{'memory': memory['memory'],
+                        'timestamp': memory['metadata']['timestamp'],
+                        'score': round(memory['score'], 2)}
+                    for memory in memories]
+            graph_memories = None
+        else:
+            semantic_memories = [{'memory': memory['memory'],
+                        'timestamp': memory['metadata']['timestamp'],
+                        'score': round(memory['score'], 2)} for memory in memories['results']]
+            graph_memories = [{"source": relation['source'], "relationship": relation['relationship'], "target": relation['target']} for relation in memories['relations']]
+        return semantic_memories, graph_memories, end_time - start_time
+
+    def answer_question(self, speaker_1_user_id, speaker_2_user_id, question, answer, category):
+        speaker_1_memories, speaker_1_graph_memories, speaker_1_memory_time = self.search_memory(speaker_1_user_id, question)
+        speaker_2_memories, speaker_2_graph_memories, speaker_2_memory_time = self.search_memory(speaker_2_user_id, question)
+
+        search_1_memory = [f"{item['timestamp']}: {item['memory']}" 
+                        for item in speaker_1_memories]
+        search_2_memory = [f"{item['timestamp']}: {item['memory']}" 
+                          for item in speaker_2_memories]
+
+        template = Template(self.ANSWER_PROMPT)
+        answer_prompt = template.render(
+            speaker_1_user_id=speaker_1_user_id.split('_')[0],
+            speaker_2_user_id=speaker_2_user_id.split('_')[0],
+            speaker_1_memories=json.dumps(search_1_memory, indent=4),
+            speaker_2_memories=json.dumps(search_2_memory, indent=4),
+            speaker_1_graph_memories=json.dumps(speaker_1_graph_memories, indent=4),
+            speaker_2_graph_memories=json.dumps(speaker_2_graph_memories, indent=4),
+            question=question
+        )
+
+        t1 = time.time()
+        response = self.openai_client.chat.completions.create(
+            model=os.getenv("MODEL"),
+            messages=[
+                {"role": "system", "content": answer_prompt}
+            ],
+            temperature=0.0
+        )
+        t2 = time.time()
+        response_time = t2 - t1
+        return response.choices[0].message.content, speaker_1_memories, speaker_2_memories, speaker_1_memory_time, speaker_2_memory_time, speaker_1_graph_memories, speaker_2_graph_memories, response_time
+
+    def process_question(self, val, speaker_a_user_id, speaker_b_user_id):
+        question = val.get('question', '')
+        answer = val.get('answer', '')
+        category = val.get('category', -1)
+        evidence = val.get('evidence', [])
+        adversarial_answer = val.get('adversarial_answer', '')
+
+        response, speaker_1_memories, speaker_2_memories, speaker_1_memory_time, speaker_2_memory_time, speaker_1_graph_memories, speaker_2_graph_memories, response_time = self.answer_question(
+            speaker_a_user_id,
+            speaker_b_user_id,
+            question,
+            answer,
+            category
+        )
+
+        result = {
+            "question": question,
+            "answer": answer,
+            "category": category,
+            "evidence": evidence,
+            "response": response,
+            "adversarial_answer": adversarial_answer,
+            "speaker_1_memories": speaker_1_memories,
+            "speaker_2_memories": speaker_2_memories,
+            'num_speaker_1_memories': len(speaker_1_memories),
+            'num_speaker_2_memories': len(speaker_2_memories),
+            'speaker_1_memory_time': speaker_1_memory_time,
+            'speaker_2_memory_time': speaker_2_memory_time,
+            "speaker_1_graph_memories": speaker_1_graph_memories,
+            "speaker_2_graph_memories": speaker_2_graph_memories,
+            "response_time": response_time
+        }
+
+        # Save results after each question is processed
+        with open(self.output_path, 'w') as f:
+            json.dump(self.results, f, indent=4)
+
+        return result
+
+    def process_data_file(self, file_path):
+        with open(file_path, 'r') as f:
+            data = json.load(f)
+
+        for idx, item in tqdm(enumerate(data), total=len(data), desc="Processing conversations"):
+            qa = item['qa']
+            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}"
+
+            for question_item in tqdm(qa, total=len(qa), desc=f"Processing questions for conversation {idx}", leave=False):
+                result = self.process_question(
+                    question_item,
+                    speaker_a_user_id,
+                    speaker_b_user_id
+                )
+                self.results[idx].append(result)
+
+                # Save results after each question is processed
+                with open(self.output_path, 'w') as f:
+                    json.dump(self.results, f, indent=4)
+
+        # Final save at the end
+        with open(self.output_path, 'w') as f:
+            json.dump(self.results, f, indent=4)
+
+    def process_questions_parallel(self, qa_list, speaker_a_user_id, speaker_b_user_id, max_workers=1):
+        def process_single_question(val):
+            result = self.process_question(val, speaker_a_user_id, speaker_b_user_id)
+            # Save results after each question is processed
+            with open(self.output_path, 'w') as f:
+                json.dump(self.results, f, indent=4)
+            return result
+
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            results = list(tqdm(
+                executor.map(process_single_question, qa_list),
+                total=len(qa_list),
+                desc="Answering Questions"
+            ))
+
+        # Final save at the end
+        with open(self.output_path, 'w') as f:
+            json.dump(self.results, f, indent=4)
+
+        return results