223 lines
8.2 KiB
Python
223 lines
8.2 KiB
Python
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, 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>{...}</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=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
|
|
n_batch=1024,
|
|
verbose=False,
|
|
)
|
|
|
|
|
|
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 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):
|
|
content = _stream_content(llm, messages, [WEB_SEARCH_TOOL], on_token)
|
|
if on_usage is not None:
|
|
on_usage({"total_tokens": llm.n_tokens})
|
|
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.
|
|
content = _stream_content(llm, messages, None, on_token)
|
|
if on_usage is not None:
|
|
on_usage({"total_tokens": llm.n_tokens})
|
|
return content
|