From b846293f686fcdb2cd9e2ff837aef38d42ce4509 Mon Sep 17 00:00:00 2001
From: WindLeaf233 <mc.windleaf@foxmail.com>
Date: Fri, 13 Sep 2024 12:46:11 +0800
Subject: [PATCH] Fixed json decode when LLM produces code block

---
 mem0/memory/main.py  | 5 ++++-
 mem0/memory/utils.py | 8 ++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/mem0/memory/main.py b/mem0/memory/main.py
index 5516227cb1..c200dbb748 100644
--- a/mem0/memory/main.py
+++ b/mem0/memory/main.py
@@ -15,7 +15,7 @@
 from mem0.memory.setup import setup_config
 from mem0.memory.storage import SQLiteManager
 from mem0.memory.telemetry import capture_event
-from mem0.memory.utils import get_fact_retrieval_messages, parse_messages
+from mem0.memory.utils import get_fact_retrieval_messages, parse_json_response, parse_messages
 from mem0.utils.factory import LlmFactory, EmbedderFactory, VectorStoreFactory
 from mem0.configs.base import MemoryItem, MemoryConfig
 
@@ -144,6 +144,8 @@ def _add_to_vector_store(self, messages, metadata, filters):
             response_format={"type": "json_object"},
         )
 
+        response = parse_json_response(response)
+
         try:
             new_retrieved_facts = json.loads(response)[
                 "facts"
@@ -170,6 +172,7 @@ def _add_to_vector_store(self, messages, metadata, filters):
             messages=[{"role": "user", "content": function_calling_prompt}],
             response_format={"type": "json_object"},
         )
+        new_memories_with_actions = parse_json_response(new_memories_with_actions)
         new_memories_with_actions = json.loads(new_memories_with_actions)
 
         returned_memories = []
diff --git a/mem0/memory/utils.py b/mem0/memory/utils.py
index a0c82fedb5..4f8a3693c7 100644
--- a/mem0/memory/utils.py
+++ b/mem0/memory/utils.py
@@ -1,3 +1,5 @@
+import re
+
 from mem0.configs.prompts import FACT_RETRIEVAL_PROMPT
 
 
@@ -14,3 +16,9 @@ def parse_messages(messages):
             if msg["role"] == "assistant":
                 response += f"assistant: {msg['content']}\n"
         return response
+
+def parse_json_response(response):
+    search_result = re.search("(```json)((.*\n)+)(```)", response)
+    if search_result:
+        response = search_result.group(2).strip()
+    return response