163 lines
6 KiB
Python
163 lines
6 KiB
Python
import asyncio
|
|
from textual.app import ComposeResult
|
|
from textual.widgets import Input, Label, Button, 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
|
|
|
|
|
|
class CognifyTUIScreen(BaseTUIScreen):
|
|
"""TUI screen for cognifying data in cognee."""
|
|
|
|
BINDINGS = [
|
|
Binding("q", "quit_app", "Quit"),
|
|
Binding("escape", "back", "Back"),
|
|
Binding("ctrl+s", "submit", "Submit"),
|
|
]
|
|
|
|
CSS = BaseTUIScreen.CSS + """
|
|
Checkbox {
|
|
margin-top: 1;
|
|
margin-bottom: 1;
|
|
}
|
|
|
|
RadioSet {
|
|
margin-top: 0;
|
|
margin-bottom: 1;
|
|
height: auto;
|
|
}
|
|
|
|
RadioButton {
|
|
height: 1;
|
|
}
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.is_processing = False
|
|
|
|
def compose_content(self) -> ComposeResult:
|
|
with Container(classes="tui-content-container"):
|
|
yield Label("⚡ Cognify Data", classes="tui-title")
|
|
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("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"
|
|
)
|
|
|
|
def on_mount(self) -> None:
|
|
"""Focus the dataset input on mount."""
|
|
dataset_input = self.query_one("#dataset-input", Input)
|
|
dataset_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_submit(self) -> None:
|
|
"""Submit the form."""
|
|
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)
|
|
chunker_radio = self.query_one("#chunker-radio", RadioSet)
|
|
background_checkbox = self.query_one("#background-checkbox", Checkbox)
|
|
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"
|
|
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
|
|
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))
|
|
|
|
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)
|
|
|
|
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
|
|
chunker_class = LangchainChunker
|
|
except ImportError:
|
|
status.update("[yellow]⚠ LangchainChunker not available, using TextChunker[/yellow]")
|
|
elif chunker_type == "CsvChunker":
|
|
try:
|
|
from cognee.modules.chunking.CsvChunker import CsvChunker
|
|
chunker_class = CsvChunker
|
|
except ImportError:
|
|
status.update("[yellow]⚠ CsvChunker not available, using TextChunker[/yellow]")
|
|
|
|
# Prepare datasets parameter
|
|
datasets = [dataset_name] if dataset_name else None
|
|
|
|
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
|
|
dataset_input = self.query_one("#dataset-input", Input)
|
|
chunker_radio = self.query_one("#chunker-radio", RadioSet)
|
|
background_checkbox = self.query_one("#background-checkbox", Checkbox)
|
|
dataset_input.disabled = False
|
|
chunker_radio.disabled = False
|
|
background_checkbox.disabled = False
|
|
self.query_one("#submit-btn", Button).disabled = False
|
|
dataset_input.focus()
|