298 lines
10 KiB
Python
298 lines
10 KiB
Python
# pyright: reportPossiblyUnboundVariable=false
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import get_config
|
|
|
|
if TYPE_CHECKING:
|
|
from textual.app import ComposeResult
|
|
|
|
try:
|
|
from textual.app import App
|
|
from textual.binding import Binding
|
|
from textual.containers import Container
|
|
from textual.widgets import Footer, Header, Input, ListItem, Static
|
|
|
|
from haiku.rag.inspector.widgets.chunk_list import ChunkList
|
|
from haiku.rag.inspector.widgets.detail_view import DetailView
|
|
from haiku.rag.inspector.widgets.document_list import DocumentList
|
|
|
|
TEXTUAL_AVAILABLE = True
|
|
except ImportError:
|
|
TEXTUAL_AVAILABLE = False
|
|
App = object # type: ignore
|
|
|
|
|
|
class InspectorApp(App): # type: ignore[misc]
|
|
"""Textual TUI for inspecting LanceDB data."""
|
|
|
|
TITLE = "haiku.rag DB Inspector"
|
|
|
|
CSS = """
|
|
Screen {
|
|
layout: grid;
|
|
grid-size: 2 2;
|
|
grid-columns: 1fr 2fr;
|
|
grid-rows: 1fr 1fr;
|
|
}
|
|
|
|
#document-list {
|
|
column-span: 1;
|
|
row-span: 2;
|
|
border: solid $primary;
|
|
}
|
|
|
|
#chunk-list {
|
|
column-span: 1;
|
|
row-span: 1;
|
|
border: solid $secondary;
|
|
}
|
|
|
|
#detail-view {
|
|
column-span: 1;
|
|
row-span: 1;
|
|
border: solid $accent;
|
|
}
|
|
|
|
ListItem {
|
|
overflow: hidden;
|
|
}
|
|
|
|
ListItem Static {
|
|
overflow: hidden;
|
|
text-overflow: ellipsis;
|
|
}
|
|
|
|
#search-container {
|
|
dock: top;
|
|
height: 3;
|
|
background: $panel;
|
|
display: none;
|
|
}
|
|
|
|
#search-container.visible {
|
|
display: block;
|
|
}
|
|
|
|
#search-input {
|
|
width: 1fr;
|
|
}
|
|
"""
|
|
|
|
BINDINGS = [
|
|
Binding("q", "quit", "Quit", show=True),
|
|
Binding("/", "search", "Search", show=True),
|
|
]
|
|
|
|
def __init__(self, db_path: Path):
|
|
super().__init__()
|
|
self.db_path = db_path
|
|
self.client: HaikuRAG | None = None
|
|
self.search_visible = False
|
|
self.search_active = False
|
|
# Track current selection for easy restoration
|
|
self.current_document_id: str | None = None
|
|
self.current_chunk_id: str | None = None
|
|
|
|
def compose(self) -> "ComposeResult":
|
|
"""Compose the UI layout."""
|
|
yield Header()
|
|
with Container(id="search-container"):
|
|
yield Input(placeholder="Search chunks...", id="search-input")
|
|
yield DocumentList(id="document-list")
|
|
yield ChunkList(id="chunk-list")
|
|
yield DetailView(id="detail-view")
|
|
yield Footer()
|
|
|
|
async def on_mount(self) -> None:
|
|
"""Initialize the app when mounted."""
|
|
config = get_config()
|
|
self.client = HaikuRAG(db_path=self.db_path, config=config, allow_create=False)
|
|
await self.client.__aenter__()
|
|
|
|
# Load initial documents
|
|
doc_list = self.query_one(DocumentList)
|
|
await doc_list.load_documents(self.client)
|
|
|
|
# Select first document and load its chunks
|
|
if doc_list.documents and doc_list.list_view:
|
|
doc_list.list_view.index = 0
|
|
first_doc = doc_list.documents[0]
|
|
self.current_document_id = first_doc.id
|
|
|
|
if first_doc.id:
|
|
chunk_list = self.query_one(ChunkList)
|
|
await chunk_list.load_chunks_for_document(self.client, first_doc.id)
|
|
|
|
# Select first chunk
|
|
if chunk_list.chunks and chunk_list.list_view:
|
|
chunk_list.list_view.index = 0
|
|
self.current_chunk_id = chunk_list.chunks[0].id
|
|
|
|
# Focus the document list view
|
|
if doc_list.list_view:
|
|
doc_list.list_view.focus()
|
|
|
|
async def on_unmount(self) -> None:
|
|
"""Clean up when unmounting."""
|
|
if self.client:
|
|
await self.client.__aexit__(None, None, None)
|
|
|
|
async def action_search(self) -> None:
|
|
"""Toggle search input visibility."""
|
|
search_container = self.query_one("#search-container", Container)
|
|
search_input = self.query_one("#search-input", Input)
|
|
|
|
if self.search_visible:
|
|
search_container.remove_class("visible")
|
|
search_input.value = ""
|
|
self.search_visible = False
|
|
self.search_active = False
|
|
|
|
# Restore full document list and reload chunks for selected document
|
|
if self.client:
|
|
doc_list = self.query_one(DocumentList)
|
|
chunk_list = self.query_one(ChunkList)
|
|
|
|
# Reload all documents
|
|
await doc_list.load_documents(self.client)
|
|
|
|
# Restore document and chunk selection
|
|
if (
|
|
self.current_document_id
|
|
and doc_list.list_view
|
|
and doc_list.documents
|
|
):
|
|
# Find and select the document
|
|
doc_found = False
|
|
for idx, doc in enumerate(doc_list.documents):
|
|
if doc.id == self.current_document_id:
|
|
doc_list.list_view.index = idx
|
|
doc_found = True
|
|
break
|
|
|
|
# Reload chunks for this document (without scores)
|
|
if doc_found:
|
|
await chunk_list.load_chunks_for_document(
|
|
self.client, self.current_document_id
|
|
)
|
|
|
|
# Restore chunk selection and show it in detail view
|
|
if (
|
|
self.current_chunk_id
|
|
and chunk_list.list_view
|
|
and chunk_list.chunks
|
|
):
|
|
for idx, chunk in enumerate(chunk_list.chunks):
|
|
if chunk.id == self.current_chunk_id:
|
|
chunk_list.list_view.index = idx
|
|
# Show the chunk in detail view
|
|
detail_view = self.query_one(DetailView)
|
|
await detail_view.show_chunk(chunk)
|
|
break
|
|
|
|
# Focus back to document list to show selection
|
|
doc_list.list_view.focus()
|
|
else:
|
|
search_container.add_class("visible")
|
|
search_input.focus()
|
|
self.search_visible = True
|
|
|
|
async def on_input_submitted(self, event: Input.Submitted) -> None:
|
|
"""Handle search input submission."""
|
|
if event.input.id == "search-input" and self.client:
|
|
query = event.value.strip()
|
|
if query:
|
|
self.search_active = True
|
|
chunk_list = self.query_one(ChunkList)
|
|
await chunk_list.load_chunks_from_search(self.client, query)
|
|
|
|
# Get unique document IDs from the search results
|
|
doc_ids = list(
|
|
{
|
|
chunk.document_id
|
|
for chunk in chunk_list.chunks
|
|
if chunk.document_id
|
|
}
|
|
)
|
|
|
|
# Update document list to show only documents with matching chunks
|
|
if doc_ids:
|
|
doc_list = self.query_one(DocumentList)
|
|
documents = []
|
|
for doc_id in doc_ids:
|
|
doc = await self.client.document_repository.get_by_id(doc_id)
|
|
if doc:
|
|
documents.append(doc)
|
|
|
|
# Update the document list with filtered documents
|
|
doc_list.documents = documents
|
|
if doc_list.list_view:
|
|
await doc_list.list_view.clear()
|
|
for doc in documents:
|
|
title = doc.title or doc.uri or doc.id or "Untitled"
|
|
item = ListItem(Static(f"{title}"))
|
|
await doc_list.list_view.append(item)
|
|
|
|
# Select first chunk and focus the chunk list
|
|
if chunk_list.list_view:
|
|
if chunk_list.chunks:
|
|
chunk_list.list_view.index = 0
|
|
chunk_list.list_view.focus()
|
|
|
|
async def on_document_list_document_selected(
|
|
self, message: DocumentList.DocumentSelected
|
|
) -> None:
|
|
"""Handle document selection from document list.
|
|
|
|
Args:
|
|
message: Message containing selected document
|
|
"""
|
|
if not self.client:
|
|
return
|
|
|
|
# Always track current document (even during search)
|
|
self.current_document_id = message.document.id
|
|
|
|
# Show document details
|
|
detail_view = self.query_one(DetailView)
|
|
await detail_view.show_document(message.document)
|
|
|
|
# Load chunks for this document (but not during search - preserve search results)
|
|
if message.document.id and not self.search_active:
|
|
chunk_list = self.query_one(ChunkList)
|
|
await chunk_list.load_chunks_for_document(self.client, message.document.id)
|
|
|
|
async def on_chunk_list_chunk_selected(
|
|
self, message: ChunkList.ChunkSelected
|
|
) -> None:
|
|
"""Handle chunk selection from chunk list.
|
|
|
|
Args:
|
|
message: Message containing selected chunk
|
|
"""
|
|
# Always track current chunk (even during search)
|
|
self.current_chunk_id = message.chunk.id
|
|
|
|
# Show chunk details
|
|
detail_view = self.query_one(DetailView)
|
|
await detail_view.show_chunk(message.chunk)
|
|
|
|
# Track the document this chunk belongs to
|
|
if message.chunk.document_id:
|
|
self.current_document_id = message.chunk.document_id
|
|
|
|
|
|
def run_inspector(db_path: Path | None = None) -> None:
|
|
"""Run the inspector TUI.
|
|
|
|
Args:
|
|
db_path: Path to the LanceDB database. If None, uses default from config.
|
|
"""
|
|
config = get_config()
|
|
if db_path is None:
|
|
db_path = config.storage.data_dir / "haiku.rag.lancedb"
|
|
|
|
app = InspectorApp(db_path)
|
|
app.run()
|