diff --git a/cognee/cli/tui/add_screen.py b/cognee/cli/tui/add_screen.py
index 5c4f31eb8..d902703b5 100644
--- a/cognee/cli/tui/add_screen.py
+++ b/cognee/cli/tui/add_screen.py
@@ -1,7 +1,8 @@
 import asyncio
+import cognee
 from textual.app import ComposeResult
-from textual.widgets import Input, Label, Button, Static, TextArea
-from textual.containers import Container, Vertical, Horizontal
+from textual.widgets import Input, Label, Static, TextArea
+from textual.containers import Container, Vertical
 from textual.binding import Binding
 
 from cognee.cli.tui.base_screen import BaseTUIScreen
@@ -76,11 +77,6 @@ class AddTUIScreen(BaseTUIScreen):
         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)
@@ -100,7 +96,6 @@ class AddTUIScreen(BaseTUIScreen):
         # 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))
@@ -110,8 +105,6 @@ class AddTUIScreen(BaseTUIScreen):
         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]")
@@ -130,5 +123,4 @@ class AddTUIScreen(BaseTUIScreen):
             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()
diff --git a/cognee/cli/tui/cognify_screen.py b/cognee/cli/tui/cognify_screen.py
index cc5a7610c..844cb345b 100644
--- a/cognee/cli/tui/cognify_screen.py
+++ b/cognee/cli/tui/cognify_screen.py
@@ -1,12 +1,25 @@
 import asyncio
+
+import cognee
+from cognee.modules.chunking.TextChunker import TextChunker
 from textual.app import ComposeResult
-from textual.widgets import Input, Label, Button, Static, Checkbox, RadioSet, RadioButton
+from textual.widgets import Input, Label, Static, Checkbox, RadioSet, RadioButton
 from textual.containers import Container, Vertical
 from textual.binding import Binding
 
 from cognee.cli.tui.base_screen import BaseTUIScreen
 from cognee.cli.config import CHUNKER_CHOICES
 
+try:
+    from cognee.modules.chunking.LangchainChunker import LangchainChunker
+except ImportError:
+    LangchainChunker = None
+
+try:
+    from cognee.modules.chunking.CsvChunker import CsvChunker
+except ImportError:
+    CsvChunker = None
+
 
 class CognifyTUIScreen(BaseTUIScreen):
     """TUI screen for cognifying data in cognee."""
@@ -83,11 +96,6 @@ class CognifyTUIScreen(BaseTUIScreen):
         if not self.is_processing:
             self._submit_cognify()
 
-    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_cognify()
-
     def _submit_cognify(self) -> None:
         """Process and submit the cognify request."""
         dataset_input = self.query_one("#dataset-input", Input)
@@ -106,7 +114,6 @@ class CognifyTUIScreen(BaseTUIScreen):
         dataset_input.disabled = True
         chunker_radio.disabled = True
         background_checkbox.disabled = True
-        self.query_one("#submit-btn", Button).disabled = True
 
         # Run async cognify operation
         asyncio.create_task(self._cognify_async(dataset_name, chunker_type, run_background))
@@ -116,22 +123,17 @@ class CognifyTUIScreen(BaseTUIScreen):
         status = self.query_one(".tui-status", Static)
         
         try:
-            import cognee
-            from cognee.modules.chunking.TextChunker import TextChunker
-            
             # Get chunker class
             chunker_class = TextChunker
             if chunker_type == "LangchainChunker":
-                try:
-                    from cognee.modules.chunking.LangchainChunker import LangchainChunker
+                if LangchainChunker is not None:
                     chunker_class = LangchainChunker
-                except ImportError:
+                else:
                     status.update("[yellow]⚠ LangchainChunker not available, using TextChunker[/yellow]")
             elif chunker_type == "CsvChunker":
-                try:
-                    from cognee.modules.chunking.CsvChunker import CsvChunker
+                if CsvChunker is not None:
                     chunker_class = CsvChunker
-                except ImportError:
+                else:
                     status.update("[yellow]⚠ CsvChunker not available, using TextChunker[/yellow]")
             
             # Prepare datasets parameter
@@ -160,5 +162,4 @@ class CognifyTUIScreen(BaseTUIScreen):
             dataset_input.disabled = False
             chunker_radio.disabled = False
             background_checkbox.disabled = False
-            self.query_one("#submit-btn", Button).disabled = False
             dataset_input.focus()
diff --git a/cognee/cli/tui/search_screen.py b/cognee/cli/tui/search_screen.py
index fa23e084d..e04165b60 100644
--- a/cognee/cli/tui/search_screen.py
+++ b/cognee/cli/tui/search_screen.py
@@ -1,4 +1,7 @@
 import asyncio
+
+import cognee
+from cognee.modules.search.types import SearchType
 from textual.app import ComposeResult
 from textual.widgets import Input, Label, Static, Select
 from textual.containers import Container, Vertical, ScrollableContainer
@@ -136,9 +139,6 @@ class SearchTUIScreen(BaseTUIScreen):
     async def _async_search(self, query_text: str, query_type: str) -> None:
         """Async search operation."""
         try:
-            import cognee
-            from cognee.modules.search.types import SearchType
-
             # Convert string to SearchType enum
             search_type = SearchType[query_type]