Dropping messages from metadata in Mem0 Storage (#3390)

* Dropped messages from metadata and added user-assistant interaction directly

* Fixed test cases for this

* Fixed static type checking issue

* Changed logic to take latest user and assistant messages

* Added default value to be string

* Linting checks

* Removed duplication of tool calling

* Fixed Linting Changes

* Ruff check

* Removed console formatter file from commit

* Linting fixed

* Linting checks

* Ignoring missing imports error

* Added suggested changes

* Fixed import untyped error

Author: Vidit-Ostwal

src/crewai/memory/storage/mem0_storage.py

@@ -1,10 +1,12 @@
 import os
-from typing import Any, Dict, List
+import re
 from collections import defaultdict
-from mem0 import Memory, MemoryClient
-from crewai.utilities.chromadb import sanitize_collection_name
+from typing import Any, Iterable
+
+from mem0 import Memory, MemoryClient  # type: ignore[import-untyped]
 
 from crewai.memory.storage.interface import Storage
+from crewai.utilities.chromadb import sanitize_collection_name
 
 MAX_AGENT_ID_LENGTH_MEM0 = 255
 
@@ -86,9 +88,28 @@ def _create_filter_for_search(self):
 
         return filter
 
-    def save(self, value: Any, metadata: Dict[str, Any]) -> None:
+    def save(self, value: Any, metadata: dict[str, Any]) -> None:
+        def _last_content(messages: Iterable[dict[str, Any]], role: str) -> str:
+            return next(
+                (m.get("content", "") for m in reversed(list(messages)) if m.get("role") == role),
+                ""
+            )
+        
+        conversations = []
+        messages = metadata.pop("messages", None)
+        if messages:
+            last_user = _last_content(messages, "user")
+            last_assistant = _last_content(messages, "assistant")
+
+            if user_msg := self._get_user_message(last_user):
+                conversations.append({"role": "user", "content": user_msg})
+                
+            if assistant_msg := self._get_assistant_message(last_assistant):
+                conversations.append({"role": "assistant", "content": assistant_msg})
+        else:
+            conversations.append({"role": "assistant", "content": value})
+
         user_id = self.config.get("user_id", "")
-        assistant_message = [{"role" : "assistant","content" : value}]
 
         base_metadata = {
             "short_term": "short_term",
@@ -119,9 +140,9 @@ def save(self, value: Any, metadata: Dict[str, Any]) -> None:
         if agent_id := self.config.get("agent_id", self._get_agent_name()):
             params["agent_id"] = agent_id
 
-        self.memory.add(assistant_message, **params)
+        self.memory.add(conversations, **params)
 
-    def search(self,query: str,limit: int = 3,score_threshold: float = 0.35) -> List[Any]:
+    def search(self,query: str,limit: int = 3,score_threshold: float = 0.35) -> list[Any]:
         params = {
             "query": query,
             "limit": limit,
@@ -160,7 +181,7 @@ def search(self,query: str,limit: int = 3,score_threshold: float = 0.35) -> List
         # This makes it compatible for Contextual Memory to retrieve
         for result in results["results"]:
             result["context"] = result["memory"]
-        
+
         return [r for r in results["results"]]
 
     def reset(self):
@@ -181,3 +202,16 @@ def _get_agent_name(self) -> str:
         agents = [self._sanitize_role(agent.role) for agent in agents]
         agents = "_".join(agents)
         return sanitize_collection_name(name=agents, max_collection_length=MAX_AGENT_ID_LENGTH_MEM0)
+
+    def _get_assistant_message(self, text: str) -> str:
+        marker = "Final Answer:"
+        if marker in text:
+            return text.split(marker, 1)[1].strip()
+        return text
+
+    def _get_user_message(self, text: str) -> str:
+        pattern = r"User message:\s*(.*)"
+        match = re.search(pattern, text)
+        if match:
+            return match.group(1).strip()
+        return text

tests/storage/test_mem0_storage.py

@@ -16,8 +16,7 @@ def __init__(self):
 @pytest.fixture
 def mock_mem0_memory():
     """Fixture to create a mock Memory instance"""
-    mock_memory = MagicMock(spec=Memory)
-    return mock_memory
+    return MagicMock(spec=Memory)
 
 
 @pytest.fixture
@@ -73,8 +72,7 @@ def test_mem0_storage_initialization(mem0_storage_with_mocked_config, mock_mem0_
 @pytest.fixture
 def mock_mem0_memory_client():
     """Fixture to create a mock MemoryClient instance"""
-    mock_memory = MagicMock(spec=MemoryClient)
-    return mock_memory
+    return MagicMock(spec=MemoryClient)
 
 
 @pytest.fixture
@@ -96,8 +94,7 @@ def mem0_storage_with_memory_client_using_config_from_crew(mock_mem0_memory_clie
                     "infer": True
                 }
 
-        mem0_storage = Mem0Storage(type="short_term", crew=crew, config=embedder_config)
-        return mem0_storage
+        return Mem0Storage(type="short_term", crew=crew, config=embedder_config)
 
 
 @pytest.fixture
@@ -111,8 +108,7 @@ def mem0_storage_with_memory_client_using_explictly_config(mock_mem0_memory_clie
         crew = MockCrew()
         new_config = {"provider": "mem0", "config": {"api_key": "new-api-key"}}
 
-        mem0_storage = Mem0Storage(type="short_term", crew=crew, config=new_config)
-        return mem0_storage
+        return Mem0Storage(type="short_term", crew=crew, config=new_config)
 
 
 def test_mem0_storage_with_memory_client_initialization(
@@ -172,14 +168,14 @@ def test_save_method_with_memory_oss(mem0_storage_with_mocked_config):
 
     # Test short_term memory type (already set in fixture)
     test_value = "This is a test memory"
-    test_metadata = {"key": "value"}
+    test_metadata = {'description': 'Respond to user conversation. User message: What do you know about me?', 'messages': [{'role': 'system', 'content': 'You are Friendly chatbot assistant. You are a kind and knowledgeable chatbot assistant. You excel at understanding user needs, providing helpful responses, and maintaining engaging conversations. You remember previous interactions to provide a personalized experience.\nYour personal goal is: Engage in useful and interesting conversations with users while remembering context.\nTo give my best complete final answer to the task respond using the exact following format:\n\nThought: I now can give a great answer\nFinal Answer: Your final answer must be the great and the most complete as possible, it must be outcome described.\n\nI MUST use these formats, my job depends on it!'}, {'role': 'user', 'content': '\nCurrent Task: Respond to user conversation. User message: What do you know about me?\n\nThis is the expected criteria for your final answer: Contextually appropriate, helpful, and friendly response.\nyou MUST return the actual complete content as the final answer, not a summary.\n\n# Useful context: \nExternal memories:\n- User is from India\n- User is interested in the solar system\n- User name is Vidit Ostwal\n- User is interested in French cuisine\n\nBegin! This is VERY important to you, use the tools available and give your best Final Answer, your job depends on it!\n\nThought:'}, {'role': 'assistant', 'content': "I now can give a great answer  \nFinal Answer: Hi Vidit! From our previous conversations, I know you're from India and have a great interest in the solar system. It's fascinating to explore the wonders of space, isn't it? Also, I remember you have a passion for French cuisine, which has so many delightful dishes to explore. If there's anything specific you'd like to discuss or learn about—whether it's about the solar system or some great French recipes—feel free to let me know! I'm here to help."}], 'agent': 'Friendly chatbot assistant'}
 
     mem0_storage.save(test_value, test_metadata)
 
     mem0_storage.memory.add.assert_called_once_with(
-        [{"role": "assistant" , "content": test_value}],
+        [{'role': 'user', 'content': 'What do you know about me?'}, {'role': 'assistant', 'content': "Hi Vidit! From our previous conversations, I know you're from India and have a great interest in the solar system. It's fascinating to explore the wonders of space, isn't it? Also, I remember you have a passion for French cuisine, which has so many delightful dishes to explore. If there's anything specific you'd like to discuss or learn about—whether it's about the solar system or some great French recipes—feel free to let me know! I'm here to help."}],
         infer=True,
-        metadata={"type": "short_term", "key": "value"},
+        metadata={'type': 'short_term', 'description': 'Respond to user conversation. User message: What do you know about me?', 'agent': 'Friendly chatbot assistant'},
         run_id="my_run_id",
         user_id="test_user",
         agent_id='Test_Agent'
@@ -191,14 +187,14 @@ def test_save_method_with_multiple_agents(mem0_storage_with_mocked_config):
     mem0_storage.memory.add = MagicMock()
 
     test_value = "This is a test memory"
-    test_metadata = {"key": "value"}
+    test_metadata = {'description': 'Respond to user conversation. User message: What do you know about me?', 'messages': [{'role': 'system', 'content': 'You are Friendly chatbot assistant. You are a kind and knowledgeable chatbot assistant. You excel at understanding user needs, providing helpful responses, and maintaining engaging conversations. You remember previous interactions to provide a personalized experience.\nYour personal goal is: Engage in useful and interesting conversations with users while remembering context.\nTo give my best complete final answer to the task respond using the exact following format:\n\nThought: I now can give a great answer\nFinal Answer: Your final answer must be the great and the most complete as possible, it must be outcome described.\n\nI MUST use these formats, my job depends on it!'}, {'role': 'user', 'content': '\nCurrent Task: Respond to user conversation. User message: What do you know about me?\n\nThis is the expected criteria for your final answer: Contextually appropriate, helpful, and friendly response.\nyou MUST return the actual complete content as the final answer, not a summary.\n\n# Useful context: \nExternal memories:\n- User is from India\n- User is interested in the solar system\n- User name is Vidit Ostwal\n- User is interested in French cuisine\n\nBegin! This is VERY important to you, use the tools available and give your best Final Answer, your job depends on it!\n\nThought:'}, {'role': 'assistant', 'content': "I now can give a great answer  \nFinal Answer: Hi Vidit! From our previous conversations, I know you're from India and have a great interest in the solar system. It's fascinating to explore the wonders of space, isn't it? Also, I remember you have a passion for French cuisine, which has so many delightful dishes to explore. If there's anything specific you'd like to discuss or learn about—whether it's about the solar system or some great French recipes—feel free to let me know! I'm here to help."}], 'agent': 'Friendly chatbot assistant'}
 
     mem0_storage.save(test_value, test_metadata)
 
     mem0_storage.memory.add.assert_called_once_with(
-        [{"role": "assistant" , "content": test_value}],
+        [{'role': 'user', 'content': 'What do you know about me?'}, {'role': 'assistant', 'content': "Hi Vidit! From our previous conversations, I know you're from India and have a great interest in the solar system. It's fascinating to explore the wonders of space, isn't it? Also, I remember you have a passion for French cuisine, which has so many delightful dishes to explore. If there's anything specific you'd like to discuss or learn about—whether it's about the solar system or some great French recipes—feel free to let me know! I'm here to help."}],
         infer=True,
-        metadata={"type": "short_term", "key": "value"},
+        metadata={'type': 'short_term', 'description': 'Respond to user conversation. User message: What do you know about me?', 'agent': 'Friendly chatbot assistant'},
         run_id="my_run_id",
         user_id="test_user",
         agent_id='Test_Agent_Test_Agent_2_Test_Agent_3'
@@ -212,14 +208,14 @@ def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_co
 
     # Test short_term memory type (already set in fixture)
     test_value = "This is a test memory"
-    test_metadata = {"key": "value"}
+    test_metadata = {'description': 'Respond to user conversation. User message: What do you know about me?', 'messages': [{'role': 'system', 'content': 'You are Friendly chatbot assistant. You are a kind and knowledgeable chatbot assistant. You excel at understanding user needs, providing helpful responses, and maintaining engaging conversations. You remember previous interactions to provide a personalized experience.\nYour personal goal is: Engage in useful and interesting conversations with users while remembering context.\nTo give my best complete final answer to the task respond using the exact following format:\n\nThought: I now can give a great answer\nFinal Answer: Your final answer must be the great and the most complete as possible, it must be outcome described.\n\nI MUST use these formats, my job depends on it!'}, {'role': 'user', 'content': '\nCurrent Task: Respond to user conversation. User message: What do you know about me?\n\nThis is the expected criteria for your final answer: Contextually appropriate, helpful, and friendly response.\nyou MUST return the actual complete content as the final answer, not a summary.\n\n# Useful context: \nExternal memories:\n- User is from India\n- User is interested in the solar system\n- User name is Vidit Ostwal\n- User is interested in French cuisine\n\nBegin! This is VERY important to you, use the tools available and give your best Final Answer, your job depends on it!\n\nThought:'}, {'role': 'assistant', 'content': "I now can give a great answer  \nFinal Answer: Hi Vidit! From our previous conversations, I know you're from India and have a great interest in the solar system. It's fascinating to explore the wonders of space, isn't it? Also, I remember you have a passion for French cuisine, which has so many delightful dishes to explore. If there's anything specific you'd like to discuss or learn about—whether it's about the solar system or some great French recipes—feel free to let me know! I'm here to help."}], 'agent': 'Friendly chatbot assistant'}
 
     mem0_storage.save(test_value, test_metadata)
 
     mem0_storage.memory.add.assert_called_once_with(
-        [{'role': 'assistant' , 'content': test_value}],
+        [{'role': 'user', 'content': 'What do you know about me?'}, {'role': 'assistant', 'content': "Hi Vidit! From our previous conversations, I know you're from India and have a great interest in the solar system. It's fascinating to explore the wonders of space, isn't it? Also, I remember you have a passion for French cuisine, which has so many delightful dishes to explore. If there's anything specific you'd like to discuss or learn about—whether it's about the solar system or some great French recipes—feel free to let me know! I'm here to help."}],
         infer=True,
-        metadata={"type": "short_term", "key": "value"},
+       metadata={'type': 'short_term', 'description': 'Respond to user conversation. User message: What do you know about me?', 'agent': 'Friendly chatbot assistant'},
         version="v2",
         run_id="my_run_id",
         includes="include1",