diff --git a/packages/crewai/tests/memory/test_external_memory.py b/packages/crewai/tests/memory/test_external_memory.py deleted file mode 100644 index 49b3a3160..000000000 --- a/packages/crewai/tests/memory/test_external_memory.py +++ /dev/null @@ -1,355 +0,0 @@ -from unittest.mock import MagicMock, patch, ANY -from collections import defaultdict -from crewai.events.event_bus import crewai_event_bus -from crewai.events.types.memory_events import ( - MemorySaveStartedEvent, - MemorySaveCompletedEvent, - MemoryQueryStartedEvent, - MemoryQueryCompletedEvent, -) -import pytest -from mem0.memory.main import Memory - -from crewai.agent import Agent -from crewai.crew import Crew, Process -from crewai.memory.external.external_memory import ExternalMemory -from crewai.memory.external.external_memory_item import ExternalMemoryItem -from crewai.memory.storage.interface import Storage -from crewai.task import Task - - -@pytest.fixture -def mock_mem0_memory(): - mock_memory = MagicMock(spec=Memory) - return mock_memory - - -@pytest.fixture -def patch_configure_mem0(mock_mem0_memory): - with patch( - "crewai.memory.external.external_memory.ExternalMemory._configure_mem0", - return_value=mock_mem0_memory, - ) as mocked: - yield mocked - - -@pytest.fixture -def external_memory_with_mocked_config(patch_configure_mem0): - embedder_config = {"provider": "mem0"} - external_memory = ExternalMemory(embedder_config=embedder_config) - return external_memory - - -@pytest.fixture -def crew_with_external_memory(external_memory_with_mocked_config, patch_configure_mem0): - agent = Agent( - role="Researcher", - goal="Search relevant data and provide results", - backstory="You are a researcher at a leading tech think tank.", - tools=[], - verbose=True, - ) - - task = Task( - description="Perform a search on specific topics.", - expected_output="A list of relevant URLs based on the search query.", - agent=agent, - ) - - crew = Crew( - agents=[agent], - tasks=[task], - verbose=True, - process=Process.sequential, - memory=True, - external_memory=external_memory_with_mocked_config, - ) - - return crew - - -@pytest.fixture -def crew_with_external_memory_without_memory_flag( - external_memory_with_mocked_config, patch_configure_mem0 -): - agent = Agent( - role="Researcher", - goal="Search relevant data and provide results", - backstory="You are a researcher at a leading tech think tank.", - tools=[], - verbose=True, - ) - - task = Task( - description="Perform a search on specific topics.", - expected_output="A list of relevant URLs based on the search query.", - agent=agent, - ) - - crew = Crew( - agents=[agent], - tasks=[task], - verbose=True, - process=Process.sequential, - external_memory=external_memory_with_mocked_config, - ) - - return crew - - -def test_external_memory_initialization(external_memory_with_mocked_config): - assert external_memory_with_mocked_config is not None - assert isinstance(external_memory_with_mocked_config, ExternalMemory) - - -def test_external_memory_save(external_memory_with_mocked_config): - memory_item = ExternalMemoryItem( - value="test value", metadata={"task": "test_task"}, agent="test_agent" - ) - - with patch.object(ExternalMemory, "save") as mock_save: - external_memory_with_mocked_config.save( - value=memory_item.value, - metadata=memory_item.metadata, - agent=memory_item.agent, - ) - - mock_save.assert_called_once_with( - value=memory_item.value, - metadata=memory_item.metadata, - agent=memory_item.agent, - ) - - -def test_external_memory_reset(external_memory_with_mocked_config): - with patch( - "crewai.memory.external.external_memory.ExternalMemory.reset" - ) as mock_reset: - external_memory_with_mocked_config.reset() - mock_reset.assert_called_once() - - -def test_external_memory_supported_storages(): - supported_storages = ExternalMemory.external_supported_storages() - assert "mem0" in supported_storages - assert callable(supported_storages["mem0"]) - - -def test_external_memory_create_storage_invalid_provider(): - embedder_config = {"provider": "invalid_provider", "config": {}} - - with pytest.raises(ValueError, match="Provider invalid_provider not supported"): - ExternalMemory.create_storage(None, embedder_config) - - -def test_external_memory_create_storage_missing_provider(): - embedder_config = {"config": {}} - - with pytest.raises( - ValueError, match="embedder_config must include a 'provider' key" - ): - ExternalMemory.create_storage(None, embedder_config) - - -def test_external_memory_create_storage_missing_config(): - with pytest.raises(ValueError, match="embedder_config is required"): - ExternalMemory.create_storage(None, None) - - -def test_crew_with_external_memory_initialization(crew_with_external_memory): - assert crew_with_external_memory._external_memory is not None - assert isinstance(crew_with_external_memory._external_memory, ExternalMemory) - assert crew_with_external_memory._external_memory.crew == crew_with_external_memory - - -@pytest.mark.parametrize("mem_type", ["external", "all"]) -def test_crew_external_memory_reset(mem_type, crew_with_external_memory): - with patch( - "crewai.memory.external.external_memory.ExternalMemory.reset" - ) as mock_reset: - crew_with_external_memory.reset_memories(mem_type) - mock_reset.assert_called_once() - - -@pytest.mark.parametrize("mem_method", ["search", "save"]) -@pytest.mark.vcr(filter_headers=["authorization"]) -def test_crew_external_memory_save_with_memory_flag( - mem_method, crew_with_external_memory -): - with patch( - f"crewai.memory.external.external_memory.ExternalMemory.{mem_method}" - ) as mock_method: - crew_with_external_memory.kickoff() - assert mock_method.call_count > 0 - - -@pytest.mark.parametrize("mem_method", ["search", "save"]) -@pytest.mark.vcr(filter_headers=["authorization"]) -def test_crew_external_memory_save_using_crew_without_memory_flag( - mem_method, crew_with_external_memory_without_memory_flag -): - with patch( - f"crewai.memory.external.external_memory.ExternalMemory.{mem_method}" - ) as mock_method: - crew_with_external_memory_without_memory_flag.kickoff() - assert mock_method.call_count > 0 - - -@pytest.fixture -def custom_storage(): - class CustomStorage(Storage): - def __init__(self): - self.memories = [] - - def save(self, value, metadata=None, agent=None): - self.memories.append({"value": value, "metadata": metadata, "agent": agent}) - - def search(self, query, limit=10, score_threshold=0.5): - return self.memories - - def reset(self): - self.memories = [] - - custom_storage = CustomStorage() - return custom_storage - - -def test_external_memory_custom_storage(custom_storage, crew_with_external_memory): - external_memory = ExternalMemory(storage=custom_storage) - - # by ensuring the crew is set, we can test that the storage is used - external_memory.set_crew(crew_with_external_memory) - - test_value = "test value" - test_metadata = {"source": "test"} - external_memory.save(value=test_value, metadata=test_metadata) - - results = external_memory.search("test") - assert len(results) == 1 - assert results[0]["value"] == test_value - assert results[0]["metadata"] == test_metadata - - external_memory.reset() - results = external_memory.search("test") - assert len(results) == 0 - - -def test_external_memory_search_events( - custom_storage, external_memory_with_mocked_config -): - events = defaultdict(list) - - external_memory_with_mocked_config.storage = custom_storage - with crewai_event_bus.scoped_handlers(): - - @crewai_event_bus.on(MemoryQueryStartedEvent) - def on_search_started(source, event): - events["MemoryQueryStartedEvent"].append(event) - - @crewai_event_bus.on(MemoryQueryCompletedEvent) - def on_search_completed(source, event): - events["MemoryQueryCompletedEvent"].append(event) - - external_memory_with_mocked_config.search( - query="test value", - limit=3, - score_threshold=0.35, - ) - - assert len(events["MemoryQueryStartedEvent"]) == 1 - assert len(events["MemoryQueryCompletedEvent"]) == 1 - - assert dict(events["MemoryQueryStartedEvent"][0]) == { - "timestamp": ANY, - "type": "memory_query_started", - "source_fingerprint": None, - "source_type": "external_memory", - "fingerprint_metadata": None, - "task_id": None, - "task_name": None, - "from_task": None, - "from_agent": None, - "agent_role": None, - "agent_id": None, - "query": "test value", - "limit": 3, - "score_threshold": 0.35, - } - - assert dict(events["MemoryQueryCompletedEvent"][0]) == { - "timestamp": ANY, - "type": "memory_query_completed", - "source_fingerprint": None, - "source_type": "external_memory", - "fingerprint_metadata": None, - "task_id": None, - "task_name": None, - "from_task": None, - "from_agent": None, - "agent_role": None, - "agent_id": None, - "query": "test value", - "results": [], - "limit": 3, - "score_threshold": 0.35, - "query_time_ms": ANY, - } - - -def test_external_memory_save_events( - custom_storage, external_memory_with_mocked_config -): - events = defaultdict(list) - - external_memory_with_mocked_config.storage = custom_storage - - with crewai_event_bus.scoped_handlers(): - - @crewai_event_bus.on(MemorySaveStartedEvent) - def on_save_started(source, event): - events["MemorySaveStartedEvent"].append(event) - - @crewai_event_bus.on(MemorySaveCompletedEvent) - def on_save_completed(source, event): - events["MemorySaveCompletedEvent"].append(event) - - external_memory_with_mocked_config.save( - value="saving value", - metadata={"task": "test_task"}, - ) - - assert len(events["MemorySaveStartedEvent"]) == 1 - assert len(events["MemorySaveCompletedEvent"]) == 1 - - assert dict(events["MemorySaveStartedEvent"][0]) == { - "timestamp": ANY, - "type": "memory_save_started", - "source_fingerprint": None, - "source_type": "external_memory", - "fingerprint_metadata": None, - "task_id": None, - "task_name": None, - "from_task": None, - "from_agent": None, - "agent_role": None, - "agent_id": None, - "value": "saving value", - "metadata": {"task": "test_task"}, - } - - assert dict(events["MemorySaveCompletedEvent"][0]) == { - "timestamp": ANY, - "type": "memory_save_completed", - "source_fingerprint": None, - "source_type": "external_memory", - "fingerprint_metadata": None, - "task_id": None, - "task_name": None, - "from_task": None, - "from_agent": None, - "agent_role": None, - "agent_id": None, - "value": "saving value", - "metadata": {"task": "test_task"}, - "save_time_ms": ANY, - } diff --git a/packages/crewai/tests/storage/test_mem0_storage.py b/packages/crewai/tests/storage/test_mem0_storage.py deleted file mode 100644 index 11cfddb3a..000000000 --- a/packages/crewai/tests/storage/test_mem0_storage.py +++ /dev/null @@ -1,505 +0,0 @@ -from unittest.mock import MagicMock, patch - -import pytest -from mem0 import Memory, MemoryClient - -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")] - - -# Test data constants -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.\n" - "Your personal goal is: Engage in useful and interesting conversations " - "with users while remembering context.\n" - "To give my best complete final answer to the task respond using the exact " - "following format:\n\n" - "Thought: I now can give a great answer\n" - "Final Answer: Your final answer must be the great and the most complete " - "as possible, it must be outcome described.\n\n" - "I MUST use these formats, my job depends on it!" -) - -USER_CONTENT = ( - "\nCurrent Task: Respond to user conversation. User message: " - "What do you know about me?\n\n" - "This is the expected criteria for your final answer: Contextually " - "appropriate, helpful, and friendly response.\n" - "you 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\n" - "Begin! This is VERY important to you, use the tools available and give " - "your best Final Answer, your job depends on it!\n\n" - "Thought:" -) - -ASSISTANT_CONTENT = ( - "I now can give a great answer \n" - "Final 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." -) - -TEST_DESCRIPTION = ( - "Respond to user conversation. User message: What do you know about me?" -) - -# Extracted content (after processing by _get_user_message and _get_assistant_message) -EXTRACTED_USER_CONTENT = "What do you know about me?" -EXTRACTED_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." -) - - -@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.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": TEST_DESCRIPTION, - "messages": [ - {"role": "system", "content": SYSTEM_CONTENT}, - {"role": "user", "content": USER_CONTENT}, - {"role": "assistant", "content": ASSISTANT_CONTENT}, - ], - "agent": "Friendly chatbot assistant", - } - - mem0_storage.save(test_value, test_metadata) - - mem0_storage.memory.add.assert_called_once_with( - [ - {"role": "user", "content": EXTRACTED_USER_CONTENT}, - { - "role": "assistant", - "content": EXTRACTED_ASSISTANT_CONTENT, - }, - ], - infer=True, - metadata={ - "type": "short_term", - "description": TEST_DESCRIPTION, - "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": TEST_DESCRIPTION, - "messages": [ - {"role": "system", "content": SYSTEM_CONTENT}, - {"role": "user", "content": USER_CONTENT}, - {"role": "assistant", "content": ASSISTANT_CONTENT}, - ], - "agent": "Friendly chatbot assistant", - } - - mem0_storage.save(test_value, test_metadata) - - mem0_storage.memory.add.assert_called_once_with( - [ - {"role": "user", "content": EXTRACTED_USER_CONTENT}, - { - "role": "assistant", - "content": EXTRACTED_ASSISTANT_CONTENT, - }, - ], - infer=True, - metadata={ - "type": "short_term", - "description": TEST_DESCRIPTION, - "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": TEST_DESCRIPTION, - "messages": [ - {"role": "system", "content": SYSTEM_CONTENT}, - {"role": "user", "content": USER_CONTENT}, - {"role": "assistant", "content": ASSISTANT_CONTENT}, - ], - "agent": "Friendly chatbot assistant", - } - - mem0_storage.save(test_value, test_metadata) - - mem0_storage.memory.add.assert_called_once_with( - [ - {"role": "user", "content": EXTRACTED_USER_CONTENT}, - { - "role": "assistant", - "content": EXTRACTED_ASSISTANT_CONTENT, - }, - ], - infer=True, - metadata={ - "type": "short_term", - "description": TEST_DESCRIPTION, - "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]["content"] == "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]["content"] == "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]["content"] == "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]["content"] == "Result 1"