Enhance memory reset functionality and JSON crew handling

- Added `reset_all` method to the `Memory` class to reset the entire memory store, ignoring `root_scope`.
- Updated the `Crew` class to utilize `reset_all` when resetting memory.
- Enhanced the `_reset_flow_memory` function to check for `Memory` instances and call `reset_all` accordingly.
- Introduced helper functions to load JSON crew configurations and handle project declarations, improving the reset command's flexibility.
- Added tests to validate the new JSON crew memory reset behavior and ensure proper handling of declared flow projects.
This commit is contained in:
Joao Moura
2026-06-16 16:27:07 -07:00
parent ee4f853270
commit 6c185d22c1
4 changed files with 94 additions and 2 deletions

View File

@@ -2275,6 +2275,8 @@ class Crew(FlowTrackable, BaseModel):
"""
def default_reset(memory: Any) -> Any:
if isinstance(memory, Memory):
return memory.reset_all()
return memory.reset()
def knowledge_reset(memory: Any) -> Any:

View File

@@ -990,6 +990,7 @@ class Memory(BaseModel):
scope: Scope to reset. If None and root_scope is set, resets only
within root_scope. If None and no root_scope, resets all.
"""
self.drain_writes()
effective_scope = scope
if effective_scope is None and self.root_scope:
effective_scope = self.root_scope
@@ -997,6 +998,11 @@ class Memory(BaseModel):
effective_scope = join_scope_paths(self.root_scope, effective_scope)
self._storage.reset(scope_prefix=effective_scope)
def reset_all(self) -> None:
"""Reset the entire backing memory store, ignoring ``root_scope``."""
self.drain_writes()
self._storage.reset(scope_prefix=None)
async def aextract_memories(self, content: str) -> list[str]:
"""Async variant of extract_memories."""
return self.extract_memories(content)

View File

@@ -6,7 +6,10 @@ from typing import Any
import click
from crewai.flow import Flow
from crewai.utilities.project_utils import get_crews, get_flows
from crewai.memory.unified_memory import Memory
from crewai.project.crew_loader import load_crew
from crewai.project.json_loader import find_crew_json_file
from crewai.utilities.project_utils import get_crews, get_flows, read_toml
def _reset_flow_memory(flow: Flow[Any]) -> None:
@@ -23,7 +26,9 @@ def _reset_flow_memory(flow: Flow[Any]) -> None:
if mem is None:
return
try:
if hasattr(mem, "reset"):
if isinstance(mem, Memory):
mem.reset_all()
elif hasattr(mem, "reset"):
mem.reset()
elif hasattr(mem, "_memory") and mem._memory is not None:
mem._memory.reset()
@@ -37,6 +42,31 @@ def _reset_flow_memory(flow: Flow[Any]) -> None:
click.echo(f"Memory reset skipped: {exc}", err=True)
def _current_project_declares_flow() -> bool:
try:
pyproject_data = read_toml()
except Exception:
return False
declared_type: str | None = (
pyproject_data.get("tool", {}).get("crewai", {}).get("type")
)
return declared_type == "flow"
def _get_json_crew() -> Any | None:
"""Load a JSON-first crew from the current project, if present."""
if _current_project_declares_flow():
return None
crew_path = find_crew_json_file()
if crew_path is None:
return None
crew, _ = load_crew(crew_path)
return crew
def reset_memories_command(
memory: bool,
knowledge: bool,
@@ -61,6 +91,8 @@ def reset_memories_command(
return
crews = get_crews()
if json_crew := _get_json_crew():
crews.append(json_crew)
flows = get_flows()
if not crews and not flows:

View File

@@ -4,6 +4,7 @@ Non-core CLI tests (train, test, version, deploy, login, flow_add_crew)
have moved to lib/cli/tests/test_cli.py.
"""
from pathlib import Path
from unittest import mock
from click.testing import CliRunner
@@ -30,6 +31,8 @@ def mock_get_crews(mock_crew):
"crewai.utilities.reset_memories.get_crews", return_value=[mock_crew]
) as mock_get_crew, mock.patch(
"crewai.utilities.reset_memories.get_flows", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories._get_json_crew", return_value=None
):
yield mock_get_crew
@@ -170,6 +173,8 @@ def mock_get_flows(mock_flow):
"crewai.utilities.reset_memories.get_flows", return_value=[mock_flow]
) as mock_get_flow, mock.patch(
"crewai.utilities.reset_memories.get_crews", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories._get_json_crew", return_value=None
):
yield mock_get_flow
@@ -197,16 +202,61 @@ def test_reset_no_crew_or_flow_found(runner):
"crewai.utilities.reset_memories.get_crews", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories.get_flows", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories._get_json_crew", return_value=None
):
result = runner.invoke(reset_memories, ["-m"])
assert "No crew or flow found." in result.output
def test_reset_json_crew_memory(mock_crew, runner, monkeypatch, tmp_path):
monkeypatch.chdir(tmp_path)
(tmp_path / "crew.jsonc").write_text("{}")
with mock.patch(
"crewai.utilities.reset_memories.get_crews", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories.get_flows", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories.load_crew",
return_value=(mock_crew, {}),
) as mock_load_crew:
result = runner.invoke(reset_memories, ["-m"])
mock_load_crew.assert_called_once_with(Path("crew.jsonc"))
mock_crew.reset_memories.assert_called_once_with(command_type="memory")
assert f"[Crew ({mock_crew.name})] Memory has been reset." in result.output
def test_reset_json_crew_skipped_for_declared_flow_project(
mock_crew, runner, monkeypatch, tmp_path
):
monkeypatch.chdir(tmp_path)
(tmp_path / "crew.jsonc").write_text("{}")
(tmp_path / "pyproject.toml").write_text('[tool.crewai]\ntype = "flow"\n')
with mock.patch(
"crewai.utilities.reset_memories.get_crews", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories.get_flows", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories.load_crew",
return_value=(mock_crew, {}),
) as mock_load_crew:
result = runner.invoke(reset_memories, ["-m"])
mock_load_crew.assert_not_called()
mock_crew.reset_memories.assert_not_called()
assert "No crew or flow found." in result.output
def test_reset_crew_and_flow_memory(mock_crew, mock_flow, runner):
with mock.patch(
"crewai.utilities.reset_memories.get_crews", return_value=[mock_crew]
), mock.patch(
"crewai.utilities.reset_memories.get_flows", return_value=[mock_flow]
), mock.patch(
"crewai.utilities.reset_memories._get_json_crew", return_value=None
):
result = runner.invoke(reset_memories, ["-m"])
mock_crew.reset_memories.assert_called_once_with(command_type="memory")
@@ -223,6 +273,8 @@ def test_reset_flow_memory_none(runner):
"crewai.utilities.reset_memories.get_crews", return_value=[]
), mock.patch(
"crewai.utilities.reset_memories.get_flows", return_value=[mock_flow]
), mock.patch(
"crewai.utilities.reset_memories._get_json_crew", return_value=None
):
result = runner.invoke(reset_memories, ["-m"])
assert "[Flow (NoMemFlow)] Memory has been reset." in result.output