Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: Handle no-crew case in reset-memories command #2124

Closed
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 88 additions & 38 deletions src/crewai/cli/reset_memories_command.py
Original file line number Diff line number Diff line change
@@ -1,64 +1,114 @@
"""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.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

_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:
"""Reset all memory types."""
ShortTermMemory().reset()
EntityMemory().reset()
LongTermMemory().reset()
TaskOutputStorageHandler().reset()
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,
short,
entity,
knowledge,
kickoff_outputs,
all,
) -> None:
long: bool,
short: bool,
entity: bool,
knowledge: bool,
kickoff_outputs: bool,
all: bool,
) -> 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 not crew:
raise ValueError("No crew found.")
if all:
crew.reset_memories(command_type="all")
if crew:
crew.reset_memories(command_type="all")
else:
# 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

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

if not crew:
click.echo("No crew found. Use --all to reset all memories.")
return 0

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)
44 changes: 36 additions & 8 deletions src/crewai/memory/storage/rag_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,18 +39,46 @@ 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.

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, allow_reset=True, embedder_config=None, crew=None, path=None
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)
agents = crew.agents if crew else []
agents = [self._sanitize_role(agent.role) for agent in agents]
agents = "_".join(agents)
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
self._initialize_app()
Expand Down
87 changes: 71 additions & 16 deletions tests/cli/cli_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
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,
Expand All @@ -12,12 +14,12 @@
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
Expand Down Expand Up @@ -55,21 +57,74 @@ 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"])

mock_crew.reset_memories.assert_called_once_with(command_type="all")
assert result.output == "All memories have been reset.\n"
class TestResetMemoriesCommand:
"""Tests for the reset-memories command."""

@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_command, ["-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(
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_command, ["-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

@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"
Expand All @@ -79,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"
Expand All @@ -89,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"
Expand All @@ -99,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"
Expand All @@ -109,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
Expand All @@ -126,15 +181,15 @@ 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"


def test_reset_no_memory_flags(runner):
result = runner.invoke(
reset_memories,
reset_memories_command,
)
assert (
result.output
Expand Down
Loading