mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 00:58:30 +00:00
refactor: Add logging and improve error handling
- Add Logger class usage - Add descriptive error messages - Add edge case test for memory reset failure - Add docstrings to RAGStorage methods Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -1,14 +1,25 @@
|
|||||||
|
"""CLI command for resetting memory storage."""
|
||||||
|
import logging
|
||||||
import subprocess
|
import subprocess
|
||||||
|
import sys
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
import click
|
import click
|
||||||
|
|
||||||
from crewai.cli.utils import get_crew
|
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.entity.entity_memory import EntityMemory
|
||||||
from crewai.memory.long_term.long_term_memory import LongTermMemory
|
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.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:
|
def _reset_all_memories() -> None:
|
||||||
@@ -20,6 +31,13 @@ def _reset_all_memories() -> None:
|
|||||||
KnowledgeStorage().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(
|
def reset_memories_command(
|
||||||
long: bool,
|
long: bool,
|
||||||
short: bool,
|
short: bool,
|
||||||
@@ -27,19 +45,18 @@ def reset_memories_command(
|
|||||||
knowledge: bool,
|
knowledge: bool,
|
||||||
kickoff_outputs: bool,
|
kickoff_outputs: bool,
|
||||||
all: bool,
|
all: bool,
|
||||||
) -> None:
|
) -> int:
|
||||||
"""
|
"""
|
||||||
Reset the crew memories.
|
Reset the crew memories.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
long (bool): Whether to reset the long-term memory.
|
long: Reset long-term memory
|
||||||
short (bool): Whether to reset the short-term memory.
|
short: Reset short-term memory
|
||||||
entity (bool): Whether to reset the entity memory.
|
entity: Reset entity memory
|
||||||
kickoff_outputs (bool): Whether to reset the latest kickoff task outputs.
|
knowledge: Reset knowledge
|
||||||
all (bool): Whether to reset all memories.
|
kickoff_outputs: Reset kickoff outputs
|
||||||
knowledge (bool): Whether to reset the knowledge.
|
all: Reset all memories
|
||||||
"""
|
"""
|
||||||
|
|
||||||
try:
|
try:
|
||||||
crew = get_crew()
|
crew = get_crew()
|
||||||
if all:
|
if all:
|
||||||
@@ -49,37 +66,49 @@ def reset_memories_command(
|
|||||||
# When no crew exists, use default storage paths
|
# When no crew exists, use default storage paths
|
||||||
_reset_all_memories()
|
_reset_all_memories()
|
||||||
click.echo("All memories have been reset.")
|
click.echo("All memories have been reset.")
|
||||||
return
|
return 0
|
||||||
|
|
||||||
if not any([long, short, entity, kickoff_outputs, knowledge]):
|
if not any([long, short, entity, kickoff_outputs, knowledge]):
|
||||||
click.echo(
|
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:
|
if not crew:
|
||||||
click.echo("No crew found. Use --all to reset all memories.")
|
click.echo("No crew found. Use --all to reset all memories.")
|
||||||
return
|
return 0
|
||||||
|
|
||||||
if long:
|
try:
|
||||||
crew.reset_memories(command_type="long")
|
if long:
|
||||||
click.echo("Long term memory has been reset.")
|
crew.reset_memories(command_type="long")
|
||||||
if short:
|
click.echo("Long term memory has been reset.")
|
||||||
crew.reset_memories(command_type="short")
|
|
||||||
click.echo("Short term memory has been reset.")
|
if short:
|
||||||
if entity:
|
crew.reset_memories(command_type="short")
|
||||||
crew.reset_memories(command_type="entity")
|
click.echo("Short term memory has been reset.")
|
||||||
click.echo("Entity memory has been reset.")
|
|
||||||
if kickoff_outputs:
|
if entity:
|
||||||
crew.reset_memories(command_type="kickoff_outputs")
|
crew.reset_memories(command_type="entity")
|
||||||
click.echo("Latest Kickoff outputs stored has been reset.")
|
click.echo("Entity memory has been reset.")
|
||||||
if knowledge:
|
|
||||||
crew.reset_memories(command_type="knowledge")
|
if kickoff_outputs:
|
||||||
click.echo("Knowledge has been reset.")
|
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:
|
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)
|
click.echo(e.output, err=True)
|
||||||
|
raise click.exceptions.Exit(code=1)
|
||||||
|
|
||||||
except Exception as e:
|
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)
|
||||||
|
|||||||
@@ -51,12 +51,30 @@ class RAGStorage(BaseRAGStorage):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
def _get_agents_string(self, crew) -> str:
|
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])
|
return "no_crew" if not crew else "_".join([self._sanitize_role(agent.role) for agent in crew.agents])
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self, type: str, allow_reset: bool = 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)
|
super().__init__(type, allow_reset, embedder_config, crew)
|
||||||
self.agents = self._get_agents_string(crew)
|
self.agents = self._get_agents_string(crew)
|
||||||
self.storage_file_name = self._build_storage_file_name(type, self.agents)
|
self.storage_file_name = self._build_storage_file_name(type, self.agents)
|
||||||
|
|||||||
@@ -4,6 +4,8 @@ from unittest import mock
|
|||||||
import pytest
|
import pytest
|
||||||
from click.testing import CliRunner
|
from click.testing import CliRunner
|
||||||
|
|
||||||
|
from crewai.cli.reset_memories_command import reset_memories_command
|
||||||
|
|
||||||
from crewai.cli.cli import (
|
from crewai.cli.cli import (
|
||||||
deploy_create,
|
deploy_create,
|
||||||
deploy_list,
|
deploy_list,
|
||||||
@@ -12,12 +14,12 @@ from crewai.cli.cli import (
|
|||||||
deploy_remove,
|
deploy_remove,
|
||||||
deply_status,
|
deply_status,
|
||||||
flow_add_crew,
|
flow_add_crew,
|
||||||
reset_memories,
|
|
||||||
signup,
|
signup,
|
||||||
test,
|
test,
|
||||||
train,
|
train,
|
||||||
version,
|
version,
|
||||||
)
|
)
|
||||||
|
from crewai.cli.reset_memories_command import reset_memories_command
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
@@ -63,7 +65,7 @@ class TestResetMemoriesCommand:
|
|||||||
"""Test resetting all memories when crew exists."""
|
"""Test resetting all memories when crew exists."""
|
||||||
mock_crew = mock.Mock()
|
mock_crew = mock.Mock()
|
||||||
mock_get_crew.return_value = mock_crew
|
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")
|
mock_crew.reset_memories.assert_called_once_with(command_type="all")
|
||||||
assert result.output == "All memories have been reset.\n"
|
assert result.output == "All memories have been reset.\n"
|
||||||
@@ -89,7 +91,7 @@ class TestResetMemoriesCommand:
|
|||||||
Should reset all memory types individually.
|
Should reset all memory types individually.
|
||||||
"""
|
"""
|
||||||
mock_get_crew.return_value = None
|
mock_get_crew.return_value = None
|
||||||
result = runner.invoke(reset_memories, ["-a"])
|
result = runner.invoke(reset_memories_command, ["-a"])
|
||||||
|
|
||||||
MockShortTermMemory().reset.assert_called_once()
|
MockShortTermMemory().reset.assert_called_once()
|
||||||
MockEntityMemory().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.output == "All memories have been reset.\n"
|
||||||
assert result.exit_code == 0
|
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")
|
@mock.patch("crewai.cli.reset_memories_command.get_crew")
|
||||||
def test_reset_short_term_memories(mock_get_crew, runner):
|
def test_reset_short_term_memories(mock_get_crew, runner):
|
||||||
mock_crew = mock.Mock()
|
mock_crew = mock.Mock()
|
||||||
mock_get_crew.return_value = mock_crew
|
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")
|
mock_crew.reset_memories.assert_called_once_with(command_type="short")
|
||||||
assert result.output == "Short term memory has been reset.\n"
|
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):
|
def test_reset_entity_memories(mock_get_crew, runner):
|
||||||
mock_crew = mock.Mock()
|
mock_crew = mock.Mock()
|
||||||
mock_get_crew.return_value = mock_crew
|
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")
|
mock_crew.reset_memories.assert_called_once_with(command_type="entity")
|
||||||
assert result.output == "Entity memory has been reset.\n"
|
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):
|
def test_reset_long_term_memories(mock_get_crew, runner):
|
||||||
mock_crew = mock.Mock()
|
mock_crew = mock.Mock()
|
||||||
mock_get_crew.return_value = mock_crew
|
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")
|
mock_crew.reset_memories.assert_called_once_with(command_type="long")
|
||||||
assert result.output == "Long term memory has been reset.\n"
|
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):
|
def test_reset_kickoff_outputs(mock_get_crew, runner):
|
||||||
mock_crew = mock.Mock()
|
mock_crew = mock.Mock()
|
||||||
mock_get_crew.return_value = mock_crew
|
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")
|
mock_crew.reset_memories.assert_called_once_with(command_type="kickoff_outputs")
|
||||||
assert result.output == "Latest Kickoff outputs stored has been reset.\n"
|
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):
|
def test_reset_multiple_memory_flags(mock_get_crew, runner):
|
||||||
mock_crew = mock.Mock()
|
mock_crew = mock.Mock()
|
||||||
mock_get_crew.return_value = mock_crew
|
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
|
# Check that reset_memories was called twice with the correct arguments
|
||||||
assert mock_crew.reset_memories.call_count == 2
|
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):
|
def test_reset_knowledge(mock_get_crew, runner):
|
||||||
mock_crew = mock.Mock()
|
mock_crew = mock.Mock()
|
||||||
mock_get_crew.return_value = mock_crew
|
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")
|
mock_crew.reset_memories.assert_called_once_with(command_type="knowledge")
|
||||||
assert result.output == "Knowledge has been reset.\n"
|
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):
|
def test_reset_no_memory_flags(runner):
|
||||||
result = runner.invoke(
|
result = runner.invoke(
|
||||||
reset_memories,
|
reset_memories_command,
|
||||||
)
|
)
|
||||||
assert (
|
assert (
|
||||||
result.output
|
result.output
|
||||||
|
|||||||
Reference in New Issue
Block a user