COG-6341 feat: Add --memory-only flag to cognee-cli forget

cognee/cli/api_dispatch.py

@@ -377,18 +377,31 @@ 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(
         dataset=dataset,
         dataset_id=dataset_id,
         data_id=data_id,
         everything=everything,
+        memory_only=memory_only,
     )
     fmt.success(f"Done: {result}")

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,13 +68,22 @@ 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(
                         data_id=data_id,
                         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

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()

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"""