diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index f5361a6e8c0..c1f87ff9369 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -10,7 +10,9 @@ import atexit import contextlib +import hashlib import json +import re import struct import structlog from loggers import get_logger @@ -2120,7 +2122,7 @@ def generate_chat_completion_with_tools( stop: Optional[list[str]] = None, cancel_event: Optional[threading.Event] = None, enable_thinking: Optional[bool] = None, - max_tool_iterations: int = 10, + max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, @@ -2172,6 +2174,29 @@ def _strip_tool_markup(text: str, *, final: bool = False) -> str: ) _MAX_BUFFER_CHARS = 32 + # ── Duplicate tool-call detection ──────────────────────── + # Track recent (tool_name, arguments) hashes to detect loops + # where the model repeats the exact same call. Retries after + # a transient failure are allowed (only block when the previous + # identical call succeeded). + _tool_call_history: list[tuple[str, bool]] = [] # (key, failed) + + def _tool_call_key(name: str, args: dict) -> str: + raw = json.dumps({"t": name, "a": args}, sort_keys = True) + return hashlib.md5(raw.encode()).hexdigest() + + def _is_duplicate_call(name: str, args: dict) -> bool: + """Block if the immediately previous call was identical and succeeded.""" + if not _tool_call_history: + return False + key = _tool_call_key(name, args) + last_key, last_failed = _tool_call_history[-1] + return last_key == key and not last_failed + + def _record_tool_call(name: str, args: dict, failed: bool) -> None: + key = _tool_call_key(name, args) + _tool_call_history.append((key, failed)) + for iteration in range(max_tool_iterations): if cancel_event is not None and cancel_event.is_set(): return @@ -2568,6 +2593,11 @@ def _strip_tool_markup(text: str, *, final: bool = False) -> str: # Merge accumulated metrics from prior tool # iterations so they are not silently dropped. yield {"type": "status", "text": ""} + if content_accum: + # Strip leaked tool-call XML before yielding + content_accum = _strip_tool_markup( + content_accum, final = True + ) if content_accum: yield {"type": "content", "text": content_accum} _fu = _iter_usage or {} @@ -2661,16 +2691,27 @@ def _strip_tool_markup(text: str, *, final: bool = False) -> str: "arguments": arguments, } - _effective_timeout = ( - None if tool_call_timeout >= 9999 else tool_call_timeout - ) - result = execute_tool( - tool_name, - arguments, - cancel_event = cancel_event, - timeout = _effective_timeout, - session_id = session_id, - ) + # ── Duplicate call detection ────────────── + if _is_duplicate_call(tool_name, arguments): + result = ( + "You already made this exact call. " + "Do not repeat the same tool call. " + "Try a different approach: fetch a URL " + "from previous results, use Python to " + "process data you already have, or " + "provide your final answer now." + ) + else: + _effective_timeout = ( + None if tool_call_timeout >= 9999 else tool_call_timeout + ) + result = execute_tool( + tool_name, + arguments, + cancel_event = cancel_event, + timeout = _effective_timeout, + session_id = session_id, + ) yield { "type": "tool_end", @@ -2679,10 +2720,32 @@ def _strip_tool_markup(text: str, *, final: bool = False) -> str: "result": result, } + # Nudge model to try a different approach on errors + _error_prefixes = ( + "Error", + "Search failed", + "Execution error", + "Blocked:", + "Exit code", + "Failed to fetch", + "Failed to resolve", + "No query provided", + ) + _is_error = isinstance(result, str) and result.lstrip().startswith( + _error_prefixes + ) + _record_tool_call(tool_name, arguments, failed = _is_error) + _result_content = result + if _is_error: + _result_content = ( + result + "\n\nThe tool call encountered an issue. " + "Please try a different approach or rephrase your request." + ) + tool_msg = { "role": "tool", "name": tool_name, - "content": result, + "content": _result_content, } tool_call_id = tc.get("id") if tool_call_id: @@ -2699,6 +2762,22 @@ def _strip_tool_markup(text: str, *, final: bool = False) -> str: return raise + # ── Tool iteration cap reached -- synthesize final answer ── + # The model used all iterations without producing a final text + # response. Inject a nudge so the final streaming pass produces + # a useful answer instead of continuing to request tools. + if max_tool_iterations > 0: + conversation.append( + { + "role": "user", + "content": ( + "You have used all available tool calls. Based on " + "everything you have found so far, provide your final " + "answer now. Do not call any more tools." + ), + } + ) + # Clear status yield {"type": "status", "text": ""} diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 55bfa095f9d..65302fe2f34 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -57,16 +57,23 @@ def _get_workdir(session_id: str | None = None) -> str: "type": "function", "function": { "name": "web_search", - "description": "Search the web for current information, recent events, or facts you are uncertain about.", + "description": ( + "Search the web and fetch page content. Returns snippets for all results. " + "Use the url parameter to fetch full page text from a specific URL." + ), "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": "The search query", - } + }, + "url": { + "type": "string", + "description": "A URL to fetch full page content from (instead of searching). Use this to read a page found in search results.", + }, }, - "required": ["query"], + "required": [], }, }, } @@ -131,7 +138,11 @@ def execute_tool( ) effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout if name == "web_search": - return _web_search(arguments.get("query", ""), timeout = effective_timeout) + return _web_search( + arguments.get("query", ""), + url = arguments.get("url"), + timeout = effective_timeout, + ) if name == "python": return _python_exec( arguments.get("code", ""), cancel_event, effective_timeout, session_id @@ -143,9 +154,180 @@ def execute_tool( return f"Unknown tool: {name}" -def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT) -> str: - """Search the web using DuckDuckGo and return formatted results.""" - if not query.strip(): +_MAX_PAGE_CHARS = 16000 # limit fetched page text +_MAX_FETCH_BYTES = _MAX_PAGE_CHARS * 4 + 1 # cap raw download size + + +def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]: + """Resolve *hostname*, reject non-public IPs, return a pinned IP string. + + Returns ``(ok, reason_or_empty, resolved_ip)``. The caller should + connect to *resolved_ip* (with a ``Host`` header) to prevent DNS + rebinding between validation and the actual fetch. + """ + import ipaddress + import socket + + try: + infos = socket.getaddrinfo(hostname, port, type = socket.SOCK_STREAM) + except OSError as e: + return False, f"Failed to resolve host: {e}", "" + + if not infos: + return False, f"Failed to resolve host: no addresses for {hostname!r}", "" + + for *_, sockaddr in infos: + ip = ipaddress.ip_address(sockaddr[0]) + if ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_multicast + or ip.is_reserved + or ip.is_unspecified + ): + return False, f"Blocked: refusing to fetch non-public address {ip}.", "" + + # Return the first resolved address for pinning + first_ip = infos[0][4][0] + return True, "", first_ip + + +def _fetch_page_text( + url: str, max_chars: int = _MAX_PAGE_CHARS, timeout: int = 30 +) -> str: + """Fetch a URL and return plain text content (HTML tags stripped). + + Blocks private/loopback/link-local targets (SSRF protection) and caps + the download size to avoid unbounded memory usage. + """ + import re as _re + from urllib.parse import urlparse + + parsed = urlparse(url) + if parsed.scheme not in ("http", "https"): + return f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r})." + if not parsed.hostname: + return "Blocked: URL is missing a hostname." + + port = parsed.port or (443 if parsed.scheme == "https" else 80) + ok, reason, pinned_ip = _validate_and_resolve_host(parsed.hostname, port) + if not ok: + return reason + + try: + import urllib.request + from urllib.error import HTTPError as _HTTPError + from urllib.parse import urljoin, urlunparse + + # Disable auto-redirect so we can validate each hop for SSRF. + # urllib raises HTTPError for 3xx when the handler returns None, + # so we catch that and extract the Location header manually. + class _NoRedirect(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): + return None + + opener = urllib.request.build_opener(_NoRedirect) + max_bytes = max_chars * 4 + 1 + current_url = url + current_host = parsed.hostname + + for _hop in range(5): + # Pin to the validated IP to prevent DNS rebinding. + # Rewrite the URL to use the IP and set the Host header. + cp = urlparse(current_url) + ip_netloc = f"{pinned_ip}:{cp.port}" if cp.port else pinned_ip + pinned_url = urlunparse(cp._replace(netloc = ip_netloc)) + + req = urllib.request.Request( + pinned_url, + headers = { + "User-Agent": "UnslothStudio/1.0", + "Host": current_host, + }, + ) + try: + resp = opener.open(req, timeout = timeout) + except _HTTPError as e: + if e.code not in (301, 302, 303, 307, 308): + return ( + f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}" + ) + location = e.headers.get("Location") + if not location: + return "Failed to fetch URL: redirect missing Location header." + current_url = urljoin(current_url, location) + rp = urlparse(current_url) + if rp.scheme not in ("http", "https") or not rp.hostname: + return "Blocked: redirect target is not a valid http/https URL." + rp_port = rp.port or (443 if rp.scheme == "https" else 80) + ok2, reason2, pinned_ip = _validate_and_resolve_host( + rp.hostname, + rp_port, + ) + if not ok2: + return reason2 + current_host = rp.hostname + continue + # Success -- read capped body + raw_bytes = resp.read(max_bytes) + break + else: + return "Failed to fetch URL: too many redirects." + + charset = resp.headers.get_content_charset() or "utf-8" + raw_html = raw_bytes.decode(charset, errors = "replace") + except _HTTPError as e: + return f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}" + except Exception as e: + return f"Failed to fetch URL: {e}" + + # Convert HTML to text -- prefer html2text for clean markdown output + try: + import html2text as _h2t + + converter = _h2t.HTML2Text() + converter.ignore_links = False + converter.ignore_images = True + converter.body_width = 0 # no wrapping + text = converter.handle(raw_html).strip() + except ImportError: + # Fallback: regex-based stripping + text = _re.sub( + r"]*>.*?]*>", + "", + raw_html, + flags = _re.DOTALL | _re.IGNORECASE, + ) + text = _re.sub( + r"]*>.*?]*>", "", text, flags = _re.DOTALL | _re.IGNORECASE + ) + text = _re.sub(r"<[^>]+>", " ", text) + text = _re.sub(r"\s+", " ", text).strip() + + if not text: + return "(page returned no readable text)" + if len(text) > max_chars: + text = text[:max_chars] + f"\n\n... (truncated, {len(text)} chars total)" + return text + + +def _web_search( + query: str, + max_results: int = 5, + timeout: int = _EXEC_TIMEOUT, + url: str | None = None, +) -> str: + """Search the web using DuckDuckGo and return formatted results. + + If ``url`` is provided, fetches that page directly instead of searching. + """ + # Direct URL fetch mode + if url and url.strip(): + fetch_timeout = 60 if timeout is None else min(timeout, 60) + return _fetch_page_text(url.strip(), timeout = fetch_timeout) + + if not query or not query.strip(): return "No query provided." try: from ddgs import DDGS @@ -160,7 +342,13 @@ def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT) f"URL: {r.get('href', '')}\n" f"Snippet: {r.get('body', '')}" ) - return "\n\n---\n\n".join(parts) + text = "\n\n---\n\n".join(parts) + text += ( + "\n\n---\n\nIMPORTANT: These are only short snippets. " + "To get the full page content, call web_search with " + 'the url parameter (e.g. {"url": ""}).' + ) + return text except Exception as e: return f"Search failed: {e}" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index aabfba9b3ad..77f70b9bd65 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -344,7 +344,7 @@ class ChatCompletionRequest(BaseModel): description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", ) max_tool_calls_per_message: Optional[int] = Field( - 10, + 25, ge = 0, description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).", ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1a942560593..9bce371775a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -86,8 +86,15 @@ def _friendly_error(exc: Exception) -> str: import wave import base64 import numpy as np +from datetime import date as _date router = APIRouter() + +# Regex for stripping leaked tool-call XML from assistant messages/stream +_TOOL_XML_RE = _re.compile( + r".*?|.*?", + _re.DOTALL, +) logger = get_logger(__name__) @@ -1078,6 +1085,68 @@ async def audio_input_stream(): else: tools_to_use = ALL_TOOLS + # ── Tool-use system prompt nudge ────────────────────── + _tool_names = {t["function"]["name"] for t in tools_to_use} + _has_web = "web_search" in _tool_names + _has_code = "python" in _tool_names or "terminal" in _tool_names + + _date_line = f"The current date is {_date.today().isoformat()}." + + _web_tips = ( + "When you search and find a relevant URL in the results, " + "fetch its full content by calling web_search with the url parameter. " + "Do not repeat the same search query. If a search returns " + "no useful results, try rephrasing or fetching a result URL directly." + ) + _code_tips = ( + "Use code execution for math, calculations, data processing, " + "or to parse and analyze information from tool results." + ) + + if _has_web and _has_code: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "tools rather than answering from memory. " + + _web_tips + + " " + + _code_tips + ) + elif _has_code: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "code execution rather than answering from memory. " + _code_tips + ) + elif _has_web: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "web search for up-to-date or uncertain factual " + "information rather than answering from memory. " + _web_tips + ) + else: + _nudge = "" + + if _nudge: + # Append nudge to system prompt (preserve user's prompt) + if system_prompt: + system_prompt = system_prompt.rstrip() + "\n\n" + _nudge + else: + system_prompt = _nudge + # Rebuild gguf_messages with updated system prompt + gguf_messages = [] + if system_prompt: + gguf_messages.append({"role": "system", "content": system_prompt}) + gguf_messages.extend(chat_messages) + + # ── Strip stale tool-call XML from conversation history ─ + for _msg in gguf_messages: + if _msg.get("role") == "assistant" and isinstance( + _msg.get("content"), str + ): + _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + def gguf_generate_with_tools(): return llama_backend.generate_chat_completion_with_tools( messages = gguf_messages, @@ -1096,7 +1165,7 @@ def gguf_generate_with_tools(): else True, max_tool_iterations = payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None - else 10, + else 25, tool_call_timeout = payload.tool_call_timeout if payload.tool_call_timeout is not None else 300, @@ -1158,9 +1227,13 @@ async def gguf_tool_stream(): continue # "content" type -- cumulative text - cumulative = event.get("text", "") - new_text = cumulative[len(prev_text) :] - prev_text = cumulative + # Sanitize the full cumulative then diff against + # the last sanitized snapshot so cross-chunk XML + # tags are handled correctly. + raw_cumulative = event.get("text", "") + clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative) + new_text = clean_cumulative[len(prev_text) :] + prev_text = clean_cumulative if not new_text: continue chunk = ChatCompletionChunk( diff --git a/studio/backend/tests/tool_calling_benchmark_results.md b/studio/backend/tests/tool_calling_benchmark_results.md new file mode 100644 index 00000000000..c2b0687895e --- /dev/null +++ b/studio/backend/tests/tool_calling_benchmark_results.md @@ -0,0 +1,62 @@ +# GGUF Tool Calling Benchmark Results + +Prompt: "List and categorize all the songs that charted #3 on the Billboard Hot 100 in 2015." +10 runs per configuration, web search + code execution + thinking enabled. +GPU: NVIDIA B200, CUDA_VISIBLE_DEVICES=2. + +Ground truth: 4 songs peaked at #3 in 2015 -- "Love Me like You Do" (Ellie Goulding), "Earned It" (The Weeknd), "Watch Me" (Silento), "Drag Me Down" (One Direction). + +## Cartesian Grid: Model x Quant x KV Cache + +| Model | Quant | KV Cache | OK/10 | Avg Time | Avg Tools | XML Leaks | URL Fetch | Peak3 Avg | All 4/4 | Best Songs | +|-------|-------|----------|-------|----------|-----------|-----------|-----------|-----------|---------|------------| +| 4B | UD-Q4_K_XL | f16 | 10/10 | 9.8s | 3.5 | 0/10 | 4/10 | 0.8/4 | 2/10 | 9 | +| 4B | UD-Q4_K_XL | bf16 | 10/10 | 10.6s | 4.5 | 0/10 | 4/10 | 0.4/4 | 1/10 | 5 | +| 4B | Q8_0 | f16 | 10/10 | 4.9s | 2.4 | 0/10 | 8/10 | 0.4/4 | 1/10 | 5 | +| 4B | Q8_0 | bf16 | 10/10 | 8.0s | 3.0 | 0/10 | 5/10 | 0.0/4 | 0/10 | 0 | +| 9B | UD-Q4_K_XL | f16 | 10/10 | 6.7s | 2.0 | 0/10 | 5/10 | 0.0/4 | 0/10 | 3 | +| 9B | UD-Q4_K_XL | bf16 | 9/10 | 49.5s | 2.4 | 0/10 | 5/10 | 0.0/4 | 0/10 | 1 | +| 9B | Q8_0 | f16 | 10/10 | 7.4s | 2.5 | 0/10 | 5/10 | 0.0/4 | 0/10 | 2 | +| 9B | Q8_0 | bf16 | 10/10 | 10.4s | 2.7 | 0/10 | 6/10 | 1.0/4 | 2/10 | 15 | +| **27B** | **UD-Q4_K_XL** | **bf16** | **9/10** | **131.1s** | **13.8** | **0/10** | **7/10** | **2.7/4** | **6/10** | **27** | +| 27B | UD-Q4_K_XL | f16 | 7/10 | 201.6s | 14.1 | 0/10 | 8/10 | 2.0/4 | 5/10 | 26 | +| 27B | Q8_0 | f16 | 4/10 | 312.5s | 16.0 | 1/10 | 10/10 | 2.4/4 | 6/10 | 28 | +| 27B | Q8_0 | bf16 | 5/10 | 258.4s | 16.5 | 2/10 | 10/10 | 0.9/4 | 1/10 | 27 | +| 35B-A3B | UD-Q4_K_XL | f16 | 3/10 | 353.6s | 14.7 | 1/10 | 6/10 | 1.2/4 | 3/10 | 27 | +| 35B-A3B | UD-Q4_K_XL | bf16 | 3/10 | 356.2s | 17.2 | 1/10 | 8/10 | 1.6/4 | 4/10 | 27 | +| 35B-A3B | Q8_0 | f16 | 2/10 | 372.1s | 17.6 | 1/10 | 7/10 | 1.2/4 | 3/10 | 26 | +| 35B-A3B | Q8_0 | bf16 | 6/10 | 267.7s | 17.5 | 1/10 | 8/10 | 2.4/4 | 6/10 | 27 | + +**Column definitions:** +- **Peak3 Avg**: Average number of correct peak-#3 songs found per run (out of 4) +- **All 4/4**: Runs where all 4 correct songs were identified +- **Best Songs**: Maximum number of Billboard 2015 songs mentioned in any single run (out of 31 tracked) +- **URL Fetch**: Runs where the model used web_search with `url` parameter to fetch full page content + +## Key Findings + +1. **27B UD-Q4_K_XL + bf16 KV is the sweet spot.** 6/10 runs found all 4 correct songs, 0 XML leaks, 131s average. Best balance of accuracy, speed, and reliability. + +2. **Larger models use tools more effectively.** 27B and 35B-A3B models used 13-17 tool calls per query (vs 2-4 for 4B/9B), performing multiple searches and URL fetches to find the answer. + +3. **27B Q8_0 had the highest raw accuracy (6/10 all-4/4) but lower reliability** -- only 4/10 OK runs due to timeouts on long agentic chains. The UD-Q4_K_XL quant is more practical. + +4. **4B models were fastest (5-10s) but least accurate.** They occasionally found all 4 songs (2/10 best case) when they happened to fetch the right Wikipedia page. + +5. **9B was surprisingly weaker than 4B on this task.** It used fewer tool calls and rarely extracted song data from fetched pages. The 9B model may need higher temperature or different prompting for this specific task type. + +6. **35B-A3B had reliability issues.** Most runs timed out or errored due to slow per-token generation with many tool iterations. When it completed (2-6/10 OK), accuracy was comparable to 27B. + +7. **bf16 KV cache had mixed effects.** For 27B it improved both speed (131s vs 202s) and accuracy (6/10 vs 5/10 all-4/4). For smaller models it had no consistent benefit. + +8. **XML leaks are nearly eliminated.** 0/10 for all 4B and 9B configs, and only 1-2/10 for the largest models (which generate much more text in complex agentic loops). + +## Before vs After (4B UD-Q4_K_XL, f16 KV) + +| Metric | Before Changes | After Changes | +|--------|---------------|---------------| +| XML leaks | 10/10 | 0/10 | +| URL fetches | 0/10 | 4/10 | +| Peak3 accuracy | 0.0/4 | 0.8/4 | +| Runs with all 4 songs | 0/10 | 2/10 | +| Avg time | 12.3s | 9.8s | diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index ca1044b3dc4..8cea234f21b 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -224,7 +224,7 @@ export const useChatRuntimeStore = create((set) => ({ toolStatus: null, generatingStatus: null, autoHealToolCalls: loadBool(AUTO_HEAL_TOOL_CALLS_KEY, true), - maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 10), + maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 25), toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5), kvCacheDtype: null, loadedKvCacheDtype: null,