unsloth/studio/backend/routes
Daniel Han e318da21a7
Fix ~1.2s TTFT penalty when tools are enabled in Studio (#4639)
* Fix ~1.2s TTFT penalty when tools are enabled in Studio

When users enable web search, Python execution, or terminal tools,
every message gets a ~1.2s delay before any text appears -- even when
the model does not call any tool. This happens because
generate_chat_completion_with_tools() does a non-streaming detection
pass (stream: False) first, waits for the complete response, then
checks for tool calls. For the ~90% of messages that don't trigger a
tool call, this blocking wait is entirely wasted.

Root cause: the detection pass payload uses stream: False, forcing
llama-server to generate the entire response before returning any
tokens.

Fix: replace the non-streaming detection pass with a streaming pass
(stream: True) and a speculative buffer state machine that detects
tool signals in the first 1-2 SSE chunks:

- BUFFERING: accumulate content tokens, check first chars for tool
  signal prefixes (<tool_call>, <function=)
- STREAMING: no tool detected, yield tokens to caller immediately
- DRAINING: tool signal found, silently accumulate rest of stream

Three detection paths:
1. Structured delta.tool_calls -- detected instantly, transition to
   DRAINING, accumulate fragments, assemble at stream end.
2. XML tool markup in content -- buffer holds up to 32 chars checking
   for <tool_call> or <function= prefix, then transitions to DRAINING.
3. No tool signal -- first non-whitespace, non-XML char triggers
   immediate transition to STREAMING (fast path, ~90% of requests).

Safety net: after any stream ends in STREAMING state, check accumulated
content for XML tool signals. Handles rare "content before tool call"
edge case.

Additional supporting changes:
- Add headers parameter to _stream_with_retry for auth forwarding
- Share _strip_tool_markup and regex patterns between the detection
  pass and the final streaming pass (removes duplication)
- Remove the iteration==0 non-streaming content shortcut (no longer
  needed since all iterations stream directly)
- Keep the final streaming pass as fallback for max_tool_iterations
  exhaustion

Benchmarked on Qwen3.5-4B Q4_K_XL:
- No tools:              TTFT ~112ms (unchanged)
- Tools enabled, no call: TTFT ~112ms (was ~1207ms)
- Decode TPS:            226 (unchanged in all cases)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add unit tests for streaming tool detection state machine

16 tests covering every tool call parsing path:
- Plain text (no tool call) streaming
- Structured delta.tool_calls detection and fragment assembly
- XML <tool_call>JSON</tool_call> detection via buffer
- XML <function=name> tag detection via buffer
- Whitespace before tool XML
- Safety net (content then tool XML)
- Parallel multi-tool calls
- Reasoning token bypass (thinking models)
- Reasoning then tool call
- Empty response handling
- Buffer prefix timeout (HTML not mistaken for tool)
- Non-XML first char instant streaming
- False positive rejection (<tool_tip> vs <tool_call>)
- Arguments split across multiple chunks
- auto_heal_tool_calls=False respects the flag
- Metrics accumulation across tool iterations

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix reasoning-only BUFFERING, pre-tool content emission, and code duplication

Addresses review feedback on the streaming tool detection:

1. Reasoning tokens are no longer yielded during BUFFERING/DRAINING
   states. The consumer in routes/inference.py tracks prev_text across
   tool iterations without resetting it, so yielding reasoning during
   a detection pass that resolves to a tool call would corrupt the
   delta computation for subsequent iterations. Reasoning is now
   silently accumulated during detection (matching the old non-streaming
   behavior) and flushed together with content when the buffer resolves
   to STREAMING.

2. Handle reasoning-only responses in the BUFFERING resolver. When a
   thinking model emits only reasoning_content with no content tokens,
   the stream ends while still in BUFFERING state. The resolver now
   detects this case and yields reasoning as plain text (without
   <think> wrapper), matching the final streaming pass behavior for
   models like Qwen3 in always-think mode.

3. Replace duplicated re.sub calls for stripping tool markup with
   the existing _strip_tool_markup(content_text, final=True) helper,
   removing ~40 lines of redundant regex code.

4. Update tests: adjust reasoning test expectations to match the new
   behavior (reasoning batched with content, not streamed individually
   during BUFFERING). Add test_reasoning_only_no_content for the
   reasoning-only edge case. 17/17 tests pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address remaining reviewer findings: late tool_call IDs and XML speculation

1. Late-arriving tool_calls.id: when a provider sends the real ID on a
   later delta chunk (after the initial one with index and function
   name), the accumulator now updates the ID instead of keeping the
   synthetic "call_{idx}" placeholder. (P2, 2/10 reviewers)

2. XML speculation respects auto_heal_tool_calls: when auto_heal is
   explicitly disabled, _TOOL_XML_SIGNALS is empty so the BUFFERING
   state never speculatively holds content for XML prefix detection.
   Content starting with literal "<tool_call>" or "<function=" text
   flows straight through without delay. (P2, 1/10 reviewers)

Skipped: finish_reason="tool_calls" without delta.tool_calls fallback
(P1, 1/10 reviewers). llama-server always sends delta.tool_calls
fragments in streaming mode. A non-streaming fallback for this edge
case would add complexity for a scenario that does not occur in
practice with the supported backend.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Check request.is_disconnected() every 20 tokens instead of every token

The disconnect check is an async round-trip that adds overhead on every
loop iteration. Since the cancel watcher in llama_cpp.py already
handles connection teardown (closes the streaming response on cancel),
this route-layer check is a secondary safety net that does not need to
run on every single token.

Check every 20 tokens across all 4 streaming paths:
- gguf_tool_stream (tool-enabled GGUF)
- gguf_stream_chunks (standard GGUF)
- audio_input_generate (audio/whisper input)
- generic backend stream (non-GGUF fallback)

* Fix safety net, DRAINING metadata, and test import path

1. Safety net no longer retroactively executes tools after visible
   content was already emitted to the user. Once _last_emitted is
   non-empty, the stream is committed to normal content mode.
   Retroactive tool execution after visible output would violate the
   streaming contract and corrupt the route-layer cumulative delta
   tracker (prev_text). The tool XML is still stripped by
   _strip_tool_markup so the user sees clean content.

2. DRAINING false-positive path now merges accumulated metrics from
   prior tool iterations instead of dropping them. Uses the same
   merge formula as the STREAMING path.

3. Test import path fixed to use repo root instead of hardcoded
   sibling directory. Works in clean checkouts and CI.

4. Renamed test_content_then_tool_xml_safety_net to
   test_content_then_tool_xml_no_retroactive_execution to reflect
   the corrected behavior.

17/17 tests pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Redact --api-key value from llama-server startup log

When UNSLOTH_DIRECT_STREAM=1, the generated bearer token was logged
verbatim in the startup command. Replace the secret with <redacted>
before logging.

* Remove test file temporarily

* Revert disconnect throttle, reset prev_text on tool_start, restore XML safety net

Addresses all P1 findings from reviewer round 3 (10 reviewers):

1. Revert disconnect check to every iteration (was every 20th).
   All 10 reviewers flagged this as a correctness regression for
   short streams and sparse tool event loops. The cancel watcher in
   llama_cpp.py is the primary mechanism but the route-layer check
   must remain per-iteration for completeness. [10/10]

2. Reset prev_text on tool_start in gguf_tool_stream. When a tool
   cycle begins after visible content was already streamed, the
   route-layer cumulative delta tracker (prev_text) must be reset
   so the post-tool synthesis response is not truncated or dropped.
   [9/10]

3. Remove the _last_emitted gate from the XML safety net. The gate
   was added to prevent retroactive tool execution after visible
   content, but with prev_text now reset on tool_start (#2), the
   root cause is fixed and the safety net can correctly handle
   content-then-tool-XML responses (matching pre-PR behavior).
   [8/10]

* Use None instead of {} for empty auth headers in TTS methods

* Include accumulated metrics in STREAMING metadata check

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-03-27 03:13:38 -07:00
..
data_recipe Fix Studio silently exiting on Windows without error output (#4527) 2026-03-22 08:23:03 -07:00
.gitkeep root studio folder 2026-02-02 09:13:49 +00:00
__init__.py feat(studio): training history persistence and past runs viewer (#4501) 2026-03-25 00:58:55 -07:00
auth.py studio: fix stale GGUF metadata, update helper model, auth improvements (#4346) 2026-03-17 01:22:08 -07:00
datasets.py fix: subprocess crash during map operation on Windows (#4507) 2026-03-22 05:21:09 -07:00
export.py Final cleanup 2026-03-12 18:28:04 +00:00
inference.py Fix ~1.2s TTFT penalty when tools are enabled in Studio (#4639) 2026-03-27 03:13:38 -07:00
models.py feat: multi-source model discovery (HF default, legacy cache, LM Studio) (#4591) 2026-03-25 07:48:04 -07:00
training.py feat(studio): training history persistence and past runs viewer (#4501) 2026-03-25 00:58:55 -07:00
training_history.py feat(studio): training history persistence and past runs viewer (#4501) 2026-03-25 00:58:55 -07:00