stable

gui.py

@@ -0,0 +1,366 @@
+import re
+import sys
+
+import pynvml
+from pygments import highlight
+from pygments.formatters import HtmlFormatter
+from pygments.lexers import TextLexer, get_lexer_by_name, guess_lexer
+from pygments.util import ClassNotFound
+from PyQt6.QtCore import QThread, QTimer, pyqtSignal
+from PyQt6.QtGui import (
+    QColor,
+    QFont,
+    QTextBlockFormat,
+    QTextCharFormat,
+    QTextCursor,
+    QTextDocument,
+    QTextDocumentFragment,
+)
+from PyQt6.QtWidgets import (
+    QApplication,
+    QHBoxLayout,
+    QLabel,
+    QLineEdit,
+    QMainWindow,
+    QPushButton,
+    QTextBrowser,
+    QVBoxLayout,
+    QWidget,
+)
+
+import chatbot
+
+# How often the bottom-right GPU/VRAM readout refreshes.
+GPU_STATUS_INTERVAL_MS = 1000
+
+# GitHub dialect adds fenced code blocks with language tags, tables, and
+# strikethrough on top of CommonMark; NoHTML stops raw HTML in the model's
+# output from being interpreted as literal markup in the chat view.
+_MARKDOWN_FEATURES = (
+    QTextDocument.MarkdownFeature.MarkdownDialectGitHub
+    | QTextDocument.MarkdownFeature.MarkdownNoHTML
+)
+
+# Matches fenced code blocks (```lang\ncode\n```) so they can be pulled out
+# and rendered through Pygments instead of QTextDocumentFragment.fromMarkdown,
+# which only gives code a monospace font and no syntax colors.
+_CODE_FENCE_RE = re.compile(r"```([\w+-]*)\n(.*?)```", re.DOTALL)
+
+# noclasses=True inlines every color as a style="..." attribute rather than
+# relying on a <style> stylesheet, which QTextBrowser's limited HTML/CSS
+# support won't pick up. nowrap=True drops the surrounding <div>/<pre>, since
+# a QTextCursor only applies that container's block styling (background,
+# monospace font) to whichever line happens to share the cursor's existing
+# block - every other line silently falls back to plain text. Instead each
+# source line is inserted as its own block, below, with that styling applied
+# explicitly and identically to all of them.
+#
+# "monokai" (dark) rather than a light Pygments style: the chat view itself
+# renders with a dark palette (Qt follows Windows' dark app theme here), so a
+# light code background would show as a glaring white box, and - worse -
+# Pygments only emits a color for tokens it specifically highlights, leaving
+# plain identifiers/punctuation uncolored to inherit whatever text color is
+# ambient; against this app's dark theme that meant near-invisible white-on-
+# near-white text. _CODE_BG/_CODE_FG give the wrapping <span> in
+# _highlight_code_lines() an explicit default color so uncolored tokens
+# inherit a color that's actually legible on a dark block, via normal CSS
+# inheritance, while Pygments' own colored spans still override it.
+_CODE_FORMATTER = HtmlFormatter(style="monokai", noclasses=True, nowrap=True)
+_CODE_BG = "#272822"
+_CODE_FG = "#f8f8f2"
+
+_CODE_BLOCK_FORMAT = QTextBlockFormat()
+_CODE_BLOCK_FORMAT.setBackground(QColor(_CODE_BG))
+_CODE_BLOCK_FORMAT.setLeftMargin(10)
+_CODE_BLOCK_FORMAT.setTopMargin(1)
+_CODE_BLOCK_FORMAT.setBottomMargin(1)
+
+_CODE_CHAR_FORMAT = QTextCharFormat()
+_CODE_CHAR_FORMAT.setFontFamilies(["Consolas", "Courier New"])
+
+
+def _highlight_code_lines(code: str, lang: str) -> list[str]:
+    lexer = None
+    if lang:
+        try:
+            lexer = get_lexer_by_name(lang)
+        except ClassNotFound:
+            lexer = None
+    if lexer is None:
+        try:
+            lexer = guess_lexer(code)
+        except ClassNotFound:
+            lexer = TextLexer()
+    highlighted = highlight(code, lexer, _CODE_FORMATTER)
+    lines = highlighted.rstrip("\n").split("\n")
+    return [f'<span style="color:{_CODE_FG};">{line}</span>' for line in lines]
+
+
+def _reset_cursor_format(cursor: QTextCursor) -> None:
+    """Start a new block with plain formatting, so a preceding code block's
+    background/monospace font can't bleed into whatever is inserted next -
+    insertBlock() otherwise inherits the current char format."""
+    cursor.insertBlock(QTextBlockFormat(), QTextCharFormat())
+
+
+class ModelLoaderThread(QThread):
+    loaded = pyqtSignal(object)
+    failed = pyqtSignal(str)
+
+    def run(self) -> None:
+        try:
+            llm = chatbot.load_model()
+        except Exception as e:
+            self.failed.emit(str(e))
+            return
+        self.loaded.emit(llm)
+
+
+class AskThread(QThread):
+    tool_call = pyqtSignal(str)
+    usage = pyqtSignal(dict)
+    finished_ok = pyqtSignal(str)
+    failed = pyqtSignal(str)
+
+    def __init__(self, llm, messages, parent=None):
+        super().__init__(parent)
+        self._llm = llm
+        self._messages = messages
+
+    def run(self) -> None:
+        try:
+            answer = chatbot.ask(
+                self._llm,
+                self._messages,
+                on_tool_call=self.tool_call.emit,
+                on_usage=self.usage.emit,
+            )
+        except Exception as e:
+            self.failed.emit(str(e))
+            return
+        self.finished_ok.emit(answer)
+
+
+class ChatWindow(QMainWindow):
+    def __init__(self):
+        super().__init__()
+        self.setWindowTitle("Dev Assist Chatbot")
+        self.resize(900, 700)
+
+        self._llm = None
+        self._history: list[tuple[str, str]] = []
+        self._pending_question = ""
+        self._ask_thread: AskThread | None = None
+
+        central = QWidget()
+        self.setCentralWidget(central)
+        layout = QVBoxLayout(central)
+
+        self.chat_view = QTextBrowser()
+        self.chat_view.setOpenExternalLinks(True)
+        self.chat_view.setFont(QFont("Segoe UI", 10))
+        layout.addWidget(self.chat_view)
+
+        input_row = QHBoxLayout()
+        self.input_box = QLineEdit()
+        self.input_box.setPlaceholderText("Ask a development question...")
+        self.input_box.returnPressed.connect(self._on_send)
+        input_row.addWidget(self.input_box)
+
+        self.send_button = QPushButton("Send")
+        self.send_button.clicked.connect(self._on_send)
+        input_row.addWidget(self.send_button)
+
+        layout.addLayout(input_row)
+
+        self._n_ctx = 0
+        self.context_status_label = QLabel("Context: --")
+        self.statusBar().addPermanentWidget(self.context_status_label)
+
+        self.gpu_status_label = QLabel()
+        self.statusBar().addPermanentWidget(self.gpu_status_label)
+        self._init_gpu_monitor()
+
+        self.statusBar().showMessage("Loading model...")
+        self._set_input_enabled(False)
+
+        self._loader_thread = ModelLoaderThread()
+        self._loader_thread.loaded.connect(self._on_model_loaded)
+        self._loader_thread.failed.connect(self._on_model_failed)
+        self._loader_thread.start()
+
+    def _init_gpu_monitor(self) -> None:
+        # llama_cpp's n_gpu_layers=-1 offloads to the first CUDA device, so
+        # that's the one whose load/VRAM is worth showing here.
+        try:
+            pynvml.nvmlInit()
+            self._gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        except pynvml.NVMLError as e:
+            self._gpu_handle = None
+            self.gpu_status_label.setText("GPU: unavailable")
+            self.gpu_status_label.setToolTip(str(e))
+            return
+
+        self._gpu_timer = QTimer(self)
+        self._gpu_timer.timeout.connect(self._update_gpu_status)
+        self._gpu_timer.start(GPU_STATUS_INTERVAL_MS)
+        self._update_gpu_status()
+
+    def _update_gpu_status(self) -> None:
+        try:
+            util = pynvml.nvmlDeviceGetUtilizationRates(self._gpu_handle)
+            mem = pynvml.nvmlDeviceGetMemoryInfo(self._gpu_handle)
+        except pynvml.NVMLError as e:
+            self._gpu_timer.stop()
+            self.gpu_status_label.setText("GPU: unavailable")
+            self.gpu_status_label.setToolTip(str(e))
+            return
+
+        used_gb = mem.used / (1024**3)
+        total_gb = mem.total / (1024**3)
+        self.gpu_status_label.setText(f"GPU {util.gpu}%  ·  VRAM {used_gb:.1f}/{total_gb:.1f} GB")
+
+    def _set_input_enabled(self, enabled: bool) -> None:
+        self.input_box.setEnabled(enabled)
+        self.send_button.setEnabled(enabled)
+        if enabled:
+            self.input_box.setFocus()
+
+    def _on_model_loaded(self, llm) -> None:
+        self._llm = llm
+        self._n_ctx = llm.n_ctx()
+        self.context_status_label.setText(f"Context: 0 / {self._n_ctx}")
+        self.statusBar().showMessage("Ready.")
+        self._append_system("Dev assist chatbot ready. Ask a development question.")
+        self._set_input_enabled(True)
+
+    def _on_model_failed(self, message: str) -> None:
+        self.statusBar().showMessage("Failed to load model.")
+        self._append_system(f"Error loading model: {message}")
+
+    def _on_send(self) -> None:
+        question = self.input_box.text().strip()
+        if not question or self._llm is None or self._ask_thread is not None:
+            return
+
+        self.input_box.clear()
+        self._append_user(question)
+        self._set_input_enabled(False)
+        self.statusBar().showMessage("Thinking...")
+
+        messages = chatbot.build_messages(self._history, question)
+        self._pending_question = question
+
+        self._ask_thread = AskThread(self._llm, messages)
+        self._ask_thread.tool_call.connect(self._on_tool_call)
+        self._ask_thread.usage.connect(self._on_usage)
+        self._ask_thread.finished_ok.connect(self._on_answer)
+        self._ask_thread.failed.connect(self._on_ask_failed)
+        self._ask_thread.start()
+
+    def _on_tool_call(self, query: str) -> None:
+        self.statusBar().showMessage(f"Searching: {query}")
+        self._append_system(f"Searching the web for: {query}")
+
+    def _on_usage(self, usage: dict) -> None:
+        total = usage["total_tokens"]
+        pct = 100 * total / self._n_ctx if self._n_ctx else 0
+        self.context_status_label.setText(f"Context: {total} / {self._n_ctx} ({pct:.0f}%)")
+
+    def _on_answer(self, answer: str) -> None:
+        self._append_assistant(answer)
+        self._history.append((self._pending_question, answer.strip()))
+        self.statusBar().showMessage("Ready.")
+        self._ask_thread = None
+        self._set_input_enabled(True)
+
+    def _on_ask_failed(self, message: str) -> None:
+        self._append_system(f"Error: {message}")
+        self.statusBar().showMessage("Ready.")
+        self._ask_thread = None
+        self._set_input_enabled(True)
+
+    def _append_user(self, text: str) -> None:
+        self._append_block("You", text, "#2b6cb0")
+
+    def _append_assistant(self, text: str) -> None:
+        self._append_markdown_block("Assistant", text, "#2f855a")
+
+    def _append_system(self, text: str) -> None:
+        self._append_block("System", text, "#718096")
+
+    def _append_block(self, label: str, text: str, color: str) -> None:
+        escaped = (
+            text.replace("&", "&amp;")
+            .replace("<", "&lt;")
+            .replace(">", "&gt;")
+            .replace("\n", "<br>")
+        )
+        self.chat_view.append(f'<p><b style="color:{color};">{label}:</b><br>{escaped}</p>')
+        self.chat_view.moveCursor(QTextCursor.MoveOperation.End)
+
+    def _append_markdown_block(self, label: str, markdown_text: str, color: str) -> None:
+        cursor = self.chat_view.textCursor()
+        cursor.movePosition(QTextCursor.MoveOperation.End)
+        # Force a fresh, format-reset paragraph before the label - moving to
+        # End only puts the cursor after the previous message's last
+        # character (inheriting its block format, e.g. a Pygments code
+        # block's background color), not on a clean new line.
+        if not self.chat_view.document().isEmpty():
+            _reset_cursor_format(cursor)
+        cursor.insertHtml(f'<b style="color:{color};">{label}:</b>')
+
+        pos = 0
+        for match in _CODE_FENCE_RE.finditer(markdown_text):
+            before = markdown_text[pos : match.start()]
+            if before.strip():
+                _reset_cursor_format(cursor)
+                fragment = QTextDocumentFragment.fromMarkdown(before, _MARKDOWN_FEATURES)
+                cursor.insertFragment(fragment)
+            lang, code = match.group(1), match.group(2)
+            for line in _highlight_code_lines(code, lang):
+                cursor.insertBlock(_CODE_BLOCK_FORMAT)
+                start = cursor.position()
+                cursor.insertHtml(line)
+                # insertHtml() builds its own char formats from the parsed
+                # HTML and ignores the cursor's format, so the monospace
+                # font has to be merged on afterwards; merging (rather than
+                # replacing) keeps Pygments' per-token colors and bold.
+                line_cursor = QTextCursor(cursor)
+                line_cursor.setPosition(start)
+                line_cursor.setPosition(cursor.position(), QTextCursor.MoveMode.KeepAnchor)
+                line_cursor.mergeCharFormat(_CODE_CHAR_FORMAT)
+            pos = match.end()
+        remainder = markdown_text[pos:]
+        if remainder.strip():
+            _reset_cursor_format(cursor)
+            fragment = QTextDocumentFragment.fromMarkdown(remainder, _MARKDOWN_FEATURES)
+            cursor.insertFragment(fragment)
+
+        # Reset the block format one more time so a code block's background/
+        # monospace font doesn't bleed into whatever message gets appended next.
+        _reset_cursor_format(cursor)
+        self.chat_view.setTextCursor(cursor)
+        self.chat_view.moveCursor(QTextCursor.MoveOperation.End)
+
+    def closeEvent(self, event) -> None:
+        # llama.cpp calls run in native code and can't be cancelled mid-call,
+        # so wait for any in-flight load/generation to finish rather than
+        # tearing down the thread while it's still running (which crashes Qt).
+        for thread in (self._loader_thread, self._ask_thread):
+            if thread is not None and thread.isRunning():
+                thread.wait()
+        if self._gpu_handle is not None:
+            pynvml.nvmlShutdown()
+        event.accept()
+
+
+def main() -> None:
+    app = QApplication(sys.argv)
+    window = ChatWindow()
+    window.show()
+    sys.exit(app.exec())
+
+
+if __name__ == "__main__":
+    main()