From 942014962e05279770933bb01627afee688bf9b3 Mon Sep 17 00:00:00 2001
From: Vidit Ostwal <110953813+Vidit-Ostwal@users.noreply.github.com>
Date: Sat, 19 Jul 2025 00:40:26 +0530
Subject: [PATCH] fixed save method, changed the test cases (#3187)

* fixed save method, changed the test cases

* Linting fixed
---
 src/crewai/memory/storage/mem0_storage.py | 4 +++-
 tests/storage/test_mem0_storage.py        | 8 ++------
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/crewai/memory/storage/mem0_storage.py b/src/crewai/memory/storage/mem0_storage.py
index 4fe75e3f5..bf1b5ef1e 100644
--- a/src/crewai/memory/storage/mem0_storage.py
+++ b/src/crewai/memory/storage/mem0_storage.py
@@ -64,6 +64,7 @@ class Mem0Storage(Storage):
     def save(self, value: Any, metadata: Dict[str, Any]) -> None:
         user_id = self._get_user_id()
         agent_name = self._get_agent_name()
+        assistant_message = [{"role" : "assistant","content" : value}] 
         params = None
         if self.memory_type == "short_term":
             params = {
@@ -93,7 +94,8 @@ class Mem0Storage(Storage):
         if params:
             if isinstance(self.memory, MemoryClient):
                 params["output_format"] = "v1.1"
-            self.memory.add(value, **params)
+            
+            self.memory.add(assistant_message, **params)
 
     def search(
         self,
diff --git a/tests/storage/test_mem0_storage.py b/tests/storage/test_mem0_storage.py
index b651521f9..c630b8298 100644
--- a/tests/storage/test_mem0_storage.py
+++ b/tests/storage/test_mem0_storage.py
@@ -1,14 +1,10 @@
-import os
 from unittest.mock import MagicMock, patch
 
 import pytest
 from mem0.client.main import MemoryClient
 from mem0.memory.main import Memory
 
-from crewai.agent import Agent
-from crewai.crew import Crew
 from crewai.memory.storage.mem0_storage import Mem0Storage
-from crewai.task import Task
 
 
 # Define the class (if not already defined)
@@ -172,7 +168,7 @@ def test_save_method_with_memory_oss(mem0_storage_with_mocked_config):
     mem0_storage.save(test_value, test_metadata)
     
     mem0_storage.memory.add.assert_called_once_with(
-        test_value,
+        [{'role': 'assistant' , 'content': test_value}],
         agent_id="Test_Agent",
         infer=False,
         metadata={"type": "short_term", "key": "value"},
@@ -191,7 +187,7 @@ def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_co
     mem0_storage.save(test_value, test_metadata)
     
     mem0_storage.memory.add.assert_called_once_with(
-        test_value,
+        [{'role': 'assistant' , 'content': test_value}],
         agent_id="Test_Agent",
         infer=False,
         metadata={"type": "short_term", "key": "value"},