From c43c8fab224422815224786ee050da5563c17afc Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Dec 2025 14:43:45 +0200 Subject: [PATCH 1/4] Switch modals, do not stack them --- haiku_rag_slim/haiku/rag/inspector/app.py | 17 ++++++++++++++--- .../haiku/rag/inspector/widgets/search_modal.py | 3 ++- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/inspector/app.py b/haiku_rag_slim/haiku/rag/inspector/app.py index dbdb62b2..48c39390 100644 --- a/haiku_rag_slim/haiku/rag/inspector/app.py +++ b/haiku_rag_slim/haiku/rag/inspector/app.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: try: from textual.app import App from textual.binding import Binding + from textual.screen import Screen from textual.widgets import Footer, Header from haiku.rag.inspector.widgets.chunk_list import ChunkList @@ -110,10 +111,20 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover chunk_list.list_view.focus() break + async def _dismiss_modals(self) -> None: + """Dismiss all modal screens, returning to the main screen.""" + while len(self.screen_stack) > 1: + self.pop_screen() + + async def _switch_modal(self, screen: Screen) -> None: + """Switch to a new modal, dismissing any existing modals first.""" + await self._dismiss_modals() + await self.push_screen(screen) + async def action_search(self) -> None: """Open search modal.""" if self.client: - await self.push_screen(SearchModal(self.client)) + await self._switch_modal(SearchModal(self.client)) async def on_search_modal_chunk_selected( self, message: SearchModal.ChunkSelected @@ -191,7 +202,7 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal - await self.push_screen(VisualGroundingModal(chunk=chunk, client=self.client)) + await self._switch_modal(VisualGroundingModal(chunk=chunk, client=self.client)) async def action_show_context(self) -> None: """Show how the currently selected chunk would be formatted for agents.""" @@ -207,7 +218,7 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover from haiku.rag.inspector.widgets.context_modal import ContextModal - await self.push_screen(ContextModal(chunk=chunk, client=self.client)) + await self._switch_modal(ContextModal(chunk=chunk, client=self.client)) def run_inspector(db_path: Path | None = None) -> None: # pragma: no cover diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py index fc4c90a8..a431c113 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py @@ -173,7 +173,8 @@ class SearchModal(Screen): # pragma: no cover from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal - await self.app.push_screen( + # Use app's _switch_modal to close this modal before opening visual + await self.app._switch_modal( # type: ignore[attr-defined] VisualGroundingModal( chunk=chunk, client=self.client, From 448d00f9b2bde6f0846741491f4fb820ef7dbd56 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Dec 2025 15:09:14 +0200 Subject: [PATCH 2/4] Add pagination to chunk repo's get_by_document_id --- .../haiku/rag/store/repositories/chunk.py | 41 +++++++++++++--- tests/test_chunk.py | 49 +++++++++++++++++++ 2 files changed, 83 insertions(+), 7 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 1d9c2181..1577c6e0 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -281,13 +281,30 @@ class ChunkRepository: results = results.limit(limit) return await self._process_search_results(results) - async def get_by_document_id(self, document_id: str) -> list[Chunk]: - """Get all chunks for a specific document.""" - results = list( - self.store.chunks_table.search() - .where(f"document_id = '{document_id}'") - .to_pydantic(self.store.ChunkRecord) - ) + async def get_by_document_id( + self, + document_id: str, + limit: int | None = None, + offset: int | None = None, + ) -> list[Chunk]: + """Get chunks for a specific document with optional pagination. + + Args: + document_id: The document ID to get chunks for. + limit: Maximum number of chunks to return. None for all. + offset: Number of chunks to skip. None for no offset. + + Returns: + List of chunks ordered by their order field. + """ + query = self.store.chunks_table.search().where(f"document_id = '{document_id}'") + + if offset is not None: + query = query.offset(offset) + if limit is not None: + query = query.limit(limit) + + results = list(query.to_pydantic(self.store.ChunkRecord)) # Get document info doc_results = list( @@ -320,6 +337,16 @@ class ChunkRepository: chunks.sort(key=lambda c: c.order) return chunks + async def count_by_document_id(self, document_id: str) -> int: + """Count the number of chunks for a specific document.""" + df = ( + self.store.chunks_table.search() + .select(["id"]) + .where(f"document_id = '{document_id}'") + .to_pandas() + ) + return len(df) + async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]: """Get adjacent chunks before and after the given chunk within the same document.""" assert chunk.document_id, "Document id is required for adjacent chunk finding" diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 4c795440..b20cbced 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -47,6 +47,55 @@ async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path): client.close() +@pytest.mark.asyncio +async def test_chunk_repository_pagination(qa_corpus: Dataset, temp_db_path): + """Test ChunkRepository pagination with get_by_document_id and count_by_document_id.""" + async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client: + # Get the first document from the corpus (should produce multiple chunks) + first_doc = qa_corpus[0] + document_text = first_doc["document_extracted"] + + # Create a document with chunks + created_document = await client.create_document( + content=document_text, metadata={"source": "test"} + ) + assert created_document.id is not None + + # Get total chunk count + total_count = await client.chunk_repository.count_by_document_id( + created_document.id + ) + assert total_count > 0 + + # Get all chunks without pagination + all_chunks = await client.chunk_repository.get_by_document_id( + created_document.id + ) + assert len(all_chunks) == total_count + + # Test pagination with limit + limit = min(2, total_count) + first_batch = await client.chunk_repository.get_by_document_id( + created_document.id, limit=limit + ) + assert len(first_batch) == limit + assert first_batch[0].id == all_chunks[0].id + + # Test pagination with offset + if total_count > limit: + second_batch = await client.chunk_repository.get_by_document_id( + created_document.id, limit=limit, offset=limit + ) + assert len(second_batch) <= limit + assert second_batch[0].id == all_chunks[limit].id + + # Test offset beyond available chunks + empty_batch = await client.chunk_repository.get_by_document_id( + created_document.id, limit=10, offset=total_count + 100 + ) + assert len(empty_batch) == 0 + + @pytest.mark.asyncio async def test_chunking_pipeline(qa_corpus: Dataset, temp_db_path): """Test document chunking using client primitives.""" From 09be31f021534409ca936c50a69d52ebec006370 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Dec 2025 15:09:55 +0200 Subject: [PATCH 3/4] Lazy load chunks in inspector --- .../haiku/rag/inspector/widgets/chunk_list.py | 49 ++++++++++++++++++- 1 file changed, 47 insertions(+), 2 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/chunk_list.py b/haiku_rag_slim/haiku/rag/inspector/widgets/chunk_list.py index 94a3d20b..d690ce5f 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/chunk_list.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/chunk_list.py @@ -7,6 +7,8 @@ from textual.widgets import ListItem, ListView, Static from haiku.rag.client import HaikuRAG from haiku.rag.store.models import Chunk +BATCH_SIZE = 50 + class ChunkList(VerticalScroll): # pragma: no cover """Widget for displaying and browsing chunks.""" @@ -24,6 +26,11 @@ class ChunkList(VerticalScroll): # pragma: no cover super().__init__(**kwargs) self.chunks: list[Chunk] = [] self.list_view = ListView() + self.has_more: bool = False + self._client: HaikuRAG | None = None + self._document_id: str | None = None + self._loading: bool = False + self._total_chunks: int = 0 def compose(self) -> ComposeResult: """Compose the chunk list.""" @@ -33,8 +40,18 @@ class ChunkList(VerticalScroll): # pragma: no cover async def load_chunks_for_document( self, client: HaikuRAG, document_id: str ) -> None: - """Load chunks for a specific document.""" - self.chunks = await client.chunk_repository.get_by_document_id(document_id) + """Load initial batch of chunks for a specific document.""" + self._client = client + self._document_id = document_id + self._total_chunks = await client.chunk_repository.count_by_document_id( + document_id + ) + + self.chunks = await client.chunk_repository.get_by_document_id( + document_id, limit=BATCH_SIZE, offset=0 + ) + self.has_more = len(self.chunks) < self._total_chunks + await self.list_view.clear() for chunk in self.chunks: first_line = chunk.content.split("\n")[0] @@ -42,6 +59,31 @@ class ChunkList(VerticalScroll): # pragma: no cover ListItem(Static(f"[{chunk.order}] {first_line}")) ) + async def load_more(self) -> None: + """Load the next batch of chunks.""" + if ( + not self.has_more + or self._loading + or not self._client + or not self._document_id + ): + return + + self._loading = True + offset = len(self.chunks) + new_chunks = await self._client.chunk_repository.get_by_document_id( + self._document_id, limit=BATCH_SIZE, offset=offset + ) + self.has_more = (offset + len(new_chunks)) < self._total_chunks + self.chunks.extend(new_chunks) + + for chunk in new_chunks: + first_line = chunk.content.split("\n")[0] + await self.list_view.append( + ListItem(Static(f"[{chunk.order}] {first_line}")) + ) + self._loading = False + @on(ListView.Highlighted) @on(ListView.Selected) async def handle_chunk_selection( @@ -53,3 +95,6 @@ class ChunkList(VerticalScroll): # pragma: no cover idx = event.list_view.index if idx is not None and 0 <= idx < len(self.chunks): self.post_message(self.ChunkSelected(self.chunks[idx])) + # Infinite scroll: load more when near the end + if self.has_more and idx >= len(self.chunks) - 10: + await self.load_more() From c0cbc22137236035d2f2936b29a56543cdf265d6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Dec 2025 15:30:31 +0200 Subject: [PATCH 4/4] Info modal --- CHANGELOG.md | 8 + haiku_rag_slim/haiku/rag/inspector/app.py | 8 + .../haiku/rag/inspector/widgets/info_modal.py | 209 ++++++++++++++++++ 3 files changed, 225 insertions(+) create mode 100644 haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e595d314..08b4fe39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,14 @@ - New `/api/documents` endpoint to list available documents - Frontend document selector component with search and multi-select - Demonstrates client-to-server state flow via AG-UI protocol +- **Inspector Info Modal**: New `i` keyboard shortcut opens a modal displaying database information + +### Changed + +- **Inspector Lazy Loading**: Chunks panel now loads chunks in batches of 50 with infinite scroll + - Fixes unresponsive UI when viewing documents with large numbers of chunks + - New `ChunkRepository.get_by_document_id()` pagination with `limit` and `offset` parameters + - New `ChunkRepository.count_by_document_id()` method ## [0.20.0] - 2025-12-10 diff --git a/haiku_rag_slim/haiku/rag/inspector/app.py b/haiku_rag_slim/haiku/rag/inspector/app.py index 48c39390..d3a9fa94 100644 --- a/haiku_rag_slim/haiku/rag/inspector/app.py +++ b/haiku_rag_slim/haiku/rag/inspector/app.py @@ -69,6 +69,7 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover BINDINGS = [ Binding("q", "quit", "Quit", show=True), Binding("/", "search", "Search", show=True), + Binding("i", "show_info", "Info", show=True), Binding("v", "show_visual", "Visual", show=True), Binding("c", "show_context", "Context", show=True), ] @@ -126,6 +127,13 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover if self.client: await self._switch_modal(SearchModal(self.client)) + async def action_show_info(self) -> None: + """Show database info modal.""" + if self.client: + from haiku.rag.inspector.widgets.info_modal import InfoModal + + await self._switch_modal(InfoModal(self.client, self.db_path)) + async def on_search_modal_chunk_selected( self, message: SearchModal.ChunkSelected ) -> None: diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py new file mode 100644 index 00000000..91715f2f --- /dev/null +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py @@ -0,0 +1,209 @@ +import json +from importlib.metadata import version as pkg_version +from pathlib import Path +from typing import TYPE_CHECKING + +from textual.app import ComposeResult +from textual.binding import Binding +from textual.containers import Vertical, VerticalScroll +from textual.screen import ModalScreen +from textual.widgets import Static + +from haiku.rag.utils import format_bytes + +if TYPE_CHECKING: + from haiku.rag.client import HaikuRAG + + +class InfoModal(ModalScreen): # pragma: no cover + """Modal screen for displaying database information.""" + + BINDINGS = [ + Binding("escape", "dismiss", "Close", show=True), + Binding("i", "dismiss", "Close", show=True), + ] + + CSS = """ + InfoModal { + align: center middle; + } + + #info-container { + width: 80; + height: auto; + max-height: 80%; + background: $surface; + border: solid $primary; + padding: 1 2; + } + + #info-header { + height: auto; + margin-bottom: 1; + } + + #info-content { + height: auto; + max-height: 100%; + } + """ + + def __init__(self, client: "HaikuRAG", db_path: Path): + super().__init__() + self.client = client + self.db_path = db_path + self._content_widget = Static("Loading...") + + def compose(self) -> ComposeResult: + with Vertical(id="info-container"): + yield Static("[bold]Database Info[/bold]", id="info-header") + with VerticalScroll(id="info-content"): + yield self._content_widget + + async def on_mount(self) -> None: + """Load and display database info.""" + import lancedb + + lines: list[str] = [] + + # Path + lines.append(f"[bold $accent]path[/bold $accent]: {self.db_path}") + + if not self.db_path.exists(): + lines.append("[red]Database path does not exist.[/red]") + self._content_widget.update("\n".join(lines)) + return + + # Connect to get table info + try: + db = lancedb.connect(self.db_path) + table_names = set(db.table_names()) + except Exception as e: + lines.append(f"[red]Failed to open database: {e}[/red]") + self._content_widget.update("\n".join(lines)) + return + + # Get versions + try: + ldb_version = pkg_version("lancedb") + except Exception: + ldb_version = "unknown" + try: + hr_version = pkg_version("haiku.rag-slim") + except Exception: + hr_version = "unknown" + try: + docling_version = pkg_version("docling") + except Exception: + docling_version = "unknown" + + # Get stats from store + table_stats = self.client.store.get_stats() + + # Read settings + stored_version = "unknown" + embed_provider: str | None = None + embed_model: str | None = None + vector_dim: int | None = None + + if "settings" in table_names: + settings_tbl = db.open_table("settings") + arrow = settings_tbl.search().where("id = 'settings'").limit(1).to_arrow() + rows = arrow.to_pylist() if arrow is not None else [] + if rows: + raw = rows[0].get("settings") or "{}" + data = json.loads(raw) if isinstance(raw, str) else (raw or {}) + stored_version = str(data.get("version", stored_version)) + embeddings = data.get("embeddings", {}) + embed_model_obj = embeddings.get("model", {}) + embed_provider = embed_model_obj.get("provider") + embed_model = embed_model_obj.get("name") + vector_dim = embed_model_obj.get("vector_dim") + + num_docs = table_stats["documents"].get("num_rows", 0) + doc_bytes = table_stats["documents"].get("total_bytes", 0) + + num_chunks = table_stats["chunks"].get("num_rows", 0) + chunk_bytes = table_stats["chunks"].get("total_bytes", 0) + + has_vector_index = table_stats["chunks"].get("has_vector_index", False) + num_indexed_rows = table_stats["chunks"].get("num_indexed_rows", 0) + num_unindexed_rows = table_stats["chunks"].get("num_unindexed_rows", 0) + + # Table versions + doc_versions = ( + len(list(db.open_table("documents").list_versions())) + if "documents" in table_names + else 0 + ) + chunk_versions = ( + len(list(db.open_table("chunks").list_versions())) + if "chunks" in table_names + else 0 + ) + + # Build output + lines.append( + f"[bold $accent]haiku.rag version (db)[/bold $accent]: {stored_version}" + ) + + if embed_provider or embed_model or vector_dim: + provider_part = embed_provider or "unknown" + model_part = embed_model or "unknown" + dim_part = f"{vector_dim}" if vector_dim is not None else "unknown" + lines.append( + f"[bold $accent]embeddings[/bold $accent]: " + f"{provider_part}/{model_part} (dim: {dim_part})" + ) + else: + lines.append("[bold $accent]embeddings[/bold $accent]: unknown") + + lines.append( + f"[bold $accent]documents[/bold $accent]: {num_docs} ({format_bytes(doc_bytes)})" + ) + lines.append( + f"[bold $accent]chunks[/bold $accent]: {num_chunks} ({format_bytes(chunk_bytes)})" + ) + + # Vector index info + if has_vector_index: + lines.append("[bold $accent]vector index[/bold $accent]: ✓ exists") + lines.append( + f"[bold $accent]indexed chunks[/bold $accent]: {num_indexed_rows}" + ) + if num_unindexed_rows > 0: + lines.append( + f"[bold $accent]unindexed chunks[/bold $accent]: [yellow]{num_unindexed_rows}[/yellow]" + ) + else: + lines.append( + f"[bold $accent]unindexed chunks[/bold $accent]: {num_unindexed_rows}" + ) + else: + if num_chunks >= 256: + lines.append( + "[bold $accent]vector index[/bold $accent]: [yellow]✗ not created[/yellow]" + ) + else: + lines.append( + f"[bold $accent]vector index[/bold $accent]: ✗ not created " + f"(need {256 - num_chunks} more chunks)" + ) + + lines.append( + f"[bold $accent]versions (documents)[/bold $accent]: {doc_versions}" + ) + lines.append( + f"[bold $accent]versions (chunks)[/bold $accent]: {chunk_versions}" + ) + + lines.append("") + lines.append("[bold]Versions[/bold]") + lines.append(f"[bold $accent]haiku.rag[/bold $accent]: {hr_version}") + lines.append(f"[bold $accent]lancedb[/bold $accent]: {ldb_version}") + lines.append(f"[bold $accent]docling[/bold $accent]: {docling_version}") + + self._content_widget.update("\n".join(lines)) + + async def action_dismiss(self, result=None) -> None: + self.app.pop_screen()