fixed save method, changed the test cases (#3187)
Some checks failed
Notify Downstream / notify-downstream (push) Has been cancelled
Mark stale issues and pull requests / stale (push) Has been cancelled

* fixed save method, changed the test cases

* Linting fixed
This commit is contained in:
Vidit Ostwal
2025-07-19 00:40:26 +05:30
committed by GitHub
parent 2ab79a7dd5
commit 942014962e
2 changed files with 5 additions and 7 deletions

View File

@@ -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,

View File

@@ -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"},