From 0e4c13a4a632c894a33a9e34af4995e719745521 Mon Sep 17 00:00:00 2001 From: Goran Date: Fri, 28 Aug 2026 07:23:51 +0200 Subject: [PATCH] COG-6341 feat: Add --memory-only flag to cognee-cli forget cognee.forget() and CogneeApiClient.forget() both accept a memory_only parameter, but cognee-cli forget never exposed it in either the local argparse configuration or the --api-url remote-dispatch path, making memory-only forget entirely unreachable from the CLI. Adds the flag to both paths. Also fixes an unrelated bug found in the same function: --api-url forget dispatch read args.dataset_name, a field the CLI parser never sets (its flag is --dataset), silently dropping --dataset in that mode. And adds a guard for --everything --memory-only, which previously silently ignored --memory-only and did a full destructive wipe instead (--everything has no confirmation prompt). Split out of #4694 per reviewer feedback that PR bundled unrelated concerns; extracted as its own change. --- cognee/cli/api_dispatch.py | 17 +++- cognee/cli/commands/forget_command.py | 23 ++++- .../cli_unit_tests/test_api_dispatch.py | 99 +++++++++++++++++++ .../cli_unit_tests/test_cli_commands.py | 93 +++++++++++++++++ 4 files changed, 229 insertions(+), 3 deletions(-) diff --git a/cognee/cli/api_dispatch.py b/cognee/cli/api_dispatch.py index 1a88b34024..393af5f806 100644 --- a/cognee/cli/api_dispatch.py +++ b/cognee/cli/api_dispatch.py @@ -377,12 +377,24 @@ def _dispatch_improve(client: CogneeApiClient, args: argparse.Namespace) -> None def _dispatch_forget(client: CogneeApiClient, args: argparse.Namespace) -> None: everything = getattr(args, "everything", False) - dataset = getattr(args, "dataset_name", None) + # ForgetCommand's own flag is --dataset (-> args.dataset), not --dataset-name; + # reading dataset_name here always returned None, silently dropping --dataset + # in --api-url mode. + dataset = getattr(args, "dataset", None) dataset_id = getattr(args, "dataset_id", None) data_id = getattr(args, "data_id", None) + memory_only = getattr(args, "memory_only", False) + if dataset and dataset_id: + fmt.error("Provide either --dataset or --dataset-id, not both.") + return if not everything and not dataset and not dataset_id and not data_id: + fmt.error("Specify --dataset or --dataset-id, --data-id with dataset, or --everything.") + return + if everything and memory_only: fmt.error( - "Specify --dataset-name or --dataset-id, --data-id with dataset, or --everything." + "--memory-only has no effect with --everything: everything deletes all " + "datasets and data outright. Specify --dataset or --dataset-id with " + "--memory-only instead." ) return result = client.forget( @@ -390,5 +402,6 @@ def _dispatch_forget(client: CogneeApiClient, args: argparse.Namespace) -> None: dataset_id=dataset_id, data_id=data_id, everything=everything, + memory_only=memory_only, ) fmt.success(f"Done: {result}") diff --git a/cognee/cli/commands/forget_command.py b/cognee/cli/commands/forget_command.py index 62167d7ce4..23716ef306 100644 --- a/cognee/cli/commands/forget_command.py +++ b/cognee/cli/commands/forget_command.py @@ -17,7 +17,9 @@ class ForgetCommand(SupportsCliCommand): Remove data from the knowledge graph. Use --everything (alias --all) to delete all user data, --dataset/--dataset-id -to delete a dataset, or dataset + --data-id to delete a single item. +to delete a dataset, or dataset + --data-id to delete a single item. Add +--memory-only to clear graph/vector memory while keeping raw files and data +records, so the dataset (or item) can be re-cognified later. """ def configure_parser(self, parser: argparse.ArgumentParser) -> None: @@ -37,6 +39,16 @@ def configure_parser(self, parser: argparse.ArgumentParser) -> None: default=False, help="Delete all datasets and data", ) + parser.add_argument( + "--memory-only", + action="store_true", + default=False, + help=( + "Delete only graph/vector memory (requires --dataset or --dataset-id); " + "raw files and data records are preserved so the dataset can be " + "re-cognified with different settings" + ), + ) def execute(self, args: argparse.Namespace) -> None: try: @@ -56,6 +68,14 @@ def execute(self, args: argparse.Namespace) -> None: ) return + if args.everything and args.memory_only: + fmt.error( + "--memory-only has no effect with --everything: everything deletes all " + "datasets and data outright. Specify --dataset or --dataset-id with " + "--memory-only instead." + ) + return + async def run_forget(): try: return await cognee.forget( @@ -63,6 +83,7 @@ async def run_forget(): dataset=dataset, dataset_id=dataset_id, everything=args.everything, + memory_only=args.memory_only, ) except Exception as e: raise CliCommandInnerException(f"Failed to forget: {str(e)}") from e diff --git a/cognee/tests/cli_tests/cli_unit_tests/test_api_dispatch.py b/cognee/tests/cli_tests/cli_unit_tests/test_api_dispatch.py index b019ed8cca..fbff504324 100644 --- a/cognee/tests/cli_tests/cli_unit_tests/test_api_dispatch.py +++ b/cognee/tests/cli_tests/cli_unit_tests/test_api_dispatch.py @@ -177,3 +177,102 @@ def test_no_user_id_no_header(self, MockClient): call_kwargs = MockClient.call_args headers = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers", {}) assert "X-User-Id" not in headers + + +class TestForgetDispatch: + """Finding 8 (COG-6335 review): --memory-only must reach the API client, + and a --dataset value must reach it too (args.dataset, not the + never-set args.dataset_name the dispatcher used to read).""" + + @patch("cognee.cli.api_dispatch.CogneeApiClient") + def test_memory_only_and_dataset_forwarded_to_client(self, MockClient): + mock_instance = MagicMock() + mock_instance.forget.return_value = { + "status": "success", + "dataset_id": "ds-id", + "data_records_reset": 0, + } + MockClient.return_value.__enter__ = MagicMock(return_value=mock_instance) + MockClient.return_value.__exit__ = MagicMock(return_value=False) + + args = argparse.Namespace( + api_url="http://localhost:8000", + command="forget", + user_id=None, + dataset="my_dataset", + dataset_id=None, + data_id=None, + everything=False, + memory_only=True, + ) + dispatch(args) + + mock_instance.forget.assert_called_once_with( + dataset="my_dataset", + dataset_id=None, + data_id=None, + everything=False, + memory_only=True, + ) + + @patch("cognee.cli.api_dispatch.CogneeApiClient") + def test_everything_with_memory_only_does_not_call_client(self, MockClient): + """--memory-only has no effect with --everything (which deletes + outright) -- must error instead of silently doing a full wipe.""" + mock_instance = MagicMock() + MockClient.return_value.__enter__ = MagicMock(return_value=mock_instance) + MockClient.return_value.__exit__ = MagicMock(return_value=False) + + args = argparse.Namespace( + api_url="http://localhost:8000", + command="forget", + user_id=None, + dataset=None, + dataset_id=None, + data_id=None, + everything=True, + memory_only=True, + ) + dispatch(args) + + mock_instance.forget.assert_not_called() + + @patch("cognee.cli.api_dispatch.CogneeApiClient") + def test_dataset_and_dataset_id_both_set_does_not_call_client(self, MockClient): + mock_instance = MagicMock() + MockClient.return_value.__enter__ = MagicMock(return_value=mock_instance) + MockClient.return_value.__exit__ = MagicMock(return_value=False) + + args = argparse.Namespace( + api_url="http://localhost:8000", + command="forget", + user_id=None, + dataset="my_dataset", + dataset_id="11111111-1111-1111-1111-111111111111", + data_id=None, + everything=False, + memory_only=False, + ) + dispatch(args) + + mock_instance.forget.assert_not_called() + + @patch("cognee.cli.api_dispatch.CogneeApiClient") + def test_missing_forget_target_does_not_call_client(self, MockClient): + mock_instance = MagicMock() + MockClient.return_value.__enter__ = MagicMock(return_value=mock_instance) + MockClient.return_value.__exit__ = MagicMock(return_value=False) + + args = argparse.Namespace( + api_url="http://localhost:8000", + command="forget", + user_id=None, + dataset=None, + dataset_id=None, + data_id=None, + everything=False, + memory_only=False, + ) + dispatch(args) + + mock_instance.forget.assert_not_called() diff --git a/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py b/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py index 08fa649350..10c7f34993 100644 --- a/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py +++ b/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py @@ -15,6 +15,7 @@ from cognee.cli.commands.recall_command import RecallCommand from cognee.cli.commands.cognify_command import CognifyCommand from cognee.cli.commands.delete_command import DeleteCommand +from cognee.cli.commands.forget_command import ForgetCommand from cognee.cli.commands.config_command import ConfigCommand from cognee.cli.exceptions import CliCommandException from cognee.modules.data.methods.get_deletion_counts import DeletionCountsPreview @@ -616,6 +617,98 @@ def test_execute_with_exception(self, mock_asyncio_run): command.execute(args) +class TestForgetCommand: + """Test the ForgetCommand class""" + + def test_command_properties(self): + command = ForgetCommand() + assert command.command_string == "forget" + assert "Remove data" in command.help_string + assert command.docs_url is not None + + def test_configure_parser(self): + command = ForgetCommand() + parser = argparse.ArgumentParser() + + command.configure_parser(parser) + + actions = {action.dest: action for action in parser._actions} + assert "dataset" in actions + assert "dataset_id" in actions + assert "data_id" in actions + assert "everything" in actions + assert "memory_only" in actions + assert actions["memory_only"].default is False + + @patch("cognee.cli.commands.forget_command.asyncio.run", side_effect=_mock_run) + def test_execute_threads_memory_only_flag(self, mock_asyncio_run): + """--memory-only must reach cognee.forget(memory_only=True).""" + mock_cognee = MagicMock() + mock_cognee.forget = AsyncMock( + return_value={"status": "success", "dataset_id": "ds", "data_records_reset": 0} + ) + + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = ForgetCommand() + args = argparse.Namespace( + dataset="my_dataset", + dataset_id=None, + data_id=None, + everything=False, + memory_only=True, + ) + command.execute(args) + + mock_cognee.forget.assert_awaited_once_with( + data_id=None, + dataset="my_dataset", + dataset_id=None, + everything=False, + memory_only=True, + ) + + def test_execute_everything_with_memory_only_errors(self): + """--memory-only has no effect with --everything (which deletes + outright) -- must error instead of silently doing a full wipe.""" + mock_cognee = MagicMock() + mock_cognee.forget = AsyncMock() + + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = ForgetCommand() + args = argparse.Namespace( + dataset=None, dataset_id=None, data_id=None, everything=True, memory_only=True + ) + # Should not raise, just print an error and return without calling forget(). + command.execute(args) + + mock_cognee.forget.assert_not_awaited() + + def test_execute_no_forget_target(self): + command = ForgetCommand() + args = argparse.Namespace( + dataset=None, dataset_id=None, data_id=None, everything=False, memory_only=False + ) + + # Should not raise, just print an error and return. + command.execute(args) + + @patch("cognee.cli.commands.forget_command.asyncio.run") + def test_execute_with_exception(self, mock_asyncio_run): + mock_asyncio_run.side_effect = Exception("Forget error") + + command = ForgetCommand() + args = argparse.Namespace( + dataset="my_dataset", + dataset_id=None, + data_id=None, + everything=False, + memory_only=False, + ) + + with pytest.raises(CliCommandException): + command.execute(args) + + class TestConfigCommand: """Test the ConfigCommand class"""