Fix not working with Gemini models (#2021)

mem0/memory/main.py

@@ -16,7 +16,11 @@ from mem0.memory.base import MemoryBase
 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_messages,
+    remove_code_blocks,
+)
 from mem0.utils.factory import EmbedderFactory, LlmFactory, VectorStoreFactory
 
 # Setup user config
@@ -152,6 +156,7 @@ class Memory(MemoryBase):
         )
 
         try:
+            response = remove_code_blocks(response)
             new_retrieved_facts = json.loads(response)["facts"]
         except Exception as e:
             logging.error(f"Error in new_retrieved_facts: {e}")
@@ -184,6 +189,8 @@ class Memory(MemoryBase):
             messages=[{"role": "user", "content": function_calling_prompt}],
             response_format={"type": "json_object"},
         )
+
+        new_memories_with_actions = remove_code_blocks(new_memories_with_actions)
         new_memories_with_actions = json.loads(new_memories_with_actions)
 
         returned_memories = []

mem0/memory/utils.py

@@ -1,3 +1,4 @@
+import re
 import json
 
 from mem0.configs.prompts import FACT_RETRIEVAL_PROMPT
@@ -21,10 +22,24 @@ def parse_messages(messages):
 def format_entities(entities):
     if not entities:
         return ""
-    
+
     formatted_lines = []
     for entity in entities:
         simplified = f"{entity['source']} -- {entity['relatationship']} -- {entity['destination']}"
         formatted_lines.append(simplified)
 
-    return "\n".join(formatted_lines)
+    return "\n".join(formatted_lines)
+
+
+def remove_code_blocks(content: str) -> str:
+    """
+    Removes enclosing code block markers ```[language] and ``` from a given string.
+
+    Remarks:
+    - The function uses a regex pattern to match code blocks that may start with ``` followed by an optional language tag (letters or numbers) and end with ```.
+    - If a code block is detected, it returns only the inner content, stripping out the markers.
+    - If no code block markers are found, the original content is returned as-is.
+    """
+    pattern = r"^```[a-zA-Z0-9]*\n([\s\S]*?)\n```$"
+    match = re.match(pattern, content.strip())
+    return match.group(1).strip() if match else content.strip()