Merge pull request #190 from ggozad/feat/chunk-lazy-load
Lazy load chunks & info modal for inspector
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
ada24eaa11
7 changed files with 371 additions and 13 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
209
haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py
Normal file
209
haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Reference in a new issue