studio: improve GGUF tool calling accuracy and reliability

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": ""}
 

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"<script[^>]*>.*?</script[^>]*>",
+            "",
+            raw_html,
+            flags = _re.DOTALL | _re.IGNORECASE,
+        )
+        text = _re.sub(
+            r"<style[^>]*>.*?</style[^>]*>", "", 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": "<URL>"}).'
+        )
+        return text
     except Exception as e:
         return f"Search failed: {e}"