mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
1b00cc71ef
* 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
347 lines
20 KiB
Python
347 lines
20 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from mem0.client.main import MemoryClient
|
|
from mem0.memory.main import Memory
|
|
|
|
from crewai.memory.storage.mem0_storage import Mem0Storage
|
|
|
|
|
|
# Define the class (if not already defined)
|
|
class MockCrew:
|
|
def __init__(self):
|
|
self.agents = [MagicMock(role="Test Agent")]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_mem0_memory():
|
|
"""Fixture to create a mock Memory instance"""
|
|
return MagicMock(spec=Memory)
|
|
|
|
|
|
@pytest.fixture
|
|
def mem0_storage_with_mocked_config(mock_mem0_memory):
|
|
"""Fixture to create a Mem0Storage instance with mocked dependencies"""
|
|
|
|
# Patch the Memory class to return our mock
|
|
with patch("mem0.memory.main.Memory.from_config", return_value=mock_mem0_memory) as mock_from_config:
|
|
config = {
|
|
"vector_store": {
|
|
"provider": "mock_vector_store",
|
|
"config": {"host": "localhost", "port": 6333},
|
|
},
|
|
"llm": {
|
|
"provider": "mock_llm",
|
|
"config": {"api_key": "mock-api-key", "model": "mock-model"},
|
|
},
|
|
"embedder": {
|
|
"provider": "mock_embedder",
|
|
"config": {"api_key": "mock-api-key", "model": "mock-model"},
|
|
},
|
|
"graph_store": {
|
|
"provider": "mock_graph_store",
|
|
"config": {
|
|
"url": "mock-url",
|
|
"username": "mock-user",
|
|
"password": "mock-password",
|
|
},
|
|
},
|
|
"history_db_path": "/mock/path",
|
|
"version": "test-version",
|
|
"custom_fact_extraction_prompt": "mock prompt 1",
|
|
"custom_update_memory_prompt": "mock prompt 2",
|
|
}
|
|
|
|
# Parameters like run_id, includes, and excludes doesn't matter in Memory OSS
|
|
crew = MockCrew()
|
|
|
|
embedder_config={"user_id": "test_user", "local_mem0_config": config, "run_id": "my_run_id", "includes": "include1","excludes": "exclude1", "infer" : True}
|
|
|
|
mem0_storage = Mem0Storage(type="short_term", crew=crew, config=embedder_config)
|
|
return mem0_storage, mock_from_config, config
|
|
|
|
|
|
def test_mem0_storage_initialization(mem0_storage_with_mocked_config, mock_mem0_memory):
|
|
"""Test that Mem0Storage initializes correctly with the mocked config"""
|
|
mem0_storage, mock_from_config, config = mem0_storage_with_mocked_config
|
|
assert mem0_storage.memory_type == "short_term"
|
|
assert mem0_storage.memory is mock_mem0_memory
|
|
mock_from_config.assert_called_once_with(config)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_mem0_memory_client():
|
|
"""Fixture to create a mock MemoryClient instance"""
|
|
return MagicMock(spec=MemoryClient)
|
|
|
|
|
|
@pytest.fixture
|
|
def mem0_storage_with_memory_client_using_config_from_crew(mock_mem0_memory_client):
|
|
"""Fixture to create a Mem0Storage instance with mocked dependencies"""
|
|
|
|
# We need to patch the MemoryClient before it's instantiated
|
|
with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client):
|
|
crew = MockCrew()
|
|
|
|
embedder_config={
|
|
"user_id": "test_user",
|
|
"api_key": "ABCDEFGH",
|
|
"org_id": "my_org_id",
|
|
"project_id": "my_project_id",
|
|
"run_id": "my_run_id",
|
|
"includes": "include1",
|
|
"excludes": "exclude1",
|
|
"infer": True
|
|
}
|
|
|
|
return Mem0Storage(type="short_term", crew=crew, config=embedder_config)
|
|
|
|
|
|
@pytest.fixture
|
|
def mem0_storage_with_memory_client_using_explictly_config(mock_mem0_memory_client, mock_mem0_memory):
|
|
"""Fixture to create a Mem0Storage instance with mocked dependencies"""
|
|
|
|
# We need to patch both MemoryClient and Memory to prevent actual initialization
|
|
with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client), \
|
|
patch.object(Memory, "__new__", return_value=mock_mem0_memory):
|
|
|
|
crew = MockCrew()
|
|
new_config = {"provider": "mem0", "config": {"api_key": "new-api-key"}}
|
|
|
|
return Mem0Storage(type="short_term", crew=crew, config=new_config)
|
|
|
|
|
|
def test_mem0_storage_with_memory_client_initialization(
|
|
mem0_storage_with_memory_client_using_config_from_crew, mock_mem0_memory_client
|
|
):
|
|
"""Test Mem0Storage initialization with MemoryClient"""
|
|
assert (
|
|
mem0_storage_with_memory_client_using_config_from_crew.memory_type
|
|
== "short_term"
|
|
)
|
|
assert (
|
|
mem0_storage_with_memory_client_using_config_from_crew.memory
|
|
is mock_mem0_memory_client
|
|
)
|
|
|
|
|
|
def test_mem0_storage_with_explict_config(
|
|
mem0_storage_with_memory_client_using_explictly_config,
|
|
):
|
|
expected_config = {"provider": "mem0", "config": {"api_key": "new-api-key"}}
|
|
assert (
|
|
mem0_storage_with_memory_client_using_explictly_config.config == expected_config
|
|
)
|
|
|
|
|
|
def test_mem0_storage_updates_project_with_custom_categories(mock_mem0_memory_client):
|
|
mock_mem0_memory_client.update_project = MagicMock()
|
|
|
|
new_categories = [
|
|
{"lifestyle_management_concerns": "Tracks daily routines, habits, hobbies and interests including cooking, time management and work-life balance"},
|
|
]
|
|
|
|
crew = MockCrew()
|
|
|
|
config={
|
|
"user_id": "test_user",
|
|
"api_key": "ABCDEFGH",
|
|
"org_id": "my_org_id",
|
|
"project_id": "my_project_id",
|
|
"custom_categories": new_categories
|
|
}
|
|
|
|
with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client):
|
|
_ = Mem0Storage(type="short_term", crew=crew, config=config)
|
|
|
|
mock_mem0_memory_client.update_project.assert_called_once_with(
|
|
custom_categories=new_categories
|
|
)
|
|
|
|
|
|
|
|
|
|
def test_save_method_with_memory_oss(mem0_storage_with_mocked_config):
|
|
"""Test save method for different memory types"""
|
|
mem0_storage, _, _ = mem0_storage_with_mocked_config
|
|
mem0_storage.memory.add = MagicMock()
|
|
|
|
# Test short_term memory type (already set in fixture)
|
|
test_value = "This is a test memory"
|
|
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': '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', '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'
|
|
)
|
|
|
|
def test_save_method_with_multiple_agents(mem0_storage_with_mocked_config):
|
|
mem0_storage, _, _ = mem0_storage_with_mocked_config
|
|
mem0_storage.crew.agents = [MagicMock(role="Test Agent"), MagicMock(role="Test Agent 2"), MagicMock(role="Test Agent 3")]
|
|
mem0_storage.memory.add = MagicMock()
|
|
|
|
test_value = "This is a test memory"
|
|
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': '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', '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'
|
|
)
|
|
|
|
|
|
def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_config_from_crew):
|
|
"""Test save method for different memory types"""
|
|
mem0_storage = mem0_storage_with_memory_client_using_config_from_crew
|
|
mem0_storage.memory.add = MagicMock()
|
|
|
|
# Test short_term memory type (already set in fixture)
|
|
test_value = "This is a test memory"
|
|
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': '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', '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",
|
|
excludes="exclude1",
|
|
output_format='v1.1',
|
|
user_id='test_user',
|
|
agent_id='Test_Agent'
|
|
)
|
|
|
|
|
|
def test_search_method_with_memory_oss(mem0_storage_with_mocked_config):
|
|
"""Test search method for different memory types"""
|
|
mem0_storage, _, _ = mem0_storage_with_mocked_config
|
|
mock_results = {"results": [{"score": 0.9, "memory": "Result 1"}, {"score": 0.4, "memory": "Result 2"}]}
|
|
mem0_storage.memory.search = MagicMock(return_value=mock_results)
|
|
|
|
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
|
|
|
|
mem0_storage.memory.search.assert_called_once_with(
|
|
query="test query",
|
|
limit=5,
|
|
user_id="test_user",
|
|
filters={'AND': [{'run_id': 'my_run_id'}]},
|
|
threshold=0.5
|
|
)
|
|
|
|
assert len(results) == 2
|
|
assert results[0]["context"] == "Result 1"
|
|
|
|
|
|
def test_search_method_with_memory_client(mem0_storage_with_memory_client_using_config_from_crew):
|
|
"""Test search method for different memory types"""
|
|
mem0_storage = mem0_storage_with_memory_client_using_config_from_crew
|
|
mock_results = {"results": [{"score": 0.9, "memory": "Result 1"}, {"score": 0.4, "memory": "Result 2"}]}
|
|
mem0_storage.memory.search = MagicMock(return_value=mock_results)
|
|
|
|
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
|
|
|
|
mem0_storage.memory.search.assert_called_once_with(
|
|
query="test query",
|
|
limit=5,
|
|
metadata={"type": "short_term"},
|
|
user_id="test_user",
|
|
version='v2',
|
|
run_id="my_run_id",
|
|
output_format='v1.1',
|
|
filters={'AND': [{'run_id': 'my_run_id'}]},
|
|
threshold=0.5
|
|
)
|
|
|
|
assert len(results) == 2
|
|
assert results[0]["context"] == "Result 1"
|
|
|
|
|
|
def test_mem0_storage_default_infer_value(mock_mem0_memory_client):
|
|
"""Test that Mem0Storage sets infer=True by default for short_term memory."""
|
|
with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client):
|
|
crew = MockCrew()
|
|
|
|
config={
|
|
"user_id": "test_user",
|
|
"api_key": "ABCDEFGH"
|
|
}
|
|
|
|
mem0_storage = Mem0Storage(type="short_term", crew=crew, config=config)
|
|
assert mem0_storage.infer is True
|
|
|
|
def test_save_memory_using_agent_entity(mock_mem0_memory_client):
|
|
config = {
|
|
"agent_id": "agent-123",
|
|
}
|
|
|
|
mock_memory = MagicMock(spec=Memory)
|
|
with patch.object(Memory, "__new__", return_value=mock_memory):
|
|
mem0_storage = Mem0Storage(type="external", config=config)
|
|
mem0_storage.save("test memory", {"key": "value"})
|
|
mem0_storage.memory.add.assert_called_once_with(
|
|
[{'role': 'assistant' , 'content': 'test memory'}],
|
|
infer=True,
|
|
metadata={"type": "external", "key": "value"},
|
|
agent_id="agent-123",
|
|
)
|
|
|
|
def test_search_method_with_agent_entity():
|
|
config = {
|
|
"agent_id": "agent-123",
|
|
}
|
|
|
|
mock_memory = MagicMock(spec=Memory)
|
|
mock_results = {"results": [{"score": 0.9, "memory": "Result 1"}, {"score": 0.4, "memory": "Result 2"}]}
|
|
|
|
with patch.object(Memory, "__new__", return_value=mock_memory):
|
|
mem0_storage = Mem0Storage(type="external", config=config)
|
|
|
|
mem0_storage.memory.search = MagicMock(return_value=mock_results)
|
|
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
|
|
|
|
mem0_storage.memory.search.assert_called_once_with(
|
|
query="test query",
|
|
limit=5,
|
|
filters={"AND": [{"agent_id": "agent-123"}]},
|
|
threshold=0.5,
|
|
)
|
|
|
|
assert len(results) == 2
|
|
assert results[0]["context"] == "Result 1"
|
|
|
|
|
|
def test_search_method_with_agent_id_and_user_id():
|
|
mock_memory = MagicMock(spec=Memory)
|
|
mock_results = {"results": [{"score": 0.9, "memory": "Result 1"}, {"score": 0.4, "memory": "Result 2"}]}
|
|
|
|
with patch.object(Memory, "__new__", return_value=mock_memory):
|
|
mem0_storage = Mem0Storage(type="external", config={"agent_id": "agent-123", "user_id": "user-123"})
|
|
|
|
mem0_storage.memory.search = MagicMock(return_value=mock_results)
|
|
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
|
|
|
|
mem0_storage.memory.search.assert_called_once_with(
|
|
query="test query",
|
|
limit=5,
|
|
user_id='user-123',
|
|
filters={"OR": [{"user_id": "user-123"}, {"agent_id": "agent-123"}]},
|
|
threshold=0.5,
|
|
)
|
|
|
|
assert len(results) == 2
|
|
assert results[0]["context"] == "Result 1"
|