haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

haiku/rag/inspector/app.py

@@ -0,0 +1,259 @@
+# pyright: reportPossiblyUnboundVariable=false
+from datetime import datetime
+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.screen import Screen
+    from textual.widgets import Footer, Header
+
+    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
+    from haiku.rag.inspector.widgets.search_modal import SearchModal
+
+    TEXTUAL_AVAILABLE = True
+except ImportError:
+    TEXTUAL_AVAILABLE = False
+    App = object  # type: ignore
+
+
+class InspectorApp(App):  # type: ignore[misc]  # pragma: no cover
+    """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;
+    }
+    """
+
+    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),
+    ]
+
+    def __init__(
+        self, db_path: Path, read_only: bool = False, before: datetime | None = None
+    ):
+        super().__init__()
+        self.db_path = db_path
+        self.read_only = read_only
+        self.before = before
+        self.client: HaikuRAG | None = None
+
+    def compose(self) -> "ComposeResult":
+        """Compose the UI layout."""
+        yield Header()
+        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,
+            read_only=self.read_only,
+            before=self.before,
+        )
+        await self.client.__aenter__()
+
+        # Load initial documents
+        doc_list = self.query_one(DocumentList)
+        await doc_list.load_documents(self.client)
+
+        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)
+
+    def _select_chunk(self, chunk_list: ChunkList, chunk_id: str) -> None:
+        """Helper to select a chunk after refresh."""
+        for idx, c in enumerate(chunk_list.chunks):
+            if c.id == chunk_id:
+                chunk_list.list_view.index = idx
+                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._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:
+        """Handle chunk selection from search modal."""
+        if not self.client:
+            return
+
+        chunk = message.chunk
+
+        # Navigate to the document containing this chunk
+        if chunk.document_id:
+            doc = await self.client.document_repository.get_by_id(chunk.document_id)
+            if doc:
+                doc_list = self.query_one(DocumentList)
+                chunk_list = self.query_one(ChunkList)
+
+                # Find and select the document
+                for idx, d in enumerate(doc_list.documents):
+                    if d.id == chunk.document_id:
+                        doc_list.list_view.index = idx
+                        break
+
+                # Load chunks for this document
+                await chunk_list.load_chunks_for_document(
+                    self.client, chunk.document_id
+                )
+
+                # Wait a tick for the ListView to process the new items
+                self.call_after_refresh(self._select_chunk, chunk_list, chunk.id)
+
+    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
+
+        # Show document details
+        detail_view = self.query_one(DetailView)
+        await detail_view.show_document(message.document)
+
+        # Load chunks for this document
+        if message.document.id:
+            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
+        """
+        # Show chunk details
+        detail_view = self.query_one(DetailView)
+        await detail_view.show_chunk(message.chunk)
+
+    async def action_show_visual(self) -> None:
+        """Show visual grounding for the currently selected chunk."""
+        if not self.client:
+            return
+
+        chunk_list = self.query_one(ChunkList)
+        idx = chunk_list.list_view.index
+        if idx is None or idx >= len(chunk_list.chunks):
+            return
+
+        chunk = chunk_list.chunks[idx]
+
+        from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
+
+        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."""
+        if not self.client:
+            return
+
+        chunk_list = self.query_one(ChunkList)
+        idx = chunk_list.list_view.index
+        if idx is None or idx >= len(chunk_list.chunks):
+            return
+
+        chunk = chunk_list.chunks[idx]
+
+        from haiku.rag.inspector.widgets.context_modal import ContextModal
+
+        await self._switch_modal(ContextModal(chunk=chunk, client=self.client))
+
+
+def run_inspector(
+    db_path: Path | None = None,
+    read_only: bool = False,
+    before: datetime | None = None,
+) -> None:  # pragma: no cover
+    """Run the inspector TUI.
+
+    Args:
+        db_path: Path to the LanceDB database. If None, uses default from config.
+        read_only: Whether to open the database in read-only mode.
+        before: Query database as it existed before this datetime.
+    """
+    config = get_config()
+    if db_path is None:
+        db_path = config.storage.data_dir / "haiku.rag.lancedb"
+
+    app = InspectorApp(db_path, read_only=read_only, before=before)
+    app.run()

haiku/rag/inspector/widgets/__init__.py

@@ -0,0 +1,6 @@
+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
+from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
+
+__all__ = ["ChunkList", "DetailView", "DocumentList", "VisualGroundingModal"]

haiku/rag/inspector/widgets/chunk_list.py

@@ -0,0 +1,100 @@
+from textual import on
+from textual.app import ComposeResult
+from textual.containers import VerticalScroll
+from textual.message import Message
+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."""
+
+    can_focus = False
+
+    class ChunkSelected(Message):
+        """Message sent when a chunk is selected."""
+
+        def __init__(self, chunk: Chunk) -> None:
+            super().__init__()
+            self.chunk = chunk
+
+    def __init__(self, **kwargs) -> None:
+        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."""
+        yield Static("[bold]Chunks[/bold]", classes="title")
+        yield self.list_view
+
+    async def load_chunks_for_document(
+        self, client: HaikuRAG, document_id: str
+    ) -> None:
+        """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]
+            await self.list_view.append(
+                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(
+        self, event: ListView.Highlighted | ListView.Selected
+    ) -> None:
+        """Handle chunk selection (arrow keys or Enter)."""
+        if event.list_view != self.list_view:
+            return
+        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/inspector/widgets/context_modal.py

@@ -0,0 +1,89 @@
+from typing import TYPE_CHECKING
+
+from textual.app import ComposeResult
+from textual.binding import Binding
+from textual.containers import VerticalScroll
+from textual.screen import Screen
+from textual.widgets import Markdown, Static
+
+from haiku.rag.store.models import SearchResult
+
+if TYPE_CHECKING:
+    from haiku.rag.client import HaikuRAG
+    from haiku.rag.store.models import Chunk
+
+
+class ContextModal(Screen):  # pragma: no cover
+    """Modal screen for displaying how a chunk appears to agents."""
+
+    BINDINGS = [
+        Binding("escape", "dismiss", "Close", show=True),
+        Binding("c", "dismiss", "Close", show=True),
+    ]
+
+    CSS = """
+    ContextModal {
+        background: $surface;
+        layout: vertical;
+    }
+
+    #context-header {
+        dock: top;
+        height: auto;
+        padding: 1;
+    }
+
+    #context-content {
+        height: 1fr;
+        width: 100%;
+        padding: 1;
+    }
+
+    #context-content Markdown {
+        width: 100%;
+    }
+    """
+
+    def __init__(self, chunk: "Chunk", client: "HaikuRAG"):
+        super().__init__()
+        self.chunk = chunk
+        self.client = client
+        self._content_widget = Markdown("Loading...")
+
+    def compose(self) -> ComposeResult:
+        yield Static("[bold]Agent Context Format[/bold]", id="context-header")
+        with VerticalScroll(id="context-content"):
+            yield self._content_widget
+
+    async def on_mount(self) -> None:
+        """Load and display the expanded context."""
+        # Create a SearchResult from the chunk
+        chunk_meta = self.chunk.get_chunk_metadata()
+        search_result = SearchResult(
+            content=self.chunk.content,
+            score=0.0,
+            chunk_id=self.chunk.id,
+            document_id=self.chunk.document_id,
+            document_uri=self.chunk.document_uri,
+            document_title=self.chunk.document_title,
+            doc_item_refs=chunk_meta.doc_item_refs,
+            page_numbers=chunk_meta.page_numbers,
+            headings=chunk_meta.headings,
+            labels=chunk_meta.labels,
+        )
+
+        # Expand context using the client (this is what agents actually receive)
+        expanded_results = await self.client.expand_context([search_result])
+        expanded = expanded_results[0] if expanded_results else search_result
+
+        formatted = expanded.format_for_agent()
+
+        content = (
+            "*This is how the chunk appears to agents after context expansion:*\n\n---\n\n"
+            f"{formatted}"
+        )
+
+        await self._content_widget.update(content)
+
+    async def action_dismiss(self, result=None) -> None:
+        self.app.pop_screen()

haiku/rag/inspector/widgets/detail_view.py

@@ -0,0 +1,130 @@
+from typing import Protocol
+
+from textual.app import ComposeResult
+from textual.containers import VerticalScroll
+from textual.widgets import Markdown, Static
+
+from haiku.rag.store.models import Chunk, Document, SearchResult
+
+
+class ProvenanceData(Protocol):
+    """Protocol for objects that have provenance metadata."""
+
+    page_numbers: list[int]
+    headings: list[str] | None
+    labels: list[str]
+    doc_item_refs: list[str]
+
+
+class DetailView(VerticalScroll):  # pragma: no cover
+    """Widget for displaying detailed content of documents or chunks."""
+
+    can_focus = True
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.title_widget = Static("[bold]Detail View[/bold]", classes="title")
+        self.content_widget = Markdown("")
+        self.content_widget.can_focus = True
+
+    def compose(self) -> ComposeResult:
+        yield self.title_widget
+        yield self.content_widget
+
+    def _format_provenance(self, prov: ProvenanceData) -> list[str]:
+        """Format provenance metadata as markdown lines."""
+        parts: list[str] = []
+        if prov.page_numbers:
+            pages_str = ", ".join(str(p) for p in prov.page_numbers)
+            parts.append(f"**Page(s):** {pages_str}")
+        if prov.headings:
+            headings_str = " > ".join(prov.headings)
+            parts.append(f"**Section:** {headings_str}")
+        if prov.labels:
+            labels_str = ", ".join(prov.labels)
+            parts.append(f"**Labels:** {labels_str}")
+        if prov.doc_item_refs:
+            refs_str = ", ".join(prov.doc_item_refs[:5])
+            if len(prov.doc_item_refs) > 5:
+                refs_str += f" ... (+{len(prov.doc_item_refs) - 5} more)"
+            parts.append(f"**DocItem Refs:** `{refs_str}`")
+        return parts
+
+    async def show_document(self, document: Document) -> None:
+        """Display document details."""
+        title = document.title or document.uri or "Untitled Document"
+        self.title_widget.update(f"[bold]Document: {title}[/bold]")
+
+        content_parts: list[str] = []
+        if document.id:
+            content_parts.append(f"**ID:** `{document.id}`")
+        if document.uri:
+            content_parts.append(f"**URI:** `{document.uri}`")
+        if document.metadata:
+            metadata_str = "\n".join(
+                f"  - {k}: {v}" for k, v in document.metadata.items()
+            )
+            content_parts.append(f"**Metadata:**\n{metadata_str}")
+        if document.created_at:
+            content_parts.append(f"**Created:** {document.created_at}")
+        if document.updated_at:
+            content_parts.append(f"**Updated:** {document.updated_at}")
+
+        content_parts.append("\n---\n")
+        content_parts.append(document.content)
+
+        await self.content_widget.update("\n\n".join(content_parts))
+
+    async def show_chunk(self, chunk: Chunk) -> None:
+        """Display chunk details."""
+        self.title_widget.update(f"[bold]Chunk {chunk.order}[/bold]")
+
+        content_parts: list[str] = []
+        if chunk.id:
+            content_parts.append(f"**ID:** `{chunk.id}`")
+        if chunk.document_id:
+            content_parts.append(f"**Document ID:** `{chunk.document_id}`")
+        if chunk.document_title:
+            content_parts.append(f"**Document Title:** {chunk.document_title}")
+        if chunk.document_uri:
+            content_parts.append(f"**Document URI:** `{chunk.document_uri}`")
+        content_parts.append(f"**Order:** {chunk.order}")
+
+        chunk_meta = chunk.get_chunk_metadata()
+        content_parts.extend(self._format_provenance(chunk_meta))
+
+        if chunk.embedding:
+            content_parts.append(f"**Embedding:** {len(chunk.embedding)} dimensions")
+
+        content_parts.append("\n---\n")
+        content_parts.append(chunk.content)
+
+        await self.content_widget.update("\n\n".join(content_parts))
+
+    async def show_search_result(
+        self, chunk: Chunk, search_result: SearchResult
+    ) -> None:
+        """Display chunk details with search result metadata."""
+        self.title_widget.update(f"[bold]Chunk {chunk.order}[/bold]")
+
+        content_parts: list[str] = []
+        if chunk.id:
+            content_parts.append(f"**ID:** `{chunk.id}`")
+        if chunk.document_id:
+            content_parts.append(f"**Document ID:** `{chunk.document_id}`")
+        if search_result.document_title:
+            content_parts.append(f"**Document Title:** {search_result.document_title}")
+        if search_result.document_uri:
+            content_parts.append(f"**Document URI:** `{search_result.document_uri}`")
+        content_parts.append(f"**Order:** {chunk.order}")
+        content_parts.append(f"**Score:** {search_result.score:.4f}")
+
+        content_parts.extend(self._format_provenance(search_result))
+
+        if chunk.embedding:
+            content_parts.append(f"**Embedding:** {len(chunk.embedding)} dimensions")
+
+        content_parts.append("\n---\n")
+        content_parts.append(chunk.content)
+
+        await self.content_widget.update("\n\n".join(content_parts))

haiku/rag/inspector/widgets/document_list.py

@@ -0,0 +1,75 @@
+from textual import on
+from textual.app import ComposeResult
+from textual.containers import VerticalScroll
+from textual.message import Message
+from textual.widgets import ListItem, ListView, Static
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.store.models import Document
+
+BATCH_SIZE = 50
+
+
+class DocumentList(VerticalScroll):  # pragma: no cover
+    """Widget for displaying and browsing documents."""
+
+    can_focus = False
+
+    class DocumentSelected(Message):
+        """Message sent when a document is selected."""
+
+        def __init__(self, document: Document) -> None:
+            super().__init__()
+            self.document = document
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.documents: list[Document] = []
+        self.list_view = ListView()
+        self.has_more: bool = True
+        self._client: HaikuRAG | None = None
+        self._loading: bool = False
+
+    def compose(self) -> ComposeResult:
+        """Compose the document list."""
+        yield Static("[bold]Documents[/bold]", classes="title")
+        yield self.list_view
+
+    async def load_documents(self, client: HaikuRAG) -> None:
+        """Load initial batch of documents from the database."""
+        self._client = client
+        self.documents = await client.list_documents(limit=BATCH_SIZE, offset=0)
+        self.has_more = len(self.documents) >= BATCH_SIZE
+        await self.list_view.clear()
+        for doc in self.documents:
+            title = doc.title or doc.uri or doc.id
+            await self.list_view.append(ListItem(Static(f"{title}")))
+
+    async def load_more(self, client: HaikuRAG) -> None:
+        """Load the next batch of documents."""
+        if not self.has_more or self._loading:
+            return
+        self._loading = True
+        offset = len(self.documents)
+        new_docs = await client.list_documents(limit=BATCH_SIZE, offset=offset)
+        self.has_more = len(new_docs) >= BATCH_SIZE
+        self.documents.extend(new_docs)
+        for doc in new_docs:
+            title = doc.title or doc.uri or doc.id
+            await self.list_view.append(ListItem(Static(f"{title}")))
+        self._loading = False
+
+    @on(ListView.Highlighted)
+    @on(ListView.Selected)
+    async def handle_document_selection(
+        self, event: ListView.Highlighted | ListView.Selected
+    ) -> None:
+        """Handle document selection (arrow keys or Enter)."""
+        if event.list_view != self.list_view:
+            return
+        idx = event.list_view.index
+        if idx is not None and 0 <= idx < len(self.documents):
+            self.post_message(self.DocumentSelected(self.documents[idx]))
+            # Infinite scroll: load more when near the end
+            if self._client and self.has_more and idx >= len(self.documents) - 10:
+                await self.load_more(self._client)