code review changes

cognee/cli/commands/tui_command.py

@@ -3,7 +3,6 @@ from cognee.cli import SupportsCliCommand
 from cognee.cli.config import DEFAULT_DOCS_URL
 import cognee.cli.echo as fmt
 from cognee.cli.exceptions import CliCommandException
-from cognee.cli.tui.home_screen import HomeScreen
 
 
 class TuiCommand(SupportsCliCommand):
@@ -26,6 +25,7 @@ class TuiCommand(SupportsCliCommand):
     def execute(self, args: argparse.Namespace) -> None:
         try:
             from textual.app import App
+            from cognee.cli.tui.home_screen import HomeScreen
             from cognee.shared.logging_utils import setup_logging
             class CogneeTUI(App):
                 """Main TUI application for cognee."""

cognee/cli/tui/delete_screen.py

@@ -126,16 +126,18 @@ class DeleteTUIScreen(BaseTUIScreen):
         """Async function to delete data."""
         status = self.query_one(".tui-status", Static)
         try:
+            if user_id is None:
+                user = await get_default_user()
+                resolved_user_id = user.id
+            else:
+                resolved_user_id = UUID(user_id)
+
             if dataset_name:
-                if user_id is None:
-                    user = await get_default_user()
-                    resolved_user_id = user.id
-                else:
-                    resolved_user_id = UUID(user_id)
                 await delete_dataset_by_name(dataset_name, resolved_user_id)
+                status.update(f"✓ Successfully deleted dataset '{dataset_name}'.")
             else:
                 await delete_data_by_user(resolved_user_id)
-            status.update(f"✓ Successfully deleted dataset '{dataset_name}'.")
+                status.update(f"✓ Successfully deleted all data for user {resolved_user_id}.")
         except Exception as e:
             status.update(f"✗ Error: {str(e)}")
         finally:

tests/reproduce_tui_search_logic.py

@@ -139,11 +139,11 @@ async def test_empty_graph_handling():
     
     # Check that notify was called with warning
     # We allow flexible matching for the exact message but check 'warning' severity
+    assert screen.notify.called, "notify was not called"
     args = screen.notify.call_args
-    if args:
-        assert args[1].get('severity') == 'warning' or 'Knowledge graph is empty' in args[0][0]
-    else:
-        print("FAIL: notify was not called")
+    assert args is not None
+    # Assert specific conditions on args
+    assert args[1].get('severity') == 'warning' or 'Knowledge graph is empty' in args[0][0]
         
     print("SUCCESS: EntityNotFoundError was caught and handled correctly.")