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 index 3c017b14..597dd1ac 100644 --- a/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py +++ b/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py @@ -6,6 +6,32 @@ from textual.screen import ModalScreen from textual.widgets import Button, Checkbox, Input, Static from haiku.rag.client import HaikuRAG +from haiku.rag.utils import escape_sql_string + +# One page of documents. A corpus is not a list to scroll: mounting a checkbox +# per document wedges the modal at tens of thousands, so the list is a page and +# the search box asks the database for the rest. +DOCUMENT_PAGE = 200 + + +def search_filter(term: str) -> str | None: + """A document filter matching `term` in a title or URI, or None for no term. + + The term is text to find, not a pattern: `LIKE` is case-sensitive and reads + `%` and `_` as wildcards, so both sides are lowered and the term's own + wildcards are escaped. + """ + term = term.strip() + if not term: + return None + literal = term.lower() + for wildcard in ("\\", "%", "_"): + literal = literal.replace(wildcard, f"\\{wildcard}") + escaped = escape_sql_string(f"%{literal}%") + return ( + f"LOWER(title) LIKE '{escaped}' ESCAPE '\\' " + f"OR LOWER(uri) LIKE '{escaped}' ESCAPE '\\'" + ) class DocumentFilterModal(ModalScreen): @@ -88,6 +114,8 @@ class DocumentFilterModal(ModalScreen): self.client = client self.initial_selected = selected or [] self._selected: set[str] = set(self.initial_selected) + self._shown = 0 + self._matching = 0 def compose(self) -> ComposeResult: with Vertical(id="filter-container"): @@ -101,38 +129,59 @@ class DocumentFilterModal(ModalScreen): yield Button("Apply", id="apply-btn", variant="primary") async def on_mount(self) -> None: - """Load documents when mounted.""" + """Load the first page of 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() + async def _load_documents(self, search: str = "") -> None: + """Show one page of documents, narrowed by `search` when given.""" + document_filter = search_filter(search) + docs = await self.client.list_documents( + limit=DOCUMENT_PAGE, filter=document_filter + ) + self._matching = await self.client.count_documents(filter=document_filter) loading = self.query_one("#loading-indicator", Static) - loading.remove() + if loading.parent is not None: + loading.remove() filter_list = self.query_one("#filter-list", VerticalScroll) - for doc in docs: - display_name = doc.title or doc.uri or str(doc.id) + await filter_list.remove_children() + + # The page is picked to represent every database; sorting is so it reads + # like a list rather than in whatever order the tables returned. + names = sorted(doc.title or doc.uri or str(doc.id) for doc in docs) + + boxes = [] + for position, display_name in enumerate(names): checkbox = Checkbox( display_name, value=display_name in self._selected, - id=f"doc-{hash(display_name)}", + # Positional, because titles repeat and a repeated id is an error. + id=f"doc-{position}", classes="doc-checkbox", ) checkbox._doc_id = display_name # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] - await filter_list.mount(checkbox) + boxes.append(checkbox) + if boxes: + await filter_list.mount_all(boxes) + self._shown = len(boxes) self._update_footer() def _update_footer(self) -> None: - """Update the footer with selection count.""" + """Update the footer with the selection and how much of the corpus is shown.""" footer = self.query_one("#filter-footer", Static) count = len(self._selected) if count == 0: - footer.update("[dim]No filter (all documents)[/dim]") + state = "[dim]No filter (all documents)[/dim]" else: - footer.update(f"[bold]{count}[/bold] document(s) selected") + state = f"[bold]{count}[/bold] document(s) selected" + if self._matching > self._shown: + state += ( + f" [dim]— showing {self._shown} of {self._matching};" + " type and press enter to search[/dim]" + ) + footer.update(state) def on_checkbox_changed(self, event: Checkbox.Changed) -> None: """Handle checkbox state changes.""" @@ -148,8 +197,12 @@ class DocumentFilterModal(ModalScreen): self._update_footer() + async def on_input_submitted(self, event: Input.Submitted) -> None: + """Ask the database for documents matching the search.""" + await self._load_documents(event.value) + def on_input_changed(self, event: Input.Changed) -> None: - """Filter document list based on search input.""" + """Narrow the page already shown, for feedback while typing.""" search_term = event.value.lower().strip() filter_list = self.query_one("#filter-list", VerticalScroll) diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 68765514..85c3f90f 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -7,6 +7,7 @@ import tempfile from collections.abc import AsyncGenerator, Sequence from enum import Enum from functools import cached_property +from itertools import zip_longest from pathlib import Path from time import monotonic from typing import TYPE_CHECKING, overload @@ -651,7 +652,11 @@ class HaikuRAG: for owner in await self.clients_covering() ) ) - merged = [doc for group in groups for doc in group] + # Round-robin, so a window shows every database. Concatenating lets + # the first one fill the whole page and hide the rest. + merged = [ + doc for row in zip_longest(*groups) for doc in row if doc is not None + ] start = offset or 0 return merged[start:] if limit is None else merged[start : start + limit] return await self.document_repository.list_all( diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index ea3f9b33..02123d04 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -539,3 +539,52 @@ async def test_visual_grounding_uses_the_database_holding_the_citation(tmp_path) owner.get_chunk_by_id.assert_awaited_once_with("c1") assert push.await_args is not None assert push.await_args.args[0].client is owner + + +class TestDocumentSearchFilter: + """The filter modal shows one page and asks the database for the rest, so the + typed term reaches SQL.""" + + def test_no_term_means_no_filter(self): + from haiku.rag.chat.widgets.document_filter_modal import search_filter + + assert search_filter(" ") is None + + def test_a_term_matches_titles_and_uris(self): + from haiku.rag.chat.widgets.document_filter_modal import search_filter + + built = search_filter("Nobel") + + assert built == ( + "LOWER(title) LIKE '%nobel%' ESCAPE '\\' " + "OR LOWER(uri) LIKE '%nobel%' ESCAPE '\\'" + ) + + def test_the_search_is_case_insensitive(self): + """LIKE is case-sensitive, so both sides are lowered.""" + from haiku.rag.chat.widgets.document_filter_modal import search_filter + + built = search_filter("NoBeL") + + assert built is not None + assert "LOWER(title) LIKE '%nobel%'" in built + assert "LOWER(uri) LIKE '%nobel%'" in built + + def test_like_wildcards_in_the_term_are_literal(self): + """A term is text to find, not a pattern: `_` matches any character.""" + from haiku.rag.chat.widgets.document_filter_modal import search_filter + + built = search_filter("100%_raw") + + assert built is not None + assert "100\\%\\_raw" in built + assert "ESCAPE '\\'" in built + + def test_a_quote_in_the_term_is_escaped(self): + """Whatever was typed is data, not syntax.""" + from haiku.rag.chat.widgets.document_filter_modal import search_filter + + built = search_filter("O'Brien") + + assert built is not None + assert "o''brien" in built diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index ca3d2aaa..d245214d 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -197,6 +197,20 @@ class TestListingAcrossDatabases: assert len(await rag.list_documents(limit=2, offset=2)) == 2 assert len(await rag.list_documents(offset=3)) == 1 + @pytest.mark.asyncio + async def test_a_page_shows_every_database(self, tmp_path): + """A window is taken across the databases, not filled from the first one: + concatenating hides every database after whichever was listed first.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", [f"alpha {i}" for i in range(5)]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + page = await rag.list_documents(limit=3) + + assert len(page) == 3 + assert {(d.uri or "").split("/")[2] for d in page} == {"alpha", "beta"} + @pytest.mark.asyncio async def test_a_filter_reaches_every_database(self, tmp_path): config = _config(tmp_path, ["alpha", "beta"])