From 190196e57280d30baf23f0c9cf39164d954659c8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 13 Feb 2025 19:10:01 +0000 Subject: [PATCH] refactor: Address code review feedback - Move imports to top level - Add helper function for resetting memories - Add type hints and docstrings - Improve test organization Co-Authored-By: Joe Moura --- src/crewai/cli/reset_memories_command.py | 39 +++++++------ src/crewai/memory/storage/rag_storage.py | 27 ++++++--- tests/cli/cli_test.py | 70 +++++++++++++----------- 3 files changed, 79 insertions(+), 57 deletions(-) diff --git a/src/crewai/cli/reset_memories_command.py b/src/crewai/cli/reset_memories_command.py index 2933a2bd4..8e4d6f02c 100644 --- a/src/crewai/cli/reset_memories_command.py +++ b/src/crewai/cli/reset_memories_command.py @@ -1,17 +1,32 @@ import subprocess +from typing import Optional import click from crewai.cli.utils import get_crew +from crewai.memory.short_term.short_term_memory import ShortTermMemory +from crewai.memory.entity.entity_memory import EntityMemory +from crewai.memory.long_term.long_term_memory import LongTermMemory +from crewai.utilities.task_output_storage_handler import TaskOutputStorageHandler +from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage + + +def _reset_all_memories() -> None: + """Reset all memory types.""" + ShortTermMemory().reset() + EntityMemory().reset() + LongTermMemory().reset() + TaskOutputStorageHandler().reset() + KnowledgeStorage().reset() def reset_memories_command( - long, - short, - entity, - knowledge, - kickoff_outputs, - all, + long: bool, + short: bool, + entity: bool, + knowledge: bool, + kickoff_outputs: bool, + all: bool, ) -> None: """ Reset the crew memories. @@ -32,17 +47,7 @@ def reset_memories_command( crew.reset_memories(command_type="all") else: # When no crew exists, use default storage paths - from crewai.memory.short_term.short_term_memory import ShortTermMemory - from crewai.memory.entity.entity_memory import EntityMemory - from crewai.memory.long_term.long_term_memory import LongTermMemory - from crewai.utilities.task_output_storage_handler import TaskOutputStorageHandler - from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage - - ShortTermMemory().reset() - EntityMemory().reset() - LongTermMemory().reset() - TaskOutputStorageHandler().reset() - KnowledgeStorage().reset() + _reset_all_memories() click.echo("All memories have been reset.") return diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py index 61e962ba4..f918f003f 100644 --- a/src/crewai/memory/storage/rag_storage.py +++ b/src/crewai/memory/storage/rag_storage.py @@ -39,18 +39,27 @@ class RAGStorage(BaseRAGStorage): app: ClientAPI | None = None + """ + RAG Storage implementation that handles both crew and no-crew scenarios. + + Args: + type: Type of storage + allow_reset: Whether storage can be reset + embedder_config: Configuration for embeddings + crew: Crew instance or None for no-crew scenario + path: Custom storage path + """ + + def _get_agents_string(self, crew) -> str: + """Get a string representation of agents for storage path.""" + return "no_crew" if not crew else "_".join([self._sanitize_role(agent.role) for agent in crew.agents]) + def __init__( - self, type, allow_reset=True, embedder_config=None, crew=None, path=None + self, type: str, allow_reset: bool = True, embedder_config=None, crew=None, path=None ): super().__init__(type, allow_reset, embedder_config, crew) - if crew: - agents = crew.agents - agents = [self._sanitize_role(agent.role) for agent in agents] - agents = "_".join(agents) - else: - agents = "no_crew" - self.agents = agents - self.storage_file_name = self._build_storage_file_name(type, agents) + self.agents = self._get_agents_string(crew) + self.storage_file_name = self._build_storage_file_name(type, self.agents) self.type = type self.allow_reset = allow_reset self.path = path diff --git a/tests/cli/cli_test.py b/tests/cli/cli_test.py index f0d5c41ca..0962ac80c 100644 --- a/tests/cli/cli_test.py +++ b/tests/cli/cli_test.py @@ -55,41 +55,49 @@ def test_train_invalid_string_iterations(train_crew, runner): ) -@mock.patch("crewai.cli.reset_memories_command.get_crew") -def test_reset_all_memories(mock_get_crew, runner): - mock_crew = mock.Mock() - mock_get_crew.return_value = mock_crew - result = runner.invoke(reset_memories, ["-a"]) +class TestResetMemoriesCommand: + """Tests for the reset-memories command.""" - mock_crew.reset_memories.assert_called_once_with(command_type="all") - assert result.output == "All memories have been reset.\n" + @mock.patch("crewai.cli.reset_memories_command.get_crew") + def test_reset_all_memories(self, mock_get_crew, runner): + """Test resetting all memories when crew exists.""" + mock_crew = mock.Mock() + mock_get_crew.return_value = mock_crew + result = runner.invoke(reset_memories, ["-a"]) + mock_crew.reset_memories.assert_called_once_with(command_type="all") + assert result.output == "All memories have been reset.\n" -@mock.patch("crewai.cli.reset_memories_command.get_crew") -@mock.patch("crewai.cli.reset_memories_command.ShortTermMemory") -@mock.patch("crewai.cli.reset_memories_command.EntityMemory") -@mock.patch("crewai.cli.reset_memories_command.LongTermMemory") -@mock.patch("crewai.cli.reset_memories_command.TaskOutputStorageHandler") -@mock.patch("crewai.cli.reset_memories_command.KnowledgeStorage") -def test_reset_all_memories_no_crew( - MockKnowledgeStorage, - MockTaskOutputStorageHandler, - MockLongTermMemory, - MockEntityMemory, - MockShortTermMemory, - mock_get_crew, - runner, -): - mock_get_crew.return_value = None - result = runner.invoke(reset_memories, ["-a"]) + @mock.patch("crewai.cli.reset_memories_command.get_crew") + @mock.patch("crewai.cli.reset_memories_command.ShortTermMemory") + @mock.patch("crewai.cli.reset_memories_command.EntityMemory") + @mock.patch("crewai.cli.reset_memories_command.LongTermMemory") + @mock.patch("crewai.cli.reset_memories_command.TaskOutputStorageHandler") + @mock.patch("crewai.cli.reset_memories_command.KnowledgeStorage") + def test_reset_all_memories_no_crew( + self, + MockKnowledgeStorage, + MockTaskOutputStorageHandler, + MockLongTermMemory, + MockEntityMemory, + MockShortTermMemory, + mock_get_crew, + runner, + ): + """ + Test resetting all memories when no crew exists. + Should reset all memory types individually. + """ + mock_get_crew.return_value = None + result = runner.invoke(reset_memories, ["-a"]) - MockShortTermMemory().reset.assert_called_once() - MockEntityMemory().reset.assert_called_once() - MockLongTermMemory().reset.assert_called_once() - MockTaskOutputStorageHandler().reset.assert_called_once() - MockKnowledgeStorage().reset.assert_called_once() - assert result.output == "All memories have been reset.\n" - assert result.exit_code == 0 + MockShortTermMemory().reset.assert_called_once() + MockEntityMemory().reset.assert_called_once() + MockLongTermMemory().reset.assert_called_once() + MockTaskOutputStorageHandler().reset.assert_called_once() + MockKnowledgeStorage().reset.assert_called_once() + assert result.output == "All memories have been reset.\n" + assert result.exit_code == 0 @mock.patch("crewai.cli.reset_memories_command.get_crew")