mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-03 14:09:24 +00:00
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:
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user