From 9c95148045a8bc1aa1fec4c6e0e13e72546a34e4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 18 Mar 2026 08:28:02 -0700 Subject: [PATCH] Fix tool call parsing, add tool outputs panel and UI improvements (#4416) * Add elapsed timer to tool status pill in Studio Show a count-up seconds timer (0s, 1s, 2s, ...) next to the tool status text in the composer area. Helps users gauge how long a tool call (web search, code execution) has been running. Timer resets when a new tool starts and disappears when all tools finish. * Fix tool call parsing, add tool outputs panel and reasoning copy button Backend: - Rewrite tool call XML parser to use balanced-brace JSON extraction instead of greedy regex, fixing truncation on nested braces in code/JSON arguments - Handle optional closing tags (, , ) that models frequently omit - Support bare tags without wrapper - Strip tool call markup from streamed content so raw XML never leaks into the chat UI - Use a persistent ~/studio_sandbox/ working directory for tool execution so files persist across calls within a session - Emit tool_start/tool_end SSE events so the frontend can display tool inputs and outputs Frontend: - Add collapsible "Tool Outputs" panel below assistant messages showing each tool call's input and output with copy buttons - Add copy button to reasoning blocks - Add elapsed timer to tool status pill - Update project URLs in pyproject.toml (http -> https, add docs link) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add interactive HTML preview with fullscreen toggle for code blocks HTML code fences now render an interactive sandboxed iframe preview below the syntax-highlighted code, similar to how SVG fences show an image preview. The iframe uses sandbox="allow-scripts" to allow JavaScript execution while blocking access to the parent page. Includes a fullscreen toggle (enlarge/minimize button) that expands the preview into a viewport overlay, dismissible via button, Escape key, or backdrop click. A streaming placeholder prevents partial HTML from rendering mid-stream. * Add tool call settings: auto-heal toggle, max iterations, timeout Add three user-configurable tool call settings to the Studio Settings panel: - Auto Heal Tool Calls: toggle to control fallback XML parsing of malformed tool calls from model output (default: on) - Max Tool Calls Per Message: slider 0-40 + Max to cap tool call iterations per message (default: 10) - Max Tool Call Duration: slider 1-30 minutes + Max to set per-tool-call execution timeout (default: 5 minutes) All settings persist to localStorage and flow through the full stack: frontend store -> API request -> Pydantic model -> route -> llama_cpp -> tools. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix tool call timeout: respect no-limit and apply to web search - Use a sentinel to distinguish timeout=None (no limit) from the default (300s). Previously None was silently replaced with _EXEC_TIMEOUT. - Pass the configured timeout to DDGS() for web searches so the setting applies uniformly to all tool types. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add input validation bounds and per-thread sandbox isolation - Add ge=0 constraint to max_tool_calls_per_message (rejects negative values) - Add ge=1 constraint to tool_call_timeout (minimum 1 second) - Thread session_id from frontend through backend to tool execution - Scope sandbox directories per conversation: ~/studio_sandbox/{thread_id}/ - Backwards compatible: API callers without session_id use ~/studio_sandbox/ * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix non-monotonic streaming and Python temp script path - Split tool markup stripping into closed-only (mid-stream) and full (final flush) to prevent cumulative text from shrinking mid-stream - Enforce monotonicity: only emit when cleaned text grows, so the proxy's delta logic (cumulative[len(prev_text):]) never breaks - Place Python temp scripts in the sandbox workdir instead of /tmp so sys.path[0] points to the sandbox and cross-call imports work * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Sanitize session_id to prevent path traversal in sandbox Strip path separators and parent-dir references from session_id before using it as a directory name. Verify the resolved path stays under ~/studio_sandbox/ as a second guard. * feat(chat): proper assistant-ui tool call UIs with sources Replace custom metadata-based ToolOutputsGroup with native assistant-ui tool-call content parts. Backend SSE tool_start/tool_end events now emit proper { type: "tool-call" } parts from the adapter, enabling per-tool UIs registered via tools.by_name in MessagePrimitive.Parts. - Web search: Globe icon, Source badges with favicons, auto-collapse when LLM starts responding - Python: Code icon, syntax-highlighted code via Streamdown/shiki, output block with copy - Terminal: Terminal icon, command in trigger, output with copy - ToolGroup wraps consecutive tool calls (skips for single calls) - Sources component renders URL badges at end of message - Flattened code block CSS (single border, no nested boxes) * fix(inference): respect empty enabled_tools allowlist `if payload.enabled_tools:` is falsy for [], falling through to ALL_TOOLS. Use `is not None` so an explicit empty list disables all tools as intended. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Shine1i --- pyproject.toml | 4 +- studio/backend/core/inference/llama_cpp.py | 245 ++++++++++++++---- studio/backend/core/inference/tools.py | 141 +++++++--- studio/backend/models/inference.py | 18 ++ studio/backend/routes/inference.py | 16 +- studio/frontend/bun.lock | 22 +- studio/frontend/package.json | 4 +- .../src/components/assistant-ui/badge.tsx | 67 +++++ .../components/assistant-ui/markdown-text.tsx | 107 +++++++- .../src/components/assistant-ui/reasoning.tsx | 55 +++- .../src/components/assistant-ui/sources.tsx | 137 ++++++++++ .../src/components/assistant-ui/thread.tsx | 33 ++- .../components/assistant-ui/tool-fallback.tsx | 38 ++- .../components/assistant-ui/tool-group.tsx | 230 ++++++++++++++++ .../assistant-ui/tool-ui-python.tsx | 136 ++++++++++ .../assistant-ui/tool-ui-terminal.tsx | 103 ++++++++ .../assistant-ui/tool-ui-web-search.tsx | 122 +++++++++ .../src/features/chat/api/chat-adapter.ts | 85 +++++- .../src/features/chat/api/chat-api.ts | 6 + .../src/features/chat/chat-settings-sheet.tsx | 70 +++++ .../chat/stores/chat-runtime-store.ts | 48 ++++ .../frontend/src/features/chat/types/api.ts | 5 + studio/frontend/src/index.css | 19 +- 23 files changed, 1587 insertions(+), 124 deletions(-) create mode 100644 studio/frontend/src/components/assistant-ui/badge.tsx create mode 100644 studio/frontend/src/components/assistant-ui/sources.tsx create mode 100644 studio/frontend/src/components/assistant-ui/tool-group.tsx create mode 100644 studio/frontend/src/components/assistant-ui/tool-ui-python.tsx create mode 100644 studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx create mode 100644 studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx diff --git a/pyproject.toml b/pyproject.toml index d9898c6a95..3e48da7c1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1154,8 +1154,8 @@ rocm711-torch2100 = [ ] [project.urls] -homepage = "http://www.unsloth.ai" -documentation = "https://github.com/unslothai/unsloth" +homepage = "https://unsloth.ai" +documentation = "https://unsloth.ai/docs" repository = "https://github.com/unslothai/unsloth" [tool.ruff] diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7bb32a81f8..8054bd2a19 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1173,50 +1173,113 @@ class LlamaCppBackend: Handles formats like: {"name":"web_search","arguments":{"query":"..."}} ... - Closing tag is optional (models sometimes omit it). + Closing tags (, , ) are all optional + since models frequently omit them. """ import re tool_calls = [] - # Pattern 1: JSON inside tags (closing tag optional) - for match in re.finditer( - r"\s*(\{.*?\})\s*(?:)?", content, re.DOTALL - ): - try: - obj = json.loads(match.group(1)) - tc = { - "id": f"call_{len(tool_calls)}", - "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, - } - if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps( - tc["function"]["arguments"] - ) - tool_calls.append(tc) - except (json.JSONDecodeError, ValueError): - pass + + # Pattern 1: JSON inside tags. + # Use balanced-brace extraction that skips braces inside JSON strings. + for m in re.finditer(r"\s*\{", content): + brace_start = m.end() - 1 # position of the opening { + depth, i = 0, brace_start + in_string = False + while i < len(content): + ch = content[i] + if in_string: + if ch == "\\" and i + 1 < len(content): + i += 2 # skip escaped character + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + break + i += 1 + if depth == 0: + json_str = content[brace_start : i + 1] + try: + obj = json.loads(json_str) + tc = { + "id": f"call_{len(tool_calls)}", + "type": "function", + "function": { + "name": obj.get("name", ""), + "arguments": obj.get("arguments", {}), + }, + } + if isinstance(tc["function"]["arguments"], dict): + tc["function"]["arguments"] = json.dumps( + tc["function"]["arguments"] + ) + tool_calls.append(tc) + except (json.JSONDecodeError, ValueError): + pass # Pattern 2: XML-style value - # Closing optional + # All closing tags optional -- models frequently omit , + # , and/or . if not tool_calls: - for match in re.finditer( - r"\s*(.*?)\s*(?:)?", - content, - re.DOTALL, - ): - func_name = match.group(1) - params_text = match.group(2) + # Step 1: Find all positions and extract their bodies. + # Body boundary: use only or next as a boundary because + # code parameter values can contain that literal string. + # After extracting, we trim a trailing if present. + func_starts = list(re.finditer(r"\s*", content)) + for idx, fm in enumerate(func_starts): + func_name = fm.group(1) + body_start = fm.end() + # Hard boundaries: next + next_func = ( + func_starts[idx + 1].start() + if idx + 1 < len(func_starts) + else len(content) + ) + end_tag = re.search(r"", content[body_start:]) + if end_tag: + body_end = body_start + end_tag.start() + else: + body_end = len(content) + body_end = min(body_end, next_func) + body = content[body_start:body_end] + # Trim trailing if present (it's the real closing tag) + body = re.sub(r"\s*\s*$", "", body) + + # Step 2: Extract parameters from body. + # For single-parameter functions (the common case: code, command, + # query), use body end as the only boundary to avoid false matches + # on inside code strings. arguments = {} - for param_match in re.finditer( - r"\s*(.*?)\s*", - params_text, - re.DOTALL, - ): - arguments[param_match.group(1)] = param_match.group(2) + param_starts = list(re.finditer(r"\s*", body)) + if len(param_starts) == 1: + # Single parameter: value is everything from after the tag + # to end of body, trimming any trailing . + pm = param_starts[0] + val = body[pm.end() :] + val = re.sub(r"\s*\s*$", "", val) + arguments[pm.group(1)] = val.strip() + else: + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + # Value ends at next if present + val = re.sub(r"\s*\s*$", "", val) + arguments[param_name] = val.strip() + tc = { "id": f"call_{len(tool_calls)}", "type": "function", @@ -1531,7 +1594,10 @@ class LlamaCppBackend: stop: Optional[list[str]] = None, cancel_event: Optional[threading.Event] = None, enable_thinking: Optional[bool] = None, - max_tool_iterations: int = 5, + max_tool_iterations: int = 10, + auto_heal_tool_calls: bool = True, + tool_call_timeout: int = 300, + session_id: Optional[str] = None, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -1596,16 +1662,45 @@ class LlamaCppBackend: tool_calls = message.get("tool_calls") # Fallback: detect tool calls embedded as XML/text in content - # Some models output XML instead of structured tool_calls + # Some models output XML instead of structured tool_calls, + # or bare tags without wrapper. content_text = message.get("content", "") or "" - if not tool_calls and "" in content_text: + if ( + auto_heal_tool_calls + and not tool_calls + and ("" in content_text or " blocks since they + # can contain arbitrary content including code. import re + # Strip ... blocks (greedy inside) content_text = re.sub( - r".*?(?:|$)", + r".*?", + "", + content_text, + flags = re.DOTALL, + ) + # Strip unterminated ... to end + content_text = re.sub( + r".*$", + "", + content_text, + flags = re.DOTALL, + ) + # Strip bare ... blocks + content_text = re.sub( + r".*?", + "", + content_text, + flags = re.DOTALL, + ) + # Strip unterminated bare to end + content_text = re.sub( + r".*$", "", content_text, flags = re.DOTALL, @@ -1632,7 +1727,10 @@ class LlamaCppBackend: try: arguments = json.loads(raw_args) except (json.JSONDecodeError, ValueError): - arguments = {"query": raw_args} + if auto_heal_tool_calls: + arguments = {"query": raw_args} + else: + arguments = {"raw": raw_args} else: arguments = raw_args @@ -1659,10 +1757,33 @@ class LlamaCppBackend: status_text = f"Calling: {tool_name}" yield {"type": "status", "text": status_text} + # Emit tool_start so the frontend can record inputs + yield { + "type": "tool_start", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "arguments": arguments, + } + # Execute the tool - result = execute_tool( - tool_name, arguments, cancel_event = cancel_event + _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, + ) + + # Emit tool_end so the frontend can record outputs + yield { + "type": "tool_end", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "result": result, + } # Append tool result to conversation tool_msg = { @@ -1714,7 +1835,30 @@ class LlamaCppBackend: if stop: stream_payload["stop"] = stop + import re as _re_final + + # Closed blocks only -- safe to strip mid-stream without shrinking later. + _TOOL_CLOSED_PATTERNS = [ + _re_final.compile(r".*?", _re_final.DOTALL), + _re_final.compile(r".*?", _re_final.DOTALL), + ] + # Open-ended patterns strip from an opening tag to end-of-string. + # Only applied on the final flush to avoid non-monotonic shrinking. + _TOOL_ALL_PATTERNS = _TOOL_CLOSED_PATTERNS + [ + _re_final.compile(r".*$", _re_final.DOTALL), + _re_final.compile(r".*$", _re_final.DOTALL), + ] + + def _strip_tool_markup(text: str, *, final: bool = False) -> str: + if not auto_heal_tool_calls: + return text + patterns = _TOOL_ALL_PATTERNS if final else _TOOL_CLOSED_PATTERNS + for pat in patterns: + text = pat.sub("", text) + return text.strip() if final else text + cumulative = "" + _last_emitted = "" in_thinking = False has_content_tokens = False reasoning_text = "" @@ -1746,7 +1890,12 @@ class LlamaCppBackend: if in_thinking: if has_content_tokens: cumulative += "" - yield {"type": "content", "text": cumulative} + yield { + "type": "content", + "text": _strip_tool_markup( + cumulative, final = True + ), + } else: cumulative = reasoning_text yield {"type": "content", "text": cumulative} @@ -1776,7 +1925,11 @@ class LlamaCppBackend: cumulative += "" in_thinking = False cumulative += token - yield {"type": "content", "text": cumulative} + cleaned = _strip_tool_markup(cumulative) + # Only emit when cleaned text grows (monotonic). + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned + yield {"type": "content", "text": cleaned} except json.JSONDecodeError: logger.debug( f"Skipping malformed SSE line: {line[:100]}" diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index ee89dd2c0b..159facabe9 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -16,12 +16,42 @@ import sys import tempfile import threading +from loggers import get_logger from unsloth_zoo.rl_environments import check_signal_escape_patterns +logger = get_logger(__name__) + _EXEC_TIMEOUT = 300 # 5 minutes _MAX_OUTPUT_CHARS = 8000 # truncate long output _BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"} +# Per-session working directories so each chat thread gets its own sandbox. +# Falls back to a shared ~/studio_sandbox/ for API callers without a session_id. +_workdirs: dict[str, str] = {} + + +def _get_workdir(session_id: str | None = None) -> str: + """Return (and lazily create) a persistent working directory for tool execution.""" + global _workdirs + key = session_id or "_default" + if key not in _workdirs or not os.path.isdir(_workdirs[key]): + home = os.path.expanduser("~") + sandbox_root = os.path.join(home, "studio_sandbox") + if session_id: + # Sanitize: strip path separators and parent-dir references + safe_id = os.path.basename(session_id.replace("..", "")) + if not safe_id: + safe_id = "_invalid" + workdir = os.path.join(sandbox_root, safe_id) + # Verify resolved path stays under sandbox root + if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)): + workdir = os.path.join(sandbox_root, "_invalid") + else: + workdir = sandbox_root + os.makedirs(workdir, exist_ok = True) + _workdirs[key] = workdir + return _workdirs[key] + WEB_SEARCH_TOOL = { "type": "function", @@ -80,25 +110,47 @@ TERMINAL_TOOL = { ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL] -def execute_tool(name: str, arguments: dict, cancel_event = None) -> str: - """Execute a tool by name with the given arguments. Returns result as a string.""" +_TIMEOUT_UNSET = object() + + +def execute_tool( + name: str, + arguments: dict, + cancel_event = None, + timeout: int | None = _TIMEOUT_UNSET, + session_id: str | None = None, +) -> str: + """Execute a tool by name with the given arguments. Returns result as a string. + + ``timeout``: int sets per-call limit in seconds, ``None`` means no limit, + unset (default) uses ``_EXEC_TIMEOUT`` (300 s). + ``session_id``: optional thread/session ID for per-conversation sandbox isolation. + """ + logger.info( + f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}" + ) + effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout if name == "web_search": - return _web_search(arguments.get("query", "")) + return _web_search(arguments.get("query", ""), timeout = effective_timeout) if name == "python": - return _python_exec(arguments.get("code", ""), cancel_event) + return _python_exec( + arguments.get("code", ""), cancel_event, effective_timeout, session_id + ) if name == "terminal": - return _bash_exec(arguments.get("command", ""), cancel_event) + return _bash_exec( + arguments.get("command", ""), cancel_event, effective_timeout, session_id + ) return f"Unknown tool: {name}" -def _web_search(query: str, max_results: int = 5) -> str: +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(): return "No query provided." try: from ddgs import DDGS - results = DDGS().text(query, max_results = max_results) + results = DDGS(timeout = timeout).text(query, max_results = max_results) if not results: return "No results found." parts = [] @@ -147,7 +199,12 @@ def _truncate(text: str, limit: int = _MAX_OUTPUT_CHARS) -> str: return text -def _python_exec(code: str, cancel_event = None) -> str: +def _python_exec( + code: str, + cancel_event = None, + timeout: int = _EXEC_TIMEOUT, + session_id: str | None = None, +) -> str: """Execute Python code in a subprocess sandbox.""" if not code or not code.strip(): return "No code provided." @@ -158,8 +215,11 @@ def _python_exec(code: str, cancel_event = None) -> str: return error tmp_path = None + workdir = _get_workdir(session_id) try: - fd, tmp_path = tempfile.mkstemp(suffix = ".py", prefix = "studio_exec_") + fd, tmp_path = tempfile.mkstemp( + suffix = ".py", prefix = "studio_exec_", dir = workdir + ) with os.fdopen(fd, "w") as f: f.write(code) @@ -168,7 +228,7 @@ def _python_exec(code: str, cancel_event = None) -> str: stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, - cwd = tempfile.gettempdir(), + cwd = workdir, ) # Spawn cancel watcher if we have a cancel event @@ -179,11 +239,11 @@ def _python_exec(code: str, cancel_event = None) -> str: watcher.start() try: - output, _ = proc.communicate(timeout = _EXEC_TIMEOUT) + output, _ = proc.communicate(timeout = timeout) except subprocess.TimeoutExpired: proc.kill() proc.communicate() - return _truncate("Execution timed out after 5 minutes.") + return _truncate(f"Execution timed out after {timeout} seconds.") if cancel_event is not None and cancel_event.is_set(): return "Execution cancelled." @@ -203,7 +263,12 @@ def _python_exec(code: str, cancel_event = None) -> str: pass -def _bash_exec(command: str, cancel_event = None) -> str: +def _bash_exec( + command: str, + cancel_event = None, + timeout: int = _EXEC_TIMEOUT, + session_id: str | None = None, +) -> str: """Execute a bash command in a subprocess sandbox.""" if not command or not command.strip(): return "No command provided." @@ -215,35 +280,35 @@ def _bash_exec(command: str, cancel_event = None) -> str: return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}" try: - with tempfile.TemporaryDirectory() as tmpdir: - proc = subprocess.Popen( - ["bash", "-c", command], - stdout = subprocess.PIPE, - stderr = subprocess.STDOUT, - text = True, - cwd = tmpdir, + workdir = _get_workdir(session_id) + proc = subprocess.Popen( + ["bash", "-c", command], + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + cwd = workdir, + ) + + if cancel_event is not None: + watcher = threading.Thread( + target = _cancel_watcher, args = (proc, cancel_event), daemon = True ) + watcher.start() - if cancel_event is not None: - watcher = threading.Thread( - target = _cancel_watcher, args = (proc, cancel_event), daemon = True - ) - watcher.start() + try: + output, _ = proc.communicate(timeout = timeout) + except subprocess.TimeoutExpired: + proc.kill() + proc.communicate() + return _truncate(f"Execution timed out after {timeout} seconds.") - try: - output, _ = proc.communicate(timeout = _EXEC_TIMEOUT) - except subprocess.TimeoutExpired: - proc.kill() - proc.communicate() - return _truncate("Execution timed out after 5 minutes.") + if cancel_event is not None and cancel_event.is_set(): + return "Execution cancelled." - if cancel_event is not None and cancel_event.is_set(): - return "Execution cancelled." - - result = output or "" - if proc.returncode != 0: - result = f"Exit code {proc.returncode}:\n{result}" - return _truncate(result) if result.strip() else "(no output)" + result = output or "" + if proc.returncode != 0: + result = f"Exit code {proc.returncode}:\n{result}" + return _truncate(result) if result.strip() else "(no output)" except Exception as e: return f"Execution error: {e}" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 5d17d2d0be..41a942d217 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -318,6 +318,24 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] List of enabled tool names (e.g. ['web_search', 'python', 'terminal']). If None, all tools are enabled.", ) + auto_heal_tool_calls: Optional[bool] = Field( + True, + description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", + ) + max_tool_calls_per_message: Optional[int] = Field( + 10, + ge = 0, + description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).", + ) + tool_call_timeout: Optional[int] = Field( + 300, + ge = 1, + description = "[x-unsloth] Timeout in seconds for each tool call execution (9999 = no limit).", + ) + session_id: Optional[str] = Field( + None, + description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.", + ) # ── Streaming response chunks ──────────────────────────────────── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d50f53ba8d..4c98dc6d24 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1028,7 +1028,7 @@ async def openai_chat_completions( if use_tools: from core.inference.tools import ALL_TOOLS - if payload.enabled_tools: + if payload.enabled_tools is not None: tools_to_use = [ t for t in ALL_TOOLS @@ -1050,6 +1050,16 @@ async def openai_chat_completions( presence_penalty = payload.presence_penalty, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, + auto_heal_tool_calls = payload.auto_heal_tool_calls + if payload.auto_heal_tool_calls is not None + else True, + max_tool_iterations = payload.max_tool_calls_per_message + if payload.max_tool_calls_per_message is not None + else 10, + tool_call_timeout = payload.tool_call_timeout + if payload.tool_call_timeout is not None + else 300, + session_id = payload.session_id, ) _tool_sentinel = object() @@ -1093,6 +1103,10 @@ async def openai_chat_completions( yield f"data: {status_data}\n\n" continue + if event["type"] in ("tool_start", "tool_end"): + yield f"data: {json.dumps(event)}\n\n" + continue + # "content" type -- cumulative text cumulative = event.get("text", "") new_text = cumulative[len(prev_text) :] diff --git a/studio/frontend/bun.lock b/studio/frontend/bun.lock index 6ac0c76470..7e3b0ac51e 100644 --- a/studio/frontend/bun.lock +++ b/studio/frontend/bun.lock @@ -5,7 +5,7 @@ "": { "name": "unsloth-theme", "dependencies": { - "@assistant-ui/react": "^0.12.17", + "@assistant-ui/react": "^0.12.19", "@assistant-ui/react-markdown": "^0.12.3", "@assistant-ui/react-streamdown": "^0.1.2", "@base-ui/react": "^1.2.0", @@ -43,7 +43,7 @@ "framer-motion": "^11.18.2", "js-yaml": "^4.1.1", "katex": "^0.16.28", - "lucide-react": "^0.575.0", + "lucide-react": "^0.577.0", "mammoth": "^1.11.0", "motion": "^12.34.0", "next": "^16.1.6", @@ -88,17 +88,17 @@ "@antfu/ni": ["@antfu/ni@25.0.0", "", { "dependencies": { "ansis": "^4.0.0", "fzf": "^0.5.2", "package-manager-detector": "^1.3.0", "tinyexec": "^1.0.1" }, "bin": { "na": "bin/na.mjs", "ni": "bin/ni.mjs", "nr": "bin/nr.mjs", "nci": "bin/nci.mjs", "nlx": "bin/nlx.mjs", "nun": "bin/nun.mjs", "nup": "bin/nup.mjs" } }, "sha512-9q/yCljni37pkMr4sPrI3G4jqdIk074+iukc5aFJl7kmDCCsiJrbZ6zKxnES1Gwg+i9RcDZwvktl23puGslmvA=="], - "@assistant-ui/core": ["@assistant-ui/core@0.1.5", "", { "dependencies": { "assistant-stream": "^0.3.5", "nanoid": "^5.1.6" }, "peerDependencies": { "@assistant-ui/store": "^0.2.2", "@assistant-ui/tap": "^0.5.2", "@types/react": "*", "assistant-cloud": "^0.1.21", "react": "^18 || ^19", "zustand": "^5.0.11" }, "optionalPeers": ["@types/react", "assistant-cloud", "react", "zustand"] }, "sha512-kLqFbRULZvE+hIwxGz705BW3QYhfwiVaVWoolfTGYkg+4xwah1PGuH0zqjXP5AMADtz+L69Lp+LX0xU9MQZ0DA=="], + "@assistant-ui/core": ["@assistant-ui/core@0.1.7", "", { "dependencies": { "assistant-stream": "^0.3.6", "nanoid": "^5.1.6" }, "peerDependencies": { "@assistant-ui/store": "^0.2.3", "@assistant-ui/tap": "^0.5.3", "@types/react": "*", "assistant-cloud": "^0.1.22", "react": "^18 || ^19", "zustand": "^5.0.11" }, "optionalPeers": ["@types/react", "assistant-cloud", "react", "zustand"] }, "sha512-219T42ihVOicbJXZLWgD2CW5Bylg9Nk7geC331X4RfJxTDYlm2zIjViGlGaqfj6URXBp6kMulO2BTUrHGmAvdw=="], - "@assistant-ui/react": ["@assistant-ui/react@0.12.17", "", { "dependencies": { "@assistant-ui/core": "^0.1.5", "@assistant-ui/store": "^0.2.2", "@assistant-ui/tap": "^0.5.2", "@radix-ui/primitive": "^1.1.3", "@radix-ui/react-compose-refs": "^1.1.2", "@radix-ui/react-context": "^1.1.3", "@radix-ui/react-primitive": "^2.1.4", "@radix-ui/react-use-callback-ref": "^1.1.1", "@radix-ui/react-use-escape-keydown": "^1.1.1", "assistant-cloud": "^0.1.21", "assistant-stream": "^0.3.4", "nanoid": "^5.1.6", "radix-ui": "^1.4.3", "react-textarea-autosize": "^8.5.9", "zod": "^4.3.6", "zustand": "^5.0.11" }, "peerDependencies": { "@types/react": "*", "@types/react-dom": "*", "react": "^18 || ^19", "react-dom": "^18 || ^19" }, "optionalPeers": ["@types/react", "@types/react-dom"] }, "sha512-t4Z8LatD3LQrtURLaYPG47r4iG7UQgkdoi5YEv+EhzvYiG8I7kAyV4SbnFH6sXPrnleV4IpBHAd8Wc7ynkQtsw=="], + "@assistant-ui/react": ["@assistant-ui/react@0.12.19", "", { "dependencies": { "@assistant-ui/core": "^0.1.7", "@assistant-ui/store": "^0.2.3", "@assistant-ui/tap": "^0.5.3", "@radix-ui/primitive": "^1.1.3", "@radix-ui/react-compose-refs": "^1.1.2", "@radix-ui/react-context": "^1.1.3", "@radix-ui/react-primitive": "^2.1.4", "@radix-ui/react-use-callback-ref": "^1.1.1", "@radix-ui/react-use-escape-keydown": "^1.1.1", "assistant-cloud": "^0.1.22", "assistant-stream": "^0.3.6", "nanoid": "^5.1.6", "radix-ui": "^1.4.3", "react-textarea-autosize": "^8.5.9", "zod": "^4.3.6", "zustand": "^5.0.11" }, "peerDependencies": { "@types/react": "*", "@types/react-dom": "*", "react": "^18 || ^19", "react-dom": "^18 || ^19" }, "optionalPeers": ["@types/react", "@types/react-dom"] }, "sha512-scAf0o8cwjuHT9Y44EFGXcE2y6BSmpeMvt0NxOn8+Y/HBlNttQMLNvrM0p2AjacXCUufagiafAnWybzBV3nKEQ=="], "@assistant-ui/react-markdown": ["@assistant-ui/react-markdown@0.12.4", "", { "dependencies": { "@radix-ui/react-primitive": "^2.1.4", "@radix-ui/react-use-callback-ref": "^1.1.1", "classnames": "^2.5.1", "react-markdown": "^10.1.0" }, "peerDependencies": { "@assistant-ui/react": "^0.12.11", "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react"] }, "sha512-6TD9guiuLJxJoOwSjNHUYAVma2ctDCG9uypUqKHE0OUhDwTDD3NsMvTnQ0n0Lh8nnCEwVglOwKKlSEYpV7SnWA=="], "@assistant-ui/react-streamdown": ["@assistant-ui/react-streamdown@0.1.3", "", { "dependencies": { "rehype-harden": "^1.1.7", "rehype-raw": "^7.0.0", "rehype-sanitize": "^6.0.0", "streamdown": "^2.1.0" }, "peerDependencies": { "@assistant-ui/react": "^0.12.11", "@streamdown/cjk": "^1.0.0", "@streamdown/code": "^1.0.0", "@streamdown/math": "^1.0.0", "@streamdown/mermaid": "^1.0.0", "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@streamdown/cjk", "@streamdown/code", "@streamdown/math", "@streamdown/mermaid", "@types/react"] }, "sha512-n1UCjXQ3svmDtJBMJj/vXqz/BqAQBuy7myrXeymz2tD9l+ENQgqu2JY5ir3J19juJTe5lsi/P3+tOJ2C1jc/nw=="], - "@assistant-ui/store": ["@assistant-ui/store@0.2.2", "", { "dependencies": { "use-effect-event": "^2.0.3" }, "peerDependencies": { "@assistant-ui/tap": "^0.5.2", "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react"] }, "sha512-JzQseWFp3UmbByBSWQmiGi/bz5jbfru04hIgb2DJBpnnTyns8Zl+8wDPnwiYGF/6SA+IzTg5M0V1wf77rwU0dA=="], + "@assistant-ui/store": ["@assistant-ui/store@0.2.3", "", { "dependencies": { "use-effect-event": "^2.0.3" }, "peerDependencies": { "@assistant-ui/tap": "^0.5.3", "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react"] }, "sha512-daStbgSQiX7+csqK6Cvo7A8p8UZkTCSMxBHxbhJvwrlVbp7BRJWTxq3U3rpTkSGIar23SXIyVRRfXU8VW7pswA=="], - "@assistant-ui/tap": ["@assistant-ui/tap@0.5.2", "", { "peerDependencies": { "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react", "react"] }, "sha512-w6gXhr+mF6cPG6ZCnkqV4kkOHzR+Fb+52S4T34PnrH0cs8l2Gqlwo/kB9BcB9fGmjwL7izdwubQ7t2VBhWpz/Q=="], + "@assistant-ui/tap": ["@assistant-ui/tap@0.5.3", "", { "peerDependencies": { "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react", "react"] }, "sha512-wy06ksqF2LfFxe4JXy31Ns89N/be1Dy3c+mG363cFHFp3CbLkRu8CrCN2SQSgCkXt628E+D8QyzqdBcl9kD4NQ=="], "@babel/code-frame": ["@babel/code-frame@7.29.0", "", { "dependencies": { "@babel/helper-validator-identifier": "^7.28.5", "js-tokens": "^4.0.0", "picocolors": "^1.1.1" } }, "sha512-9NhCeYjq9+3uxgdtp20LSiJXJvN0FeCtNGpJxuMFZ1Kv3cWUNb6DOhJwUvcVCzKGR66cw4njwM6hrJLqgOwbcw=="], @@ -846,7 +846,7 @@ "aria-hidden": ["aria-hidden@1.2.6", "", { "dependencies": { "tslib": "^2.0.0" } }, "sha512-ik3ZgC9dY/lYVVM++OISsaYDeg1tb0VtP5uL3ouh1koGOaUMDPpbFIei4JkFimWUFPn90sbMNMXQAIVOlnYKJA=="], - "assistant-cloud": ["assistant-cloud@0.1.21", "", { "dependencies": { "assistant-stream": "^0.3.4" } }, "sha512-KZ9ZsF1i1zMhozvD4m8TsmTdtufqULMaqgOoSLRyVtnhwvxkDufL87tSjv7epddZ4kbebe31biWSg7KIlgzvQA=="], + "assistant-cloud": ["assistant-cloud@0.1.22", "", { "dependencies": { "assistant-stream": "^0.3.6" } }, "sha512-AEE9shV+oFrGDv/MRTRERctNKpIYS0n34UpAQXXICiOkSWD6QZnS1ljLqruFko7fJoT5CIWq8dNeJWdzQLTBLg=="], "assistant-stream": ["assistant-stream@0.3.3", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-Ne/uTseMIiZx740dTbr/SWxONM8nYj4Z5BRmUfqQN+TNgtOCgWOlC/oTUQ+A7LIUHtmGbcoyZwDf8yd2RASnDA=="], @@ -1450,7 +1450,7 @@ "lru-cache": ["lru-cache@5.1.1", "", { "dependencies": { "yallist": "^3.0.2" } }, "sha512-KpNARQA3Iwv+jTA0utUVVbrh+Jlrr1Fv0e56GGzAFOXN7dk/FviaDW8LHmK52DlcH4WP2n6gI8vN1aesBFgo9w=="], - "lucide-react": ["lucide-react@0.575.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-VuXgKZrk0uiDlWjGGXmKV6MSk9Yy4l10qgVvzGn2AWBx1Ylt0iBexKOAoA6I7JO3m+M9oeovJd3yYENfkUbOeg=="], + "lucide-react": ["lucide-react@0.577.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-4LjoFv2eEPwYDPg/CUdBJQSDfPyzXCRrVW1X7jrx/trgxnxkHFjnVZINbzvzxjN70dxychOfg+FTYwBiS3pQ5A=="], "magic-string": ["magic-string@0.30.21", "", { "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.5" } }, "sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ=="], @@ -2082,9 +2082,9 @@ "zwitch": ["zwitch@2.0.4", "", {}, "sha512-bXE4cR/kVZhKZX/RjPEflHaKVhUVl85noU3v6b8apfQEc1x4A+zBxjZ4lN8LqGd6WZ3dl98pY4o717VFmoPp+A=="], - "@assistant-ui/core/assistant-stream": ["assistant-stream@0.3.5", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-OGxVClfpEOoSsJDraPoe+GYTwh9TJX1wxK3hT5Qs7gIOyD/MZbqwyWwabRO2KTnNU4w7usvmC/vneUzxSk4bBg=="], + "@assistant-ui/core/assistant-stream": ["assistant-stream@0.3.6", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-NdtSRrQfWCDA/aqQ1xhobf/xnhuMZkhFAw9xzAt5iAoL3ouxVXOowSRN87OL4MYBQEvqtcjw9/CE6YcsXoBtuw=="], - "@assistant-ui/react/assistant-stream": ["assistant-stream@0.3.5", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-OGxVClfpEOoSsJDraPoe+GYTwh9TJX1wxK3hT5Qs7gIOyD/MZbqwyWwabRO2KTnNU4w7usvmC/vneUzxSk4bBg=="], + "@assistant-ui/react/assistant-stream": ["assistant-stream@0.3.6", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-NdtSRrQfWCDA/aqQ1xhobf/xnhuMZkhFAw9xzAt5iAoL3ouxVXOowSRN87OL4MYBQEvqtcjw9/CE6YcsXoBtuw=="], "@assistant-ui/react/zod": ["zod@4.3.6", "", {}, "sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg=="], @@ -2282,7 +2282,7 @@ "ajv-formats/ajv": ["ajv@8.18.0", "", { "dependencies": { "fast-deep-equal": "^3.1.3", "fast-uri": "^3.0.1", "json-schema-traverse": "^1.0.0", "require-from-string": "^2.0.2" } }, "sha512-PlXPeEWMXMZ7sPYOHqmDyCJzcfNrUr3fGNKtezX14ykXOEIvyK81d+qydx89KY5O71FKMPaQ2vBfBFI5NHR63A=="], - "assistant-cloud/assistant-stream": ["assistant-stream@0.3.5", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-OGxVClfpEOoSsJDraPoe+GYTwh9TJX1wxK3hT5Qs7gIOyD/MZbqwyWwabRO2KTnNU4w7usvmC/vneUzxSk4bBg=="], + "assistant-cloud/assistant-stream": ["assistant-stream@0.3.6", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-NdtSRrQfWCDA/aqQ1xhobf/xnhuMZkhFAw9xzAt5iAoL3ouxVXOowSRN87OL4MYBQEvqtcjw9/CE6YcsXoBtuw=="], "chalk/ansi-styles": ["ansi-styles@4.3.0", "", { "dependencies": { "color-convert": "^2.0.1" } }, "sha512-zbB9rCJAT1rbjiVDb2hqKFHNYLxgtk8NURxZ3IZwD3F6NtxbXZQCnnSi1Lkx+IDohdPlFp222wVALIheZJQSEg=="], diff --git a/studio/frontend/package.json b/studio/frontend/package.json index b3acb6468f..4b40759d62 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -13,7 +13,7 @@ "biome:fix": "biome check . --write" }, "dependencies": { - "@assistant-ui/react": "^0.12.17", + "@assistant-ui/react": "^0.12.19", "@assistant-ui/react-markdown": "^0.12.3", "@assistant-ui/react-streamdown": "^0.1.2", "@base-ui/react": "^1.2.0", @@ -51,7 +51,7 @@ "framer-motion": "^11.18.2", "js-yaml": "^4.1.1", "katex": "^0.16.28", - "lucide-react": "^0.575.0", + "lucide-react": "^0.577.0", "mammoth": "^1.11.0", "motion": "^12.34.0", "next": "^16.1.6", diff --git a/studio/frontend/src/components/assistant-ui/badge.tsx b/studio/frontend/src/components/assistant-ui/badge.tsx new file mode 100644 index 0000000000..7b189b349c --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/badge.tsx @@ -0,0 +1,67 @@ +"use client"; + +import type { ComponentProps } from "react"; +import { Slot } from "radix-ui"; +import { cva, type VariantProps } from "class-variance-authority"; +import { cn } from "@/lib/utils"; + +const badgeVariants = cva( + "inline-flex items-center justify-center gap-1 rounded-md font-medium text-xs transition-colors [&_svg]:size-3 [&_svg]:shrink-0", + { + variants: { + variant: { + outline: + "border border-input bg-transparent text-muted-foreground hover:bg-accent hover:text-accent-foreground", + secondary: + "bg-secondary text-secondary-foreground hover:bg-secondary/80", + muted: + "bg-muted text-muted-foreground hover:bg-muted/80 hover:text-foreground", + ghost: + "bg-transparent text-muted-foreground hover:bg-accent hover:text-accent-foreground", + info: "bg-blue-100 text-blue-700 hover:bg-blue-100/80 dark:bg-blue-900/50 dark:text-blue-300", + warning: + "bg-amber-100 text-amber-700 hover:bg-amber-100/80 dark:bg-amber-900/50 dark:text-amber-300", + success: + "bg-emerald-100 text-emerald-700 hover:bg-emerald-100/80 dark:bg-emerald-900/50 dark:text-emerald-300", + destructive: + "bg-red-100 text-red-700 hover:bg-red-100/80 dark:bg-red-900/50 dark:text-red-300", + }, + size: { + sm: "px-1.5 py-0.5", + default: "px-2 py-1", + lg: "px-2.5 py-1.5 text-sm", + }, + }, + defaultVariants: { + variant: "outline", + size: "default", + }, + }, +); + +export type BadgeProps = ComponentProps<"span"> & + VariantProps & { + asChild?: boolean; + }; + +function Badge({ + className, + variant, + size, + asChild = false, + ...props +}: BadgeProps) { + const Comp = asChild ? Slot.Root : "span"; + + return ( + + ); +} + +export { Badge, badgeVariants }; diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index e41ec10dfe..0bbbe94fdc 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -10,7 +10,7 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { code } from "@streamdown/code"; import { math } from "@streamdown/math"; import { mermaid } from "@streamdown/mermaid"; -import { DownloadIcon } from "lucide-react"; +import { DownloadIcon, Maximize2Icon, Minimize2Icon } from "lucide-react"; import { useEffect, useRef, useState } from "react"; import { Block, type BlockProps, Streamdown } from "streamdown"; import "katex/dist/katex.min.css"; @@ -84,6 +84,11 @@ function isSvgFence(codeFence: CodeFence): boolean { return false; } +function isHtmlFence(codeFence: CodeFence): boolean { + const lang = codeFence.language?.toLowerCase() ?? ""; + return lang === "html" && !codeFence.source.trimStart().startsWith("]|on\w+\s*=|javascript:|]|]|]|]/i; function sanitizeSvg(source: string): string | null { @@ -104,6 +109,96 @@ function SvgPreview({ source }: { source: string }) { ); } +const HTML_PREVIEW_DEFAULT_HEIGHT = 400; +const HTML_PREVIEW_MAX_HEIGHT = 800; + +function HtmlPreview({ source }: { source: string }) { + const iframeRef = useRef(null); + const [height, setHeight] = useState(HTML_PREVIEW_DEFAULT_HEIGHT); + const [enlarged, setEnlarged] = useState(false); + + useEffect(() => { + const handler = (e: MessageEvent) => { + if (e.source !== iframeRef.current?.contentWindow) return; + if (typeof e.data?.htmlPreviewHeight === "number") { + setHeight(Math.min(Math.max(e.data.htmlPreviewHeight, 100), HTML_PREVIEW_MAX_HEIGHT)); + } + }; + window.addEventListener("message", handler); + return () => window.removeEventListener("message", handler); + }, []); + + useEffect(() => { + if (!enlarged) return; + const handler = (e: KeyboardEvent) => { + if (e.key === "Escape") setEnlarged(false); + }; + window.addEventListener("keydown", handler); + return () => window.removeEventListener("keydown", handler); + }, [enlarged]); + + const resizeScript = ``; + + const srcDoc = source + resizeScript; + + if (enlarged) { + return ( + <> +
+ {/* Placeholder keeps layout stable while overlay is shown */} +
+
{ if (e.target === e.currentTarget) setEnlarged(false); }} + > +
+ +
+
+