Lazy load chunks in inspector
This commit is contained in:
parent
448d00f9b2
commit
09be31f021
1 changed files with 47 additions and 2 deletions
|
|
@ -7,6 +7,8 @@ from textual.widgets import ListItem, ListView, Static
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.store.models import Chunk
|
from haiku.rag.store.models import Chunk
|
||||||
|
|
||||||
|
BATCH_SIZE = 50
|
||||||
|
|
||||||
|
|
||||||
class ChunkList(VerticalScroll): # pragma: no cover
|
class ChunkList(VerticalScroll): # pragma: no cover
|
||||||
"""Widget for displaying and browsing chunks."""
|
"""Widget for displaying and browsing chunks."""
|
||||||
|
|
@ -24,6 +26,11 @@ class ChunkList(VerticalScroll): # pragma: no cover
|
||||||
super().__init__(**kwargs)
|
super().__init__(**kwargs)
|
||||||
self.chunks: list[Chunk] = []
|
self.chunks: list[Chunk] = []
|
||||||
self.list_view = ListView()
|
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:
|
def compose(self) -> ComposeResult:
|
||||||
"""Compose the chunk list."""
|
"""Compose the chunk list."""
|
||||||
|
|
@ -33,8 +40,18 @@ class ChunkList(VerticalScroll): # pragma: no cover
|
||||||
async def load_chunks_for_document(
|
async def load_chunks_for_document(
|
||||||
self, client: HaikuRAG, document_id: str
|
self, client: HaikuRAG, document_id: str
|
||||||
) -> None:
|
) -> 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()
|
await self.list_view.clear()
|
||||||
for chunk in self.chunks:
|
for chunk in self.chunks:
|
||||||
first_line = chunk.content.split("\n")[0]
|
first_line = chunk.content.split("\n")[0]
|
||||||
|
|
@ -42,6 +59,31 @@ class ChunkList(VerticalScroll): # pragma: no cover
|
||||||
ListItem(Static(f"[{chunk.order}] {first_line}"))
|
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.Highlighted)
|
||||||
@on(ListView.Selected)
|
@on(ListView.Selected)
|
||||||
async def handle_chunk_selection(
|
async def handle_chunk_selection(
|
||||||
|
|
@ -53,3 +95,6 @@ class ChunkList(VerticalScroll): # pragma: no cover
|
||||||
idx = event.list_view.index
|
idx = event.list_view.index
|
||||||
if idx is not None and 0 <= idx < len(self.chunks):
|
if idx is not None and 0 <= idx < len(self.chunks):
|
||||||
self.post_message(self.ChunkSelected(self.chunks[idx]))
|
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()
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue