367 lines
14 KiB
Python
367 lines
14 KiB
Python
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}; white-space:pre;">{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("&", "&")
|
|
.replace("<", "<")
|
|
.replace(">", ">")
|
|
.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()
|