stable

chatbot.py

@@ -0,0 +1,198 @@
+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 / "Qwen2.5-Coder-14B-Instruct-Q5_K_M.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. 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>{...}</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)
+
+
+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 _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)
+    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),
+        # the model supports up to 32k context natively; give plenty of room
+        # for chat history plus a long generated response
+        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 ask(
+    llm: Llama,
+    messages: list[dict],
+    on_tool_call: Callable[[str], None] | None = None,
+    on_usage: Callable[[dict], 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."""
+    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,
+        )
+        if on_usage is not None:
+            on_usage(response["usage"])
+        content = response["choices"][0]["message"]["content"] or ""
+        calls = extract_tool_calls(content)
+        if not calls:
+            return 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.
+    response = llm.create_chat_completion(
+        messages,
+        max_tokens=8192,
+        temperature=0.4,
+        top_p=0.9,
+        repeat_penalty=1.1,
+    )
+    if on_usage is not None:
+        on_usage(response["usage"])
+    return response["choices"][0]["message"]["content"] or ""