diff --git a/src/crewai/cli/reset_memories_command.py b/src/crewai/cli/reset_memories_command.py index 8e4d6f02c..f55862cdf 100644 --- a/src/crewai/cli/reset_memories_command.py +++ b/src/crewai/cli/reset_memories_command.py @@ -1,14 +1,25 @@ +"""CLI command for resetting memory storage.""" +import logging import subprocess +import sys 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.knowledge.storage.knowledge_storage import KnowledgeStorage from crewai.memory.entity.entity_memory import EntityMemory from crewai.memory.long_term.long_term_memory import LongTermMemory +from crewai.memory.short_term.short_term_memory import ShortTermMemory from crewai.utilities.task_output_storage_handler import TaskOutputStorageHandler -from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage + +_logger = logging.getLogger(__name__) + + +def _log_error(message: str) -> None: + """Log an error message.""" + _logger.error(message) + click.echo(message, err=True) def _reset_all_memories() -> None: @@ -20,6 +31,13 @@ def _reset_all_memories() -> None: KnowledgeStorage().reset() +@click.command() +@click.option("-l", "--long", is_flag=True, help="Reset long-term memory") +@click.option("-s", "--short", is_flag=True, help="Reset short-term memory") +@click.option("-e", "--entity", is_flag=True, help="Reset entity memory") +@click.option("--knowledge", is_flag=True, help="Reset knowledge") +@click.option("-k", "--kickoff-outputs", is_flag=True, help="Reset kickoff outputs") +@click.option("-a", "--all", is_flag=True, help="Reset all memories") def reset_memories_command( long: bool, short: bool, @@ -27,19 +45,18 @@ def reset_memories_command( knowledge: bool, kickoff_outputs: bool, all: bool, -) -> None: +) -> int: """ Reset the crew memories. Args: - long (bool): Whether to reset the long-term memory. - short (bool): Whether to reset the short-term memory. - entity (bool): Whether to reset the entity memory. - kickoff_outputs (bool): Whether to reset the latest kickoff task outputs. - all (bool): Whether to reset all memories. - knowledge (bool): Whether to reset the knowledge. + long: Reset long-term memory + short: Reset short-term memory + entity: Reset entity memory + knowledge: Reset knowledge + kickoff_outputs: Reset kickoff outputs + all: Reset all memories """ - try: crew = get_crew() if all: @@ -49,37 +66,49 @@ def reset_memories_command( # When no crew exists, use default storage paths _reset_all_memories() click.echo("All memories have been reset.") - return + return 0 if not any([long, short, entity, kickoff_outputs, knowledge]): click.echo( - "No memory type specified. Please specify at least one type to reset." + "Please specify at least one memory type to reset using the appropriate flags." ) - return + return 0 if not crew: click.echo("No crew found. Use --all to reset all memories.") - return + return 0 - if long: - crew.reset_memories(command_type="long") - click.echo("Long term memory has been reset.") - if short: - crew.reset_memories(command_type="short") - click.echo("Short term memory has been reset.") - if entity: - crew.reset_memories(command_type="entity") - click.echo("Entity memory has been reset.") - if kickoff_outputs: - crew.reset_memories(command_type="kickoff_outputs") - click.echo("Latest Kickoff outputs stored has been reset.") - if knowledge: - crew.reset_memories(command_type="knowledge") - click.echo("Knowledge has been reset.") + try: + if long: + crew.reset_memories(command_type="long") + click.echo("Long term memory has been reset.") + + if short: + crew.reset_memories(command_type="short") + click.echo("Short term memory has been reset.") + + if entity: + crew.reset_memories(command_type="entity") + click.echo("Entity memory has been reset.") + + if kickoff_outputs: + crew.reset_memories(command_type="kickoff_outputs") + click.echo("Latest Kickoff outputs stored has been reset.") + + if knowledge: + crew.reset_memories(command_type="knowledge") + click.echo("Knowledge has been reset.") + + return 0 + except Exception as e: + _log_error(f"An unexpected error occurred: {e}") + raise click.exceptions.Exit(code=1) except subprocess.CalledProcessError as e: - click.echo(f"An error occurred while resetting the memories: {e}", err=True) + _log_error(f"An error occurred while resetting the memories: {e}") click.echo(e.output, err=True) + raise click.exceptions.Exit(code=1) except Exception as e: - click.echo(f"An unexpected error occurred: {e}", err=True) + _log_error(f"An unexpected error occurred: {e}") + raise click.exceptions.Exit(code=1) diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py index f918f003f..284c65619 100644 --- a/src/crewai/memory/storage/rag_storage.py +++ b/src/crewai/memory/storage/rag_storage.py @@ -51,12 +51,30 @@ class RAGStorage(BaseRAGStorage): """ def _get_agents_string(self, crew) -> str: - """Get a string representation of agents for storage path.""" + """ + Get a string representation of agents for storage path. + + Args: + crew: Optional crew instance. If None, returns "no_crew". + + Returns: + str: String representation of agents or "no_crew" if no crew exists. + """ return "no_crew" if not crew else "_".join([self._sanitize_role(agent.role) for agent in crew.agents]) def __init__( self, type: str, allow_reset: bool = True, embedder_config=None, crew=None, path=None ): + """ + Initialize 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 + """ super().__init__(type, allow_reset, embedder_config, crew) self.agents = self._get_agents_string(crew) self.storage_file_name = self._build_storage_file_name(type, self.agents) diff --git a/tests/cli/cli_test.py b/tests/cli/cli_test.py index 0962ac80c..a50a2b45a 100644 --- a/tests/cli/cli_test.py +++ b/tests/cli/cli_test.py @@ -4,6 +4,8 @@ from unittest import mock import pytest from click.testing import CliRunner +from crewai.cli.reset_memories_command import reset_memories_command + from crewai.cli.cli import ( deploy_create, deploy_list, @@ -12,12 +14,12 @@ from crewai.cli.cli import ( deploy_remove, deply_status, flow_add_crew, - reset_memories, signup, test, train, version, ) +from crewai.cli.reset_memories_command import reset_memories_command @pytest.fixture @@ -63,7 +65,7 @@ class TestResetMemoriesCommand: """Test resetting all memories when crew exists.""" mock_crew = mock.Mock() mock_get_crew.return_value = mock_crew - result = runner.invoke(reset_memories, ["-a"]) + result = runner.invoke(reset_memories_command, ["-a"]) mock_crew.reset_memories.assert_called_once_with(command_type="all") assert result.output == "All memories have been reset.\n" @@ -89,7 +91,7 @@ class TestResetMemoriesCommand: Should reset all memory types individually. """ mock_get_crew.return_value = None - result = runner.invoke(reset_memories, ["-a"]) + result = runner.invoke(reset_memories_command, ["-a"]) MockShortTermMemory().reset.assert_called_once() MockEntityMemory().reset.assert_called_once() @@ -99,12 +101,30 @@ class TestResetMemoriesCommand: assert result.output == "All memories have been reset.\n" assert result.exit_code == 0 + @mock.patch("crewai.cli.reset_memories_command.get_crew") + def test_reset_memories_handles_failure( + self, + mock_get_crew, + runner, + ): + """ + Test handling of memory reset failures. + Should handle exceptions gracefully and return appropriate error code. + """ + mock_crew = mock.Mock() + mock_get_crew.return_value = mock_crew + mock_crew.reset_memories.side_effect = Exception("Failed to reset") + result = runner.invoke(reset_memories_command, ["-s"], catch_exceptions=True) + + assert result.exit_code == 1 + assert "An unexpected error occurred: Failed to reset" in result.output + @mock.patch("crewai.cli.reset_memories_command.get_crew") def test_reset_short_term_memories(mock_get_crew, runner): mock_crew = mock.Mock() mock_get_crew.return_value = mock_crew - result = runner.invoke(reset_memories, ["-s"]) + result = runner.invoke(reset_memories_command, ["-s"]) mock_crew.reset_memories.assert_called_once_with(command_type="short") assert result.output == "Short term memory has been reset.\n" @@ -114,7 +134,7 @@ def test_reset_short_term_memories(mock_get_crew, runner): def test_reset_entity_memories(mock_get_crew, runner): mock_crew = mock.Mock() mock_get_crew.return_value = mock_crew - result = runner.invoke(reset_memories, ["-e"]) + result = runner.invoke(reset_memories_command, ["-e"]) mock_crew.reset_memories.assert_called_once_with(command_type="entity") assert result.output == "Entity memory has been reset.\n" @@ -124,7 +144,7 @@ def test_reset_entity_memories(mock_get_crew, runner): def test_reset_long_term_memories(mock_get_crew, runner): mock_crew = mock.Mock() mock_get_crew.return_value = mock_crew - result = runner.invoke(reset_memories, ["-l"]) + result = runner.invoke(reset_memories_command, ["-l"]) mock_crew.reset_memories.assert_called_once_with(command_type="long") assert result.output == "Long term memory has been reset.\n" @@ -134,7 +154,7 @@ def test_reset_long_term_memories(mock_get_crew, runner): def test_reset_kickoff_outputs(mock_get_crew, runner): mock_crew = mock.Mock() mock_get_crew.return_value = mock_crew - result = runner.invoke(reset_memories, ["-k"]) + result = runner.invoke(reset_memories_command, ["-k"]) mock_crew.reset_memories.assert_called_once_with(command_type="kickoff_outputs") assert result.output == "Latest Kickoff outputs stored has been reset.\n" @@ -144,7 +164,7 @@ def test_reset_kickoff_outputs(mock_get_crew, runner): def test_reset_multiple_memory_flags(mock_get_crew, runner): mock_crew = mock.Mock() mock_get_crew.return_value = mock_crew - result = runner.invoke(reset_memories, ["-s", "-l"]) + result = runner.invoke(reset_memories_command, ["-s", "-l"]) # Check that reset_memories was called twice with the correct arguments assert mock_crew.reset_memories.call_count == 2 @@ -161,7 +181,7 @@ def test_reset_multiple_memory_flags(mock_get_crew, runner): def test_reset_knowledge(mock_get_crew, runner): mock_crew = mock.Mock() mock_get_crew.return_value = mock_crew - result = runner.invoke(reset_memories, ["--knowledge"]) + result = runner.invoke(reset_memories_command, ["--knowledge"]) mock_crew.reset_memories.assert_called_once_with(command_type="knowledge") assert result.output == "Knowledge has been reset.\n" @@ -169,7 +189,7 @@ def test_reset_knowledge(mock_get_crew, runner): def test_reset_no_memory_flags(runner): result = runner.invoke( - reset_memories, + reset_memories_command, ) assert ( result.output