diff --git a/chatbot.py b/chatbot.py index 6610687..62f2a1f 100644 --- a/chatbot.py +++ b/chatbot.py @@ -135,7 +135,7 @@ def load_model() -> 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=32768, + n_ctx=24000, # -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,31 +144,56 @@ def load_model() -> Llama: ) +def _stream_content( + llm: Llama, + messages: list[dict], + tools: list[dict] | None, + on_token: Callable[[str], 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.""" + kwargs = {} + if tools is not None: + kwargs["tools"] = tools + stream = llm.create_chat_completion( + messages, + max_tokens=8192, + temperature=0.4, + top_p=0.9, + repeat_penalty=1.1, + stream=True, + **kwargs, + ) + content = "" + for chunk in stream: + piece = chunk["choices"][0]["delta"].get("content") + if piece: + content += piece + if on_token is not None: + on_token(piece) + return content + + 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, ) -> 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. If given, on_tool_call is invoked with the search query each time a tool is about to run, so a caller (CLI, GUI, ...) can surface progress. - If given, on_usage is invoked with the {prompt_tokens, completion_tokens, - total_tokens} usage dict after each model call, so a caller can track how - much of the context window is filled.""" + If given, on_usage is invoked with a {total_tokens} usage dict after each + 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.""" for _ in range(MAX_TOOL_ROUNDS): - response = llm.create_chat_completion( - messages, - tools=[WEB_SEARCH_TOOL], - max_tokens=8192, - temperature=0.4, - top_p=0.9, - repeat_penalty=1.1, - ) + content = _stream_content(llm, messages, [WEB_SEARCH_TOOL], on_token) if on_usage is not None: - on_usage(response["usage"]) - content = response["choices"][0]["message"]["content"] or "" + on_usage({"total_tokens": llm.n_tokens}) calls = extract_tool_calls(content) if not calls: return content @@ -191,13 +216,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. - response = llm.create_chat_completion( - messages, - max_tokens=8192, - temperature=0.4, - top_p=0.9, - repeat_penalty=1.1, - ) + content = _stream_content(llm, messages, None, on_token) if on_usage is not None: - on_usage(response["usage"]) - return response["choices"][0]["message"]["content"] or "" + on_usage({"total_tokens": llm.n_tokens}) + return content diff --git a/gui.py b/gui.py index 953d540..fb59115 100644 --- a/gui.py +++ b/gui.py @@ -362,6 +362,7 @@ class ModelLoaderThread(QThread): class AskThread(QThread): tool_call = pyqtSignal(str) usage = pyqtSignal(dict) + token = pyqtSignal(str) finished_ok = pyqtSignal(str) failed = pyqtSignal(str) @@ -377,6 +378,7 @@ class AskThread(QThread): self._messages, on_tool_call=self.tool_call.emit, on_usage=self.usage.emit, + on_token=self.token.emit, ) except Exception as e: self.failed.emit(str(e)) @@ -394,6 +396,11 @@ class ChatWindow(QMainWindow): self._history: list[tuple[str, str]] = [] self._pending_question = "" self._ask_thread: AskThread | None = None + # Position in chat_view's document where the in-progress streamed + # answer card starts, and a cursor into its body cell for appending + # incoming tokens; both None when no answer is currently streaming. + self._stream_start_pos: int | None = None + self._stream_body_cursor: QTextCursor | None = None central = QWidget() self.setCentralWidget(central) @@ -526,6 +533,8 @@ class ChatWindow(QMainWindow): self._append_user(question) self._set_input_enabled(False) self.statusBar().showMessage("Thinking...") + self._stream_start_pos = None + self._stream_body_cursor = None messages = chatbot.build_messages(self._history, question) self._pending_question = question @@ -533,11 +542,43 @@ class ChatWindow(QMainWindow): 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.token.connect(self._on_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.statusBar().showMessage("Generating...") + cursor = self.chat_view.textCursor() + cursor.movePosition(QTextCursor.MoveOperation.End) + self._stream_start_pos = cursor.position() + accent, bg = ROLE_STYLES["assistant"] + body = _insert_message_card(cursor, accent, bg) + body.insertHtml(f'ASSISTANT') + body.insertBlock(QTextBlockFormat(), QTextCharFormat()) + self._stream_body_cursor = body + + char_format = QTextCharFormat() + char_format.setForeground(QColor(TEXT_PRIMARY)) + self._stream_body_cursor.insertText(piece, char_format) + self.chat_view.moveCursor(QTextCursor.MoveOperation.End) + + def _discard_stream_card(self) -> None: + """Remove the in-progress streamed answer card, if any, so a + tool-call round's raw output or a failed generation doesn't linger + in the chat view.""" + if self._stream_start_pos is None: + return + cursor = self.chat_view.textCursor() + cursor.setPosition(self._stream_start_pos) + cursor.movePosition(QTextCursor.MoveOperation.End, QTextCursor.MoveMode.KeepAnchor) + cursor.removeSelectedText() + self._stream_start_pos = None + self._stream_body_cursor = None + def _on_tool_call(self, query: str) -> None: + self._discard_stream_card() self.statusBar().showMessage(f"Searching: {query}") self._append_system(f"Searching the web for: {query}") @@ -547,6 +588,10 @@ class ChatWindow(QMainWindow): self.context_status_label.setText(f"Context: {total} / {self._n_ctx} ({pct:.0f}%)") def _on_answer(self, answer: str) -> None: + # 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. + self._discard_stream_card() self._append_assistant(answer) self._history.append((self._pending_question, answer.strip())) self.statusBar().showMessage("Ready.") @@ -554,6 +599,7 @@ class ChatWindow(QMainWindow): self._set_input_enabled(True) def _on_ask_failed(self, message: str) -> None: + self._discard_stream_card() self._append_system(f"Error: {message}") self.statusBar().showMessage("Ready.") self._ask_thread = None