From af786596b1e288100cd9f5c6e8cb2dc839ec49a4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 26 Jan 2026 13:33:23 +0200 Subject: [PATCH] Document filter in the TUI, use command palette instead of shortcuts --- haiku_rag_slim/haiku/rag/chat/app.py | 63 +++++- .../rag/chat/widgets/document_filter_modal.py | 179 ++++++++++++++++++ 2 files changed, 235 insertions(+), 7 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 7091ead8..463be357 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -1,7 +1,7 @@ # pyright: reportPossiblyUnboundVariable=false import asyncio import uuid -from collections.abc import AsyncIterable +from collections.abc import AsyncIterable, Iterable from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Any @@ -39,7 +39,7 @@ except ImportError: try: import textual_image.widget # noqa: F401 - import early for renderer detection - from textual.app import App + from textual.app import App, SystemCommand from textual.binding import Binding from textual.widgets import Footer, Header, Input from textual.worker import Worker @@ -50,6 +50,7 @@ try: except ImportError: TEXTUAL_AVAILABLE = False App = object # type: ignore + SystemCommand = object # type: ignore class ChatApp(App): @@ -79,10 +80,6 @@ class ChatApp(App): """ BINDINGS = [ - Binding("ctrl+l", "clear_chat", "Clear", show=True), - Binding("ctrl+g", "show_visual", "Visual", show=True), - Binding("ctrl+i", "show_info", "Info", show=True), - Binding("ctrl+o", "show_context", "Context", show=True), Binding("escape", "focus_input", "Focus Input", show=False), ] @@ -105,6 +102,7 @@ class ChatApp(App): self._last_citations: list[Citation] = [] self._current_worker: Worker[None] | None = None self._message_history: list[ModelMessage] = [] + self._document_filter: list[str] = [] def compose(self) -> "ComposeResult": """Compose the UI layout.""" @@ -113,6 +111,35 @@ class ChatApp(App): yield Input(placeholder="Ask a question...", id="chat-input") yield Footer() + def get_system_commands(self, screen: Any) -> Iterable[SystemCommand]: + """Add commands to the command palette.""" + yield from super().get_system_commands(screen) + yield SystemCommand( + "Clear chat", + "Clear the chat history and reset session", + self.action_clear_chat, + ) + yield SystemCommand( + "Filter documents", + "Select documents to filter searches", + self.action_show_filter, + ) + yield SystemCommand( + "Show visual grounding", + "Show visual grounding for selected citation", + self.action_show_visual, + ) + yield SystemCommand( + "Database info", + "Show database information", + self.action_show_info, + ) + yield SystemCommand( + "Session context", + "Show current session context", + self.action_show_context, + ) + async def on_mount(self) -> None: """Initialize the app when mounted.""" self.client = HaikuRAG( @@ -127,6 +154,7 @@ class ChatApp(App): self.agent = create_chat_agent(self.config) self.session_state = ChatSessionState( session_id=str(uuid.uuid4()), + document_filter=self._document_filter, ) # Focus the input field @@ -270,9 +298,10 @@ class ChatApp(App): await chat_history.clear_messages() self._last_citations.clear() self._message_history.clear() - # Reset session state for fresh conversation + # Reset session state for fresh conversation (preserve document filter) self.session_state = ChatSessionState( session_id=str(uuid.uuid4()), + document_filter=self._document_filter, ) def action_focus_input(self) -> None: @@ -339,3 +368,23 @@ class ChatApp(App): citation_widgets = list(chat_history.query(CitationWidget)) if 0 <= event.citation_index < len(citation_widgets): citation_widgets[event.citation_index].add_class("selected") + + async def action_show_filter(self) -> None: + """Show document filter modal.""" + if not self.client: + return + + from haiku.rag.chat.widgets.document_filter_modal import DocumentFilterModal + + await self.push_screen( + DocumentFilterModal( + client=self.client, + selected=self._document_filter, + ) + ) + + def on_document_filter_modal_filter_changed(self, event: Any) -> None: + """Handle document filter changes from modal.""" + self._document_filter = event.selected + if self.session_state: + self.session_state.document_filter = self._document_filter diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py b/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py new file mode 100644 index 00000000..6da412bf --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py @@ -0,0 +1,179 @@ +from textual.app import ComposeResult +from textual.binding import Binding +from textual.containers import Horizontal, Vertical, VerticalScroll +from textual.message import Message +from textual.screen import ModalScreen +from textual.widgets import Button, Checkbox, Input, Static + +from haiku.rag.client import HaikuRAG + + +class DocumentFilterModal(ModalScreen): # pragma: no cover + """Modal screen for selecting documents to filter searches.""" + + BINDINGS = [ + Binding("escape", "cancel", "Cancel", show=False), + ] + + CSS = """ + DocumentFilterModal { + align: center middle; + background: rgba(0, 0, 0, 0.5); + } + + #filter-container { + width: 60; + height: auto; + max-height: 28; + background: $surface; + border: tall $primary; + padding: 1 2; + } + + #filter-header { + height: auto; + margin-bottom: 1; + } + + #filter-search { + margin-bottom: 1; + } + + #filter-list { + height: 1fr; + min-height: 8; + max-height: 16; + scrollbar-gutter: stable; + } + + #filter-footer { + height: auto; + margin-top: 1; + color: $text-muted; + } + + #button-row { + height: auto; + margin-top: 1; + align: right middle; + } + + #button-row Button { + margin-left: 1; + } + + .doc-checkbox { + height: auto; + padding: 0 1; + } + + .doc-checkbox:hover { + background: $surface-lighten-1; + } + """ + + class FilterChanged(Message): + """Emitted when the document filter selection changes.""" + + def __init__(self, selected: list[str]) -> None: + super().__init__() + self.selected = selected + + def __init__( + self, + client: HaikuRAG, + selected: list[str] | None = None, + ) -> None: + super().__init__() + self.client = client + self.initial_selected = selected or [] + self._selected: set[str] = set(self.initial_selected) + + def compose(self) -> ComposeResult: + with Vertical(id="filter-container"): + yield Static("[bold]Filter Documents[/bold]", id="filter-header") + yield Input(placeholder="Search documents...", id="filter-search") + with VerticalScroll(id="filter-list"): + yield Static("Loading...", id="loading-indicator") + yield Static("", id="filter-footer") + with Horizontal(id="button-row"): + yield Button("Cancel", id="cancel-btn", variant="default") + yield Button("Apply", id="apply-btn", variant="primary") + + async def on_mount(self) -> None: + """Load documents when mounted.""" + await self._load_documents() + + async def _load_documents(self) -> None: + """Load all documents from the client.""" + docs = await self.client.list_documents() + + # Remove loading indicator + loading = self.query_one("#loading-indicator", Static) + loading.remove() + + # Add checkboxes for all documents + filter_list = self.query_one("#filter-list", VerticalScroll) + for doc in docs: + display_name = doc.title or doc.uri or str(doc.id) + checkbox = Checkbox( + display_name, + value=display_name in self._selected, + id=f"doc-{hash(display_name)}", + classes="doc-checkbox", + ) + checkbox._doc_id = display_name # type: ignore[attr-defined] + await filter_list.mount(checkbox) + + self._update_footer() + + def _update_footer(self) -> None: + """Update the footer with selection count.""" + footer = self.query_one("#filter-footer", Static) + count = len(self._selected) + if count == 0: + footer.update("[dim]No filter (all documents)[/dim]") + else: + footer.update(f"[bold]{count}[/bold] document(s) selected") + + def on_checkbox_changed(self, event: Checkbox.Changed) -> None: + """Handle checkbox state changes.""" + checkbox = event.checkbox + doc_id = getattr(checkbox, "_doc_id", None) + if doc_id is None: + return + + if event.value: + self._selected.add(doc_id) + else: + self._selected.discard(doc_id) + + self._update_footer() + + def on_input_changed(self, event: Input.Changed) -> None: + """Filter document list based on search input.""" + search_term = event.value.lower().strip() + filter_list = self.query_one("#filter-list", VerticalScroll) + + for checkbox in filter_list.query(Checkbox): + doc_id = getattr(checkbox, "_doc_id", "") + if search_term == "" or search_term in doc_id.lower(): + checkbox.display = True + else: + checkbox.display = False + + def on_button_pressed(self, event: Button.Pressed) -> None: + """Handle button presses.""" + if event.button.id == "cancel-btn": + self.action_cancel() + elif event.button.id == "apply-btn": + self.action_confirm() + + def action_cancel(self) -> None: + """Cancel and close without saving.""" + self.app.pop_screen() + + def action_confirm(self) -> None: + """Confirm selection and close.""" + self.post_message(self.FilterChanged(list(self._selected))) + self.app.pop_screen()