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 <joao@crewai.com>

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
 

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

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")