From 73ef9c2c27335bd9ea59bd4b6912e7bb73718520 Mon Sep 17 00:00:00 2001 From: rajeevrajeshuni Date: Sat, 29 Nov 2025 18:54:39 +0530 Subject: [PATCH] WIP - adding a add data screen --- cognee/cli/tui/add_screen.py | 191 ++++++++++++++++++++++++++++++++++ cognee/cli/tui/home_screen.py | 5 +- 2 files changed, 195 insertions(+), 1 deletion(-) create mode 100644 cognee/cli/tui/add_screen.py diff --git a/cognee/cli/tui/add_screen.py b/cognee/cli/tui/add_screen.py new file mode 100644 index 000000000..bfcf3b007 --- /dev/null +++ b/cognee/cli/tui/add_screen.py @@ -0,0 +1,191 @@ +import asyncio +from textual.app import ComposeResult +from textual.widgets import Input, Label, Button, Static, TextArea +from textual.containers import Container, Vertical +from textual.binding import Binding + +from cognee.cli.tui.base_screen import BaseTUIScreen + + +class AddTUIScreen(BaseTUIScreen): + """TUI screen for adding data to cognee.""" + + BINDINGS = [ + Binding("q", "quit_app", "Quit"), + Binding("escape", "back", "Back"), + Binding("ctrl+s", "submit", "Submit"), + Binding("ctrl+v", "paste", "Paste", show=False), + ] + + CSS = BaseTUIScreen.CSS + """ + AddTUIScreen { + background: $surface; + } + + #add-container { + height: auto; + padding: 1; + } + + #add-form { + width: 100%; + height: auto; + border: solid $primary; + padding: 2; + background: $surface; + } + + #form-title { + text-align: center; + width: 100%; + text-style: bold; + color: $accent; + margin-bottom: 2; + } + + .field-label { + color: $text-muted; + margin-top: 1; + margin-bottom: 0; + } + + Input { + width: 100%; + margin-bottom: 1; + } + + #data-input { + height: 8; + min-height: 8; + } + + #submit-btn { + width: 20; + margin-top: 2; + } + + #status-message { + margin-top: 2; + text-align: center; + height: auto; + } + + #add-footer { + dock: bottom; + height: 3; + background: $boost; + color: $text-muted; + content-align: center middle; + border: solid $primary; + } + """ + + def __init__(self): + super().__init__() + self.is_processing = False + + def compose_content(self) -> ComposeResult: + with Container(id="add-container"): + yield Label("Add Data to Cognee", id="form-title") + with Vertical(id="add-form"): + yield Label("Data (text, file path, URL, or S3 path):", classes="field-label") + yield TextArea( + "", + id="data-input", + ) + + yield Label("Dataset Name:", classes="field-label") + yield Input( + placeholder="main_dataset", + value="main_dataset", + id="dataset-input" + ) + + yield Button("Add Data [Ctrl+S]", variant="primary", id="submit-btn") + yield Static("", id="status-message") + + def compose_footer(self) -> ComposeResult: + yield Static( + "Ctrl+S: Submit • Esc: Back • q: Quit", + id="add-footer" + ) + + def on_mount(self) -> None: + """Focus the data input on mount.""" + data_input = self.query_one("#data-input", TextArea) + data_input.focus() + + def action_back(self) -> None: + """Go back to home screen.""" + if not self.is_processing: + self.app.pop_screen() + + def action_quit_app(self) -> None: + """Quit the entire application.""" + self.app.exit() + + def action_paste(self) -> None: + """Handle paste action - Textual handles this automatically.""" + pass + + def action_submit(self) -> None: + """Submit the form.""" + if not self.is_processing: + self._submit_data() + + def on_button_pressed(self, event: Button.Pressed) -> None: + """Handle button press.""" + if event.button.id == "submit-btn" and not self.is_processing: + self._submit_data() + + def _submit_data(self) -> None: + """Process and submit the data.""" + data_input = self.query_one("#data-input", TextArea) + dataset_input = self.query_one("#dataset-input", Input) + status = self.query_one("#status-message", Static) + + data = data_input.text.strip() + dataset_name = dataset_input.value.strip() or "main_dataset" + + if not data: + status.update("[red]✗ Please enter data to add[/red]") + return + + self.is_processing = True + status.update("[yellow]⏳ Processing...[/yellow]") + + # Disable inputs during processing + data_input.disabled = True + dataset_input.disabled = True + self.query_one("#submit-btn", Button).disabled = True + + # Run async add operation + asyncio.create_task(self._add_data_async(data, dataset_name)) + + async def _add_data_async(self, data: str, dataset_name: str) -> None: + """Async function to add data to cognee.""" + status = self.query_one("#status-message", Static) + + try: + import cognee + + await cognee.add(data=data, dataset_name=dataset_name) + + status.update(f"[green]✓ Successfully added data to dataset '{dataset_name}'[/green]") + + # Clear the data input after successful add + data_input = self.query_one("#data-input", TextArea) + data_input.clear() + + except Exception as e: + status.update(f"[red]✗ Failed to add data: {str(e)}[/red]") + + finally: + # Re-enable inputs + self.is_processing = False + data_input = self.query_one("#data-input", TextArea) + dataset_input = self.query_one("#dataset-input", Input) + data_input.disabled = False + dataset_input.disabled = False + self.query_one("#submit-btn", Button).disabled = False + data_input.focus() \ No newline at end of file diff --git a/cognee/cli/tui/home_screen.py b/cognee/cli/tui/home_screen.py index f8f75fa7b..0c3cc8d18 100644 --- a/cognee/cli/tui/home_screen.py +++ b/cognee/cli/tui/home_screen.py @@ -5,6 +5,7 @@ from textual.binding import Binding from cognee.cli.tui.base_screen import BaseTUIScreen from cognee.cli.tui.config_screen import ConfigTUIScreen +from cognee.cli.tui.add_screen import AddTUIScreen def make_item(icon: str, command: str, description: str) -> ListItem: @@ -167,7 +168,9 @@ class HomeScreen(BaseTUIScreen): def on_list_view_selected(self, event: ListView.Selected) -> None: selected_index = event.index - if selected_index == 4: + if selected_index == 0: # add + self.app.push_screen(AddTUIScreen()) + elif selected_index == 4: # config self.app.push_screen(ConfigTUIScreen()) else: self.app.exit()