From 37e862c0535071985ef2cc8b0286000dbe2a972e Mon Sep 17 00:00:00 2001 From: rajeevrajeshuni Date: Sun, 30 Nov 2025 14:20:25 +0530 Subject: [PATCH] ruff changes --- cognee/cli/commands/delete_command.py | 2 +- cognee/cli/commands/tui_command.py | 1 + cognee/cli/tui/add_screen.py | 35 +++++++------ cognee/cli/tui/base_screen.py | 5 +- cognee/cli/tui/cognify_screen.py | 50 +++++++++++-------- cognee/cli/tui/common_styles.py | 2 +- cognee/cli/tui/config_screen.py | 34 ++++++++----- cognee/cli/tui/delete_screen.py | 35 +++++++------ cognee/cli/tui/home_screen.py | 10 ++-- cognee/cli/tui/search_screen.py | 19 ++++--- .../sqlalchemy/SqlAlchemyAdapter.py | 9 ++-- .../data/methods/delete_data_by_user.py | 2 +- .../data/methods/delete_dataset_by_name.py | 6 +-- 13 files changed, 120 insertions(+), 90 deletions(-) diff --git a/cognee/cli/commands/delete_command.py b/cognee/cli/commands/delete_command.py index 1a52d5e1b..8400d3b0f 100644 --- a/cognee/cli/commands/delete_command.py +++ b/cognee/cli/commands/delete_command.py @@ -109,4 +109,4 @@ Be careful with deletion operations as they are irreversible. except Exception as e: if isinstance(e, CliCommandInnerException): raise CliCommandException(str(e), error_code=1) from e - raise CliCommandException(f"Error deleting data: {str(e)}", error_code=1) from e \ No newline at end of file + raise CliCommandException(f"Error deleting data: {str(e)}", error_code=1) from e diff --git a/cognee/cli/commands/tui_command.py b/cognee/cli/commands/tui_command.py index 08ba54353..19f46a18e 100644 --- a/cognee/cli/commands/tui_command.py +++ b/cognee/cli/commands/tui_command.py @@ -27,6 +27,7 @@ class TuiCommand(SupportsCliCommand): def execute(self, args: argparse.Namespace) -> None: try: + class CogneeTUI(App): """Main TUI application for cognee.""" diff --git a/cognee/cli/tui/add_screen.py b/cognee/cli/tui/add_screen.py index c74b539a6..371023baa 100644 --- a/cognee/cli/tui/add_screen.py +++ b/cognee/cli/tui/add_screen.py @@ -17,12 +17,15 @@ class AddTUIScreen(BaseTUIScreen): Binding("ctrl+v", "paste", "Paste", show=False), ] - CSS = BaseTUIScreen.CSS + """ + CSS = ( + BaseTUIScreen.CSS + + """ #data-input { height: 8; min-height: 8; } """ + ) def __init__(self): super().__init__() @@ -33,25 +36,21 @@ class AddTUIScreen(BaseTUIScreen): with Container(classes="tui-title-wrapper"): yield Static("📥 Add Data to Cognee", classes="tui-title-bordered") with Vertical(classes="tui-form"): - yield Label("Data (text, file path (/path/to/doc), URL, or S3 path (s3://bucket)):", classes="tui-label-spaced") + yield Label( + "Data (text, file path (/path/to/doc), URL, or S3 path (s3://bucket)):", + classes="tui-label-spaced", + ) yield TextArea( "", id="data-input", ) - + yield Label("Dataset Name:", classes="tui-label-spaced") - yield Input( - placeholder="main_dataset", - value="main_dataset", - id="dataset-input" - ) + yield Input(placeholder="main_dataset", value="main_dataset", id="dataset-input") yield Static("", classes="tui-status") def compose_footer(self) -> ComposeResult: - yield Static( - "Ctrl+S: Add • Esc: Back • q: Quit", - classes="tui-footer" - ) + yield Static("Ctrl+S: Add • Esc: Back • q: Quit", classes="tui-footer") def on_mount(self) -> None: """Focus the data input on mount.""" @@ -91,7 +90,7 @@ class AddTUIScreen(BaseTUIScreen): self.is_processing = True status.update("[yellow]⏳ Processing...[/yellow]") - + # Disable inputs during processing data_input.disabled = True dataset_input.disabled = True @@ -102,21 +101,21 @@ class AddTUIScreen(BaseTUIScreen): async def _add_data_async(self, data: str, dataset_name: str) -> None: """Async function to add data to cognee.""" status = self.query_one(".tui-status", 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 diff --git a/cognee/cli/tui/base_screen.py b/cognee/cli/tui/base_screen.py index 7256e0353..f03688ff6 100644 --- a/cognee/cli/tui/base_screen.py +++ b/cognee/cli/tui/base_screen.py @@ -10,7 +10,9 @@ class BaseTUIScreen(Screen): """Base screen class with constant header for all TUI screens.""" # Subclasses should override this CSS and add their own styles - CSS = COMMON_STYLES + """ + CSS = ( + COMMON_STYLES + + """ #header { dock: top; background: $boost; @@ -21,6 +23,7 @@ class BaseTUIScreen(Screen): padding: 1; } """ + ) def compose_header(self) -> ComposeResult: """Compose the constant header widget.""" diff --git a/cognee/cli/tui/cognify_screen.py b/cognee/cli/tui/cognify_screen.py index 15035a912..c939da02e 100644 --- a/cognee/cli/tui/cognify_screen.py +++ b/cognee/cli/tui/cognify_screen.py @@ -16,7 +16,9 @@ class CognifyTUIScreen(BaseTUIScreen): Binding("ctrl+s", "submit", "Submit"), ] - CSS = BaseTUIScreen.CSS + """ + CSS = ( + BaseTUIScreen.CSS + + """ Checkbox { margin-top: 1; margin-bottom: 1; @@ -32,6 +34,7 @@ class CognifyTUIScreen(BaseTUIScreen): height: 1; } """ + ) def __init__(self): super().__init__() @@ -42,26 +45,23 @@ class CognifyTUIScreen(BaseTUIScreen): with Container(classes="tui-title-wrapper"): yield Static("⚡ Cognify Data", classes="tui-title-bordered") with Vertical(classes="tui-form"): - yield Label("Dataset Name (optional, leave empty for all):", classes="tui-label-spaced") - yield Input( - placeholder="Leave empty to process all datasets", - value="", - id="dataset-input" + yield Label( + "Dataset Name (optional, leave empty for all):", classes="tui-label-spaced" ) - + yield Input( + placeholder="Leave empty to process all datasets", value="", id="dataset-input" + ) + yield Label("Chunker Type:", classes="tui-label-spaced") with RadioSet(id="chunker-radio"): for chunker in CHUNKER_CHOICES: yield RadioButton(chunker, value=(chunker == "TextChunker")) - + yield Checkbox("Run in background", id="background-checkbox") yield Static("", classes="tui-status") def compose_footer(self) -> ComposeResult: - yield Static( - "Ctrl+S: Start • Esc: Back • q: Quit", - classes="tui-footer" - ) + yield Static("Ctrl+S: Start • Esc: Back • q: Quit", classes="tui-footer") def on_mount(self) -> None: """Focus the dataset input on mount.""" @@ -90,12 +90,16 @@ class CognifyTUIScreen(BaseTUIScreen): status = self.query_one(".tui-status", Static) dataset_name = dataset_input.value.strip() or None - chunker_type = str(chunker_radio.pressed_button.label) if chunker_radio.pressed_button else "TextChunker" + chunker_type = ( + str(chunker_radio.pressed_button.label) + if chunker_radio.pressed_button + else "TextChunker" + ) run_background = background_checkbox.value self.is_processing = True status.update("[yellow]⏳ Starting cognification...[/yellow]") - + # Disable inputs during processing dataset_input.disabled = True chunker_radio.disabled = True @@ -104,10 +108,13 @@ class CognifyTUIScreen(BaseTUIScreen): # Run async cognify operation asyncio.create_task(self._cognify_async(dataset_name, chunker_type, run_background)) - async def _cognify_async(self, dataset_name: str | None, chunker_type: str, run_background: bool) -> None: + async def _cognify_async( + self, dataset_name: str | None, chunker_type: str, run_background: bool + ) -> None: """Async function to cognify data.""" status = self.query_one(".tui-status", Static) from cognee.modules.chunking.TextChunker import TextChunker + try: # Get chunker class chunker_class = TextChunker @@ -119,7 +126,9 @@ class CognifyTUIScreen(BaseTUIScreen): if LangchainChunker is not None: chunker_class = LangchainChunker else: - status.update("[yellow]⚠ LangchainChunker not available, using TextChunker[/yellow]") + status.update( + "[yellow]⚠ LangchainChunker not available, using TextChunker[/yellow]" + ) elif chunker_type == "CsvChunker": try: from cognee.modules.chunking.CsvChunker import CsvChunker @@ -129,24 +138,25 @@ class CognifyTUIScreen(BaseTUIScreen): chunker_class = CsvChunker else: status.update("[yellow]⚠ CsvChunker not available, using TextChunker[/yellow]") - + # Prepare datasets parameter datasets = [dataset_name] if dataset_name else None import cognee + await cognee.cognify( datasets=datasets, chunker=chunker_class, run_in_background=run_background, ) - + if run_background: status.update("[green]✓ Cognification started in background![/green]") else: status.update("[green]✓ Cognification completed successfully![/green]") - + except Exception as e: status.update(f"[red]✗ Failed to cognify: {str(e)}[/red]") - + finally: # Re-enable inputs self.is_processing = False diff --git a/cognee/cli/tui/common_styles.py b/cognee/cli/tui/common_styles.py index 95a4ddd8e..4cd97ce7c 100644 --- a/cognee/cli/tui/common_styles.py +++ b/cognee/cli/tui/common_styles.py @@ -132,4 +132,4 @@ Button { height: auto; margin-bottom: 2; } -""" \ No newline at end of file +""" diff --git a/cognee/cli/tui/config_screen.py b/cognee/cli/tui/config_screen.py index f062a0f01..508d6d9b3 100644 --- a/cognee/cli/tui/config_screen.py +++ b/cognee/cli/tui/config_screen.py @@ -97,7 +97,9 @@ class ConfigTUIScreen(BaseTUIScreen): Binding("down", "cursor_down", "Down", show=False), ] - CSS = BaseTUIScreen.CSS + """ + CSS = ( + BaseTUIScreen.CSS + + """ DataTable { height: 1fr; text-align: center; @@ -129,6 +131,7 @@ class ConfigTUIScreen(BaseTUIScreen): margin-bottom: 1; } """ + ) # Config key mappings: Key -> (Reset Method Name, Default Value) CONFIG_MAP = { @@ -166,7 +169,7 @@ class ConfigTUIScreen(BaseTUIScreen): def compose_footer(self) -> ComposeResult: yield Static( "↑↓: Navigate • e: Edit • Enter: Save • r: Reset • Esc: Back • q: Quit", - classes="tui-footer" + classes="tui-footer", ) def on_mount(self) -> None: @@ -184,6 +187,7 @@ class ConfigTUIScreen(BaseTUIScreen): try: import cognee + # Check if get method exists, otherwise warn has_get = hasattr(cognee.config, "get") except ImportError: @@ -205,11 +209,13 @@ class ConfigTUIScreen(BaseTUIScreen): table.add_row(key, value_display, key=key) def action_cursor_up(self) -> None: - if self.editing_key: return + if self.editing_key: + return self.query_one(DataTable).action_cursor_up() def action_cursor_down(self) -> None: - if self.editing_key: return + if self.editing_key: + return self.query_one(DataTable).action_cursor_down() def action_cancel_or_back(self) -> None: @@ -223,10 +229,12 @@ class ConfigTUIScreen(BaseTUIScreen): def action_edit(self) -> None: """Start inline editing for the selected config value.""" - if self.editing_key: return + if self.editing_key: + return table = self.query_one(DataTable) - if table.cursor_row < 0: return + if table.cursor_row < 0: + return # Get row data using the cursor row_key = table.coordinate_to_cell_key(table.cursor_coordinate).row_key @@ -253,7 +261,8 @@ class ConfigTUIScreen(BaseTUIScreen): def action_confirm_edit(self) -> None: """Confirm the inline edit and save the value.""" - if not self.editing_key: return + if not self.editing_key: + return input_widget = self.query_one("#inline-input", Input) value = input_widget.value.strip() @@ -276,7 +285,8 @@ class ConfigTUIScreen(BaseTUIScreen): def action_reset(self) -> None: """Reset the selected config to default.""" table = self.query_one(DataTable) - if table.cursor_row < 0: return + if table.cursor_row < 0: + return row_key_obj = table.coordinate_to_cell_key(table.cursor_coordinate).row_key key = str(row_key_obj.value) @@ -292,10 +302,7 @@ class ConfigTUIScreen(BaseTUIScreen): if confirmed: self._reset_config(key) - self.app.push_screen( - ConfirmModal(key, display_default), - handle_confirm_result - ) + self.app.push_screen(ConfirmModal(key, display_default), handle_confirm_result) def _save_config(self, key: str, value: str) -> None: """Save config value and update UI.""" @@ -373,6 +380,7 @@ class ConfigTUICommand(SupportsCliCommand): class ConfigTUIApp(App): """Simple app wrapper for config TUI.""" + CSS = """ Screen { background: $surface; } """ @@ -393,4 +401,4 @@ class ConfigTUICommand(SupportsCliCommand): f"Failed to launch config TUI: {str(ex)}", docs_url=self.docs_url, raiseable_exception=ex, - ) \ No newline at end of file + ) diff --git a/cognee/cli/tui/delete_screen.py b/cognee/cli/tui/delete_screen.py index db278ca24..37959cf4a 100644 --- a/cognee/cli/tui/delete_screen.py +++ b/cognee/cli/tui/delete_screen.py @@ -21,13 +21,16 @@ class DeleteTUIScreen(BaseTUIScreen): Binding("ctrl+a", "delete_all", "Delete All", priority=True), ] - CSS = BaseTUIScreen.CSS + """ + CSS = ( + BaseTUIScreen.CSS + + """ #button-group { height: auto; align: center middle; margin-top: 2; } """ + ) def __init__(self): super().__init__() @@ -42,26 +45,22 @@ class DeleteTUIScreen(BaseTUIScreen): yield Label("Dataset Name (optional):", classes="tui-label") yield Input( placeholder="Enter dataset name to delete specific dataset", - id="dataset-input" + id="dataset-input", ) - + with Vertical(classes="tui-input-group"): yield Label("User ID (optional):", classes="tui-label") - yield Input( - placeholder="Enter user ID to delete user's data", - id="user-input" - ) - + yield Input(placeholder="Enter user ID to delete user's data", id="user-input") + with Horizontal(id="button-group"): yield Button("Delete", variant="error", id="delete-btn") yield Button("Delete All", variant="error", id="delete-all-btn") - + yield Static("", classes="tui-status") def compose_footer(self) -> ComposeResult: yield Static( - "Ctrl+s: Delete • Ctrl+a: Delete All • Esc: Back • q: Quit", - classes="tui-footer" + "Ctrl+s: Delete • Ctrl+a: Delete All • Esc: Back • q: Quit", classes="tui-footer" ) def on_mount(self) -> None: @@ -133,17 +132,19 @@ class DeleteTUIScreen(BaseTUIScreen): user = await get_default_user() user_id = user.id result = await delete_dataset_by_name(dataset_name, user_id) - + if result["not_found"]: status.update(f"⚠️ Dataset '{dataset_name}' not found") self.is_processing = False return - + status.update(f"✓ Successfully deleted {result['deleted_count']} dataset(s)") else: # For user_id deletion, use the new delete_data_by_user method result = await delete_data_by_user(UUID(user_id)) - status.update(f"✓ Successfully deleted {result['datasets_deleted']} datasets and {result['data_entries_deleted']} data entries for user '{user_id}'") + status.update( + f"✓ Successfully deleted {result['datasets_deleted']} datasets and {result['data_entries_deleted']} data entries for user '{user_id}'" + ) except Exception as e: status.update(f"✗ Error: {str(e)}") @@ -191,6 +192,7 @@ class DeleteTUIScreen(BaseTUIScreen): # Perform deletion - delete all uses the original cognee.delete import cognee + await cognee.delete(dataset_name=None, user_id=None) status.update("✓ Successfully deleted all data") @@ -214,7 +216,9 @@ class DeleteAllConfirmModal(BaseTUIScreen): Binding("escape", "cancel", "Cancel"), ] - CSS = BaseTUIScreen.CSS + """ + CSS = ( + BaseTUIScreen.CSS + + """ DeleteAllConfirmModal { align: center middle; } @@ -241,6 +245,7 @@ class DeleteAllConfirmModal(BaseTUIScreen): margin-bottom: 2; } """ + ) def compose_content(self) -> ComposeResult: with Container(id="confirm-dialog"): diff --git a/cognee/cli/tui/home_screen.py b/cognee/cli/tui/home_screen.py index e841a7f66..d23061175 100644 --- a/cognee/cli/tui/home_screen.py +++ b/cognee/cli/tui/home_screen.py @@ -34,7 +34,9 @@ class HomeScreen(BaseTUIScreen): Binding("down", "nav_down", "Down", priority=True), ] - CSS = BaseTUIScreen.CSS + """ + CSS = ( + BaseTUIScreen.CSS + + """ ListView > ListItem { width: 100%; padding: 0; @@ -108,6 +110,7 @@ class HomeScreen(BaseTUIScreen): color: $text-muted; } """ + ) def __init__(self): super().__init__() @@ -130,10 +133,7 @@ class HomeScreen(BaseTUIScreen): ) def compose_footer(self) -> ComposeResult: - yield Static( - "↑↓: Navigate • Enter: Select • q/Esc: Quit", - classes="tui-footer" - ) + yield Static("↑↓: Navigate • Enter: Select • q/Esc: Quit", classes="tui-footer") def on_mount(self) -> None: """Focus the list view on mount.""" diff --git a/cognee/cli/tui/search_screen.py b/cognee/cli/tui/search_screen.py index cb9152aff..97a292762 100644 --- a/cognee/cli/tui/search_screen.py +++ b/cognee/cli/tui/search_screen.py @@ -5,6 +5,7 @@ from textual.containers import Container, Vertical, ScrollableContainer from textual.binding import Binding from cognee.cli.tui.base_screen import BaseTUIScreen + class SearchTUIScreen(BaseTUIScreen): """Simple search screen with query input and results display.""" @@ -14,7 +15,9 @@ class SearchTUIScreen(BaseTUIScreen): Binding("ctrl+s", "search", "Search"), ] - CSS = BaseTUIScreen.CSS + """ + CSS = ( + BaseTUIScreen.CSS + + """ #search-form { height: auto; border: solid $primary; @@ -48,6 +51,7 @@ class SearchTUIScreen(BaseTUIScreen): overflow-y: auto; } """ + ) def __init__(self): super().__init__() @@ -75,13 +79,12 @@ class SearchTUIScreen(BaseTUIScreen): with Container(id="results-container"): yield Static("Results", id="results-title") with ScrollableContainer(id="results-content"): - yield Static("Enter a query and click Search to see results.", id="results-text") + yield Static( + "Enter a query and click Search to see results.", id="results-text" + ) def compose_footer(self) -> ComposeResult: - yield Static( - "Ctrl+S: Search • Esc: Back • q: Quit", - classes="tui-footer" - ) + yield Static("Ctrl+S: Search • Esc: Back • q: Quit", classes="tui-footer") def on_mount(self) -> None: """Focus the query input on mount.""" @@ -157,7 +160,9 @@ class SearchTUIScreen(BaseTUIScreen): if query_type in ["GRAPH_COMPLETION", "RAG_COMPLETION"]: formatted = "\n\n".join([f"📝 {result}" for result in results]) elif query_type == "CHUNKS": - formatted = "\n\n".join([f"📄 Chunk {i + 1}:\n{result}" for i, result in enumerate(results)]) + formatted = "\n\n".join( + [f"📄 Chunk {i + 1}:\n{result}" for i, result in enumerate(results)] + ) else: formatted = "\n\n".join([f"• {result}" for result in results]) diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py index c02c486a4..eeeaefdde 100644 --- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py +++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py @@ -235,10 +235,10 @@ class SQLAlchemyAdapter: return [] async def delete_entities_by_id( - self, - table_name: str, - data_id: Union[UUID, List[UUID]], # Supports a single UUID or a List of UUIDs - schema_name: Optional[str] = "public" + self, + table_name: str, + data_id: Union[UUID, List[UUID]], # Supports a single UUID or a List of UUIDs + schema_name: Optional[str] = "public", ): """ Delete one or more entities from the specified table based on their ID(s). @@ -265,6 +265,7 @@ class SQLAlchemyAdapter: # Handle SQLite's foreign key requirement if self.engine.dialect.name == "sqlite": from sqlalchemy import text + await session.execute(text("PRAGMA foreign_keys = ON;")) # Construct the DELETE statement using the 'in_()' operator diff --git a/cognee/modules/data/methods/delete_data_by_user.py b/cognee/modules/data/methods/delete_data_by_user.py index 0f94df852..ab4a5aad7 100644 --- a/cognee/modules/data/methods/delete_data_by_user.py +++ b/cognee/modules/data/methods/delete_data_by_user.py @@ -33,4 +33,4 @@ async def delete_data_by_user(user_id: UUID): datasets_query = select(Dataset.id).where(Dataset.owner_id == user_id) user_datasets_ids = (await session.execute(datasets_query)).scalars().all() if user_datasets_ids: - await db_engine.delete_entities_by_id(Dataset.__table__.name, user_datasets_ids) \ No newline at end of file + await db_engine.delete_entities_by_id(Dataset.__table__.name, user_datasets_ids) diff --git a/cognee/modules/data/methods/delete_dataset_by_name.py b/cognee/modules/data/methods/delete_dataset_by_name.py index ae8f7704f..7c5ed06b8 100644 --- a/cognee/modules/data/methods/delete_dataset_by_name.py +++ b/cognee/modules/data/methods/delete_dataset_by_name.py @@ -4,9 +4,7 @@ from cognee.infrastructure.databases.relational import get_relational_engine from ..models import Dataset -async def delete_dataset_by_name( - dataset_name: str, user_id: UUID -): +async def delete_dataset_by_name(dataset_name: str, user_id: UUID): """ Delete a single dataset by name for a specific user. @@ -25,6 +23,6 @@ async def delete_dataset_by_name( .filter(Dataset.name == dataset_name) ) ).first() - #Keeping this out of the first session, since delete_entities_by_id creates another session. + # Keeping this out of the first session, since delete_entities_by_id creates another session. if dataset_id: await db_engine.delete_entities_by_id(Dataset.__table__.name, dataset_id)