feat streaming the inference in realtime

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

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'<span style="color:{accent}; font-weight:600; font-size:8.5pt;">ASSISTANT</span>')
+            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