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