Formatting (#2750)

2025-05-22 01:17:29 +05:30
parent dff91154a7
commit d85fcda037
71 changed files with 1391 additions and 1823 deletions
@@ -13,7 +13,7 @@ load_dotenv()


 # Update custom instructions
-custom_instructions ="""
+custom_instructions = """
 Generate personal memories that follow these guidelines:

 1. Each memory should be self-contained with complete context, including:
@@ -47,7 +47,7 @@ class MemoryADD:
        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")
+            project_id=os.getenv("MEM0_PROJECT_ID"),
        )

        self.mem0_client.update_project(custom_instructions=custom_instructions)
@@ -59,15 +59,16 @@ class MemoryADD:
            self.load_data()

    def load_data(self):
-        with open(self.data_path, 'r') as f:
+        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)
+                _ = 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:
@@ -78,13 +79,13 @@ class MemoryADD:

    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]
+            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']
+        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}"
@@ -94,7 +95,7 @@ class MemoryADD:
        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:
+            if key in ["speaker_a", "speaker_b"] or "date" in key or "timestamp" in key:
                continue

            date_time_key = key + "_date_time"
@@ -104,10 +105,10 @@ class MemoryADD:
            messages = []
            messages_reverse = []
            for chat in chats:
-                if chat['speaker'] == speaker_a:
+                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:
+                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:
@@ -116,11 +117,11 @@ class MemoryADD:
            # 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")
+                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")
+                args=(speaker_b_user_id, messages_reverse, timestamp, "Adding Memories for Speaker B"),
            )

            thread_a.start()
@@ -134,10 +135,7 @@ class MemoryADD:
        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)
-            ]
+            futures = [executor.submit(self.process_conversation, item, idx) for idx, item in enumerate(self.data)]

            for future in futures:
-                future.result()
+                future.result()
@@ -16,12 +16,11 @@ load_dotenv()


 class MemorySearch:
-
-    def __init__(self, output_path='results.json', top_k=10, filter_memories=False, is_graph=False):
+    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")
+            project_id=os.getenv("MEM0_PROJECT_ID"),
        )
        self.top_k = top_k
        self.openai_client = OpenAI()
@@ -42,11 +41,18 @@ class MemorySearch:
            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')
+                    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)
+                    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...")
@@ -57,64 +63,86 @@ class MemorySearch:

        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]
+            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']]
+            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)
+        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]
+        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_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
+            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
+            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
+        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', '')
+        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
-        )
+        (
+            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,
@@ -125,67 +153,63 @@ class MemorySearch:
            "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,
+            "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
+            "response_time": response_time,
        }

        # Save results after each question is processed
-        with open(self.output_path, 'w') as f:
+        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:
+        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']
+            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
-                )
+            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:
+                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:
+        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:
+            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"
-            ))
+            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:
+        with open(self.output_path, "w") as f:
            json.dump(self.results, f, indent=4)

        return results