refactor to modularize to multiple screens

cognee/cli/commands/tui_command.py

@@ -3,25 +3,10 @@ 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.config_tui import ConfigTUIScreen
+from cognee.cli.tui.home_screen import HomeScreen
-from cognee.version import get_cognee_version
+from textual.app import App
-from textual.app import App, ComposeResult
-from textual.widgets import ListView, ListItem, Static
-from textual.containers import Container, Horizontal
-from textual.binding import Binding
 
 
-def make_item(icon, command, description):
-    # Compose a ListItem that contains a Horizontal container with 3 children
-    return ListItem(
-        Horizontal(
-            Static(icon, classes="cmd-icon"),
-            Static(command, classes="cmd-name"),
-            Static(description, classes="cmd-desc"),
-            classes="cmd-row",
-        )
-    )
-
 
 class TuiCommand(SupportsCliCommand):
     @property
@@ -42,190 +27,18 @@ class TuiCommand(SupportsCliCommand):
 
     def execute(self, args: argparse.Namespace) -> None:
         try:
-            class CommandItem(Static):
-                """A custom widget for command items with icon and description."""
-
-                def __init__(self, icon: str, command: str, description: str):
-                    self.icon = icon
-                    self.command = command
-                    self.description = description
-                    super().__init__()
-
-                def render(self) -> str:
-                    return f"{self.icon}  {self.command:<12} {self.description}"
-
             class CogneeTUI(App):
+                """Main TUI application for cognee."""
+
                 CSS = """
                 Screen {
                     background: $surface;
                 }
-
-                #header {
-                    dock: top;
-                    background: $boost;
-                    color: $text;
-                    content-align: center middle;
-                    border: solid $primary;
-                    text-style: bold;
-                    padding: 1;
-                }
-
-                #main-container {
-                    height: 100%;
-                    border: solid $primary;
-                    background: $surface;
-                    padding: 1;
-                }
-                #title-wrapper {
-                    width: 100%;
-                    height: auto;
-                    align: center middle; 
-                }
-
-                #title {
-                    text-align: center;
-                    width: auto;
-                    color: $accent;
-                    text-style: bold;
-                    padding: 0 10;
-                    border: solid $accent;
-                    margin-bottom: 2;
-                }
-                
-                ListView > ListItem {
-                    width: 100%;
-                    padding: 0;
-                    margin: 0;
-                }
-
-                ListView {
-                    height: auto;
-                    background: $surface;
-                    border: none;
-                    padding: 0 0;
-                }
-
-                ListItem {
-                    background: $surface;
-                    color: $text;
-                    padding: 0 1;
-                    height: auto;
-                    width: 100%;
-                }
-                
-                ListItem.highlighted {
-                    background: $primary-darken-2;
-                }
-                
-                CommandItem {
-                    width: 100%;
-                    background: transparent;
-                }
-
-                #footer-info {
-                    dock: bottom;
-                    height: 3;
-                    background: $boost;
-                    color: $text-muted;
-                    content-align: center middle;
-                    border: solid $primary;
-                }
-                
-                .cmd-row {
-                    width: 100%;
-                    height: auto;
-                    align-horizontal: left;
-                    padding: 0 1;
-                }
-                
-                .cmd-icon {
-                    width: 4;
-                    text-align: center;
-                }
-                
-                .cmd-name {
-                    width: 14;
-                    padding-left: 1;
-                }
-                
-                .cmd-desc {
-                    width: 1fr;
-                    overflow: auto;
-                    padding-left: 1;
-                }
-                
                 """
 
-                BINDINGS = [
-                    Binding("q", "quit", "Quit", priority=True),
-                    Binding("escape", "quit", "Quit", priority=True),
-                    Binding("enter", "select", "Select", priority=True),
-                    Binding("up", "nav_up", "Up", priority=True),
-                    Binding("down", "nav_down", "Down", priority=True),
-                ]
-
-                def __init__(self):
-                    super().__init__()
-                    self.lv = None
-                    self.current_index = 0
-
-                def compose(self) -> ComposeResult:
-                    version = get_cognee_version()
-                    yield Static(f"🧠 cognee v{version}", id="header")
-
-                    with Container(id="main-container"):
-                        with Container(id="title-wrapper"):
-                            yield Static("Select Command", id="title")
-                        yield ListView(
-                            make_item("📥", "add", "Add data to cognee"),
-                            make_item("🔍", "search", "Search data in cognee"),
-                            make_item("⚡", "cognify", "Process data in cognee"),
-                            make_item("🗑️", "delete", "Delete data from cognee"),
-                            make_item("⚙️", "config", "Configure cognee settings"),
-                        )
-
-                    yield Static(
-                        "↑↓: Navigate  •  Enter: Select  •  q/Esc: Quit",
-                        id="footer-info"
-                    )
-
                 def on_mount(self) -> None:
-                    """Focus the list view on mount."""
+                    """Push the home screen on mount."""
-                    self.lv = self.query_one(ListView)
+                    self.push_screen(HomeScreen())
-                    self.current_index = 0
-                    self.set_focus(self.lv)
-                    self._apply_highlight()
-
-                def _apply_highlight(self) -> None:
-                    lv = self.lv
-                    children = list(lv.children)
-                    self.lv.index = self.current_index
-                    for idx, item in enumerate(children):
-                        if idx == self.current_index:
-                            item.add_class("highlighted")
-                        else:
-                            item.remove_class("highlighted")
-
-                def action_nav_up(self) -> None:
-                    self.current_index = max(0, self.current_index - 1)
-                    self._apply_highlight()
-
-                def action_nav_down(self) -> None:
-                    children = list(self.lv.children)
-                    self.current_index = min(len(children) - 1, self.current_index + 1)
-                    self._apply_highlight()
-
-                def on_list_view_selected(self, event: ListView.Selected) -> None:
-                    selected_index = event.index
-                    if selected_index == 4:
-                        self.app.push_screen(ConfigTUIScreen())
-                    else:
-                        self.exit()
-
-                def action_select(self) -> None:
-                    """Select the current item."""
-                    list_view = self.query_one(ListView)
-                    list_view.action_select_cursor()
 
             app = CogneeTUI()
             app.run()

cognee/cli/tui/base_screen.py

@@ -0,0 +1,41 @@
+from textual.screen import Screen
+from textual.app import ComposeResult
+from textual.widgets import Static
+
+from cognee.version import get_cognee_version
+
+
+class BaseTUIScreen(Screen):
+    """Base screen class with constant header for all TUI screens."""
+
+    # Subclasses should override this CSS and add their own styles
+    CSS = """
+    #header {
+        dock: top;
+        background: $boost;
+        color: $text;
+        content-align: center middle;
+        border: solid $primary;
+        text-style: bold;
+        padding: 1;
+    }
+    """
+
+    def compose_header(self) -> ComposeResult:
+        """Compose the constant header widget."""
+        version = get_cognee_version()
+        yield Static(f"🧠 cognee v{version}", id="header")
+
+    def compose_content(self) -> ComposeResult:
+        """Override this method in subclasses to provide screen-specific content."""
+        yield from ()
+
+    def compose_footer(self) -> ComposeResult:
+        """Override this method in subclasses to provide screen-specific footer."""
+        yield from ()
+
+    def compose(self) -> ComposeResult:
+        """Compose the screen with header, content, and footer."""
+        yield from self.compose_header()
+        yield from self.compose_content()
+        yield from self.compose_footer()

cognee/cli/tui/config_screen.py

@@ -1,6 +1,6 @@
 import argparse
 import json
-from typing import Optional, Tuple
+from typing import Optional
 
 from cognee.cli.reference import SupportsCliCommand
 from cognee.cli import DEFAULT_DOCS_URL
@@ -9,10 +9,11 @@ from cognee.cli.exceptions import CliCommandException
 
 from textual.app import App, ComposeResult
 from textual.screen import Screen
-from textual.widgets import DataTable, Footer, Header, Input, Label, Button
+from textual.widgets import DataTable, Input, Label, Button
-from textual.containers import Container, Vertical, Horizontal
+from textual.containers import Container, Horizontal
 from textual.binding import Binding
-from textual.coordinate import Coordinate
+
+from cognee.cli.tui.base_screen import BaseTUIScreen
 
 
 class EditModal(Screen):
@@ -158,7 +159,7 @@ class ConfirmModal(Screen):
         self.dismiss(False)
 
 
-class ConfigTUIScreen(Screen):
+class ConfigTUIScreen(BaseTUIScreen):
     """Main config TUI screen."""
 
     BINDINGS = [
@@ -170,21 +171,11 @@ class ConfigTUIScreen(Screen):
         Binding("down", "cursor_down", "Down", show=False),
     ]
 
-    CSS = """
+    CSS = BaseTUIScreen.CSS + """
     ConfigTUIScreen {
         background: $surface;
     }
 
-    #config-header {
-        dock: top;
-        background: $boost;
-        color: $text;
-        content-align: center middle;
-        text-style: bold;
-        padding: 1;
-        border: solid $primary;
-    }
-
     #config-container {
         height: 100%;
         padding: 1;
@@ -217,15 +208,14 @@ class ConfigTUIScreen(Screen):
         "chunk_overlap": ("set_chunk_overlap", "10"),
     }
 
-    def compose(self) -> ComposeResult:
+    def compose_content(self) -> ComposeResult:
-        yield Label("🧠 cognee Config Manager", id="config-header")
-
         with Container(id="config-container"):
             table = DataTable()
             table.cursor_type = "row"
             table.zebra_stripes = True
             yield table
 
+    def compose_footer(self) -> ComposeResult:
         yield Label(
             "[↑↓] Navigate  [e] Edit  [r] Reset  [Esc] Back  [q] Quit",
             id="config-footer"

cognee/cli/tui/home_screen.py

@@ -0,0 +1,182 @@
+from textual.app import ComposeResult
+from textual.widgets import ListView, ListItem, Static
+from textual.containers import Container, Horizontal
+from textual.binding import Binding
+
+from cognee.cli.tui.base_screen import BaseTUIScreen
+from cognee.cli.tui.config_screen import ConfigTUIScreen
+
+
+def make_item(icon: str, command: str, description: str) -> ListItem:
+    """Compose a ListItem that contains a Horizontal container with 3 children."""
+    return ListItem(
+        Horizontal(
+            Static(icon, classes="cmd-icon"),
+            Static(command, classes="cmd-name"),
+            Static(description, classes="cmd-desc"),
+            classes="cmd-row",
+        )
+    )
+
+
+class HomeScreen(BaseTUIScreen):
+    """Home screen with command selection menu."""
+
+    BINDINGS = [
+        Binding("q", "quit_app", "Quit"),
+        Binding("escape", "quit_app", "Quit"),
+        Binding("enter", "select", "Select"),
+        Binding("up", "nav_up", "Up", priority=True),
+        Binding("down", "nav_down", "Down", priority=True),
+    ]
+
+    CSS = BaseTUIScreen.CSS + """
+    HomeScreen {
+        background: $surface;
+    }
+
+    #main-container {
+        height: 100%;
+        border: solid $primary;
+        background: $surface;
+        padding: 1;
+    }
+
+    #title-wrapper {
+        width: 100%;
+        height: auto;
+        align: center middle;
+    }
+
+    #title {
+        text-align: center;
+        width: auto;
+        color: $accent;
+        text-style: bold;
+        padding: 0 10;
+        border: solid $accent;
+        margin-bottom: 2;
+    }
+
+    ListView > ListItem {
+        width: 100%;
+        padding: 0;
+        margin: 0;
+    }
+
+    ListView {
+        height: auto;
+        background: $surface;
+        border: none;
+        padding: 0 0;
+    }
+
+    ListItem {
+        background: $surface;
+        color: $text;
+        padding: 0 1;
+        height: auto;
+        width: 100%;
+    }
+
+    ListItem.highlighted {
+        background: $primary-darken-2;
+    }
+
+    .cmd-row {
+        width: 100%;
+        height: auto;
+        align-horizontal: left;
+        padding: 0 1;
+    }
+
+    .cmd-icon {
+        width: 4;
+        text-align: center;
+    }
+
+    .cmd-name {
+        width: 14;
+        padding-left: 1;
+    }
+
+    .cmd-desc {
+        width: 1fr;
+        overflow: auto;
+        padding-left: 1;
+    }
+
+    #home-footer {
+        dock: bottom;
+        height: 3;
+        background: $boost;
+        color: $text-muted;
+        content-align: center middle;
+        border: solid $primary;
+    }
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.lv = None
+        self.current_index = 0
+
+    def compose_content(self) -> ComposeResult:
+        with Container(id="main-container"):
+            with Container(id="title-wrapper"):
+                yield Static("Select Command", id="title")
+            yield ListView(
+                make_item("📥", "add", "Add data to cognee"),
+                make_item("🔍", "search", "Search data in cognee"),
+                make_item("⚡", "cognify", "Process data in cognee"),
+                make_item("🗑️", "delete", "Delete data from cognee"),
+                make_item("⚙️", "config", "Configure cognee settings"),
+            )
+
+    def compose_footer(self) -> ComposeResult:
+        yield Static(
+            "↑↓: Navigate  •  Enter: Select  •  q/Esc: Quit",
+            id="home-footer"
+        )
+
+    def on_mount(self) -> None:
+        """Focus the list view on mount."""
+        self.lv = self.query_one(ListView)
+        self.current_index = 0
+        self.set_focus(self.lv)
+        self._apply_highlight()
+
+    def _apply_highlight(self) -> None:
+        lv = self.lv
+        children = list(lv.children)
+        self.lv.index = self.current_index
+        for idx, item in enumerate(children):
+            if idx == self.current_index:
+                item.add_class("highlighted")
+            else:
+                item.remove_class("highlighted")
+
+    def action_nav_up(self) -> None:
+        self.current_index = max(0, self.current_index - 1)
+        self._apply_highlight()
+
+    def action_nav_down(self) -> None:
+        children = list(self.lv.children)
+        self.current_index = min(len(children) - 1, self.current_index + 1)
+        self._apply_highlight()
+
+    def on_list_view_selected(self, event: ListView.Selected) -> None:
+        selected_index = event.index
+        if selected_index == 4:
+            self.app.push_screen(ConfigTUIScreen())
+        else:
+            self.app.exit()
+
+    def action_select(self) -> None:
+        """Select the current item."""
+        list_view = self.query_one(ListView)
+        list_view.action_select_cursor()
+
+    def action_quit_app(self) -> None:
+        """Quit the entire application."""
+        self.app.exit()