diff --git a/chatbot.py b/chatbot.py index 62f2a1f..f147476 100644 --- a/chatbot.py +++ b/chatbot.py @@ -16,7 +16,7 @@ if _cuda_path: from ddgs import DDGS # noqa: E402 from llama_cpp import Llama # noqa: E402 -MODEL_PATH = Path(__file__).parent / "Qwen2.5-Coder-14B-Instruct-Q5_K_M.gguf" +MODEL_PATH = Path(__file__).parent / "Qwen3.5-27B-IQ4_XS.gguf" SYSTEM_PROMPT = ( "You are a knowledgeable developer assistant. Answer questions about " @@ -56,9 +56,20 @@ WEB_SEARCH_TOOL = { }, } -# Matches the {...} blocks Qwen's chat template trains -# it to emit when it wants to invoke a tool. -TOOL_CALL_RE = re.compile(r"\s*(\{.*?\})\s*", re.DOTALL) +# Matches the ... blocks +# Qwen3.5's chat template trains it to emit when it wants to invoke a tool - +# an XML-style syntax, not Qwen2.5's JSON-in- format. +TOOL_CALL_RE = re.compile( + r"\s*]+)>(.*?)\s*", re.DOTALL +) +# Matches each value inside a tool call's body. +PARAMETER_RE = re.compile(r"]+)>\n?(.*?)\n?", re.DOTALL) + +# Qwen3.5 thinks by default, wrapping its reasoning trace in these tags +# before the actual answer (or tool call). _route_stream_piece splits a +# streamed response between the two so a caller can render them separately. +THINK_OPEN = "" +THINK_CLOSE = "" def web_search(query: str, max_results: int = 5) -> str: @@ -82,42 +93,15 @@ def web_search(query: str, max_results: int = 5) -> str: TOOLS = {"web_search": web_search} -def _iter_balanced_braces(text: str): - """Yield each top-level {...} substring of text, respecting nesting.""" - depth = 0 - start = None - for i, ch in enumerate(text): - if ch == "{": - if depth == 0: - start = i - depth += 1 - elif ch == "}" and depth > 0: - depth -= 1 - if depth == 0 and start is not None: - yield text[start : i + 1] - - def extract_tool_calls(text: str) -> list[dict]: calls = [] - for match in TOOL_CALL_RE.finditer(text): - try: - calls.append(json.loads(match.group(1))) - except json.JSONDecodeError: - continue - if calls: - return calls - - # This model often skips the wrapper it's instructed to use - # and just emits a bare JSON object instead - sometimes with leading - # prose alongside it (e.g. "Let's try again. {...}") - so scan for any - # JSON object in the text rather than requiring the whole message to be one. - for candidate in _iter_balanced_braces(text): - try: - call = json.loads(candidate) - except json.JSONDecodeError: - continue - if isinstance(call, dict) and call.get("name") in TOOLS: - calls.append(call) + for call_match in TOOL_CALL_RE.finditer(text): + name = call_match.group(1).strip() + arguments = { + param_match.group(1).strip(): param_match.group(2) + for param_match in PARAMETER_RE.finditer(call_match.group(2)) + } + calls.append({"name": name, "arguments": arguments}) return calls @@ -133,9 +117,11 @@ def build_messages(history: list[tuple[str, str]], question: str) -> list[dict[s def load_model() -> Llama: return Llama( model_path=str(MODEL_PATH), - # the model supports up to 32k context natively; give plenty of room - # for chat history plus a long generated response - n_ctx=24000, + # Qwen3.5-27B's hybrid attention keeps KV cache small (most layers + # carry a fixed-size state rather than one that grows with context), + # so this fits alongside the IQ4_XS weights on a 16GB card with room + # to spare - unlike a same-sized dense-attention model. + n_ctx=32768, # -1 offloads every layer to the GPU n_gpu_layers=-1, # number of tokens in the prompt that are fed into the model at a time @@ -144,14 +130,55 @@ def load_model() -> Llama: ) +def _route_stream_piece( + piece: str, + state: dict, + on_token: Callable[[str], None] | None, + on_thinking_token: Callable[[str], None] | None, +) -> None: + """Forward piece to on_thinking_token while inside a ... + block and to on_token otherwise, using state (a per-call {"in_think", + "buffer"} dict) to track which side of a tag boundary we're on and to + hold back text that might be the start of a tag split across pieces.""" + buffer = state["buffer"] + piece + while True: + tag = THINK_CLOSE if state["in_think"] else THINK_OPEN + idx = buffer.find(tag) + if idx == -1: + break + before, buffer = buffer[:idx], buffer[idx + len(tag) :] + callback = on_thinking_token if state["in_think"] else on_token + if before and callback is not None: + callback(before) + state["in_think"] = not state["in_think"] + + # The buffer may end with a prefix of the tag we're now watching for + # (e.g. a chunk boundary landing mid-""), so hold that part back + # rather than misrouting it as ordinary content. + tag = THINK_CLOSE if state["in_think"] else THINK_OPEN + keep = 0 + for length in range(min(len(tag), len(buffer)), 0, -1): + if tag.startswith(buffer[-length:]): + keep = length + break + emit, state["buffer"] = buffer[: len(buffer) - keep], buffer[len(buffer) - keep :] + callback = on_thinking_token if state["in_think"] else on_token + if emit and callback is not None: + callback(emit) + + def _stream_content( llm: Llama, messages: list[dict], tools: list[dict] | None, on_token: Callable[[str], None] | None, + on_thinking_token: Callable[[str], None] | None = None, ) -> str: """Run one streaming create_chat_completion call, forwarding each piece - of generated text to on_token as it arrives, and return the full text.""" + of generated text to on_token (or on_thinking_token, while inside a + block) as it arrives, and return the full raw text, thinking + tags included, so the caller's message history matches what the model + actually produced.""" kwargs = {} if tools is not None: kwargs["tools"] = tools @@ -165,21 +192,40 @@ def _stream_content( **kwargs, ) content = "" + # Qwen3.5's chat template appends "\n" to the prompt itself before + # generation starts (see its tokenizer.chat_template's add_generation_prompt + # block), so the model's own streamed output never contains that opening + # tag - it starts already "inside" one, and only emits "" once + # reasoning is done. + route_state = {"in_think": True, "buffer": ""} for chunk in stream: piece = chunk["choices"][0]["delta"].get("content") if piece: content += piece - if on_token is not None: - on_token(piece) + _route_stream_piece(piece, route_state, on_token, on_thinking_token) + if route_state["buffer"]: + callback = on_thinking_token if route_state["in_think"] else on_token + if callback is not None: + callback(route_state["buffer"]) return content +def _strip_thinking(text: str) -> str: + """Drop a leading reasoning trace from raw model output, the same way + the chat template splits it back out when re-rendering a past turn: + everything before the first is reasoning (there's no opening + tag to look for - see _stream_content), everything after is the answer.""" + before, sep, after = text.partition(THINK_CLOSE) + return after.lstrip("\n") if sep else before + + def ask( llm: Llama, messages: list[dict], on_tool_call: Callable[[str], None] | None = None, on_usage: Callable[[dict], None] | None = None, on_token: Callable[[str], None] | None = None, + on_thinking_token: Callable[[str], None] | None = None, ) -> str: """Run the chat loop for one question, letting the model call web_search up to MAX_TOOL_ROUNDS times before returning its final answer text. @@ -189,14 +235,16 @@ def ask( model call, so a caller can track how much of the context window is filled. If given, on_token is invoked with each piece of generated text as it streams from the model, so a caller can render the answer live - instead of waiting for it to finish.""" + instead of waiting for it to finish. If given, on_thinking_token is + invoked the same way but for text inside the model's block, + letting a caller render its reasoning trace separately from the answer.""" for _ in range(MAX_TOOL_ROUNDS): - content = _stream_content(llm, messages, [WEB_SEARCH_TOOL], on_token) + content = _stream_content(llm, messages, [WEB_SEARCH_TOOL], on_token, on_thinking_token) if on_usage is not None: on_usage({"total_tokens": llm.n_tokens}) calls = extract_tool_calls(content) if not calls: - return content + return _strip_thinking(content) messages.append({"role": "assistant", "content": content}) for call in calls: @@ -216,7 +264,7 @@ def ask( # Ran out of tool-call rounds without a plain answer; ask one more time # without offering the tool, forcing it to answer from what it has # instead of returning a leftover tool-call as if it were the answer. - content = _stream_content(llm, messages, None, on_token) + content = _stream_content(llm, messages, None, on_token, on_thinking_token) if on_usage is not None: on_usage({"total_tokens": llm.n_tokens}) - return content + return _strip_thinking(content) diff --git a/gui.py b/gui.py index 48bc345..5060ffd 100644 --- a/gui.py +++ b/gui.py @@ -6,9 +6,10 @@ 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 Qt, QThread, QTimer, pyqtSignal +from PyQt6.QtCore import Qt, QThread, QTimer, QUrl, pyqtSignal from PyQt6.QtGui import ( QColor, + QDesktopServices, QFont, QIcon, QLinearGradient, @@ -87,8 +88,14 @@ ROLE_STYLES = { "user": (CORAL_GLOW, "#3a2c26"), "assistant": (SILVER, "#2f3342"), "system": (BLUE_SLATE, "#262a36"), + "thinking": (TEXT_SECONDARY, "#282b38"), } +# Prefix for the fake anchor href used to make a thinking block's header +# clickable (e.g. "think:think-3"), intercepted in _on_chat_anchor_clicked +# rather than let QTextBrowser try to navigate to it. +_THINK_ANCHOR_PREFIX = "think:" + _STYLESHEET = f""" QMainWindow {{ background-color: {BG_APP}; @@ -376,6 +383,7 @@ class AskThread(QThread): tool_call = pyqtSignal(str) usage = pyqtSignal(dict) token = pyqtSignal(str) + thinking_token = pyqtSignal(str) finished_ok = pyqtSignal(str) failed = pyqtSignal(str) @@ -392,6 +400,7 @@ class AskThread(QThread): on_tool_call=self.tool_call.emit, on_usage=self.usage.emit, on_token=self.token.emit, + on_thinking_token=self.thinking_token.emit, ) except Exception as e: self.failed.emit(str(e)) @@ -414,6 +423,22 @@ class ChatWindow(QMainWindow): # incoming tokens; both None when no answer is currently streaming. self._stream_start_pos: int | None = None self._stream_body_cursor: QTextCursor | None = None + # Whether a real (non-newline) answer character has been inserted + # yet for the in-progress card - the model's text after + # typically opens with a blank line or two, which _on_token trims + # rather than showing as a gap at the top of the answer bubble. + self._answer_started = False + + # Thinking blocks (one per model reasoning round), keyed by an id + # referenced from the clickable "think:" anchor in their header. + # Each holds the full reasoning text, whether it's currently shown + # expanded, a fixed cursor marking where its content region starts, + # and a live cursor tracking where that region currently ends - see + # _render_thinking(). self._stream_thinking_id names whichever block + # (if any) is still actively streaming in the current round. + self._thinking_blocks: dict[str, dict] = {} + self._thinking_seq = 0 + self._stream_thinking_id: str | None = None central = QWidget() self.setCentralWidget(central) @@ -448,7 +473,12 @@ class ChatWindow(QMainWindow): self.chat_view = QTextBrowser() self.chat_view.setObjectName("chatView") - self.chat_view.setOpenExternalLinks(True) + # Links are handled by hand in _on_chat_anchor_clicked (real URLs get + # opened externally there too) rather than left to QTextBrowser, + # since its own link handling doesn't know about the fake + # "think:" anchors used to toggle a thinking block. + self.chat_view.setOpenLinks(False) + self.chat_view.anchorClicked.connect(self._on_chat_anchor_clicked) self.chat_view.setFont(QFont("Segoe UI Variable Text", 10)) layout.addWidget(self.chat_view) @@ -551,6 +581,8 @@ class ChatWindow(QMainWindow): self.statusBar().showMessage("Thinking...") self._stream_start_pos = None self._stream_body_cursor = None + self._answer_started = False + self._stream_thinking_id = None messages = chatbot.build_messages(self._history, question) self._pending_question = question @@ -559,12 +591,14 @@ class ChatWindow(QMainWindow): self._ask_thread.tool_call.connect(self._on_tool_call) self._ask_thread.usage.connect(self._on_usage) self._ask_thread.token.connect(self._on_token) + self._ask_thread.thinking_token.connect(self._on_thinking_token) self._ask_thread.finished_ok.connect(self._on_answer) self._ask_thread.failed.connect(self._on_ask_failed) self._ask_thread.start() def _on_token(self, piece: str) -> None: if self._stream_body_cursor is None: + self._collapse_open_thinking() self.statusBar().showMessage("Generating...") cursor = self.chat_view.textCursor() cursor.movePosition(QTextCursor.MoveOperation.End) @@ -574,6 +608,13 @@ class ChatWindow(QMainWindow): body.insertHtml(f'ASSISTANT') body.insertBlock(QTextBlockFormat(), QTextCharFormat()) self._stream_body_cursor = body + self._answer_started = False + + if not self._answer_started: + piece = piece.lstrip("\n") + if not piece: + return + self._answer_started = True char_format = QTextCharFormat() char_format.setForeground(QColor(TEXT_PRIMARY)) @@ -593,7 +634,95 @@ class ChatWindow(QMainWindow): self._stream_start_pos = None self._stream_body_cursor = None + def _on_thinking_token(self, piece: str) -> None: + if self._stream_thinking_id is None: + cursor = self.chat_view.textCursor() + cursor.movePosition(QTextCursor.MoveOperation.End) + accent, bg = ROLE_STYLES["thinking"] + body = _insert_message_card(cursor, accent, bg) + body.insertHtml(f'THINKING') + body.insertBlock(QTextBlockFormat(), QTextCharFormat()) + + self._thinking_seq += 1 + block_id = f"think-{self._thinking_seq}" + content_start = QTextCursor(body) + # Without this, inserting the header right at content_start's own + # position (both here and on every later toggle) would drag it + # forward along with that insertion, per QTextCursor's default - + # it needs to stay pinned at the true start of the region instead. + content_start.setKeepPositionOnInsert(True) + self._thinking_blocks[block_id] = { + "text": "", + "expanded": True, + "content_start": content_start, + "cursor": body, + } + self._insert_thinking_header(body, block_id, expanded=True) + body.insertBlock(QTextBlockFormat(), QTextCharFormat()) + self._stream_thinking_id = block_id + + block = self._thinking_blocks[self._stream_thinking_id] + block["text"] += piece + italic_format = QTextCharFormat() + italic_format.setForeground(QColor(TEXT_SECONDARY)) + italic_format.setFontItalic(True) + block["cursor"].insertText(piece, italic_format) + self.chat_view.moveCursor(QTextCursor.MoveOperation.End) + + def _insert_thinking_header(self, cursor: QTextCursor, block_id: str, expanded: bool) -> None: + arrow = "▼" if expanded else "▶" + cursor.insertHtml( + f'{arrow} Thinking' + ) + + def _render_thinking(self, block_id: str, expanded: bool) -> None: + """Rewrite a thinking block's header and body to match expanded, + replacing everything from its fixed content_start cursor through + wherever its content currently ends (tracked by its own cursor, + which - like any QTextCursor - keeps up with edits made elsewhere in + the document, so this stays correct even if other messages have + since been appended below it).""" + block = self._thinking_blocks[block_id] + cursor = block["cursor"] + old_end = cursor.position() + cursor.setPosition(block["content_start"].position()) + cursor.setPosition(old_end, QTextCursor.MoveMode.KeepAnchor) + cursor.removeSelectedText() + + self._insert_thinking_header(cursor, block_id, expanded) + if expanded: + cursor.insertBlock(QTextBlockFormat(), QTextCharFormat()) + italic_format = QTextCharFormat() + italic_format.setForeground(QColor(TEXT_SECONDARY)) + italic_format.setFontItalic(True) + cursor.insertText(block["text"], italic_format) + block["expanded"] = expanded + self.chat_view.moveCursor(QTextCursor.MoveOperation.End) + + def _collapse_open_thinking(self) -> None: + if self._stream_thinking_id is not None: + self._render_thinking(self._stream_thinking_id, expanded=False) + self._stream_thinking_id = None + + def _on_chat_anchor_clicked(self, url: QUrl) -> None: + href = url.toString() + if not href.startswith(_THINK_ANCHOR_PREFIX): + QDesktopServices.openUrl(url) + return + # A generation currently in flight may be inserting content through + # this same block's cursor further down; leave the block alone + # rather than risk racing that live cursor. + if self._ask_thread is not None: + return + block_id = href[len(_THINK_ANCHOR_PREFIX) :] + block = self._thinking_blocks.get(block_id) + if block is None: + return + self._render_thinking(block_id, expanded=not block["expanded"]) + def _on_tool_call(self, query: str) -> None: + self._collapse_open_thinking() self._discard_stream_card() self.statusBar().showMessage(f"Searching: {query}") self._append_system(f"Searching the web for: {query}") @@ -604,6 +733,10 @@ class ChatWindow(QMainWindow): self.context_status_label.setText(f"Context: {total} / {self._n_ctx} ({pct:.0f}%)") def _on_answer(self, answer: str) -> None: + # Normally already collapsed by the first answer token (_on_token), + # but a response that hit max_tokens mid- would leave this + # block open with nothing to trigger the collapse - catch that here. + self._collapse_open_thinking() # The streamed card holds the raw, plain-text answer; swap it out for # a fully markdown-rendered and syntax-highlighted one now that the # complete text is available. @@ -615,6 +748,7 @@ class ChatWindow(QMainWindow): self._set_input_enabled(True) def _on_ask_failed(self, message: str) -> None: + self._collapse_open_thinking() self._discard_stream_card() self._append_system(f"Error: {message}") self.statusBar().showMessage("Ready.")