From 6d535951dfdb9bbc5c287c2bad3456b0074e7869 Mon Sep 17 00:00:00 2001
From: Mayank <83648453+spike-spiegel-21@users.noreply.github.com>
Date: Sat, 9 Nov 2024 03:13:24 +0530
Subject: [PATCH] [graph_improvement]: Unique Id removed from update prompt
 (#2020)

---
 mem0/memory/graph_memory.py |  6 ++++--
 mem0/memory/utils.py        | 17 +++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/mem0/memory/graph_memory.py b/mem0/memory/graph_memory.py
index 7fcfdd7b..5fad61a4 100644
--- a/mem0/memory/graph_memory.py
+++ b/mem0/memory/graph_memory.py
@@ -1,5 +1,7 @@
 import logging
 
+from mem0.memory.utils import format_entities
+
 try:
     from langchain_community.graphs import Neo4jGraph
 except ImportError:
@@ -94,8 +96,8 @@ class MemoryGraph:
             extracted_entities = []
 
         logger.debug(f"Extracted entities: {extracted_entities}")
-
-        update_memory_prompt = get_update_memory_messages(search_output, extracted_entities)
+        search_output_string = format_entities(search_output)
+        update_memory_prompt = get_update_memory_messages(search_output_string, extracted_entities)
 
         _tools = [UPDATE_MEMORY_TOOL_GRAPH, ADD_MEMORY_TOOL_GRAPH, NOOP_TOOL]
         if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
diff --git a/mem0/memory/utils.py b/mem0/memory/utils.py
index a7e7bc35..5b0a2a1c 100644
--- a/mem0/memory/utils.py
+++ b/mem0/memory/utils.py
@@ -1,3 +1,5 @@
+import json
+
 from mem0.configs.prompts import FACT_RETRIEVAL_PROMPT
 
 
@@ -15,3 +17,18 @@ def parse_messages(messages):
         if msg["role"] == "assistant":
             response += f"assistant: {msg['content']}\n"
     return response
+
+def format_entities(entities):
+    if not entities:
+        return ""
+    
+    formatted_lines = []
+    for entity in entities:
+        simplified = {
+            "source": entity["source"],
+            "relation": entity["relation"],
+            "destination": entity["destination"]
+        }
+        formatted_lines.append(json.dumps(simplified))
+
+    return "\n".join(formatted_lines)
\ No newline at end of file