diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 164f202681..064be30c06 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -828,6 +828,7 @@ class InferenceBackend: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -877,6 +878,7 @@ class InferenceBackend: execute_tool = execute_tool, cancel_event = cancel_event, auto_heal_tool_calls = auto_heal_tool_calls, + nudge_tool_calls = nudge_tool_calls, max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index aab3470193..5e6287f528 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -75,6 +75,12 @@ from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs +from core.inference.tool_call_parser import ( + MAX_ACT_REPROMPTS as _MAX_REPROMPTS, + REPROMPT_MAX_CHARS as _REPROMPT_MAX_CHARS, + is_short_intent_without_action as _is_short_intent_without_action, + reprompt_to_act_message as _reprompt_to_act_message, +) from core.inference.tool_loop_controller import ( ToolLoopController, tool_event_provenance, @@ -223,25 +229,8 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]": return out -# ── Pre-compiled patterns for plan-without-action re-prompt ── -# Forward-looking intent signals: the model is describing what it *will* -# do rather than giving a final answer. -_INTENT_SIGNAL = re.compile( - r"(?i)(" - # Direct intent ("I'll ...", "Let me ...", straight + curly apostrophes). - # Excludes "I can"/"I should"/"I want to"/"let's" (common in answers). - # Negative lookahead drops negated forms ("I will not") so a refusal - # doesn't trigger a re-prompt. - r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" - r"|" - # Step/plan framing: "First ...", "Step 1:", "Here's my plan" - r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))" - r"|" - # "Now I" / "Next I" patterns - r"\b(?:now i|next i)\b" - r")" -) -_MAX_REPROMPTS = 3 +# Plan-without-action re-prompt state (intent signal, caps, message) now lives +# in tool_call_parser, imported above under its old aliases. # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. @@ -252,7 +241,6 @@ _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min # is exempt because it needs immediate artifact feedback. _PROVISIONAL_ARGS_MIN_CHARS = 256 _DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min -_REPROMPT_MAX_CHARS = 2000 # Cap tool calls from a single TEXTUAL-fallback turn (mirrors the safetensors # loop). Structured delta.tool_calls are grammar-bounded by llama-server; text # parsed from content is not, so one runaway turn could fan out unbounded. @@ -333,11 +321,6 @@ def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int: return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0 -def _is_short_intent_without_action(text: str) -> bool: - stripped = text.strip() - return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None - - def _should_suppress_forced_no_tool_output(text: str) -> bool: """Suppress only repeated forced-turn planning text, not final answers.""" stripped = text.strip() @@ -8456,6 +8439,7 @@ class LlamaCppBackend: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -8626,11 +8610,9 @@ class LlamaCppBackend: _kb_search_count = 0 # ── Re-prompt on plan-without-action ───────────────── - # When the model describes what it intends to do (forward-looking - # language) without calling a tool, re-prompt once. Only triggers on - # responses signaling intent/planning -- a direct answer like "4" or - # "Hello!" won't match. Pattern compiled at module level - # (_INTENT_SIGNAL). + # Model describes intent without calling a tool: re-prompt once. A + # direct answer ("4", "Hello!") won't match. Pattern shared with the + # safetensors loop (tool_call_parser.INTENT_SIGNAL). _reprompt_count = 0 # Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved # re-prompt slots don't extend the budget. Mirrors the safetensors guard. @@ -9153,8 +9135,10 @@ class LlamaCppBackend: r"(?i)\brender[_\s-]?html\b", _stripped, ) + # None keeps the default-on re-prompt; False disables it. if ( auto_heal_tool_calls + and (nudge_tool_calls is None or nudge_tool_calls) and active_tools and not _render_html_already_done_intent and _reprompt_count < _MAX_REPROMPTS @@ -9183,12 +9167,7 @@ class LlamaCppBackend: conversation.append( { "role": "user", - "content": ( - "You have access to enabled tools. If a tool is needed to satisfy " - "the user's request or complete the action you described, call " - f"{tool_hint} now. If no tool is needed, provide the final answer " - "and follow the user's requested format." - ), + "content": _reprompt_to_act_message(tool_hint), } ) # Accumulate tokens and timing from this iteration. diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 5dbd5fb479..675bd9f3ea 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -945,6 +945,7 @@ class InferenceOrchestrator: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -1007,6 +1008,7 @@ class InferenceOrchestrator: execute_tool = execute_tool, cancel_event = cancel_event, auto_heal_tool_calls = auto_heal_tool_calls, + nudge_tool_calls = nudge_tool_calls, max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index aa732f47e4..81c25b777e 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -33,10 +33,13 @@ from core.inference.tool_call_parser import ( _strip_mistral_closed_calls, _strip_mistral_reasoning, BUDGET_EXHAUSTED_NUDGE, + MAX_ACT_REPROMPTS, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, + is_short_intent_without_action, parse_tool_calls_from_text, + reprompt_to_act_message, strip_leading_bare_json_call, strip_llama3_leading_sentinels, strip_tool_markup, @@ -75,21 +78,6 @@ _MAX_BUFFER_CHARS = 32 # Memory bound for holding a leading bare-JSON object whose top-level "{" never balances. _MAX_BARE_JSON_BUFFER = 16384 -# Forward-looking intent ("I'll", "First,", "Step 1:") = planning, not answering; nudge a call. -# Negative lookahead drops negated forms ("I will not") so a refusal doesn't trigger it. Mirrors GGUF. -_INTENT_SIGNAL = re.compile( - r"(?i)(" - r"\b(i['’](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" - r"|\b(?:first\b|step \d+:?|here['’]?s (?:my |the |a )?(?:plan|approach))" - r"|\b(?:now i|next i)\b" - r")" -) -_MAX_REPROMPTS = 3 -_REPROMPT_MAX_CHARS = 2000 -# Templated so the nudge names the caller's enabled tools, not a hardcoded set. Mirrors GGUF tool_hint. -_REPROMPT_INSTRUCTION_TEMPLATE = ( - "STOP. Do NOT write code or explain. You MUST call a tool NOW. Call {tool_hint} immediately." -) # No grammar constraint here (unlike llama-server's lazy grammar): collapse # exact-duplicate calls and cap the count so a runaway turn cannot fan out. @@ -432,6 +420,7 @@ def run_safetensors_tool_loop( execute_tool: Callable[..., str], cancel_event: Optional[threading.Event] = None, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, @@ -471,6 +460,9 @@ def run_safetensors_tool_loop( for _ev in _auto["events"]: yield _ev conversation.extend(_auto["messages"]) + # Autoinject ran a KB search outside the controller, so it counts as an + # executed tool for the plan-without-action gate. + rag_autoinjected = bool(_auto) unrestricted_tools = not tools # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the @@ -488,6 +480,9 @@ def run_safetensors_tool_loop( final_attempt_done = False next_call_id = 0 reprompt_count = 0 + # A denied tool confirmation must not be answered with a plan-without-action + # re-prompt (which would raise the confirmation gate again). + tool_denied = False # Real tool-call turns completed. Only turns that actually executed a tool count # against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a # plan-without-action re-prompt) must not consume budget, matching the GGUF loop. @@ -510,7 +505,7 @@ def run_safetensors_tool_loop( _state_draining = 2 # Reserve re-prompt slots so they don't eat the caller's tool budget. - _extra_iters = _MAX_REPROMPTS if max_tool_iterations > 0 else 0 + _extra_iters = MAX_ACT_REPROMPTS if max_tool_iterations > 0 else 0 for iteration in range(max_tool_iterations + _extra_iters + 1): if cancel_event is not None and cancel_event.is_set(): return @@ -869,33 +864,39 @@ def run_safetensors_tool_loop( enabled_tool_names = _enabled_tool_names, ) if not safety_tc: - # Re-prompt only when the model planned without acting (intent - # signal); "4" / "Hello!" never trigger. Mirrors GGUF. - _stripped = content_accum.strip() + # Re-prompt once on plan-without-action, before any tool runs + # (GGUF loop parity). The retry is gated on nudge_tool_calls so + # Studio callers (which send True) always nudge, while API callers + # who omit the flag keep today's no-reprompt behavior (opt-in). + stripped_answer = content_accum.strip() if ( - tools - and auto_heal_tool_calls - and reprompt_count < _MAX_REPROMPTS - and 0 < len(_stripped) < _REPROMPT_MAX_CHARS - and _INTENT_SIGNAL.search(_stripped) - and not final_attempt_done + auto_heal_tool_calls + and nudge_tool_calls + and active_tools + and reprompt_count < MAX_ACT_REPROMPTS + and not rag_autoinjected + and not tool_denied + and not any(record.executed for record in tool_controller.history) + and is_short_intent_without_action(stripped_answer) ): reprompt_count += 1 logger.info( - "Safetensors re-prompt %d/%d: model planned without " + "Safetensors re-prompt %d/%d: model responded without " "calling tools (%d chars)", reprompt_count, - _MAX_REPROMPTS, - len(_stripped), + MAX_ACT_REPROMPTS, + len(stripped_answer), ) + conversation.append({"role": "assistant", "content": stripped_answer}) tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" - conversation.append({"role": "assistant", "content": _stripped}) conversation.append( { "role": "user", - "content": _REPROMPT_INSTRUCTION_TEMPLATE.format(tool_hint = tool_hint), + "content": reprompt_to_act_message(tool_hint), } ) + # Empty status clears the badge and resets the route's + # per-turn text cursor before the re-prompted turn streams. yield {"type": "status", "text": ""} continue @@ -1085,6 +1086,7 @@ def run_safetensors_tool_loop( "result": TOOL_REJECTED_MESSAGE, "provenance": decision.provenance, } + tool_denied = True denied_message = { "role": "tool", "name": decision.tool_name, diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 70115e5744..1ab1142eba 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -163,6 +163,41 @@ RAG_SEARCH_CAP_NUDGE = ( ) +# ── Plan-without-action re-prompt (shared by the GGUF and safetensors loops) ── +# Forward-looking intent: the model says what it *will* do, not a final answer. +INTENT_SIGNAL = re.compile( + r"(?i)(" + # Direct intent ("I'll", "Let me"); lookahead drops negated forms + # ("I will not") so a refusal does not re-prompt. + r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" + r"|" + # Step/plan framing: "First ...", "Step 1:", "Here's my plan" + r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))" + r"|" + r"\b(?:now i|next i)\b" + r")" +) +# Matches GGUF's established default (llama_cpp.py has re-prompted up to 3 +# times since #5620); safetensors and MLX inherit the same cap from here. +MAX_ACT_REPROMPTS = 3 +REPROMPT_MAX_CHARS = 2000 + + +def is_short_intent_without_action(text: str) -> bool: + stripped = text.strip() + return 0 < len(stripped) < REPROMPT_MAX_CHARS and INTENT_SIGNAL.search(stripped) is not None + + +def reprompt_to_act_message(tool_hint: str) -> str: + """The user message appended when re-prompting a plan-without-action turn.""" + return ( + "You have access to enabled tools. If a tool is needed to satisfy " + "the user's request or complete the action you described, call " + f"{tool_hint} now. If no tool is needed, provide the final answer " + "and follow the user's requested format." + ) + + # Qwen / Hermes ``{json}``. _TC_JSON_START_RE = re.compile(r"\s*\{") # Qwen3.5 ```` and the attribute form ```` diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 2c87ce8c6e..963d584303 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -177,6 +177,7 @@ class ChatSettingsPayload(BaseModel): collapseHtmlArtifacts: Optional[bool] = None allowArtifactNetworkAccess: Optional[bool] = None autoHealToolCalls: Optional[bool] = None + nudgeToolCalls: Optional[bool] = None maxToolCallsPerMessage: Optional[int] = Field(default = None, ge = 1) toolCallTimeout: Optional[int] = Field(default = None, ge = 1) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1042dda004..032a6e874a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6162,6 +6162,7 @@ async def openai_chat_completions( reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, max_tool_iterations = payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25, @@ -6888,6 +6889,7 @@ async def openai_chat_completions( reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, auto_heal_tool_calls = _sf_auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, max_tool_iterations = _sf_tool_budget, tool_call_timeout = payload.tool_call_timeout if payload.tool_call_timeout is not None @@ -10074,6 +10076,7 @@ async def anthropic_messages( cancel_event = cancel_event, max_tool_iterations = 25, auto_heal_tool_calls = True, + nudge_tool_calls = payload.nudge_tool_calls, tool_call_timeout = 300, session_id = payload.session_id, # Anthropic passthrough has no rag_scope field (RAG is local-only). diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index 2a6ebe244f..a60ac700bf 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -91,6 +91,17 @@ def test_chat_settings_payload_accepts_fast_mode_presets(): assert dumped["customPresets"][0]["params"]["fastMode"] is True +def test_chat_settings_payload_accepts_nudge_tool_calls(): + # extra="forbid" 400s PUT /api/chat/settings on unknown keys, so the + # frontend's persisted nudgeToolCalls needs a payload field (like + # autoHealToolCalls). + payload = chat_history.ChatSettingsPayload.model_validate( + {"autoHealToolCalls": True, "nudgeToolCalls": False} + ) + dumped = payload.model_dump(exclude_unset = True) + assert dumped == {"autoHealToolCalls": True, "nudgeToolCalls": False} + + def test_chat_inference_settings_covers_frontend_persisted_fields(): # Drift guard: every InferenceParams field the UI persists (all but # checkpoint) must exist on ChatInferenceSettings, else extra="forbid" diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 9e16be2160..fb1b0e52b7 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -1168,6 +1168,48 @@ def test_internal_reprompt_disabled_when_auto_heal_disabled(monkeypatch): assert len(payloads) == 1 +def test_internal_reprompt_disabled_when_nudge_tool_calls_false(monkeypatch): + # Explicit nudge_tool_calls=False disables the plan-without-action + # re-prompt even with Auto-Heal on (None keeps the default-on behavior). + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + auto_heal_tool_calls = True, + nudge_tool_calls = False, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now."] + assert len(payloads) == 1 + + def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatch): streams = [ [ diff --git a/studio/backend/tests/test_nudge_tool_calls_wiring.py b/studio/backend/tests/test_nudge_tool_calls_wiring.py new file mode 100644 index 0000000000..2c27b220ba --- /dev/null +++ b/studio/backend/tests/test_nudge_tool_calls_wiring.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Wiring guard for the plan-without-action ``nudge_tool_calls`` policy. + +Decided policy: the re-prompt is ALWAYS ON for the Studio inference paths +(safetensors, GGUF/llama_cpp, MLX) and OPT-IN for the API (/v1 OpenAI-compat + +Anthropic-compat, controlled by the request's ``nudge_tool_calls``, default off). + +Mechanism (verified here without loading a model): + + * every backend tool-loop entry point accepts and forwards ``nudge_tool_calls`` + (safetensors -> ``InferenceBackend``; MLX -> ``InferenceOrchestrator``; both + call the shared ``run_safetensors_tool_loop``; GGUF -> ``LlamaCppBackend``); + * the safetensors/MLX loop gates the retry on a truthy flag (new retry -> + opt-in), while the GGUF loop keeps its pre-existing default-on behaviour + (``None`` keeps nudging) so an omitted flag never disables GGUF; + * the API request models default the flag to ``None`` (opt-in / off); + * the Studio-facing routes forward the request's flag, and the Studio frontend + sends ``nudge_tool_calls: true`` -- exercised behaviourally in + ``test_safetensors_tool_loop.py`` and ``test_llama_cpp_tool_loop.py``. +""" + +import inspect + +from core.inference.inference import InferenceBackend +from core.inference.llama_cpp import LlamaCppBackend +from core.inference.orchestrator import InferenceOrchestrator +from core.inference.safetensors_agentic import run_safetensors_tool_loop + + +def _params(fn): + return inspect.signature(fn).parameters + + +def test_shared_loop_accepts_nudge_flag(): + assert "nudge_tool_calls" in _params(run_safetensors_tool_loop) + + +def test_all_three_backends_accept_the_flag(): + for method in ( + InferenceBackend.generate_chat_completion_with_tools, + InferenceOrchestrator.generate_chat_completion_with_tools, + LlamaCppBackend.generate_chat_completion_with_tools, + ): + assert "nudge_tool_calls" in _params(method), method.__qualname__ + + +def test_delegating_backends_forward_the_flag_to_the_shared_loop(): + # safetensors (in-process transformers) and MLX (parent-process orchestrator) + # both delegate to run_safetensors_tool_loop; GGUF runs its own in-file loop + # and consumes the flag directly (asserted separately by the gate test). + for method in ( + InferenceBackend.generate_chat_completion_with_tools, + InferenceOrchestrator.generate_chat_completion_with_tools, + ): + src = inspect.getsource(method) + assert "nudge_tool_calls = nudge_tool_calls" in src, method.__qualname__ + + +def test_safetensors_loop_is_opt_in_while_gguf_stays_default_on(): + # Safetensors/MLX: the retry is new here, so it requires a truthy flag. + sf_src = inspect.getsource(run_safetensors_tool_loop) + assert "and nudge_tool_calls" in sf_src + # GGUF: pre-existing nudge must not be accidentally disabled -- an omitted + # (None) flag keeps nudging; only an explicit False turns it off. + gguf_src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert "nudge_tool_calls is None or nudge_tool_calls" in gguf_src + + +def test_api_request_models_default_the_flag_off(): + from models.inference import AnthropicMessagesRequest, ChatCompletionRequest + for model in (ChatCompletionRequest, AnthropicMessagesRequest): + field = model.model_fields["nudge_tool_calls"] + assert field.default is None, model.__name__ + + +def test_studio_routes_forward_the_request_flag(): + # The Studio chat frontend posts to /v1/chat/completions and /v1/messages + # with nudge_tool_calls=true; the route handlers forward the request value + # (external API clients that omit it fall back to the opt-in default). + from routes import inference as routes_inference + for handler in ( + routes_inference.openai_chat_completions, + routes_inference.anthropic_messages, + ): + src = inspect.getsource(handler) + assert "nudge_tool_calls = payload.nudge_tool_calls" in src, handler.__name__ diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index f826f3cddf..a8546b82c4 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -2229,6 +2229,9 @@ def _reprompt_loop(*, auto_heal_tool_calls): tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], execute_tool = exec_fn, auto_heal_tool_calls = auto_heal_tool_calls, + # Studio always nudges (always-on for the Studio inference paths); the + # API opts in per request. Model the Studio caller here. + nudge_tool_calls = True, max_tool_iterations = 3, ) ) @@ -3100,7 +3103,7 @@ class TestLoopBehaviour: class TestLoopRePrompt: - """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``_MAX_REPROMPTS`` extra slots.""" + """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Studio always nudges, so these drive the loop with ``nudge_tool_calls=True``.""" def test_intent_signal_triggers_reprompt(self): # Turn 1: intent signal, no tool call. @@ -3116,6 +3119,7 @@ class TestLoopRePrompt: ["The sky is blue."], ], exec_results = ["Blue (Rayleigh scattering)"], + nudge_tool_calls = True, ) events = _collect_events(loop) # web_search must have been called once (after the re-prompt). @@ -3159,13 +3163,14 @@ class TestLoopRePrompt: contents = [e for e in events if e["type"] == "content"] assert contents and contents[-1]["text"].strip() == "4" - def test_max_reprompts_capped_at_three(self): - # Model keeps stalling with intent -- after 3 re-prompts the - # loop must give up rather than burn forever. + def test_max_reprompts_capped(self): + # Model keeps stalling with intent -- after MAX_ACT_REPROMPTS re-prompts + # the loop must give up rather than burn forever. turns = [["Let me search for that."]] * 6 # well over the cap loop, exec_fn = _make_loop( turns = turns, exec_results = [], + nudge_tool_calls = True, ) events = _collect_events(loop, max_events = 500) # No tool ever ran, but the loop terminated cleanly. @@ -3184,6 +3189,7 @@ class TestLoopRePrompt: ["found"], ], exec_results = ["..."], + nudge_tool_calls = True, ) events = _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "x"})] @@ -3194,7 +3200,7 @@ class TestLoopRePrompt: # re-prompt ate the slot the tool call would never run. loop, exec_fn = _make_loop( turns = [ - # 1. Intent stall (re-prompt 1/3). + # 1. Intent stall (re-prompt). ["Let me search for that."], # 2. Real tool call (uses the budget slot). ['{"name":"web_search","arguments":{"query":"weather"}}'], @@ -3203,6 +3209,7 @@ class TestLoopRePrompt: ], exec_results = ["sunny"], max_tool_iterations = 1, + nudge_tool_calls = True, ) events = _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "weather"})] @@ -3305,13 +3312,22 @@ class TestGGUFSafetensorsHealingParity: } def test_intent_regex_matches_same_phrases_as_gguf(self): - # The intent re-prompt regex must match the SAME forward-looking - # phrases on both backends so behaviour is the same on Mac (MLX - # / safetensors) and on Linux (GGUF). - from core.inference.llama_cpp import _INTENT_SIGNAL as gguf_re - from core.inference.safetensors_agentic import ( - _INTENT_SIGNAL as sf_re, + # The intent re-prompt regex is now a single shared source of truth + # (tool_call_parser.INTENT_SIGNAL) consumed by both the GGUF and the + # safetensors/MLX loops, so behaviour is identical on Mac and Linux. + # Both backends must resolve to that one shared helper. + from core.inference.llama_cpp import ( + _is_short_intent_without_action as gguf_fn, ) + from core.inference.safetensors_agentic import ( + is_short_intent_without_action as sf_fn, + ) + from core.inference.tool_call_parser import ( + INTENT_SIGNAL as shared_re, + is_short_intent_without_action as shared_fn, + ) + + assert gguf_fn is shared_fn and sf_fn is shared_fn for phrase in ( "I'll search for that", @@ -3322,8 +3338,8 @@ class TestGGUFSafetensorsHealingParity: "Here's my plan", "Now I need to call web_search", ): - assert gguf_re.search(phrase), f"GGUF missed {phrase!r}" - assert sf_re.search(phrase), f"safetensors missed {phrase!r}" + assert shared_re.search(phrase), f"missed {phrase!r}" + assert shared_fn(phrase), f"helper missed {phrase!r}" for plain in ( "4", @@ -3337,13 +3353,16 @@ class TestGGUFSafetensorsHealingParity: "I will not search the web for that.", "I'll never call that tool.", ): - assert not gguf_re.search(plain), f"GGUF wrongly fired on {plain!r}" - assert not sf_re.search(plain), f"safetensors wrongly fired on {plain!r}" + assert not shared_re.search(plain), f"wrongly fired on {plain!r}" + assert not shared_fn(plain), f"helper wrongly fired on {plain!r}" def test_max_reprompts_equal_on_both_backends(self): + # Both loops draw the cap from the shared constant, so they stay equal. from core.inference.llama_cpp import _MAX_REPROMPTS as gguf_cap - from core.inference.safetensors_agentic import _MAX_REPROMPTS as sf_cap - assert gguf_cap == sf_cap == 3 + from core.inference.safetensors_agentic import MAX_ACT_REPROMPTS as sf_cap + from core.inference.tool_call_parser import MAX_ACT_REPROMPTS as shared_cap + + assert gguf_cap == sf_cap == shared_cap class TestLoopControl: @@ -3822,6 +3841,193 @@ class TestGptOssNameDetection: assert is_gpt_oss_model_name(cast(str, None)) is False +# ──────────────────────────────────────────────────────────────────── +# Plan-without-action re-prompt (GGUF loop parity) +# ──────────────────────────────────────────────────────────────────── + + +class TestPlanWithoutActionReprompt: + def test_short_intent_is_reprompted_and_tool_executes(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the final answer."], + ], + exec_results = ["result-1"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search"] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("Here is the final answer." in t for t in texts) + + def test_reprompt_fires_up_to_the_cap(self): + # GGUF parity: a persistently stalling model is re-prompted up to + # MAX_ACT_REPROMPTS times, then the last stall is surrendered as the + # final answer and no further turn is generated. + from core.inference.tool_call_parser import MAX_ACT_REPROMPTS + + stall = "Let me look into it first." + turns = [["I'll search the web for that."]] + turns += [[stall]] * MAX_ACT_REPROMPTS + turns += [["SHOULD NOT APPEAR"]] + + generations = {"count": 0} + turn_iter = iter(turns) + + def _gen(_messages): + generations["count"] += 1 + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + # One initial turn plus exactly MAX_ACT_REPROMPTS re-prompted turns. + assert generations["count"] == MAX_ACT_REPROMPTS + 1 + texts = [e["text"] for e in events if e["type"] == "content"] + assert any(stall in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_long_prose_answer_is_not_reprompted(self): + long_answer = "I'll keep explaining the details of the topic. " * 60 + loop, exec_fn = _make_loop( + turns = [ + [long_answer], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_disabled_auto_heal_is_not_reprompted(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + auto_heal_tool_calls = False, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_explicit_nudge_off_is_not_reprompted(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = False, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_omitted_nudge_flag_is_not_reprompted(self): + # The retry is new on this loop: API callers who do not send the flag + # must keep today's behavior. Studio opts in explicitly. + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_rag_autoinject_counts_as_executed_tool(self, monkeypatch): + # Autoinject already ran a KB search outside the controller; a short + # post-retrieval intent must not trigger a spurious re-prompt. + import core.inference.tools as tools_mod + + def fake_autoinject(conversation, rag_scope): + return { + "events": [ + {"type": "tool_start", "tool_name": "search_knowledge_base"}, + {"type": "tool_end", "tool_name": "search_knowledge_base"}, + ], + "messages": [{"role": "tool", "content": "kb result"}], + } + + monkeypatch.setattr(tools_mod, "build_rag_autoinject", fake_autoinject) + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the docs."], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + assert any(e.get("type") == "tool_start" for e in events) + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the docs." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_no_reprompt_after_a_denied_tool_confirmation(self, monkeypatch): + # An explicit user denial must not be answered with a nudge to call + # the tool again (which would raise another confirmation prompt). + monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: "appr-1") + monkeypatch.setattr(safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object()) + monkeypatch.setattr(safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "deny") + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["I'll search again."], + ["SHOULD NOT APPEAR"], + ], + confirm_tool_calls = True, + session_id = "sess", + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search again." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_no_reprompt_after_a_tool_already_executed(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Now I'll refine the search."], + ["SHOULD NOT APPEAR"], + ], + exec_results = ["result-1"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search"] + texts = [e["text"] for e in events if e["type"] == "content"] + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + # Routes-level python_tag strip (multi-line; stop on next sentinel) class TestRoutesPythonTagStrip: """``_TOOL_XML_RE`` must consume multi-line code, embedded JSON, and bare ``<`` (earlier ``[^\n<]*`` / ``[^\n]*`` revisions leaked tails); the streaming route-level strip is the regression-prone path.""" diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 929c0385b0..df266fd749 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -2739,6 +2739,7 @@ export function createOpenAIStreamAdapter( : {}), auto_heal_tool_calls: useChatRuntimeStore.getState().autoHealToolCalls, + nudge_tool_calls: useChatRuntimeStore.getState().nudgeToolCalls, max_tool_calls_per_message: useChatRuntimeStore.getState().maxToolCallsPerMessage, tool_call_timeout: (() => { diff --git a/studio/frontend/src/features/chat/api/chat-settings-api.ts b/studio/frontend/src/features/chat/api/chat-settings-api.ts index 1ac569efc7..1e00357ea4 100644 --- a/studio/frontend/src/features/chat/api/chat-settings-api.ts +++ b/studio/frontend/src/features/chat/api/chat-settings-api.ts @@ -26,6 +26,7 @@ export interface PersistedChatSettings { collapseHtmlArtifacts?: boolean; allowArtifactNetworkAccess?: boolean; autoHealToolCalls?: boolean; + nudgeToolCalls?: boolean; maxToolCallsPerMessage?: number; toolCallTimeout?: number; } diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 74578b9438..ea6c409b40 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1732,6 +1732,7 @@ export function ChatSettingsPanel({
+ @@ -2006,6 +2007,30 @@ function AutoHealToolCallsToggle() { ); } +function NudgeToolCallsToggle() { + const nudgeToolCalls = useChatRuntimeStore((s) => s.nudgeToolCalls); + const setNudgeToolCalls = useChatRuntimeStore((s) => s.setNudgeToolCalls); + + return ( +
+
+ + Nudge Tool Calls + + + When a tool call cannot be repaired, re-ask the model once so the + intended tool still runs. API requests stay opt-in. + +
+ +
+ ); +} + function ConfirmToolCallsToggle() { const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); 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 7c9685d6d0..ad50c4ece8 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -646,6 +646,7 @@ type ChatRuntimeStore = { toolStatus: string | null; generatingStatus: string | null; autoHealToolCalls: boolean; + nudgeToolCalls: boolean; maxToolCallsPerMessage: number; toolCallTimeout: number; kvCacheDtype: string | null; @@ -780,6 +781,7 @@ type ChatRuntimeStore = { setGeneratingStatus: (status: string | null) => void; setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void; setAutoHealToolCalls: (enabled: boolean) => void; + setNudgeToolCalls: (enabled: boolean) => void; setMaxToolCallsPerMessage: (value: number) => void; setToolCallTimeout: (value: number) => void; setKvCacheDtype: (dtype: string | null) => void; @@ -832,6 +834,7 @@ type ScalarSettingKey = | "collapseHtmlArtifacts" | "allowArtifactNetworkAccess" | "autoHealToolCalls" + | "nudgeToolCalls" | "maxToolCallsPerMessage" | "toolCallTimeout"; @@ -869,6 +872,7 @@ const SCALAR_SETTING_KEYS = [ "collapseHtmlArtifacts", "allowArtifactNetworkAccess", "autoHealToolCalls", + "nudgeToolCalls", "maxToolCallsPerMessage", "toolCallTimeout", ] as const satisfies readonly ScalarSettingKey[]; @@ -1103,6 +1107,7 @@ export const useChatRuntimeStore = create((set, get) => ({ generatingStatus: null, activeDiffusionCanvas: null, autoHealToolCalls: true, + nudgeToolCalls: true, maxToolCallsPerMessage: 25, toolCallTimeout: 5, kvCacheDtype: null, @@ -1544,6 +1549,15 @@ export const useChatRuntimeStore = create((set, get) => ({ ); return { autoHealToolCalls }; }), + setNudgeToolCalls: (nudgeToolCalls) => + set((state) => { + setScalarSettingVersion( + "nudgeToolCalls", + nudgeToolCalls, + state.nudgeToolCalls, + ); + return { nudgeToolCalls }; + }), setMaxToolCallsPerMessage: (maxToolCallsPerMessage) => set((state) => { setScalarSettingVersion( diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 94a7e1c56b..954e88e86b 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -357,6 +357,7 @@ export interface OpenAIChatCompletionsRequest { context_length?: number; }; auto_heal_tool_calls?: boolean; + nudge_tool_calls?: boolean; max_tool_calls_per_message?: number; tool_call_timeout?: number; session_id?: string; diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts index 4e457920d1..4899cb9c83 100644 --- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts @@ -21,6 +21,7 @@ import type { ReasoningEffort } from "../stores/chat-runtime-store"; const AUTO_TITLE_KEY = "unsloth_chat_auto_title"; const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls"; +const NUDGE_TOOL_CALLS_KEY = "unsloth_nudge_tool_calls"; const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message"; const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout"; const INFERENCE_PARAMS_KEY = "unsloth_chat_inference_params"; @@ -223,6 +224,7 @@ function sanitizeChatSettings(value: unknown): PersistedChatSettings { value.allowArtifactNetworkAccess, ); const autoHealToolCalls = sanitizeBool(value.autoHealToolCalls); + const nudgeToolCalls = sanitizeBool(value.nudgeToolCalls); const maxToolCallsPerMessage = sanitizeInt(value.maxToolCallsPerMessage, 1); const toolCallTimeout = sanitizeInt(value.toolCallTimeout, 1); @@ -245,6 +247,9 @@ function sanitizeChatSettings(value: unknown): PersistedChatSettings { if (autoHealToolCalls !== undefined) { settings.autoHealToolCalls = autoHealToolCalls; } + if (nudgeToolCalls !== undefined) { + settings.nudgeToolCalls = nudgeToolCalls; + } if (maxToolCallsPerMessage !== undefined) { settings.maxToolCallsPerMessage = maxToolCallsPerMessage; } @@ -305,6 +310,7 @@ export function isEmptyChatSettings(settings: PersistedChatSettings): boolean { settings.collapseHtmlArtifacts === undefined && settings.allowArtifactNetworkAccess === undefined && settings.autoHealToolCalls === undefined && + settings.nudgeToolCalls === undefined && settings.maxToolCallsPerMessage === undefined && settings.toolCallTimeout === undefined ); @@ -335,6 +341,7 @@ export function loadLegacyChatSettings(): PersistedChatSettings { const collapseHtmlArtifacts = loadBool(COLLAPSE_HTML_ARTIFACTS_KEY); const allowArtifactNetworkAccess = loadBool(ALLOW_ARTIFACT_NETWORK_ACCESS_KEY); const autoHealToolCalls = loadBool(AUTO_HEAL_TOOL_CALLS_KEY); + const nudgeToolCalls = loadBool(NUDGE_TOOL_CALLS_KEY); const maxToolCallsPerMessage = loadInt(MAX_TOOL_CALLS_KEY, 1); const toolCallTimeout = loadInt(TOOL_CALL_TIMEOUT_KEY, 1); const allCustomPresets = sanitizeCustomPresets([ @@ -361,6 +368,9 @@ export function loadLegacyChatSettings(): PersistedChatSettings { if (autoHealToolCalls !== undefined) { settings.autoHealToolCalls = autoHealToolCalls; } + if (nudgeToolCalls !== undefined) { + settings.nudgeToolCalls = nudgeToolCalls; + } if (maxToolCallsPerMessage !== undefined) { settings.maxToolCallsPerMessage = maxToolCallsPerMessage; } diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index df684b7752..7670aae5fa 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -78,6 +78,7 @@ const PREFS_KEYS: string[] = [ "unsloth_chat_auto_title", "unsloth_hf_token", "unsloth_auto_heal_tool_calls", + "unsloth_nudge_tool_calls", "unsloth_max_tool_calls_per_message", "unsloth_tool_call_timeout", "unsloth_chat_inference_params",