Merge pull request #190 from ggozad/feat/chunk-lazy-load

Lazy load chunks & info modal for inspector

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
 

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
@@ -68,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),
     ]
@@ -110,10 +112,27 @@ 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 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
@@ -191,7 +210,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 +226,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

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."""
+        """Load initial batch of chunks for a specific document."""
-        self.chunks = await client.chunk_repository.get_by_document_id(document_id)
+        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()

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

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,

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]:
+    async def get_by_document_id(
-        """Get all chunks for a specific document."""
+        self,
-        results = list(
+        document_id: str,
-            self.store.chunks_table.search()
+        limit: int | None = None,
-            .where(f"document_id = '{document_id}'")
+        offset: int | None = None,
-            .to_pydantic(self.store.ChunkRecord)
+    ) -> 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"

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."""