feat reasoning

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 <tool_call>{...}</tool_call> blocks Qwen's chat template trains
-# it to emit when it wants to invoke a tool.
-TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
+# Matches the <tool_call><function=name>...</function></tool_call> 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-<tool_call> format.
+TOOL_CALL_RE = re.compile(
+    r"<tool_call>\s*<function=([^>]+)>(.*?)</function>\s*</tool_call>", re.DOTALL
+)
+# Matches each <parameter=name>value</parameter> inside a tool call's body.
+PARAMETER_RE = re.compile(r"<parameter=([^>]+)>\n?(.*?)\n?</parameter>", 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>"
+THINK_CLOSE = "</think>"
 
 
 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 <tool_call> 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 <think>...</think>
+    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-"</think>"), 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
+    <think> 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 "<think>\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 "</think>" 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 </think> 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 <think> 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)