zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/tui/screens/settings.py

@@ -0,0 +1,219 @@
+"""Settings: manage downloaded models, quantized outputs, and the HF token."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+
+from rich.text import Text
+from textual import work
+from textual.app import ComposeResult
+from textual.binding import Binding
+from textual.containers import Vertical
+from textual.screen import Screen
+from textual.widgets import DataTable, Footer, Static, TabbedContent, TabPane
+
+from zeroquantz.core import artifacts
+from zeroquantz.models import hf_auth
+from zeroquantz.tui.screens.confirm import ConfirmScreen
+from zeroquantz.tui.screens.token import TokenScreen
+from zeroquantz.utils import units
+
+
+class SettingsScreen(Screen):
+    BINDINGS = [
+        Binding("escape", "back", "Back"),
+        Binding("1", "tab_downloaded", "Downloaded", show=False),
+        Binding("2", "tab_quantized", "Quantized", show=False),
+        Binding("3", "tab_token", "Token", show=False),
+        Binding("r", "refresh", "Refresh"),
+        Binding("d", "delete", "Delete"),
+        Binding("o", "open", "Open"),
+        Binding("e", "token_edit", "Set token"),
+        Binding("c", "token_clear", "Clear token"),
+    ]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._downloaded: list[artifacts.DownloadedModel] = []
+        self._quantized: list[artifacts.QuantizedArtifact] = []
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="settings-page"):
+            yield Static("Settings", classes="page-title")
+            with TabbedContent(id="settings-tabs", initial="tab-downloaded"):
+                with TabPane("Downloaded models", id="tab-downloaded"):
+                    yield DataTable(id="dl-table", cursor_type="row", zebra_stripes=True)
+                with TabPane("Quantized models", id="tab-quantized"):
+                    yield DataTable(id="q-table", cursor_type="row", zebra_stripes=True)
+                with TabPane("Hugging Face token", id="tab-token"):
+                    yield Static("", id="tok-status")
+        yield Footer()
+
+    def on_mount(self) -> None:
+        self.query_one("#dl-table", DataTable).add_columns("Model", "Size", "Files")
+        self.query_one("#q-table", DataTable).add_columns("Name", "Base model", "Format", "Size")
+        self._refresh_quantized()
+        self._refresh_token()
+        self._refresh_downloaded()
+
+    # ---- refresh ------------------------------------------------------------
+
+    @work(thread=True, exclusive=True)
+    def _refresh_downloaded(self) -> None:
+        data = artifacts.list_downloaded_models()
+        self.app.call_from_thread(self._fill_downloaded, data)
+
+    def _fill_downloaded(self, data: list[artifacts.DownloadedModel]) -> None:
+        self._downloaded = data
+        table = self.query_one("#dl-table", DataTable)
+        table.clear()
+        total = 0
+        for m in data:
+            total += m.size_bytes
+            table.add_row(m.repo_id, units.humanize_bytes(m.size_bytes), str(m.nb_files))
+        if not data:
+            table.add_row("(cache empty)", "", "")
+        self._set_tab_title("tab-downloaded", f"Downloaded models · {units.humanize_bytes(total)}")
+
+    def _refresh_quantized(self) -> None:
+        self._quantized = artifacts.default_quantized_registry().list()
+        table = self.query_one("#q-table", DataTable)
+        table.clear()
+        total = 0
+        for a in self._quantized:
+            total += a.size_bytes
+            name = os.path.basename(a.path.rstrip("/\\"))
+            table.add_row(name, a.base_model or "?", a.format_id or "?", units.humanize_bytes(a.size_bytes))
+        if not self._quantized:
+            table.add_row("(none yet)", "", "", "")
+        self._set_tab_title("tab-quantized", f"Quantized models · {units.humanize_bytes(total)}")
+
+    def _refresh_token(self) -> None:
+        token = hf_auth.current_token()
+        status = self.query_one("#tok-status", Static)
+        if token:
+            status.update(Text.assemble(
+                ("● ", "#3fb950"), (f"Signed in — token {hf_auth.masked(token)}\n", "#3fb950"),
+                ("Gated models and higher Hub rate limits are enabled.", "#8b949e"),
+            ))
+        else:
+            status.update(Text.assemble(
+                ("○ ", "#8b949e"), ("No token set\n", "#8b949e"),
+                ("Add one to access gated models (Llama, …) and raise rate limits.", "#8b949e"),
+            ))
+
+    def _set_tab_title(self, pane_id: str, title: str) -> None:
+        try:
+            self.query_one("#settings-tabs", TabbedContent).get_tab(pane_id).label = title
+        except Exception:
+            pass
+
+    # ---- tab switching ------------------------------------------------------
+
+    def _active(self) -> str:
+        return self.query_one("#settings-tabs", TabbedContent).active
+
+    def action_tab_downloaded(self) -> None:
+        self.query_one("#settings-tabs", TabbedContent).active = "tab-downloaded"
+
+    def action_tab_quantized(self) -> None:
+        self.query_one("#settings-tabs", TabbedContent).active = "tab-quantized"
+
+    def action_tab_token(self) -> None:
+        self.query_one("#settings-tabs", TabbedContent).active = "tab-token"
+
+    # ---- actions ------------------------------------------------------------
+
+    def action_refresh(self) -> None:
+        active = self._active()
+        if active == "tab-downloaded":
+            self._refresh_downloaded()
+        elif active == "tab-quantized":
+            self._refresh_quantized()
+        else:
+            self._refresh_token()
+
+    def action_delete(self) -> None:
+        active = self._active()
+        if active == "tab-downloaded":
+            self._delete_downloaded()
+        elif active == "tab-quantized":
+            self._delete_quantized()
+
+    def _delete_downloaded(self) -> None:
+        row = self.query_one("#dl-table", DataTable).cursor_row
+        if row is None or row >= len(self._downloaded):
+            return
+        model = self._downloaded[row]
+        msg = f"Delete cached model '{model.repo_id}'?"
+        detail = f"Frees {units.humanize_bytes(model.size_bytes)} from the Hugging Face cache."
+        self.app.push_screen(
+            ConfirmScreen(msg, detail=detail),
+            lambda ok, r=model.repo_id: self._do_delete_downloaded(r) if ok else None,
+        )
+
+    @work(thread=True, exclusive=True)
+    def _do_delete_downloaded(self, repo_id: str) -> None:
+        try:
+            artifacts.delete_downloaded_model(repo_id)
+        except Exception:
+            pass
+        data = artifacts.list_downloaded_models()
+        self.app.call_from_thread(self._fill_downloaded, data)
+
+    def _delete_quantized(self) -> None:
+        row = self.query_one("#q-table", DataTable).cursor_row
+        if row is None or row >= len(self._quantized):
+            return
+        artifact = self._quantized[row]
+        msg = f"Delete quantized model at '{artifact.path}'?"
+        detail = f"Permanently removes the directory ({units.humanize_bytes(artifact.size_bytes)})."
+        self.app.push_screen(
+            ConfirmScreen(msg, detail=detail),
+            lambda ok, p=artifact.path: self._do_delete_quantized(p) if ok else None,
+        )
+
+    def _do_delete_quantized(self, path: str) -> None:
+        artifacts.default_quantized_registry().delete(path)
+        self._refresh_quantized()
+
+    def action_open(self) -> None:
+        if self._active() != "tab-quantized":
+            return
+        row = self.query_one("#q-table", DataTable).cursor_row
+        if row is None or row >= len(self._quantized):
+            return
+        _open_folder(self._quantized[row].path)
+
+    def action_token_edit(self) -> None:
+        if self._active() == "tab-token":
+            self.app.push_screen(TokenScreen(), lambda _changed: self._refresh_token())
+
+    def action_token_clear(self) -> None:
+        if self._active() != "tab-token" or not hf_auth.current_token():
+            return
+        self.app.push_screen(
+            ConfirmScreen("Clear the saved Hugging Face token?"),
+            lambda ok: self._do_clear_token() if ok else None,
+        )
+
+    def _do_clear_token(self) -> None:
+        hf_auth.clear_token()
+        self._refresh_token()
+
+    def action_back(self) -> None:
+        self.app.pop_screen()
+
+
+def _open_folder(path: str) -> None:
+    try:
+        if sys.platform.startswith("win"):
+            os.startfile(path)  # noqa: S606
+        elif sys.platform == "darwin":
+            subprocess.Popen(["open", path])
+        else:
+            subprocess.Popen(["xdg-open", path])
+    except Exception:
+        pass

zeroquantz/tui/screens/token.py

@@ -0,0 +1,94 @@
+"""A modal dialog for entering / clearing a Hugging Face token.
+
+Keyboard-driven: Enter validates & saves, Ctrl+D clears, Esc closes.
+"""
+
+from __future__ import annotations
+
+from rich.text import Text
+from textual import on, work
+from textual.app import ComposeResult
+from textual.binding import Binding
+from textual.containers import Vertical
+from textual.screen import ModalScreen
+from textual.widgets import Input, Static
+
+from zeroquantz.models import hf_auth
+
+
+class TokenScreen(ModalScreen[bool]):
+    """Enter, validate, and persist a Hugging Face access token."""
+
+    BINDINGS = [
+        Binding("escape", "close", "Close"),
+        Binding("ctrl+d", "clear", "Clear token"),
+    ]
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="token-dialog"):
+            yield Static("Hugging Face token", classes="dialog-title")
+            yield Static(
+                "Add a token to access gated models (Llama, etc.) and raise Hub rate limits.\n"
+                "Create one at  https://huggingface.co/settings/tokens  (read scope is enough).",
+                classes="dialog-sub",
+            )
+            yield Static("", id="token-status")
+            yield Input(placeholder="hf_…", password=True, id="token-input")
+            yield Static(
+                "[white]Enter[/white] save    [white]Ctrl+D[/white] clear    [white]Esc[/white] close",
+                id="token-hint",
+                classes="page-hint",
+            )
+
+    def on_mount(self) -> None:
+        self._refresh_status()
+        self.query_one("#token-input", Input).focus()
+
+    def _refresh_status(self) -> None:
+        token = hf_auth.current_token()
+        status = self.query_one("#token-status", Static)
+        if token:
+            status.update(Text.assemble(("● ", "#3fb950"), (f"token set ({hf_auth.masked(token)})", "#3fb950")))
+        else:
+            status.update(Text("○ no token set", style="#8b949e"))
+
+    # ---- actions ------------------------------------------------------------
+
+    @on(Input.Submitted, "#token-input")
+    def _submit(self, event: Input.Submitted) -> None:
+        token = event.value.strip()
+        if not token:
+            self.query_one("#token-status", Static).update(Text("Enter a token first.", style="#d29922"))
+            return
+        self.query_one("#token-status", Static).update(Text("checking token …", style="#d29922"))
+        self._validate(token)
+
+    @work(thread=True, exclusive=True)
+    def _validate(self, token: str) -> None:
+        try:
+            name = hf_auth.save_token(token)
+            self.app.call_from_thread(self._saved, name)
+        except Exception as exc:
+            message = exc.format() if hasattr(exc, "format") else str(exc)
+            self.app.call_from_thread(self._failed, message)
+
+    def _saved(self, name: str | None) -> None:
+        who = f" as {name}" if name else ""
+        self.query_one("#token-status", Static).update(
+            Text(f"✓ signed in{who} — gated models & higher limits enabled", style="#3fb950")
+        )
+        self.query_one("#token-input", Input).value = ""
+        self._changed = True
+
+    def _failed(self, message: str) -> None:
+        self.query_one("#token-status", Static).update(Text(f"✗ {message.splitlines()[0]}", style="#f85149"))
+
+    def action_clear(self) -> None:
+        hf_auth.clear_token()
+        self._changed = True
+        self.query_one("#token-input", Input).value = ""
+        self._refresh_status()
+        self.query_one("#token-status", Static).update(Text("token removed", style="#8b949e"))
+
+    def action_close(self) -> None:
+        self.dismiss(getattr(self, "_changed", False))

zeroquantz/tui/screens/welcome.py

@@ -0,0 +1,128 @@
+"""The landing screen: ZEROQUANTZ wordmark, hardware status, arrow-key menu."""
+
+from __future__ import annotations
+
+from rich.text import Text
+from textual import on
+from textual.app import ComposeResult
+from textual.containers import Vertical
+from textual.screen import Screen
+from textual.widgets import OptionList, Static
+from textual.widgets.option_list import Option
+
+from zeroquantz.models import hf_auth
+from zeroquantz.tui.banner import banner_text
+
+
+class WelcomeScreen(Screen):
+    """First screen — mirrors the requested launcher aesthetic."""
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="welcome"):
+            yield Static("", classes="banner", id="banner")
+            yield Static(
+                "Interactive Model Optimization Environment", classes="tagline"
+            )
+            yield Static("", classes="status-line", id="status")
+            yield Static(
+                "What would you like to do?  (Use arrow keys)", classes="prompt-line"
+            )
+            yield OptionList(id="menu")
+
+    def on_mount(self) -> None:
+        self._refresh_banner()
+        self._refresh_status()
+        self._build_menu()
+        self.query_one("#menu", OptionList).focus()
+
+    def on_resize(self) -> None:
+        self._refresh_banner()
+
+    # ---- content ------------------------------------------------------------
+
+    def _refresh_banner(self) -> None:
+        width = self.app.size.width or 80
+        self.query_one("#banner", Static).update(
+            Text(banner_text(width - 4), style="bold #c9d1d9", justify="center")
+        )
+
+    def _refresh_status(self) -> None:
+        hw = self.app.ctx.hardware
+        text = Text()
+        if hw.cuda_available:
+            text.append("● ", style="#3fb950")
+            text.append(hw.gpu_name, style="#3fb950")
+            text.append(f"   VRAM {hw.total_vram_gb:g} GB", style="#8b949e")
+            text.append(f"   CUDA {hw.cuda_version or 'detected'}", style="#8b949e")
+        else:
+            text.append("⚠ No CUDA GPU detected — planning-only mode", style="#d29922")
+        if hf_auth.current_token():
+            text.append("     HF token ✓", style="#3fb950")
+        self.query_one("#status", Static).update(text)
+
+    def _build_menu(self) -> None:
+        menu = self.query_one("#menu", OptionList)
+        menu.clear_options()
+        self._items = [
+            ("optimize", "Optimize a model", "guided: model → quantization → export", "●", "bold #58a6ff"),
+            ("hardware", "Inspect hardware", "detected GPU & precision support", "●", "bold #58a6ff"),
+            ("command", "Command mode & help", "slash commands, natural language, /help", "●", "bold #58a6ff"),
+        ]
+        latest = self.app.ctx.repo.latest()
+        if latest is not None:
+            self._items.append(
+                ("resume", "Resume last session", f"continue “{latest.name}”", "●", "bold #58a6ff")
+            )
+        self._items.append((
+            "settings", "Settings",
+            "manage downloaded / quantized models & HF token", "●", "bold #58a6ff",
+        ))
+        self._items.append(("exit", "Exit", "", "○", "#8b949e"))
+        for idx, item in enumerate(self._items):
+            menu.add_option(Option(_option_prompt(item, pointed=idx == 0), id=item[0]))
+
+    @on(OptionList.OptionHighlighted)
+    def _highlighted(self, event: OptionList.OptionHighlighted) -> None:
+        if not getattr(self, "_items", None):
+            return
+        menu = event.option_list
+        try:  # move the → pointer to the highlighted row (graceful if API differs)
+            for idx, item in enumerate(self._items):
+                menu.replace_option_prompt_at_index(
+                    idx, _option_prompt(item, pointed=idx == event.option_index)
+                )
+        except Exception:
+            pass
+
+    # ---- actions ------------------------------------------------------------
+
+    @on(OptionList.OptionSelected)
+    def _selected(self, event: OptionList.OptionSelected) -> None:
+        choice = event.option.id
+        if choice == "exit":
+            self.app.exit()
+        elif choice == "hardware":
+            self.app.open_workspace("/hardware")
+        elif choice == "command":
+            self.app.open_workspace("/help")  # enter command mode with help shown
+        elif choice == "settings":
+            self.app.push_settings()
+        elif choice == "resume":
+            self.app.resume_last_session()
+        else:  # optimize -> guided wizard
+            self.app.start_wizard()
+
+    def on_screen_resume(self) -> None:
+        # Returning from Settings may have changed the token — refresh the chip.
+        self._refresh_status()
+
+
+def _option_prompt(item: tuple[str, str, str, str, str], *, pointed: bool) -> Text:
+    _id, label, desc, bullet, label_style = item
+    text = Text()
+    text.append("→ " if pointed else "  ", style="bold #3fb950")
+    text.append(f"{bullet} ", style="#3fb950" if bullet == "●" else "#6e7681")
+    text.append(label, style=label_style)
+    if desc:
+        text.append(f"   —  {desc}", style="#6e7681")
+    return text

zeroquantz/tui/screens/workspace.py

@@ -0,0 +1,175 @@
+"""The main workspace: conversation log, input, and the state sidebar.
+
+User input and slash commands are parsed to intents and dispatched on a worker
+thread (so a long quantize/benchmark never freezes the UI), with progress shown
+live. Results render through the shared :mod:`zeroquantz.render` layer.
+"""
+
+from __future__ import annotations
+
+from rich.text import Text
+from textual import work
+from textual.app import ComposeResult
+from textual.binding import Binding
+from textual.containers import Horizontal
+from textual.screen import Screen
+from textual.widgets import Input, RichLog, Static
+
+from zeroquantz.agent.dispatcher import COMMANDS
+from zeroquantz.agent.intents import IntentKind
+from zeroquantz.agent.parser import IntentParser
+from zeroquantz.render import render_result
+from zeroquantz.tui.widgets import Sidebar, StatusBar
+
+_SLASH_COMMANDS = [cmd.split()[0] for cmd, _ in COMMANDS]
+
+
+class WorkspaceScreen(Screen):
+    BINDINGS = [
+        Binding("ctrl+l", "clear", "Clear"),
+        Binding("ctrl+b", "toggle_sidebar", "Sidebar"),
+        Binding("ctrl+c", "quit", "Quit", priority=True),
+        Binding("escape", "home", "Home"),
+    ]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._history: list[str] = []
+        self._hist_idx = 0
+
+    def compose(self) -> ComposeResult:
+        yield StatusBar(id="topbar")
+        with Horizontal(id="body"):
+            yield RichLog(id="conversation", wrap=True, highlight=False, markup=False)
+            yield Sidebar(id="sidebar")
+        yield Static("", id="progress")
+        yield Input(
+            placeholder="Type a command or ask in plain English…   (/help, /exit)",
+            id="prompt",
+        )
+
+    def on_mount(self) -> None:
+        log = self.query_one("#conversation", RichLog)
+        log.write(
+            Text.assemble(
+                ("ZeroQuantz ready. ", "bold #3fb950"),
+                ("Try ", "#8b949e"),
+                ("/model Qwen/Qwen3-8B", "#58a6ff"),
+                (" or “fit under 8GB for vLLM”.", "#8b949e"),
+            )
+        )
+        self._refresh_panels()
+        self.query_one("#prompt", Input).focus()
+        pending = getattr(self.app, "pending_command", None)
+        if pending:
+            self.app.pending_command = None
+            self._echo(pending)
+            self._dispatch(pending)
+
+    # ---- input handling -----------------------------------------------------
+
+    def on_input_submitted(self, event: Input.Submitted) -> None:
+        text = event.value.strip()
+        event.input.value = ""
+        if not text:
+            return
+        self._history.append(text)
+        self._hist_idx = len(self._history)
+        self._echo(text)
+        # EXIT is handled synchronously so the app closes immediately.
+        if IntentParser.parse(text).kind is IntentKind.EXIT:
+            self.app.exit()
+            return
+        self._dispatch(text)
+
+    def on_key(self, event) -> None:  # noqa: ANN001
+        inp = self.query_one("#prompt", Input)
+        if not inp.has_focus:
+            return
+        if event.key == "up":
+            self._history_step(-1)
+            event.stop()
+        elif event.key == "down":
+            self._history_step(1)
+            event.stop()
+        elif event.key == "tab":
+            self._autocomplete()
+            event.stop()
+            event.prevent_default()
+
+    def _history_step(self, delta: int) -> None:
+        if not self._history:
+            return
+        self._hist_idx = max(0, min(len(self._history), self._hist_idx + delta))
+        inp = self.query_one("#prompt", Input)
+        inp.value = self._history[self._hist_idx] if self._hist_idx < len(self._history) else ""
+        inp.cursor_position = len(inp.value)
+
+    def _autocomplete(self) -> None:
+        inp = self.query_one("#prompt", Input)
+        value = inp.value
+        if not value.startswith("/"):
+            return
+        matches = [c for c in _SLASH_COMMANDS if c.startswith(value)]
+        if len(matches) == 1:
+            inp.value = matches[0] + " "
+            inp.cursor_position = len(inp.value)
+        elif len(matches) > 1:
+            self.query_one("#conversation", RichLog).write(
+                Text("  ".join(matches), style="#6e7681")
+            )
+
+    # ---- dispatch -----------------------------------------------------------
+
+    @work(thread=True, exclusive=True)
+    def _dispatch(self, text: str) -> None:
+        intent = IntentParser.parse(text)
+
+        def progress(stage: str, fraction: float) -> None:
+            self.app.call_from_thread(self._set_progress, stage, fraction)
+
+        result = self.app.dispatcher.dispatch(intent, self.app.ctx, progress)
+        self.app.call_from_thread(self._show_result, result)
+
+    def _show_result(self, result) -> None:  # noqa: ANN001
+        self._set_progress(None, 0.0)
+        self.query_one("#conversation", RichLog).write(render_result(result))
+        self._refresh_panels()
+        if result.should_exit:
+            self.app.exit()
+
+    def _set_progress(self, stage: str | None, fraction: float) -> None:
+        prog = self.query_one("#progress", Static)
+        if stage is None:
+            prog.remove_class("active")
+            prog.update("")
+            return
+        filled = int(round(fraction * 20))
+        bar = "█" * filled + "░" * (20 - filled)
+        prog.update(Text(f"  {bar}  {stage} ({fraction*100:.0f}%)", style="#d29922"))
+        prog.add_class("active")
+
+    # ---- helpers ------------------------------------------------------------
+
+    def _echo(self, text: str) -> None:
+        self.query_one("#conversation", RichLog).write(
+            Text.assemble(("› ", "bold #3fb950"), (text, "bold white"))
+        )
+
+    def _refresh_panels(self) -> None:
+        self.query_one("#topbar", StatusBar).update_context(self.app.ctx)
+        self.query_one("#sidebar", Sidebar).update_context(self.app.ctx)
+
+    # ---- actions ------------------------------------------------------------
+
+    def action_clear(self) -> None:
+        self.query_one("#conversation", RichLog).clear()
+
+    def action_toggle_sidebar(self) -> None:
+        self.query_one("#sidebar", Sidebar).toggle_class("hidden")
+
+    def action_home(self) -> None:
+        self.app.pop_screen()
+
+    def action_quit(self) -> None:
+        self.app.exit()