66 lines
1.9 KiB
Python
66 lines
1.9 KiB
Python
"""Custom Textual widgets for SneakyCode TUI."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from rich.text import Text
|
|
from textual.widgets import Static
|
|
|
|
|
|
class StatusBar(Static):
|
|
"""Single-line status bar showing token usage and iteration count."""
|
|
|
|
DEFAULT_CSS = """
|
|
StatusBar {
|
|
dock: bottom;
|
|
height: 1;
|
|
background: $surface;
|
|
color: $text-muted;
|
|
padding: 0 2;
|
|
}
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__("")
|
|
self._tokens: int = 0
|
|
self._budget: int = 0
|
|
self._iteration: int = 0
|
|
self._max_iterations: int = 0
|
|
|
|
def update_tokens(self, tokens: int, budget: int) -> None:
|
|
"""Update the token usage display."""
|
|
self._tokens = tokens
|
|
self._budget = budget
|
|
self._refresh_display()
|
|
|
|
def update_iteration(self, iteration: int, max_iterations: int) -> None:
|
|
"""Update the iteration count display."""
|
|
self._iteration = iteration
|
|
self._max_iterations = max_iterations
|
|
self._refresh_display()
|
|
|
|
def _refresh_display(self) -> None:
|
|
"""Rebuild the status bar text."""
|
|
parts: list[str] = []
|
|
if self._budget > 0:
|
|
parts.append(f"Tokens: ~{self._tokens:,} / {self._budget:,}")
|
|
if self._max_iterations > 0:
|
|
parts.append(f"Iteration {self._iteration}/{self._max_iterations}")
|
|
self.update(Text(" │ ".join(parts), style="dim"))
|
|
|
|
|
|
class StreamingStatic(Static):
|
|
"""A Static widget that stays mounted but hidden during non-streaming periods.
|
|
|
|
During streaming, call show() to make visible and update() with partial content.
|
|
When streaming ends, call hide() to conceal.
|
|
"""
|
|
|
|
def show_streaming(self) -> None:
|
|
"""Make the widget visible for streaming."""
|
|
self.add_class("visible")
|
|
|
|
def hide_streaming(self) -> None:
|
|
"""Hide the widget and clear content."""
|
|
self.remove_class("visible")
|
|
self.update("")
|