Document filter in the TUI, use command palette instead of shortcuts

This commit is contained in:
Yiorgis Gozadinos 2026-01-26 13:33:23 +02:00
parent 30e9ff3038
commit af786596b1
No known key found for this signature in database
2 changed files with 235 additions and 7 deletions

View file

@ -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

View file

@ -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()