mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 08:38:30 +00:00
fixed save method, changed the test cases (#3187)
* fixed save method, changed the test cases * Linting fixed
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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"},
|
||||
|
||||
Reference in New Issue
Block a user