studio: improve GGUF tool calling accuracy and reliability #4700
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||
|---|---|---|---|---|---|---|---|---|
|
|
@@ -57,16 +57,23 @@ def _get_workdir(session_id: str | None = None) -> str: | |||||||
| "type": "function", | ||||||||
| "function": { | ||||||||
| "name": "web_search", | ||||||||
| "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": [], | ||||||||
| }, | ||||||||
| }, | ||||||||
| } | ||||||||
|
|
@@ -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", ""), | ||||||||
| 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,148 @@ def execute_tool( | |||||||
| return f"Unknown tool: {name}" | ||||||||
|
|
||||||||
|
|
||||||||
| _MAX_PAGE_CHARS = 16000 # limit fetched page text | ||||||||
| _MAX_FETCH_BYTES = _MAX_PAGE_CHARS * 4 + 1 # cap raw download size | ||||||||
|
|
||||||||
|
|
||||||||
| def _is_public_host(hostname: str, port: int) -> tuple[bool, str]: | ||||||||
| """Resolve *hostname* and reject private/loopback/link-local addresses.""" | ||||||||
| 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}" | ||||||||
|
|
||||||||
| 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 True, "" | ||||||||
|
|
||||||||
|
|
||||||||
| def _fetch_page_text( | ||||||||
|
Contributor
There was a problem hiding this comment. The To improve clarity and avoid potential confusion, consider making the default |
||||||||
| 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})." | ||||||||
|
Comment on lines
+207
to
+209
There was a problem hiding this comment.
Useful? React with 👍 / 👎. |
||||||||
| if not parsed.hostname: | ||||||||
| return "Blocked: URL is missing a hostname." | ||||||||
|
|
||||||||
| ok, reason = _is_public_host( | ||||||||
| parsed.hostname, | ||||||||
| parsed.port or (443 if parsed.scheme == "https" else 80), | ||||||||
| ) | ||||||||
| if not ok: | ||||||||
| return reason | ||||||||
|
|
||||||||
| try: | ||||||||
| import urllib.request | ||||||||
| from urllib.parse import urljoin | ||||||||
|
|
||||||||
| # Disable auto-redirect so we can validate each hop for SSRF | ||||||||
| class _NoRedirect(urllib.request.HTTPRedirectHandler): | ||||||||
| def redirect_request(self, req, fp, code, msg, headers, newurl): | ||||||||
| return None | ||||||||
|
Comment on lines
+227
to
+228
There was a problem hiding this comment.
Returning Useful? React with 👍 / 👎. |
||||||||
|
|
||||||||
| opener = urllib.request.build_opener(_NoRedirect) | ||||||||
| max_bytes = max_chars * 4 + 1 | ||||||||
| current_url = url | ||||||||
|
|
||||||||
| for _hop in range(5): | ||||||||
| req = urllib.request.Request( | ||||||||
| current_url, | ||||||||
| headers = {"User-Agent": "UnslothStudio/1.0"}, | ||||||||
| ) | ||||||||
| resp = opener.open(req, timeout = timeout) | ||||||||
|
There was a problem hiding this comment.
The private-range check is done on a preflight DNS lookup, but the subsequent Useful? React with 👍 / 👎. |
||||||||
| if resp.status in (301, 302, 303, 307, 308): | ||||||||
| location = resp.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." | ||||||||
| ok2, reason2 = _is_public_host( | ||||||||
| rp.hostname, | ||||||||
| rp.port or (443 if rp.scheme == "https" else 80), | ||||||||
| ) | ||||||||
| if not ok2: | ||||||||
| return reason2 | ||||||||
| continue | ||||||||
| # Success -- read capped body | ||||||||
| raw_bytes = resp.read(max_bytes) | ||||||||
| break | ||||||||
| else: | ||||||||
| return "Failed to fetch URL: too many redirects." | ||||||||
|
|
||||||||
| raw_html = raw_bytes.decode("utf-8", errors = "replace") | ||||||||
|
Contributor
There was a problem hiding this comment. Hardcoding
Suggested change
|
||||||||
| 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() | ||||||||
|
Comment on lines
+296
to
+306
Contributor
There was a problem hiding this comment. Using regular expressions to strip HTML is fragile and can be easily bypassed or produce poor results with complex or malformed HTML. While this is a fallback, consider making |
||||||||
|
|
||||||||
| 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 +310,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', '')}" | ||||||||
| ) | ||||||||
| 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}" | ||||||||
|
|
||||||||
|
|
||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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( | ||
| 25, | ||
| ge = 0, | ||
| description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).", | ||
| ) | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Importing
hashlibinside the function on every call is inefficient. This import should be moved to the top of the file to follow standard Python practices and improve performance.References