python-inference/chatbot.py

import json
import os
import re
import time
from collections.abc import Callable
from pathlib import Path

# CUDA 13.x moved its runtime DLLs to bin/x64, but llama_cpp's loader only
# checks CUDA_PATH/bin, so point it at the right folder before importing.
_cuda_path = os.environ.get("CUDA_PATH")
if _cuda_path:
    _cuda_bin_x64 = Path(_cuda_path) / "bin" / "x64"
    if _cuda_bin_x64.is_dir():
        os.environ["PATH"] = str(_cuda_bin_x64) + os.pathsep + os.environ["PATH"]

from ddgs import DDGS  # noqa: E402
from llama_cpp import Llama  # noqa: E402

MODEL_PATH = Path(__file__).parent / "Qwen3.5-27B-IQ4_XS.gguf"

SYSTEM_PROMPT = (
    "You are a knowledgeable developer assistant. Answer questions about "
    "programming languages, libraries, frameworks, tools, debugging, and "
    "software design with clear, accurate, and concise explanations. Use "
    "code snippets when they help illustrate an answer, and always place "
    "them in fenced Markdown code blocks with a language tag right after "
    "the opening backticks (e.g. ```python), never inline in a paragraph "
    "or left unfenced. Use Markdown throughout: numbered or bulleted "
    "lists for steps, **bold** for emphasis, and inline `code` spans for "
    "identifiers, filenames, and commands mentioned in prose. Use the "
    "web_search tool when a question depends on current information you "
    "may not know (recent releases, changelogs, current best practices, "
    "error messages you don't recognize)."
)

# Number of prior user/assistant exchanges to keep in the prompt. Bounds the
# prompt size so a long chat session doesn't overflow the model's context.
MAX_HISTORY_TURNS = 10

# Caps how many times the model can call a tool while answering a single
# question, so a confused model can't loop searches forever.
MAX_TOOL_ROUNDS = 3

WEB_SEARCH_TOOL = {
    "type": "function",
    "function": {
        "name": "web_search",
        "description": "Search the web and return brief result snippets.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The search query."},
            },
            "required": ["query"],
        },
    },
}

# 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:
    # ddgs is a metasearch client and defaults to aggregating several backends
    # (Mojeek, Brave, ...); pin it to duckduckgo specifically, since some of
    # the others time out from this machine. The duckduckgo backend itself
    # occasionally returns zero results transiently, so retry once.
    for attempt in range(2):
        try:
            results = DDGS(timeout=10).text(query, max_results=max_results, backend="duckduckgo")
            break
        except Exception as e:
            if attempt == 1:
                return f"Search failed: {e}"
            time.sleep(1)
    if not results:
        return "No results found."
    return "\n\n".join(f"{r['title']}\n{r['href']}\n{r['body']}" for r in results)


TOOLS = {"web_search": web_search}


def extract_tool_calls(text: str) -> list[dict]:
    calls = []
    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


def build_messages(history: list[tuple[str, str]], question: str) -> list[dict[str, str]]:
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user_turn, assistant_turn in history[-MAX_HISTORY_TURNS:]:
        messages.append({"role": "user", "content": user_turn})
        messages.append({"role": "assistant", "content": assistant_turn})
    messages.append({"role": "user", "content": question})
    return messages


def load_model() -> Llama:
    return Llama(
        model_path=str(MODEL_PATH),
        # 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
        n_batch=1024,
        verbose=False,
    )


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 (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
    stream = llm.create_chat_completion(
        messages,
        max_tokens=8192,
        temperature=0.4,
        top_p=0.9,
        repeat_penalty=1.1,
        stream=True,
        **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
            _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.
    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 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. 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, 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 _strip_thinking(content)

        messages.append({"role": "assistant", "content": content})
        for call in calls:
            name = call.get("name")
            args = call.get("arguments") or {}
            if isinstance(args, str):
                args = json.loads(args)
            func = TOOLS.get(name)
            if func is None:
                result = f"Unknown tool: {name}"
            else:
                if on_tool_call is not None:
                    on_tool_call(args.get("query", ""))
                result = func(**args)
            messages.append({"role": "tool", "content": result})

    # 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, on_thinking_token)
    if on_usage is not None:
        on_usage({"total_tokens": llm.n_tokens})
    return _strip_thinking(content)