From a067609fc36ae2f35640bb2212b3d3979f16fbad Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 06:16:09 +0000 Subject: [PATCH 01/18] 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 (, or .*?", _re_tool.DOTALL), + _re_tool.compile(r".*?", _re_tool.DOTALL), + ] + _TOOL_ALL_PATTERNS = _TOOL_CLOSED_PATTERNS + [ + _re_tool.compile(r".*$", _re_tool.DOTALL), + _re_tool.compile(r".*$", _re_tool.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 + + # XML prefixes that signal a tool call in content + _TOOL_XML_SIGNALS = ("", "= 0 else 0, @@ -1706,65 +1747,327 @@ class LlamaCppBackend: payload["stop"] = stop try: - with httpx.Client(timeout = None) as client: - resp = client.post(url, json = payload) - if resp.status_code != 200: - raise RuntimeError( - f"llama-server returned {resp.status_code}: {resp.text}" + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} + if self._api_key else None + ) + + # ── Speculative buffer state machine ────────────────── + # BUFFERING: accumulating content, checking for tool signals + # STREAMING: no tool detected, yielding tokens to caller + # DRAINING: tool signal found, silently consuming rest + _S_BUFFERING = 0 + _S_STREAMING = 1 + _S_DRAINING = 2 + + detect_state = _S_BUFFERING + content_buffer = "" # Raw content held during BUFFERING + content_accum = "" # All content tokens (for tool parsing) + reasoning_accum = "" + cumulative_display = "" # Cumulative text yielded (with ) + in_thinking = False + has_content_tokens = False + tool_calls_acc = {} # Structured delta.tool_calls fragments + has_structured_tc = False + _iter_usage = None + _iter_timings = None + _stream_done = False + _last_emitted = "" + + stream_timeout = httpx.Timeout( + connect = 10, read = 0.5, write = 10, pool = 10, + ) + with httpx.Client(timeout = stream_timeout) as client: + with self._stream_with_retry( + client, url, payload, cancel_event, + headers = _auth_headers, + ) as response: + if response.status_code != 200: + error_body = response.read().decode() + raise RuntimeError( + f"llama-server returned {response.status_code}: " + f"{error_body}" + ) + + raw_buf = "" + for raw_chunk in self._iter_text_cancellable( + response, cancel_event, + ): + raw_buf += raw_chunk + while "\n" in raw_buf: + line, raw_buf = raw_buf.split("\n", 1) + line = line.strip() + + if not line: + continue + if line == "data: [DONE]": + # Flush thinking state for STREAMING + if detect_state == _S_STREAMING and in_thinking: + if has_content_tokens: + cumulative_display += "" + yield { + "type": "content", + "text": _strip_tool_markup( + cumulative_display, + final = True, + ), + } + else: + cumulative_display = reasoning_accum + yield { + "type": "content", + "text": cumulative_display, + } + _stream_done = True + break # exit inner while + if not line.startswith("data: "): + continue + + try: + chunk_data = json.loads(line[6:]) + _ct = chunk_data.get("timings") + if _ct: + _iter_timings = _ct + _cu = chunk_data.get("usage") + if _cu: + _iter_usage = _cu + + choices = chunk_data.get("choices", []) + if not choices: + continue + + delta = choices[0].get("delta", {}) + + # ── Structured tool_calls ── + tc_deltas = delta.get("tool_calls") + if tc_deltas: + has_structured_tc = True + detect_state = _S_DRAINING + for tc_d in tc_deltas: + idx = tc_d.get("index", 0) + if idx not in tool_calls_acc: + tool_calls_acc[idx] = { + "id": tc_d.get( + "id", f"call_{idx}" + ), + "type": "function", + "function": { + "name": "", + "arguments": "", + }, + } + func = tc_d.get("function", {}) + if func.get("name"): + tool_calls_acc[idx][ + "function" + ]["name"] += func["name"] + if func.get("arguments"): + tool_calls_acc[idx][ + "function" + ]["arguments"] += func[ + "arguments" + ] + continue + + # ── Reasoning tokens (bypass buffer) ── + reasoning = delta.get( + "reasoning_content", "" + ) + if reasoning: + reasoning_accum += reasoning + if detect_state != _S_DRAINING: + if not in_thinking: + cumulative_display += "" + in_thinking = True + cumulative_display += reasoning + yield { + "type": "content", + "text": cumulative_display, + } + + # ── Content tokens ── + token = delta.get("content", "") + if token: + has_content_tokens = True + content_accum += token + + if detect_state == _S_DRAINING: + pass # accumulate silently + + elif detect_state == _S_STREAMING: + if in_thinking: + cumulative_display += "" + in_thinking = False + cumulative_display += token + cleaned = _strip_tool_markup( + cumulative_display, + ) + if len(cleaned) > len( + _last_emitted + ): + _last_emitted = cleaned + yield { + "type": "content", + "text": cleaned, + } + + elif detect_state == _S_BUFFERING: + content_buffer += token + stripped_buf = ( + content_buffer.lstrip() + ) + if not stripped_buf: + continue + + # Check tool signal prefixes + is_prefix = False + is_match = False + for sig in _TOOL_XML_SIGNALS: + if stripped_buf.startswith( + sig + ): + is_match = True + break + if sig.startswith( + stripped_buf + ): + is_prefix = True + break + + if is_match: + detect_state = _S_DRAINING + elif ( + is_prefix + and len(stripped_buf) + < _MAX_BUFFER_CHARS + ): + pass # keep buffering + else: + # Not a tool -- flush buffer + detect_state = _S_STREAMING + if in_thinking: + cumulative_display += ( + "" + ) + in_thinking = False + cumulative_display += ( + content_buffer + ) + cleaned = _strip_tool_markup( + cumulative_display, + ) + if len(cleaned) > len( + _last_emitted + ): + _last_emitted = cleaned + yield { + "type": "content", + "text": cleaned, + } + + except json.JSONDecodeError: + logger.debug( + f"Skipping malformed SSE line: " + f"{line[:100]}" + ) + if _stream_done: + break # exit outer for + + # ── Resolve BUFFERING at stream end ── + if detect_state == _S_BUFFERING: + stripped_buf = content_buffer.lstrip() + if stripped_buf and auto_heal_tool_calls and any( + s in stripped_buf for s in _TOOL_XML_SIGNALS + ): + detect_state = _S_DRAINING + elif content_accum or reasoning_accum: + detect_state = _S_STREAMING + if content_buffer: + if in_thinking: + cumulative_display += "" + in_thinking = False + cumulative_display += content_buffer + yield { + "type": "content", + "text": _strip_tool_markup( + cumulative_display, final = True, + ), + } + else: + return + + # ── STREAMING path: no tool call ── + if detect_state == _S_STREAMING: + # Safety net: check for XML tool signals in content + _safety_tc = None + if auto_heal_tool_calls and any( + s in content_accum for s in _TOOL_XML_SIGNALS + ): + _safety_tc = self._parse_tool_calls_from_text( + content_accum, ) - data = resp.json() - except httpx.ConnectError: - raise RuntimeError("Lost connection to llama-server") + if not _safety_tc: + # Content was already streamed. Yield metadata. + yield {"type": "status", "text": ""} + _fu = _iter_usage or {} + _fc = _fu.get("completion_tokens", 0) + _fp = _fu.get("prompt_tokens", 0) + _tc = _fc + _accumulated_completion_tokens + if _iter_usage or _iter_timings: + _mt = ( + dict(_iter_timings) if _iter_timings else {} + ) + if ( + _accumulated_predicted_ms + or _accumulated_predicted_n + ): + _mt["predicted_ms"] = ( + _mt.get("predicted_ms", 0) + + _accumulated_predicted_ms + ) + _tn = ( + _mt.get("predicted_n", 0) + + _accumulated_predicted_n + ) + _mt["predicted_n"] = _tn + _tms = _mt["predicted_ms"] + if _tms > 0: + _mt["predicted_per_second"] = ( + _tn / (_tms / 1000.0) + ) + yield { + "type": "metadata", + "usage": { + "prompt_tokens": _fp, + "completion_tokens": _tc, + "total_tokens": _fp + _tc, + }, + "timings": _mt, + } + return - choices = data.get("choices", []) - if not choices: - return - - choice = choices[0] - finish_reason = choice.get("finish_reason", "") - message = choice.get("message", {}) - - # If model wants to call tools - 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, - # or bare tags without wrapper. - content_text = message.get("content", "") or "" - if ( - auto_heal_tool_calls - and not tool_calls - and ("" in content_text or " blocks since they - # can contain arbitrary content including code. + # Safety net caught tool XML -- treat as tool call + tool_calls = _safety_tc + content_text = content_accum import re - - # Strip ... blocks (greedy inside) content_text = re.sub( 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".*$", "", @@ -1772,30 +2075,86 @@ class LlamaCppBackend: flags = re.DOTALL, ).strip() logger.info( - f"Parsed {len(tool_calls)} tool call(s) from content text" + f"Safety net: parsed {len(tool_calls)} tool call(s) " + f"from streamed content" ) + else: + # ── DRAINING path: assemble tool_calls ── + tool_calls = None + content_text = content_accum + if has_structured_tc: + tool_calls = [ + tool_calls_acc[i] + for i in sorted(tool_calls_acc) + ] + if not tool_calls and auto_heal_tool_calls and any( + s in content_accum for s in _TOOL_XML_SIGNALS + ): + tool_calls = self._parse_tool_calls_from_text( + content_accum, + ) + if tool_calls and not has_structured_tc: + import re + content_text = re.sub( + r".*?", + "", + content_text, + flags = re.DOTALL, + ) + content_text = re.sub( + r".*$", + "", + content_text, + flags = re.DOTALL, + ) + content_text = re.sub( + r".*?", + "", + content_text, + flags = re.DOTALL, + ) + content_text = re.sub( + r".*$", + "", + content_text, + flags = re.DOTALL, + ).strip() + if tool_calls: + logger.info( + f"Parsed {len(tool_calls)} tool call(s) from " + f"{'structured delta' if has_structured_tc else 'content text'}" + ) + if not tool_calls: + # DRAINING but no tool calls (false positive) + yield {"type": "status", "text": ""} + if content_accum: + yield {"type": "content", "text": content_accum} + if _iter_usage or _iter_timings: + yield { + "type": "metadata", + "usage": _iter_usage, + "timings": _iter_timings, + } + return - if finish_reason == "tool_calls" or (tool_calls and len(tool_calls) > 0): - # Only accumulate metrics for responses that are actually used - _accumulated_completion_tokens += data.get("usage", {}).get( - "completion_tokens", 0 + # ── Execute tool calls ── + _accumulated_completion_tokens += ( + (_iter_usage or {}).get("completion_tokens", 0) ) - _iter_timings = data.get("timings", {}) - _accumulated_predicted_ms += _iter_timings.get("predicted_ms", 0) - _accumulated_predicted_n += _iter_timings.get("predicted_n", 0) - # Append the assistant message with tool_calls to conversation + _it = _iter_timings or {} + _accumulated_predicted_ms += _it.get("predicted_ms", 0) + _accumulated_predicted_n += _it.get("predicted_n", 0) + assistant_msg = {"role": "assistant", "content": content_text} if tool_calls: assistant_msg["tool_calls"] = tool_calls conversation.append(assistant_msg) - # Execute each tool call for tc in tool_calls or []: func = tc.get("function", {}) tool_name = func.get("name", "") raw_args = func.get("arguments", {}) - # Handle arguments as either string or dict if isinstance(raw_args, str): try: arguments = json.loads(raw_args) @@ -1807,12 +2166,13 @@ class LlamaCppBackend: else: arguments = raw_args - # Yield status update if tool_name == "web_search": status_text = f"Searching: {arguments.get('query', '')}" elif tool_name == "python": preview = ( - (arguments.get("code") or "").strip().split("\n")[0][:60] + (arguments.get("code") or "") + .strip() + .split("\n")[0][:60] ) status_text = ( f"Running Python: {preview}" @@ -1830,7 +2190,6 @@ 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, @@ -1838,7 +2197,6 @@ class LlamaCppBackend: "arguments": arguments, } - # Execute the tool _effective_timeout = ( None if tool_call_timeout >= 9999 else tool_call_timeout ) @@ -1850,7 +2208,6 @@ class LlamaCppBackend: session_id = session_id, ) - # Emit tool_end so the frontend can record outputs yield { "type": "tool_end", "tool_name": tool_name, @@ -1858,7 +2215,6 @@ class LlamaCppBackend: "result": result, } - # Append tool result to conversation tool_msg = { "role": "tool", "name": tool_name, @@ -1872,26 +2228,13 @@ class LlamaCppBackend: # Continue the loop to let model respond with context continue - # No tool calls -- model answered directly. - # If no tools were executed at all, just yield the content - # from this response instead of making a redundant second request. - if iteration == 0 and content_text: - yield {"type": "status", "text": ""} - yield {"type": "content", "text": content_text} - _direct_usage = data.get("usage") - _direct_timings = data.get("timings") - if _direct_usage or _direct_timings: - yield { - "type": "metadata", - "usage": _direct_usage, - "timings": _direct_timings, - } - return + except httpx.ConnectError: + raise RuntimeError("Lost connection to llama-server") + except Exception as e: + if cancel_event is not None and cancel_event.is_set(): + return + raise - # Tools were called in previous iterations; do a final - # streaming pass so the model can synthesize a response - # incorporating the tool results. - break # Clear status yield {"type": "status", "text": ""} @@ -1917,27 +2260,6 @@ class LlamaCppBackend: stream_payload["stop"] = stop stream_payload["stream_options"] = {"include_usage": True} - 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 = "" @@ -1950,9 +2272,10 @@ class LlamaCppBackend: try: stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, stream_payload, cancel_event + client, url, stream_payload, cancel_event, headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -2078,7 +2401,8 @@ class LlamaCppBackend: if not self.is_loaded: return None try: - with httpx.Client(timeout = 10) as client: + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + with httpx.Client(timeout = 10, headers = _auth_headers) as client: def _detok(tid: int) -> str: r = client.post( @@ -2196,7 +2520,8 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 - with httpx.Client(timeout = httpx.Timeout(300, connect = 10)) as client: + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: raise RuntimeError( From 190124d4acd392b6ccaf46dd922f33930a75c7ea Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 27 Mar 2026 06:17:46 +0000 Subject: [PATCH 02/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 162 ++++++++++----------- 1 file changed, 81 insertions(+), 81 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index eac54c4a53..c4901d4df2 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1486,7 +1486,10 @@ class LlamaCppBackend: pool = 10, ) with client.stream( - "POST", url, json = payload, timeout = prefill_timeout, + "POST", + url, + json = payload, + timeout = prefill_timeout, headers = headers, ) as response: _response_ref[0] = response @@ -1561,10 +1564,16 @@ class LlamaCppBackend: # can finish. Cancel during streaming is handled by the # watcher thread (closes the response on cancel_event). stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, payload, cancel_event, headers = _auth_headers, + client, + url, + payload, + cancel_event, + headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -1749,7 +1758,8 @@ class LlamaCppBackend: try: _auth_headers = ( {"Authorization": f"Bearer {self._api_key}"} - if self._api_key else None + if self._api_key + else None ) # ── Speculative buffer state machine ────────────────── @@ -1758,16 +1768,16 @@ class LlamaCppBackend: # DRAINING: tool signal found, silently consuming rest _S_BUFFERING = 0 _S_STREAMING = 1 - _S_DRAINING = 2 + _S_DRAINING = 2 detect_state = _S_BUFFERING - content_buffer = "" # Raw content held during BUFFERING - content_accum = "" # All content tokens (for tool parsing) + content_buffer = "" # Raw content held during BUFFERING + content_accum = "" # All content tokens (for tool parsing) reasoning_accum = "" - cumulative_display = "" # Cumulative text yielded (with ) + cumulative_display = "" # Cumulative text yielded (with ) in_thinking = False has_content_tokens = False - tool_calls_acc = {} # Structured delta.tool_calls fragments + tool_calls_acc = {} # Structured delta.tool_calls fragments has_structured_tc = False _iter_usage = None _iter_timings = None @@ -1775,11 +1785,17 @@ class LlamaCppBackend: _last_emitted = "" stream_timeout = httpx.Timeout( - connect = 10, read = 0.5, write = 10, pool = 10, + connect = 10, + read = 0.5, + write = 10, + pool = 10, ) with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, payload, cancel_event, + client, + url, + payload, + cancel_event, headers = _auth_headers, ) as response: if response.status_code != 200: @@ -1791,7 +1807,8 @@ class LlamaCppBackend: raw_buf = "" for raw_chunk in self._iter_text_cancellable( - response, cancel_event, + response, + cancel_event, ): raw_buf += raw_chunk while "\n" in raw_buf: @@ -1847,9 +1864,7 @@ class LlamaCppBackend: idx = tc_d.get("index", 0) if idx not in tool_calls_acc: tool_calls_acc[idx] = { - "id": tc_d.get( - "id", f"call_{idx}" - ), + "id": tc_d.get("id", f"call_{idx}"), "type": "function", "function": { "name": "", @@ -1858,21 +1873,17 @@ class LlamaCppBackend: } func = tc_d.get("function", {}) if func.get("name"): - tool_calls_acc[idx][ - "function" - ]["name"] += func["name"] + tool_calls_acc[idx]["function"][ + "name" + ] += func["name"] if func.get("arguments"): - tool_calls_acc[idx][ - "function" - ]["arguments"] += func[ + tool_calls_acc[idx]["function"][ "arguments" - ] + ] += func["arguments"] continue # ── Reasoning tokens (bypass buffer) ── - reasoning = delta.get( - "reasoning_content", "" - ) + reasoning = delta.get("reasoning_content", "") if reasoning: reasoning_accum += reasoning if detect_state != _S_DRAINING: @@ -1902,9 +1913,7 @@ class LlamaCppBackend: cleaned = _strip_tool_markup( cumulative_display, ) - if len(cleaned) > len( - _last_emitted - ): + if len(cleaned) > len(_last_emitted): _last_emitted = cleaned yield { "type": "content", @@ -1913,9 +1922,7 @@ class LlamaCppBackend: elif detect_state == _S_BUFFERING: content_buffer += token - stripped_buf = ( - content_buffer.lstrip() - ) + stripped_buf = content_buffer.lstrip() if not stripped_buf: continue @@ -1923,14 +1930,10 @@ class LlamaCppBackend: is_prefix = False is_match = False for sig in _TOOL_XML_SIGNALS: - if stripped_buf.startswith( - sig - ): + if stripped_buf.startswith(sig): is_match = True break - if sig.startswith( - stripped_buf - ): + if sig.startswith(stripped_buf): is_prefix = True break @@ -1946,19 +1949,13 @@ class LlamaCppBackend: # Not a tool -- flush buffer detect_state = _S_STREAMING if in_thinking: - cumulative_display += ( - "" - ) + cumulative_display += "" in_thinking = False - cumulative_display += ( - content_buffer - ) + cumulative_display += content_buffer cleaned = _strip_tool_markup( cumulative_display, ) - if len(cleaned) > len( - _last_emitted - ): + if len(cleaned) > len(_last_emitted): _last_emitted = cleaned yield { "type": "content", @@ -1967,8 +1964,7 @@ class LlamaCppBackend: except json.JSONDecodeError: logger.debug( - f"Skipping malformed SSE line: " - f"{line[:100]}" + f"Skipping malformed SSE line: " f"{line[:100]}" ) if _stream_done: break # exit outer for @@ -1976,8 +1972,10 @@ class LlamaCppBackend: # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: stripped_buf = content_buffer.lstrip() - if stripped_buf and auto_heal_tool_calls and any( - s in stripped_buf for s in _TOOL_XML_SIGNALS + if ( + stripped_buf + and auto_heal_tool_calls + and any(s in stripped_buf for s in _TOOL_XML_SIGNALS) ): detect_state = _S_DRAINING elif content_accum or reasoning_accum: @@ -1990,7 +1988,8 @@ class LlamaCppBackend: yield { "type": "content", "text": _strip_tool_markup( - cumulative_display, final = True, + cumulative_display, + final = True, ), } else: @@ -2014,27 +2013,19 @@ class LlamaCppBackend: _fp = _fu.get("prompt_tokens", 0) _tc = _fc + _accumulated_completion_tokens if _iter_usage or _iter_timings: - _mt = ( - dict(_iter_timings) if _iter_timings else {} - ) - if ( - _accumulated_predicted_ms - or _accumulated_predicted_n - ): + _mt = dict(_iter_timings) if _iter_timings else {} + if _accumulated_predicted_ms or _accumulated_predicted_n: _mt["predicted_ms"] = ( _mt.get("predicted_ms", 0) + _accumulated_predicted_ms ) _tn = ( - _mt.get("predicted_n", 0) - + _accumulated_predicted_n + _mt.get("predicted_n", 0) + _accumulated_predicted_n ) _mt["predicted_n"] = _tn _tms = _mt["predicted_ms"] if _tms > 0: - _mt["predicted_per_second"] = ( - _tn / (_tms / 1000.0) - ) + _mt["predicted_per_second"] = _tn / (_tms / 1000.0) yield { "type": "metadata", "usage": { @@ -2050,6 +2041,7 @@ class LlamaCppBackend: tool_calls = _safety_tc content_text = content_accum import re + content_text = re.sub( r".*?", "", @@ -2083,18 +2075,18 @@ class LlamaCppBackend: tool_calls = None content_text = content_accum if has_structured_tc: - tool_calls = [ - tool_calls_acc[i] - for i in sorted(tool_calls_acc) - ] - if not tool_calls and auto_heal_tool_calls and any( - s in content_accum for s in _TOOL_XML_SIGNALS + tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] + if ( + not tool_calls + and auto_heal_tool_calls + and any(s in content_accum for s in _TOOL_XML_SIGNALS) ): tool_calls = self._parse_tool_calls_from_text( content_accum, ) if tool_calls and not has_structured_tc: import re + content_text = re.sub( r".*?", "", @@ -2138,8 +2130,8 @@ class LlamaCppBackend: return # ── Execute tool calls ── - _accumulated_completion_tokens += ( - (_iter_usage or {}).get("completion_tokens", 0) + _accumulated_completion_tokens += (_iter_usage or {}).get( + "completion_tokens", 0 ) _it = _iter_timings or {} _accumulated_predicted_ms += _it.get("predicted_ms", 0) @@ -2170,9 +2162,7 @@ class LlamaCppBackend: status_text = f"Searching: {arguments.get('query', '')}" elif tool_name == "python": preview = ( - (arguments.get("code") or "") - .strip() - .split("\n")[0][:60] + (arguments.get("code") or "").strip().split("\n")[0][:60] ) status_text = ( f"Running Python: {preview}" @@ -2235,7 +2225,6 @@ class LlamaCppBackend: return raise - # Clear status yield {"type": "status", "text": ""} @@ -2260,7 +2249,6 @@ class LlamaCppBackend: stream_payload["stop"] = stop stream_payload["stream_options"] = {"include_usage": True} - cumulative = "" _last_emitted = "" in_thinking = False @@ -2272,10 +2260,16 @@ class LlamaCppBackend: try: stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, stream_payload, cancel_event, headers = _auth_headers, + client, + url, + stream_payload, + cancel_event, + headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -2401,7 +2395,9 @@ class LlamaCppBackend: if not self.is_loaded: return None try: - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + ) with httpx.Client(timeout = 10, headers = _auth_headers) as client: def _detok(tid: int) -> str: @@ -2520,8 +2516,12 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} - with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client: + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + ) + with httpx.Client( + timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers + ) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: raise RuntimeError( From 531229811d76d1a20c49536ccb2f2be5f1e0fbff Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 06:22:34 +0000 Subject: [PATCH 03/18] 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 JSON detection via buffer - XML 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 ( vs ) - Arguments split across multiple chunks - auto_heal_tool_calls=False respects the flag - Metrics accumulation across tool iterations --- tests/test_streaming_tool_detection.py | 712 +++++++++++++++++++++++++ 1 file changed, 712 insertions(+) create mode 100644 tests/test_streaming_tool_detection.py diff --git a/tests/test_streaming_tool_detection.py b/tests/test_streaming_tool_detection.py new file mode 100644 index 0000000000..b2429f0f63 --- /dev/null +++ b/tests/test_streaming_tool_detection.py @@ -0,0 +1,712 @@ +""" +Exhaustive tests for the speculative-buffer streaming tool detection in +generate_chat_completion_with_tools(). + +We mock the HTTP layer so llama-server is not required. Each test constructs +the exact SSE byte stream that llama-server would emit, feeds it through the +real method, and asserts on the yielded events. +""" + +import json +import threading +import types +import contextlib +from unittest.mock import MagicMock, patch, PropertyMock +import sys, os + +# ── helpers ────────────────────────────────────────────────────────────── + +def _sse_line(data: dict) -> str: + """One SSE data line (no trailing blank line -- we add those in the stream).""" + return f"data: {json.dumps(data)}" + + +def _sse_done() -> str: + return "data: [DONE]" + + +def _make_chunk(delta: dict, finish_reason=None, usage=None, timings=None): + """Build a chat-completions streaming chunk.""" + choice = {"index": 0, "delta": delta} + if finish_reason: + choice["finish_reason"] = finish_reason + chunk = {"choices": [choice]} + if usage: + chunk["usage"] = usage + if timings: + chunk["timings"] = timings + return chunk + + +def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) -> str: + """ + Build a complete SSE text stream from a list of chunk dicts. + Includes the role chunk, content/tool chunks, and [DONE]. + """ + lines = [] + for c in chunks: + lines.append(_sse_line(c)) + lines.append("") # blank line separator + # Final usage chunk (if provided) + if final_usage or final_timings: + meta = {} + if final_usage: + meta["usage"] = final_usage + if final_timings: + meta["timings"] = final_timings + meta["choices"] = [] + lines.append(_sse_line(meta)) + lines.append("") + lines.append(_sse_done()) + lines.append("") + return "\n".join(lines) + + +class FakeResponse: + """Mimics httpx.Response for streaming.""" + def __init__(self, text: str, status_code: int = 200): + self._text = text + self.status_code = status_code + self._closed = False + + def iter_text(self): + # Yield the whole thing in one shot (simplest case) + yield self._text + + def read(self): + return self._text.encode() + + def close(self): + self._closed = True + + +class FakeClient: + """Mimics httpx.Client context manager.""" + def __init__(self, response: FakeResponse): + self._response = response + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + @contextlib.contextmanager + def stream(self, method, url, json=None, timeout=None, headers=None): + yield self._response + + +# ── Build a minimal LlamaCppBackend for testing ───────────────────────── + +def _make_backend(): + """Create a minimal mock backend with just enough to run the method.""" + # We need the real class but only care about generate_chat_completion_with_tools + # Import the real module + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "unsloth_studio_src")) + + # Instead of importing the full module (which has other deps), we'll + # build a lightweight object that has the method and its dependencies. + from studio.backend.core.inference.llama_cpp import LlamaCppBackend + + backend = object.__new__(LlamaCppBackend) + backend._process = True # is_loaded checks _process is not None + backend._healthy = True # is_loaded checks _healthy + backend._port = 9999 # base_url property reads _port + backend._api_key = None + backend._supports_reasoning = False + return backend + + +def _synthesis_sse(): + """Build a simple text SSE response for post-tool synthesis.""" + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"content": "Done."}), + _make_chunk({}, finish_reason="stop"), + ] + usage = {"prompt_tokens": 20, "completion_tokens": 1} + return _build_sse_stream(chunks, final_usage=usage) + + +def _collect_events(backend, sse_text, tools=None, **kwargs): + """ + Run generate_chat_completion_with_tools with a fake SSE stream + and collect all yielded events. + + After the first iteration (tool detection), subsequent iterations + return a plain text synthesis response so the agentic loop terminates. + """ + if tools is None: + tools = [{"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}] + + call_count = [0] + synth_sse = _synthesis_sse() + + @contextlib.contextmanager + def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): + idx = call_count[0] + call_count[0] += 1 + # First call: use the provided SSE. Subsequent: plain text synthesis. + text = sse_text if idx == 0 else synth_sse + yield FakeResponse(text) + + # Patch execute_tool to return a dummy result + def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): + return f"Tool {tool_name} result: OK" + + original_stream = backend._stream_with_retry + backend._stream_with_retry = fake_stream_with_retry + + events = [] + with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): + try: + for event in backend.generate_chat_completion_with_tools( + messages=[{"role": "user", "content": "Hello"}], + tools=tools, + **kwargs, + ): + events.append(event) + except Exception as e: + import traceback + traceback.print_exc() + events.append({"type": "error", "error": str(e)}) + + backend._stream_with_retry = original_stream + return events + + +# ── The actual tests ───────────────────────────────────────────────────── + +def test_no_tool_call_plain_text(): + """90% case: model responds with plain text, no tool call. + Should stream content immediately without delay.""" + backend = _make_backend() + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"content": "Hello"}), + _make_chunk({"content": " there"}), + _make_chunk({"content": "!"}, finish_reason="stop"), + ] + usage = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13} + timings = {"predicted_ms": 100, "predicted_n": 3, "predicted_per_second": 30.0} + sse = _build_sse_stream(chunks, final_usage=usage, final_timings=timings) + + events = _collect_events(backend, sse) + + # Should have content events with cumulative text + content_events = [e for e in events if e["type"] == "content"] + assert len(content_events) >= 1, f"Expected content events, got: {events}" + + # Final content should contain the full text + final_content = content_events[-1]["text"] + assert "Hello there!" in final_content, f"Missing text in: {final_content}" + + # Should have metadata + meta_events = [e for e in events if e["type"] == "metadata"] + assert len(meta_events) == 1, f"Expected 1 metadata event, got: {meta_events}" + + # Should have status clear + status_events = [e for e in events if e["type"] == "status"] + assert any(e["text"] == "" for e in status_events), "Missing status clear" + + print("PASS: test_no_tool_call_plain_text") + + +def test_structured_tool_calls(): + """Model emits structured delta.tool_calls (the standard path). + Should detect instantly and execute.""" + backend = _make_backend() + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": ""}}]}), + _make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": '{"query":'}}]}), + _make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": ' "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), + ] + usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + sse = _build_sse_stream(chunks, final_usage=usage) + + events = _collect_events(backend, sse) + + # Should have status update for tool execution + status_events = [e for e in events if e["type"] == "status" and "Searching" in e.get("text", "")] + assert len(status_events) >= 1, f"Expected search status, got: {events}" + + # Should have tool_start event + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 1, f"Expected 1 tool_start, got: {tool_starts}" + assert tool_starts[0]["tool_name"] == "web_search" + assert tool_starts[0]["arguments"] == {"query": "test"} + + # Should have tool_end event + tool_ends = [e for e in events if e["type"] == "tool_end"] + assert len(tool_ends) == 1, f"Expected 1 tool_end, got: {tool_ends}" + + print("PASS: test_structured_tool_calls") + + +def test_xml_tool_call_at_start(): + """Model emits JSON instead of structured tool_calls. + Buffer should detect prefix and drain.""" + backend = _make_backend() + + tc_json = json.dumps({"name": "web_search", "arguments": {"query": "hello"}}) + content = f"{tc_json}" + + # Stream the XML content token by token to simulate real streaming + chunks = [_make_chunk({"role": "assistant"})] + for char in content: + chunks.append(_make_chunk({"content": char})) + chunks.append(_make_chunk({}, finish_reason="stop")) + + usage = {"prompt_tokens": 10, "completion_tokens": len(content), "total_tokens": 10 + len(content)} + sse = _build_sse_stream(chunks, final_usage=usage) + + events = _collect_events(backend, sse) + + # Should detect tool call and execute it + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 1, f"Expected 1 tool_start, got: {events}" + assert tool_starts[0]["tool_name"] == "web_search" + + tool_ends = [e for e in events if e["type"] == "tool_end"] + assert len(tool_ends) == 1, f"Expected 1 tool_end" + + print("PASS: test_xml_tool_call_at_start") + + +def test_xml_function_tag_at_start(): + """Model emits tag. + Buffer should detect . Buffer should strip + leading whitespace before prefix check.""" + backend = _make_backend() + + tc_json = json.dumps({"name": "web_search", "arguments": {"query": "test"}}) + content = f" \n {tc_json}" + + chunks = [_make_chunk({"role": "assistant"})] + # Send whitespace as one chunk, then the rest + chunks.append(_make_chunk({"content": " \n "})) + rest = f"{tc_json}" + for char in rest: + chunks.append(_make_chunk({"content": char})) + chunks.append(_make_chunk({}, finish_reason="stop")) + + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse) + + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 1, f"Expected 1 tool_start after whitespace, got: {events}" + + print("PASS: test_whitespace_before_tool_xml") + + +def test_content_then_tool_xml_safety_net(): + """Rare case: model emits normal content first, then tool XML later. + Safety net at [DONE] should catch the tool call.""" + backend = _make_backend() + + tc_json = json.dumps({"name": "web_search", "arguments": {"query": "q"}}) + # Start with normal text (triggers STREAMING), then tool XML + # Send as separate content chunks + chunks = [_make_chunk({"role": "assistant"})] + chunks.append(_make_chunk({"content": "Let me search for that. "})) + chunks.append(_make_chunk({"content": f"{tc_json}"})) + chunks.append(_make_chunk({}, finish_reason="stop")) + + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse) + + # The safety net should catch the tool call + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) >= 1, f"Safety net should catch tool, got: {events}" + assert tool_starts[0]["tool_name"] == "web_search" + + print("PASS: test_content_then_tool_xml_safety_net") + + +def test_multiple_structured_tool_calls(): + """Model calls two tools in one response (parallel tool calls).""" + backend = _make_backend() + + tools = [ + {"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}, + {"type": "function", "function": {"name": "python", + "parameters": {"type": "object", "properties": {"code": {"type": "string"}}}}}, + ] + + chunks = [ + _make_chunk({"role": "assistant"}), + # Two tool calls streamed with different indices + _make_chunk({"tool_calls": [ + {"index": 0, "id": "call_0", "function": {"name": "web_search", "arguments": ""}}, + {"index": 1, "id": "call_1", "function": {"name": "python", "arguments": ""}}, + ]}), + _make_chunk({"tool_calls": [ + {"index": 0, "function": {"arguments": '{"query": "test"}'}}, + ]}), + _make_chunk({"tool_calls": [ + {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, + ]}), + _make_chunk({}, finish_reason="tool_calls"), + ] + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse, tools=tools) + + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 2, f"Expected 2 tool_start events, got: {tool_starts}" + + names = {ts["tool_name"] for ts in tool_starts} + assert names == {"web_search", "python"}, f"Wrong tool names: {names}" + + print("PASS: test_multiple_structured_tool_calls") + + +def test_reasoning_tokens_stream_immediately(): + """Thinking model: reasoning_content tokens should stream to user + immediately, even during BUFFERING state.""" + backend = _make_backend() + backend._supports_reasoning = True + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"reasoning_content": "Let me think..."}), + _make_chunk({"reasoning_content": " about this."}), + _make_chunk({"content": "The answer is 42."}), + _make_chunk({}, finish_reason="stop"), + ] + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse, enable_thinking=True) + + content_events = [e for e in events if e["type"] == "content"] + assert len(content_events) >= 3, f"Expected at least 3 content events (2 reasoning + 1 content), got: {content_events}" + + # First content events should contain tag + assert "" in content_events[0]["text"], "First content should have tag" + # Last content should have the actual answer + final = content_events[-1]["text"] + assert "42" in final, f"Final content should have answer: {final}" + + print("PASS: test_reasoning_tokens_stream_immediately") + + +def test_reasoning_then_tool_call(): + """Thinking model that reasons then calls a tool. + Reasoning should stream, then tool detected and executed.""" + backend = _make_backend() + backend._supports_reasoning = True + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"reasoning_content": "I need to search for this."}), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), + ] + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse, enable_thinking=True) + + # Reasoning should have been yielded + content_events = [e for e in events if e["type"] == "content"] + assert len(content_events) >= 1, "Reasoning should be yielded" + assert "" in content_events[0]["text"] + assert "I need to search" in content_events[0]["text"] + + # Tool should be executed + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 1, f"Expected tool_start: {events}" + assert tool_starts[0]["tool_name"] == "web_search" + + print("PASS: test_reasoning_then_tool_call") + + +def test_empty_response(): + """Model returns empty stream (just role + [DONE]). Should not crash.""" + backend = _make_backend() + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({}, finish_reason="stop"), + ] + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse) + + # Should not crash, just return with no content + error_events = [e for e in events if e.get("type") == "error"] + assert len(error_events) == 0, f"Should not error: {error_events}" + + print("PASS: test_empty_response") + + +def test_buffer_prefix_timeout(): + """Content starts with '<' but is not a tool call (e.g., '

Hello

'). + Buffer should hold briefly then flush when no prefix match at 32 chars.""" + backend = _make_backend() + + content = "

This is a paragraph of HTML content that is not a tool call

" + chunks = [_make_chunk({"role": "assistant"})] + # Stream char by char + for char in content: + chunks.append(_make_chunk({"content": char})) + chunks.append(_make_chunk({}, finish_reason="stop")) + + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse) + + content_events = [e for e in events if e["type"] == "content"] + assert len(content_events) >= 1, f"Should have content events: {events}" + + # No tool calls should be detected + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 0, f"Should not detect tools in HTML: {tool_starts}" + + # Final content should contain the HTML + final = content_events[-1]["text"] + assert "

" in final, f"HTML content should pass through: {final}" + + print("PASS: test_buffer_prefix_timeout") + + +def test_buffer_resolves_to_streaming_on_non_xml_first_char(): + """First content char is not '<' and not whitespace. + Should immediately transition to STREAMING.""" + backend = _make_backend() + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"content": "H"}), # 'H' is not '<', instant STREAMING + _make_chunk({"content": "ello"}), + ] + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse) + + content_events = [e for e in events if e["type"] == "content"] + # First content event should appear immediately with just "H" + assert len(content_events) >= 1 + assert "H" in content_events[0]["text"] + + print("PASS: test_buffer_resolves_to_streaming_on_non_xml_first_char") + + +def test_draining_false_positive(): + """Buffer detects 'Use a screwdriver"}), + _make_chunk({}, finish_reason="stop"), + ] + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse) + + # "" so it enters BUFFERING. + # Then "_tip>" does NOT match "" since the buffer becomes + # "..." which doesn't start with "" or "32 chars the buffer should flush. + # No tool should be executed. + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 0, f"Should not detect tool in : {tool_starts}" + + print("PASS: test_draining_false_positive") + + +def test_structured_tool_args_json_parsing(): + """Verify that arguments streamed across multiple chunks get reassembled + and parsed correctly as JSON.""" + backend = _make_backend() + + # Arguments split across 4 chunks + arg_parts = ['{"qu', 'ery":', ' "wha', 't is python?"}'] + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_abc", + "function": {"name": "web_search", "arguments": ""}}]}), + ] + for part in arg_parts: + chunks.append(_make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": part}}]})) + chunks.append(_make_chunk({}, finish_reason="tool_calls")) + + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse) + + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 1 + assert tool_starts[0]["arguments"] == {"query": "what is python?"}, \ + f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" + assert tool_starts[0]["tool_call_id"] == "call_abc" + + print("PASS: test_structured_tool_args_json_parsing") + + +def test_auto_heal_disabled(): + """When auto_heal_tool_calls=False, XML tool calls in content should NOT + be parsed -- only structured tool_calls are honored.""" + backend = _make_backend() + + tc_json = json.dumps({"name": "web_search", "arguments": {"query": "test"}}) + content = f"{tc_json}" + + chunks = [_make_chunk({"role": "assistant"})] + # Send as one big content chunk + chunks.append(_make_chunk({"content": content})) + chunks.append(_make_chunk({}, finish_reason="stop")) + + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse, auto_heal_tool_calls=False) + + # With auto_heal disabled, the XML should NOT be parsed as a tool call + tool_starts = [e for e in events if e["type"] == "tool_start"] + assert len(tool_starts) == 0, \ + f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" + + print("PASS: test_auto_heal_disabled") + + +def test_metrics_accumulation_across_tool_iterations(): + """When tools are called, metrics from the tool iteration should be + accumulated and included in the final metadata.""" + backend = _make_backend() + + # First iteration: tool call + tool_chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), + ] + tool_usage = {"prompt_tokens": 10, "completion_tokens": 5} + tool_timings = {"predicted_ms": 50, "predicted_n": 5} + tool_sse = _build_sse_stream(tool_chunks, final_usage=tool_usage, final_timings=tool_timings) + + # Second iteration: plain text response (synthesis) + synth_chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"content": "Based on my search, the answer is X."}), + _make_chunk({}, finish_reason="stop"), + ] + synth_usage = {"prompt_tokens": 20, "completion_tokens": 8} + synth_timings = {"predicted_ms": 100, "predicted_n": 8} + synth_sse = _build_sse_stream(synth_chunks, final_usage=synth_usage, final_timings=synth_timings) + + # We need to return different SSE streams for each iteration + call_count = [0] + original_sse = [tool_sse, synth_sse] + + fake_responses = [FakeResponse(tool_sse), FakeResponse(synth_sse)] + + @contextlib.contextmanager + def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): + idx = min(call_count[0], len(fake_responses) - 1) + call_count[0] += 1 + yield fake_responses[idx] + + def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): + return "Search result: success" + + backend._stream_with_retry = fake_stream_with_retry + + events = [] + with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): + for event in backend.generate_chat_completion_with_tools( + messages=[{"role": "user", "content": "Search for test"}], + tools=[{"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}], + ): + events.append(event) + + meta_events = [e for e in events if e["type"] == "metadata"] + assert len(meta_events) == 1, f"Expected exactly 1 metadata event, got: {meta_events}" + + meta = meta_events[0] + # completion_tokens should be accumulated: 5 (tool iter) + 8 (synthesis) = 13 + assert meta["usage"]["completion_tokens"] == 13, \ + f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" + + # predicted_ms and predicted_n should also accumulate + assert meta["timings"]["predicted_n"] == 13, \ + f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" + + print("PASS: test_metrics_accumulation_across_tool_iterations") + + +# ── Run all tests ──────────────────────────────────────────────────────── + +if __name__ == "__main__": + tests = [ + test_no_tool_call_plain_text, + test_structured_tool_calls, + test_xml_tool_call_at_start, + test_xml_function_tag_at_start, + test_whitespace_before_tool_xml, + test_content_then_tool_xml_safety_net, + test_multiple_structured_tool_calls, + test_reasoning_tokens_stream_immediately, + test_reasoning_then_tool_call, + test_empty_response, + test_buffer_prefix_timeout, + test_buffer_resolves_to_streaming_on_non_xml_first_char, + test_draining_false_positive, + test_structured_tool_args_json_parsing, + test_auto_heal_disabled, + test_metrics_accumulation_across_tool_iterations, + ] + + passed = 0 + failed = 0 + errors = [] + + for test_fn in tests: + try: + test_fn() + passed += 1 + except Exception as e: + failed += 1 + errors.append((test_fn.__name__, str(e))) + import traceback + print(f"FAIL: {test_fn.__name__}: {e}") + traceback.print_exc() + print() + + print(f"\n{'='*60}") + print(f"Results: {passed} passed, {failed} failed, {len(tests)} total") + if errors: + print(f"\nFailed tests:") + for name, err in errors: + print(f" - {name}: {err}") + print(f"{'='*60}") From 6d6b28db3ec27fb07450ab155f76b6e8e55a263c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 27 Mar 2026 06:22:50 +0000 Subject: [PATCH 04/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_streaming_tool_detection.py | 322 ++++++++++++++++++------- 1 file changed, 229 insertions(+), 93 deletions(-) diff --git a/tests/test_streaming_tool_detection.py b/tests/test_streaming_tool_detection.py index b2429f0f63..bdda6ed834 100644 --- a/tests/test_streaming_tool_detection.py +++ b/tests/test_streaming_tool_detection.py @@ -16,6 +16,7 @@ import sys, os # ── helpers ────────────────────────────────────────────────────────────── + def _sse_line(data: dict) -> str: """One SSE data line (no trailing blank line -- we add those in the stream).""" return f"data: {json.dumps(data)}" @@ -25,7 +26,7 @@ def _sse_done() -> str: return "data: [DONE]" -def _make_chunk(delta: dict, finish_reason=None, usage=None, timings=None): +def _make_chunk(delta: dict, finish_reason = None, usage = None, timings = None): """Build a chat-completions streaming chunk.""" choice = {"index": 0, "delta": delta} if finish_reason: @@ -38,7 +39,7 @@ def _make_chunk(delta: dict, finish_reason=None, usage=None, timings=None): return chunk -def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) -> str: +def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = None) -> str: """ Build a complete SSE text stream from a list of chunk dicts. Includes the role chunk, content/tool chunks, and [DONE]. @@ -46,7 +47,7 @@ def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) lines = [] for c in chunks: lines.append(_sse_line(c)) - lines.append("") # blank line separator + lines.append("") # blank line separator # Final usage chunk (if provided) if final_usage or final_timings: meta = {} @@ -64,6 +65,7 @@ def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) class FakeResponse: """Mimics httpx.Response for streaming.""" + def __init__(self, text: str, status_code: int = 200): self._text = text self.status_code = status_code @@ -82,6 +84,7 @@ class FakeResponse: class FakeClient: """Mimics httpx.Client context manager.""" + def __init__(self, response: FakeResponse): self._response = response @@ -92,26 +95,29 @@ class FakeClient: pass @contextlib.contextmanager - def stream(self, method, url, json=None, timeout=None, headers=None): + def stream(self, method, url, json = None, timeout = None, headers = None): yield self._response # ── Build a minimal LlamaCppBackend for testing ───────────────────────── + def _make_backend(): """Create a minimal mock backend with just enough to run the method.""" # We need the real class but only care about generate_chat_completion_with_tools # Import the real module - sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "unsloth_studio_src")) + sys.path.insert( + 0, os.path.join(os.path.dirname(__file__), "..", "unsloth_studio_src") + ) # Instead of importing the full module (which has other deps), we'll # build a lightweight object that has the method and its dependencies. from studio.backend.core.inference.llama_cpp import LlamaCppBackend backend = object.__new__(LlamaCppBackend) - backend._process = True # is_loaded checks _process is not None - backend._healthy = True # is_loaded checks _healthy - backend._port = 9999 # base_url property reads _port + backend._process = True # is_loaded checks _process is not None + backend._healthy = True # is_loaded checks _healthy + backend._port = 9999 # base_url property reads _port backend._api_key = None backend._supports_reasoning = False return backend @@ -122,13 +128,13 @@ def _synthesis_sse(): chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"content": "Done."}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] usage = {"prompt_tokens": 20, "completion_tokens": 1} - return _build_sse_stream(chunks, final_usage=usage) + return _build_sse_stream(chunks, final_usage = usage) -def _collect_events(backend, sse_text, tools=None, **kwargs): +def _collect_events(backend, sse_text, tools = None, **kwargs): """ Run generate_chat_completion_with_tools with a fake SSE stream and collect all yielded events. @@ -137,14 +143,24 @@ def _collect_events(backend, sse_text, tools=None, **kwargs): return a plain text synthesis response so the agentic loop terminates. """ if tools is None: - tools = [{"type": "function", "function": {"name": "web_search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}] + tools = [ + { + "type": "function", + "function": { + "name": "web_search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } + ] call_count = [0] synth_sse = _synthesis_sse() @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): + def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): idx = call_count[0] call_count[0] += 1 # First call: use the provided SSE. Subsequent: plain text synthesis. @@ -152,23 +168,26 @@ def _collect_events(backend, sse_text, tools=None, **kwargs): yield FakeResponse(text) # Patch execute_tool to return a dummy result - def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): + def fake_execute_tool( + tool_name, arguments, cancel_event = None, timeout = None, session_id = None + ): return f"Tool {tool_name} result: OK" original_stream = backend._stream_with_retry backend._stream_with_retry = fake_stream_with_retry events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): + with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): try: for event in backend.generate_chat_completion_with_tools( - messages=[{"role": "user", "content": "Hello"}], - tools=tools, + messages = [{"role": "user", "content": "Hello"}], + tools = tools, **kwargs, ): events.append(event) except Exception as e: import traceback + traceback.print_exc() events.append({"type": "error", "error": str(e)}) @@ -178,6 +197,7 @@ def _collect_events(backend, sse_text, tools=None, **kwargs): # ── The actual tests ───────────────────────────────────────────────────── + def test_no_tool_call_plain_text(): """90% case: model responds with plain text, no tool call. Should stream content immediately without delay.""" @@ -187,11 +207,11 @@ def test_no_tool_call_plain_text(): _make_chunk({"role": "assistant"}), _make_chunk({"content": "Hello"}), _make_chunk({"content": " there"}), - _make_chunk({"content": "!"}, finish_reason="stop"), + _make_chunk({"content": "!"}, finish_reason = "stop"), ] usage = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13} timings = {"predicted_ms": 100, "predicted_n": 3, "predicted_per_second": 30.0} - sse = _build_sse_stream(chunks, final_usage=usage, final_timings=timings) + sse = _build_sse_stream(chunks, final_usage = usage, final_timings = timings) events = _collect_events(backend, sse) @@ -221,21 +241,34 @@ def test_structured_tool_calls(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", - "function": {"name": "web_search", "arguments": ""}}]}), - _make_chunk({"tool_calls": [{"index": 0, - "function": {"arguments": '{"query":'}}]}), - _make_chunk({"tool_calls": [{"index": 0, - "function": {"arguments": ' "test"}'}}]}), - _make_chunk({}, finish_reason="tool_calls"), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "function": {"name": "web_search", "arguments": ""}, + } + ] + } + ), + _make_chunk( + {"tool_calls": [{"index": 0, "function": {"arguments": '{"query":'}}]} + ), + _make_chunk( + {"tool_calls": [{"index": 0, "function": {"arguments": ' "test"}'}}]} + ), + _make_chunk({}, finish_reason = "tool_calls"), ] usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - sse = _build_sse_stream(chunks, final_usage=usage) + sse = _build_sse_stream(chunks, final_usage = usage) events = _collect_events(backend, sse) # Should have status update for tool execution - status_events = [e for e in events if e["type"] == "status" and "Searching" in e.get("text", "")] + status_events = [ + e for e in events if e["type"] == "status" and "Searching" in e.get("text", "") + ] assert len(status_events) >= 1, f"Expected search status, got: {events}" # Should have tool_start event @@ -263,10 +296,14 @@ def test_xml_tool_call_at_start(): chunks = [_make_chunk({"role": "assistant"})] for char in content: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) - usage = {"prompt_tokens": 10, "completion_tokens": len(content), "total_tokens": 10 + len(content)} - sse = _build_sse_stream(chunks, final_usage=usage) + usage = { + "prompt_tokens": 10, + "completion_tokens": len(content), + "total_tokens": 10 + len(content), + } + sse = _build_sse_stream(chunks, final_usage = usage) events = _collect_events(backend, sse) @@ -286,15 +323,15 @@ def test_xml_function_tag_at_start(): Buffer should detect {tc_json}" for char in rest: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 1, f"Expected 1 tool_start after whitespace, got: {events}" + assert ( + len(tool_starts) == 1 + ), f"Expected 1 tool_start after whitespace, got: {events}" print("PASS: test_whitespace_before_tool_xml") @@ -341,7 +380,7 @@ def test_content_then_tool_xml_safety_net(): chunks = [_make_chunk({"role": "assistant"})] chunks.append(_make_chunk({"content": "Let me search for that. "})) chunks.append(_make_chunk({"content": f"{tc_json}"})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -359,29 +398,65 @@ def test_multiple_structured_tool_calls(): backend = _make_backend() tools = [ - {"type": "function", "function": {"name": "web_search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}, - {"type": "function", "function": {"name": "python", - "parameters": {"type": "object", "properties": {"code": {"type": "string"}}}}}, + { + "type": "function", + "function": { + "name": "web_search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + }, + { + "type": "function", + "function": { + "name": "python", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + }, + }, + }, ] chunks = [ _make_chunk({"role": "assistant"}), # Two tool calls streamed with different indices - _make_chunk({"tool_calls": [ - {"index": 0, "id": "call_0", "function": {"name": "web_search", "arguments": ""}}, - {"index": 1, "id": "call_1", "function": {"name": "python", "arguments": ""}}, - ]}), - _make_chunk({"tool_calls": [ - {"index": 0, "function": {"arguments": '{"query": "test"}'}}, - ]}), - _make_chunk({"tool_calls": [ - {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, - ]}), - _make_chunk({}, finish_reason="tool_calls"), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "function": {"name": "web_search", "arguments": ""}, + }, + { + "index": 1, + "id": "call_1", + "function": {"name": "python", "arguments": ""}, + }, + ] + } + ), + _make_chunk( + { + "tool_calls": [ + {"index": 0, "function": {"arguments": '{"query": "test"}'}}, + ] + } + ), + _make_chunk( + { + "tool_calls": [ + {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, + ] + } + ), + _make_chunk({}, finish_reason = "tool_calls"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, tools=tools) + events = _collect_events(backend, sse, tools = tools) tool_starts = [e for e in events if e["type"] == "tool_start"] assert len(tool_starts) == 2, f"Expected 2 tool_start events, got: {tool_starts}" @@ -403,16 +478,20 @@ def test_reasoning_tokens_stream_immediately(): _make_chunk({"reasoning_content": "Let me think..."}), _make_chunk({"reasoning_content": " about this."}), _make_chunk({"content": "The answer is 42."}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking=True) + events = _collect_events(backend, sse, enable_thinking = True) content_events = [e for e in events if e["type"] == "content"] - assert len(content_events) >= 3, f"Expected at least 3 content events (2 reasoning + 1 content), got: {content_events}" + assert ( + len(content_events) >= 3 + ), f"Expected at least 3 content events (2 reasoning + 1 content), got: {content_events}" # First content events should contain tag - assert "" in content_events[0]["text"], "First content should have tag" + assert ( + "" in content_events[0]["text"] + ), "First content should have tag" # Last content should have the actual answer final = content_events[-1]["text"] assert "42" in final, f"Final content should have answer: {final}" @@ -429,12 +508,24 @@ def test_reasoning_then_tool_call(): chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"reasoning_content": "I need to search for this."}), - _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", - "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), - _make_chunk({}, finish_reason="tool_calls"), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "function": { + "name": "web_search", + "arguments": '{"query": "test"}', + }, + } + ] + } + ), + _make_chunk({}, finish_reason = "tool_calls"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking=True) + events = _collect_events(backend, sse, enable_thinking = True) # Reasoning should have been yielded content_events = [e for e in events if e["type"] == "content"] @@ -456,7 +547,7 @@ def test_empty_response(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -478,7 +569,7 @@ def test_buffer_prefix_timeout(): # Stream char by char for char in content: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -528,7 +619,7 @@ def test_draining_false_positive(): _make_chunk({"role": "assistant"}), _make_chunk({"content": "Use a screwdriver"}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -554,21 +645,32 @@ def test_structured_tool_args_json_parsing(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({"tool_calls": [{"index": 0, "id": "call_abc", - "function": {"name": "web_search", "arguments": ""}}]}), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_abc", + "function": {"name": "web_search", "arguments": ""}, + } + ] + } + ), ] for part in arg_parts: - chunks.append(_make_chunk({"tool_calls": [{"index": 0, - "function": {"arguments": part}}]})) - chunks.append(_make_chunk({}, finish_reason="tool_calls")) + chunks.append( + _make_chunk({"tool_calls": [{"index": 0, "function": {"arguments": part}}]}) + ) + chunks.append(_make_chunk({}, finish_reason = "tool_calls")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) tool_starts = [e for e in events if e["type"] == "tool_start"] assert len(tool_starts) == 1 - assert tool_starts[0]["arguments"] == {"query": "what is python?"}, \ - f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" + assert tool_starts[0]["arguments"] == { + "query": "what is python?" + }, f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" assert tool_starts[0]["tool_call_id"] == "call_abc" print("PASS: test_structured_tool_args_json_parsing") @@ -585,15 +687,16 @@ def test_auto_heal_disabled(): chunks = [_make_chunk({"role": "assistant"})] # Send as one big content chunk chunks.append(_make_chunk({"content": content})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, auto_heal_tool_calls=False) + events = _collect_events(backend, sse, auto_heal_tool_calls = False) # With auto_heal disabled, the XML should NOT be parsed as a tool call tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 0, \ - f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" + assert ( + len(tool_starts) == 0 + ), f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" print("PASS: test_auto_heal_disabled") @@ -606,23 +709,39 @@ def test_metrics_accumulation_across_tool_iterations(): # First iteration: tool call tool_chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", - "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), - _make_chunk({}, finish_reason="tool_calls"), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "function": { + "name": "web_search", + "arguments": '{"query": "test"}', + }, + } + ] + } + ), + _make_chunk({}, finish_reason = "tool_calls"), ] tool_usage = {"prompt_tokens": 10, "completion_tokens": 5} tool_timings = {"predicted_ms": 50, "predicted_n": 5} - tool_sse = _build_sse_stream(tool_chunks, final_usage=tool_usage, final_timings=tool_timings) + tool_sse = _build_sse_stream( + tool_chunks, final_usage = tool_usage, final_timings = tool_timings + ) # Second iteration: plain text response (synthesis) synth_chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"content": "Based on my search, the answer is X."}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] synth_usage = {"prompt_tokens": 20, "completion_tokens": 8} synth_timings = {"predicted_ms": 100, "predicted_n": 8} - synth_sse = _build_sse_stream(synth_chunks, final_usage=synth_usage, final_timings=synth_timings) + synth_sse = _build_sse_stream( + synth_chunks, final_usage = synth_usage, final_timings = synth_timings + ) # We need to return different SSE streams for each iteration call_count = [0] @@ -631,36 +750,52 @@ def test_metrics_accumulation_across_tool_iterations(): fake_responses = [FakeResponse(tool_sse), FakeResponse(synth_sse)] @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): + def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): idx = min(call_count[0], len(fake_responses) - 1) call_count[0] += 1 yield fake_responses[idx] - def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): + def fake_execute_tool( + tool_name, arguments, cancel_event = None, timeout = None, session_id = None + ): return "Search result: success" backend._stream_with_retry = fake_stream_with_retry events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): + with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): for event in backend.generate_chat_completion_with_tools( - messages=[{"role": "user", "content": "Search for test"}], - tools=[{"type": "function", "function": {"name": "web_search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}], + messages = [{"role": "user", "content": "Search for test"}], + tools = [ + { + "type": "function", + "function": { + "name": "web_search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } + ], ): events.append(event) meta_events = [e for e in events if e["type"] == "metadata"] - assert len(meta_events) == 1, f"Expected exactly 1 metadata event, got: {meta_events}" + assert ( + len(meta_events) == 1 + ), f"Expected exactly 1 metadata event, got: {meta_events}" meta = meta_events[0] # completion_tokens should be accumulated: 5 (tool iter) + 8 (synthesis) = 13 - assert meta["usage"]["completion_tokens"] == 13, \ - f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" + assert ( + meta["usage"]["completion_tokens"] == 13 + ), f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" # predicted_ms and predicted_n should also accumulate - assert meta["timings"]["predicted_n"] == 13, \ - f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" + assert ( + meta["timings"]["predicted_n"] == 13 + ), f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" print("PASS: test_metrics_accumulation_across_tool_iterations") @@ -699,6 +834,7 @@ if __name__ == "__main__": failed += 1 errors.append((test_fn.__name__, str(e))) import traceback + print(f"FAIL: {test_fn.__name__}: {e}") traceback.print_exc() print() From 620b1522108631c3fbc6f4d7cb642b500d0cbdcd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 06:44:56 +0000 Subject: [PATCH 05/18] 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 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. --- studio/backend/core/inference/llama_cpp.py | 91 ++++++++-------------- 1 file changed, 33 insertions(+), 58 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c4901d4df2..d1f878fed8 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1882,11 +1882,18 @@ class LlamaCppBackend: ] += func["arguments"] continue - # ── Reasoning tokens (bypass buffer) ── - reasoning = delta.get("reasoning_content", "") + # ── Reasoning tokens ── + # Only yield in STREAMING state. In BUFFERING + # and DRAINING, accumulate silently so we don't + # corrupt the consumer's prev_text tracker + # (routes/inference.py never resets prev_text + # between tool iterations). + reasoning = delta.get( + "reasoning_content", "" + ) if reasoning: reasoning_accum += reasoning - if detect_state != _S_DRAINING: + if detect_state == _S_STREAMING: if not in_thinking: cumulative_display += "" in_thinking = True @@ -1948,9 +1955,12 @@ class LlamaCppBackend: else: # Not a tool -- flush buffer detect_state = _S_STREAMING - if in_thinking: + # Flush any reasoning accumulated + # during BUFFERING phase + if reasoning_accum: + cumulative_display += "" + cumulative_display += reasoning_accum cumulative_display += "" - in_thinking = False cumulative_display += content_buffer cleaned = _strip_tool_markup( cumulative_display, @@ -1981,9 +1991,11 @@ class LlamaCppBackend: elif content_accum or reasoning_accum: detect_state = _S_STREAMING if content_buffer: - if in_thinking: + # Flush any reasoning accumulated first + if reasoning_accum: + cumulative_display += "" + cumulative_display += reasoning_accum cumulative_display += "" - in_thinking = False cumulative_display += content_buffer yield { "type": "content", @@ -1992,6 +2004,16 @@ class LlamaCppBackend: final = True, ), } + elif reasoning_accum and not has_content_tokens: + # Reasoning-only response (no content tokens): + # show reasoning as plain text, matching + # the final streaming pass behavior for + # models that put everything in reasoning. + cumulative_display = reasoning_accum + yield { + "type": "content", + "text": cumulative_display, + } else: return @@ -2039,33 +2061,9 @@ class LlamaCppBackend: # Safety net caught tool XML -- treat as tool call tool_calls = _safety_tc - content_text = content_accum - import re - - content_text = re.sub( - r".*?", - "", - content_text, - flags = re.DOTALL, + content_text = _strip_tool_markup( + content_accum, final = True, ) - content_text = re.sub( - r".*$", - "", - content_text, - flags = re.DOTALL, - ) - content_text = re.sub( - r".*?", - "", - content_text, - flags = re.DOTALL, - ) - content_text = re.sub( - r".*$", - "", - content_text, - flags = re.DOTALL, - ).strip() logger.info( f"Safety net: parsed {len(tool_calls)} tool call(s) " f"from streamed content" @@ -2085,32 +2083,9 @@ class LlamaCppBackend: content_accum, ) if tool_calls and not has_structured_tc: - import re - - content_text = re.sub( - r".*?", - "", - content_text, - flags = re.DOTALL, + content_text = _strip_tool_markup( + content_text, final = True, ) - content_text = re.sub( - r".*$", - "", - content_text, - flags = re.DOTALL, - ) - content_text = re.sub( - r".*?", - "", - content_text, - flags = re.DOTALL, - ) - content_text = re.sub( - r".*$", - "", - content_text, - flags = re.DOTALL, - ).strip() if tool_calls: logger.info( f"Parsed {len(tool_calls)} tool call(s) from " From e75c2013d5de59a97181216277710cce7df881f7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 27 Mar 2026 06:45:11 +0000 Subject: [PATCH 06/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index d1f878fed8..746ab26abb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1888,9 +1888,7 @@ class LlamaCppBackend: # corrupt the consumer's prev_text tracker # (routes/inference.py never resets prev_text # between tool iterations). - reasoning = delta.get( - "reasoning_content", "" - ) + reasoning = delta.get("reasoning_content", "") if reasoning: reasoning_accum += reasoning if detect_state == _S_STREAMING: @@ -1959,7 +1957,9 @@ class LlamaCppBackend: # during BUFFERING phase if reasoning_accum: cumulative_display += "" - cumulative_display += reasoning_accum + cumulative_display += ( + reasoning_accum + ) cumulative_display += "" cumulative_display += content_buffer cleaned = _strip_tool_markup( @@ -2062,7 +2062,8 @@ class LlamaCppBackend: # Safety net caught tool XML -- treat as tool call tool_calls = _safety_tc content_text = _strip_tool_markup( - content_accum, final = True, + content_accum, + final = True, ) logger.info( f"Safety net: parsed {len(tool_calls)} tool call(s) " @@ -2084,7 +2085,8 @@ class LlamaCppBackend: ) if tool_calls and not has_structured_tc: content_text = _strip_tool_markup( - content_text, final = True, + content_text, + final = True, ) if tool_calls: logger.info( From aa9bea12b9f077b23edd8d46e042da71196480fb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 06:46:23 +0000 Subject: [PATCH 07/18] 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 "" or "", "", " Date: Fri, 27 Mar 2026 06:46:42 +0000 Subject: [PATCH 08/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b3ed3124db..7b4f123ba7 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1878,9 +1878,7 @@ class LlamaCppBackend: elif tc_d.get("id"): # Update ID if real one # arrives on a later delta - tool_calls_acc[idx]["id"] = ( - tc_d["id"] - ) + tool_calls_acc[idx]["id"] = tc_d["id"] func = tc_d.get("function", {}) if func.get("name"): tool_calls_acc[idx]["function"][ From 09399ab877468cc60572b0718eca3518519f53db Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 07:19:44 +0000 Subject: [PATCH 09/18] 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) --- studio/backend/routes/inference.py | 38 ++++++++++++++++++++---------- 1 file changed, 25 insertions(+), 13 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 78d95fedbd..cde29246b9 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -930,10 +930,13 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" + _dc_counter = 0 for chunk_text in audio_input_generate(): - if await request.is_disconnected(): - cancel_event.set() - return + _dc_counter += 1 + if _dc_counter % 20 == 0: + if await request.is_disconnected(): + cancel_event.set() + return if chunk_text: chunk = ChatCompletionChunk( id = completion_id, @@ -1107,10 +1110,13 @@ async def openai_chat_completions( prev_text = "" _stream_usage = None _stream_timings = None + _dc_counter = 0 while True: - if await request.is_disconnected(): - cancel_event.set() - return + _dc_counter += 1 + if _dc_counter % 20 == 0: + if await request.is_disconnected(): + cancel_event.set() + return event = await asyncio.to_thread(next, gen, _tool_sentinel) if event is _tool_sentinel: @@ -1256,10 +1262,13 @@ async def openai_chat_completions( prev_text = "" _stream_usage = None _stream_timings = None + _dc_counter = 0 while True: - if await request.is_disconnected(): - cancel_event.set() - return + _dc_counter += 1 + if _dc_counter % 20 == 0: + if await request.is_disconnected(): + cancel_event.set() + return cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) if cumulative is _gguf_sentinel: break @@ -1465,6 +1474,7 @@ async def openai_chat_completions( _DONE = object() # sentinel for generator exhaustion loop = asyncio.get_event_loop() gen = generate() + _dc_counter = 0 while True: # next(gen, _DONE) returns _DONE instead of raising # StopIteration — StopIteration cannot propagate @@ -1472,10 +1482,12 @@ async def openai_chat_completions( cumulative = await loop.run_in_executor(None, next, gen, _DONE) if cumulative is _DONE: break - if await request.is_disconnected(): - cancel_event.set() - backend.reset_generation_state() - return + _dc_counter += 1 + if _dc_counter % 20 == 0: + if await request.is_disconnected(): + cancel_event.set() + backend.reset_generation_state() + return new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: From 32ce2324f04daec306de0c10249b7370d5c360da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 08:22:24 +0000 Subject: [PATCH 10/18] 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. --- studio/backend/core/inference/llama_cpp.py | 51 ++- tests/test_streaming_tool_detection.py | 401 ++++++++------------- 2 files changed, 197 insertions(+), 255 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7b4f123ba7..6ce7a9b51c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2027,10 +2027,16 @@ class LlamaCppBackend: # ── STREAMING path: no tool call ── if detect_state == _S_STREAMING: - # Safety net: check for XML tool signals in content + # Safety net: check for XML tool signals in content. + # Only if we have NOT already emitted visible text -- + # retroactively switching to tool mode after the user + # has seen content violates the streaming contract and + # corrupts the route-layer cumulative delta tracker. _safety_tc = None - if auto_heal_tool_calls and any( - s in content_accum for s in _TOOL_XML_SIGNALS + if ( + auto_heal_tool_calls + and not _last_emitted + and any(s in content_accum for s in _TOOL_XML_SIGNALS) ): _safety_tc = self._parse_tool_calls_from_text( content_accum, @@ -2102,15 +2108,46 @@ class LlamaCppBackend: f"{'structured delta' if has_structured_tc else 'content text'}" ) if not tool_calls: - # DRAINING but no tool calls (false positive) + # DRAINING but no tool calls (false positive). + # Merge accumulated metrics from prior tool + # iterations so they are not silently dropped. yield {"type": "status", "text": ""} if content_accum: yield {"type": "content", "text": content_accum} - if _iter_usage or _iter_timings: + _fu = _iter_usage or {} + _fc = _fu.get("completion_tokens", 0) + _fp = _fu.get("prompt_tokens", 0) + _tc = _fc + _accumulated_completion_tokens + if _iter_usage or _iter_timings or _accumulated_completion_tokens: + _mt = ( + dict(_iter_timings) if _iter_timings else {} + ) + if ( + _accumulated_predicted_ms + or _accumulated_predicted_n + ): + _mt["predicted_ms"] = ( + _mt.get("predicted_ms", 0) + + _accumulated_predicted_ms + ) + _tn = ( + _mt.get("predicted_n", 0) + + _accumulated_predicted_n + ) + _mt["predicted_n"] = _tn + _tms = _mt["predicted_ms"] + if _tms > 0: + _mt["predicted_per_second"] = ( + _tn / (_tms / 1000.0) + ) yield { "type": "metadata", - "usage": _iter_usage, - "timings": _iter_timings, + "usage": { + "prompt_tokens": _fp, + "completion_tokens": _tc, + "total_tokens": _fp + _tc, + }, + "timings": _mt, } return diff --git a/tests/test_streaming_tool_detection.py b/tests/test_streaming_tool_detection.py index bdda6ed834..53b8adb87c 100644 --- a/tests/test_streaming_tool_detection.py +++ b/tests/test_streaming_tool_detection.py @@ -16,7 +16,6 @@ import sys, os # ── helpers ────────────────────────────────────────────────────────────── - def _sse_line(data: dict) -> str: """One SSE data line (no trailing blank line -- we add those in the stream).""" return f"data: {json.dumps(data)}" @@ -26,7 +25,7 @@ def _sse_done() -> str: return "data: [DONE]" -def _make_chunk(delta: dict, finish_reason = None, usage = None, timings = None): +def _make_chunk(delta: dict, finish_reason=None, usage=None, timings=None): """Build a chat-completions streaming chunk.""" choice = {"index": 0, "delta": delta} if finish_reason: @@ -39,7 +38,7 @@ def _make_chunk(delta: dict, finish_reason = None, usage = None, timings = None) return chunk -def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = None) -> str: +def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) -> str: """ Build a complete SSE text stream from a list of chunk dicts. Includes the role chunk, content/tool chunks, and [DONE]. @@ -47,7 +46,7 @@ def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = No lines = [] for c in chunks: lines.append(_sse_line(c)) - lines.append("") # blank line separator + lines.append("") # blank line separator # Final usage chunk (if provided) if final_usage or final_timings: meta = {} @@ -65,7 +64,6 @@ def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = No class FakeResponse: """Mimics httpx.Response for streaming.""" - def __init__(self, text: str, status_code: int = 200): self._text = text self.status_code = status_code @@ -84,7 +82,6 @@ class FakeResponse: class FakeClient: """Mimics httpx.Client context manager.""" - def __init__(self, response: FakeResponse): self._response = response @@ -95,29 +92,28 @@ class FakeClient: pass @contextlib.contextmanager - def stream(self, method, url, json = None, timeout = None, headers = None): + def stream(self, method, url, json=None, timeout=None, headers=None): yield self._response # ── Build a minimal LlamaCppBackend for testing ───────────────────────── - def _make_backend(): """Create a minimal mock backend with just enough to run the method.""" # We need the real class but only care about generate_chat_completion_with_tools # Import the real module - sys.path.insert( - 0, os.path.join(os.path.dirname(__file__), "..", "unsloth_studio_src") - ) + _repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) + if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) # Instead of importing the full module (which has other deps), we'll # build a lightweight object that has the method and its dependencies. from studio.backend.core.inference.llama_cpp import LlamaCppBackend backend = object.__new__(LlamaCppBackend) - backend._process = True # is_loaded checks _process is not None - backend._healthy = True # is_loaded checks _healthy - backend._port = 9999 # base_url property reads _port + backend._process = True # is_loaded checks _process is not None + backend._healthy = True # is_loaded checks _healthy + backend._port = 9999 # base_url property reads _port backend._api_key = None backend._supports_reasoning = False return backend @@ -128,13 +124,13 @@ def _synthesis_sse(): chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"content": "Done."}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] usage = {"prompt_tokens": 20, "completion_tokens": 1} - return _build_sse_stream(chunks, final_usage = usage) + return _build_sse_stream(chunks, final_usage=usage) -def _collect_events(backend, sse_text, tools = None, **kwargs): +def _collect_events(backend, sse_text, tools=None, **kwargs): """ Run generate_chat_completion_with_tools with a fake SSE stream and collect all yielded events. @@ -143,24 +139,14 @@ def _collect_events(backend, sse_text, tools = None, **kwargs): return a plain text synthesis response so the agentic loop terminates. """ if tools is None: - tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - } - ] + tools = [{"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}] call_count = [0] synth_sse = _synthesis_sse() @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): + def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): idx = call_count[0] call_count[0] += 1 # First call: use the provided SSE. Subsequent: plain text synthesis. @@ -168,26 +154,23 @@ def _collect_events(backend, sse_text, tools = None, **kwargs): yield FakeResponse(text) # Patch execute_tool to return a dummy result - def fake_execute_tool( - tool_name, arguments, cancel_event = None, timeout = None, session_id = None - ): + def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): return f"Tool {tool_name} result: OK" original_stream = backend._stream_with_retry backend._stream_with_retry = fake_stream_with_retry events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): + with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): try: for event in backend.generate_chat_completion_with_tools( - messages = [{"role": "user", "content": "Hello"}], - tools = tools, + messages=[{"role": "user", "content": "Hello"}], + tools=tools, **kwargs, ): events.append(event) except Exception as e: import traceback - traceback.print_exc() events.append({"type": "error", "error": str(e)}) @@ -197,7 +180,6 @@ def _collect_events(backend, sse_text, tools = None, **kwargs): # ── The actual tests ───────────────────────────────────────────────────── - def test_no_tool_call_plain_text(): """90% case: model responds with plain text, no tool call. Should stream content immediately without delay.""" @@ -207,11 +189,11 @@ def test_no_tool_call_plain_text(): _make_chunk({"role": "assistant"}), _make_chunk({"content": "Hello"}), _make_chunk({"content": " there"}), - _make_chunk({"content": "!"}, finish_reason = "stop"), + _make_chunk({"content": "!"}, finish_reason="stop"), ] usage = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13} timings = {"predicted_ms": 100, "predicted_n": 3, "predicted_per_second": 30.0} - sse = _build_sse_stream(chunks, final_usage = usage, final_timings = timings) + sse = _build_sse_stream(chunks, final_usage=usage, final_timings=timings) events = _collect_events(backend, sse) @@ -241,34 +223,21 @@ def test_structured_tool_calls(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": {"name": "web_search", "arguments": ""}, - } - ] - } - ), - _make_chunk( - {"tool_calls": [{"index": 0, "function": {"arguments": '{"query":'}}]} - ), - _make_chunk( - {"tool_calls": [{"index": 0, "function": {"arguments": ' "test"}'}}]} - ), - _make_chunk({}, finish_reason = "tool_calls"), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": ""}}]}), + _make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": '{"query":'}}]}), + _make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": ' "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), ] usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - sse = _build_sse_stream(chunks, final_usage = usage) + sse = _build_sse_stream(chunks, final_usage=usage) events = _collect_events(backend, sse) # Should have status update for tool execution - status_events = [ - e for e in events if e["type"] == "status" and "Searching" in e.get("text", "") - ] + status_events = [e for e in events if e["type"] == "status" and "Searching" in e.get("text", "")] assert len(status_events) >= 1, f"Expected search status, got: {events}" # Should have tool_start event @@ -296,14 +265,10 @@ def test_xml_tool_call_at_start(): chunks = [_make_chunk({"role": "assistant"})] for char in content: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) - usage = { - "prompt_tokens": 10, - "completion_tokens": len(content), - "total_tokens": 10 + len(content), - } - sse = _build_sse_stream(chunks, final_usage = usage) + usage = {"prompt_tokens": 10, "completion_tokens": len(content), "total_tokens": 10 + len(content)} + sse = _build_sse_stream(chunks, final_usage=usage) events = _collect_events(backend, sse) @@ -323,15 +288,15 @@ def test_xml_function_tag_at_start(): Buffer should detect {tc_json}" for char in rest: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) tool_starts = [e for e in events if e["type"] == "tool_start"] - assert ( - len(tool_starts) == 1 - ), f"Expected 1 tool_start after whitespace, got: {events}" + assert len(tool_starts) == 1, f"Expected 1 tool_start after whitespace, got: {events}" print("PASS: test_whitespace_before_tool_xml") -def test_content_then_tool_xml_safety_net(): +def test_content_then_tool_xml_no_retroactive_execution(): """Rare case: model emits normal content first, then tool XML later. - Safety net at [DONE] should catch the tool call.""" + Once visible content has been emitted to the user, we must NOT + retroactively switch to tool execution -- that would violate the + streaming contract and corrupt the route-layer cumulative delta + tracker. The tool XML is stripped by _strip_tool_markup, and the + user sees the cleaned content as a normal response.""" backend = _make_backend() tc_json = json.dumps({"name": "web_search", "arguments": {"query": "q"}}) # Start with normal text (triggers STREAMING), then tool XML - # Send as separate content chunks chunks = [_make_chunk({"role": "assistant"})] chunks.append(_make_chunk({"content": "Let me search for that. "})) chunks.append(_make_chunk({"content": f"{tc_json}"})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) - # The safety net should catch the tool call + # Tool should NOT be executed (visible content was already emitted) tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) >= 1, f"Safety net should catch tool, got: {events}" - assert tool_starts[0]["tool_name"] == "web_search" + assert len(tool_starts) == 0, ( + f"Should NOT retroactively execute tools after visible content: {tool_starts}" + ) - print("PASS: test_content_then_tool_xml_safety_net") + # Content should be present (tool XML stripped by _strip_tool_markup) + content_events = [e for e in events if e["type"] == "content"] + assert len(content_events) >= 1, f"Should have content: {events}" + assert "Let me search" in content_events[0]["text"] + + print("PASS: test_content_then_tool_xml_no_retroactive_execution") def test_multiple_structured_tool_calls(): @@ -398,65 +370,29 @@ def test_multiple_structured_tool_calls(): backend = _make_backend() tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - }, - { - "type": "function", - "function": { - "name": "python", - "parameters": { - "type": "object", - "properties": {"code": {"type": "string"}}, - }, - }, - }, + {"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}, + {"type": "function", "function": {"name": "python", + "parameters": {"type": "object", "properties": {"code": {"type": "string"}}}}}, ] chunks = [ _make_chunk({"role": "assistant"}), # Two tool calls streamed with different indices - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": {"name": "web_search", "arguments": ""}, - }, - { - "index": 1, - "id": "call_1", - "function": {"name": "python", "arguments": ""}, - }, - ] - } - ), - _make_chunk( - { - "tool_calls": [ - {"index": 0, "function": {"arguments": '{"query": "test"}'}}, - ] - } - ), - _make_chunk( - { - "tool_calls": [ - {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), + _make_chunk({"tool_calls": [ + {"index": 0, "id": "call_0", "function": {"name": "web_search", "arguments": ""}}, + {"index": 1, "id": "call_1", "function": {"name": "python", "arguments": ""}}, + ]}), + _make_chunk({"tool_calls": [ + {"index": 0, "function": {"arguments": '{"query": "test"}'}}, + ]}), + _make_chunk({"tool_calls": [ + {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, + ]}), + _make_chunk({}, finish_reason="tool_calls"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, tools = tools) + events = _collect_events(backend, sse, tools=tools) tool_starts = [e for e in events if e["type"] == "tool_start"] assert len(tool_starts) == 2, f"Expected 2 tool_start events, got: {tool_starts}" @@ -468,8 +404,9 @@ def test_multiple_structured_tool_calls(): def test_reasoning_tokens_stream_immediately(): - """Thinking model: reasoning_content tokens should stream to user - immediately, even during BUFFERING state.""" + """Thinking model: reasoning_content is accumulated during BUFFERING + and flushed together with content when transitioning to STREAMING. + The final output includes ... wrapping.""" backend = _make_backend() backend._supports_reasoning = True @@ -478,60 +415,46 @@ def test_reasoning_tokens_stream_immediately(): _make_chunk({"reasoning_content": "Let me think..."}), _make_chunk({"reasoning_content": " about this."}), _make_chunk({"content": "The answer is 42."}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking = True) + events = _collect_events(backend, sse, enable_thinking=True) content_events = [e for e in events if e["type"] == "content"] - assert ( - len(content_events) >= 3 - ), f"Expected at least 3 content events (2 reasoning + 1 content), got: {content_events}" + assert len(content_events) >= 1, f"Expected content events, got: {content_events}" - # First content events should contain tag - assert ( - "" in content_events[0]["text"] - ), "First content should have tag" - # Last content should have the actual answer + # Content should contain both tags and the answer final = content_events[-1]["text"] - assert "42" in final, f"Final content should have answer: {final}" + assert "" in final, f"Should have tag: {final}" + assert "Let me think" in final, f"Should have reasoning: {final}" + assert "42" in final, f"Should have answer: {final}" + assert "" in final, f"Should have closing : {final}" print("PASS: test_reasoning_tokens_stream_immediately") def test_reasoning_then_tool_call(): """Thinking model that reasons then calls a tool. - Reasoning should stream, then tool detected and executed.""" + Reasoning is silently accumulated during tool detection (matching + old non-streaming behavior) so the consumer's prev_text is not + corrupted for subsequent iterations. Tool is still detected.""" backend = _make_backend() backend._supports_reasoning = True chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"reasoning_content": "I need to search for this."}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": { - "name": "web_search", - "arguments": '{"query": "test"}', - }, - } - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking = True) + events = _collect_events(backend, sse, enable_thinking=True) - # Reasoning should have been yielded - content_events = [e for e in events if e["type"] == "content"] - assert len(content_events) >= 1, "Reasoning should be yielded" - assert "" in content_events[0]["text"] - assert "I need to search" in content_events[0]["text"] + # Reasoning should NOT be yielded during tool detection + # (prevents prev_text corruption in consumer). Instead it's + # accumulated silently, matching old non-streaming behavior. + # After tool execution, the synthesis pass handles display. # Tool should be executed tool_starts = [e for e in events if e["type"] == "tool_start"] @@ -541,13 +464,39 @@ def test_reasoning_then_tool_call(): print("PASS: test_reasoning_then_tool_call") +def test_reasoning_only_no_content(): + """Thinking model produces only reasoning_content with no content tokens. + Should yield reasoning as plain text (no wrapper), matching + the final streaming pass behavior for models like Qwen3 always-think.""" + backend = _make_backend() + backend._supports_reasoning = True + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"reasoning_content": "The answer is simply 42."}), + _make_chunk({}, finish_reason="stop"), + ] + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse, enable_thinking=True) + + content_events = [e for e in events if e["type"] == "content"] + assert len(content_events) >= 1, f"Should yield reasoning as content: {events}" + + final = content_events[-1]["text"] + assert "42" in final, f"Should contain reasoning text: {final}" + # Should NOT have wrapper (reasoning-only fallback) + assert "" not in final, f"Reasoning-only should not have wrapper: {final}" + + print("PASS: test_reasoning_only_no_content") + + def test_empty_response(): """Model returns empty stream (just role + [DONE]). Should not crash.""" backend = _make_backend() chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -569,7 +518,7 @@ def test_buffer_prefix_timeout(): # Stream char by char for char in content: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -619,7 +568,7 @@ def test_draining_false_positive(): _make_chunk({"role": "assistant"}), _make_chunk({"content": "Use a screwdriver"}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -645,32 +594,21 @@ def test_structured_tool_args_json_parsing(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_abc", - "function": {"name": "web_search", "arguments": ""}, - } - ] - } - ), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_abc", + "function": {"name": "web_search", "arguments": ""}}]}), ] for part in arg_parts: - chunks.append( - _make_chunk({"tool_calls": [{"index": 0, "function": {"arguments": part}}]}) - ) - chunks.append(_make_chunk({}, finish_reason = "tool_calls")) + chunks.append(_make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": part}}]})) + chunks.append(_make_chunk({}, finish_reason="tool_calls")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) tool_starts = [e for e in events if e["type"] == "tool_start"] assert len(tool_starts) == 1 - assert tool_starts[0]["arguments"] == { - "query": "what is python?" - }, f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" + assert tool_starts[0]["arguments"] == {"query": "what is python?"}, \ + f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" assert tool_starts[0]["tool_call_id"] == "call_abc" print("PASS: test_structured_tool_args_json_parsing") @@ -687,16 +625,15 @@ def test_auto_heal_disabled(): chunks = [_make_chunk({"role": "assistant"})] # Send as one big content chunk chunks.append(_make_chunk({"content": content})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, auto_heal_tool_calls = False) + events = _collect_events(backend, sse, auto_heal_tool_calls=False) # With auto_heal disabled, the XML should NOT be parsed as a tool call tool_starts = [e for e in events if e["type"] == "tool_start"] - assert ( - len(tool_starts) == 0 - ), f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" + assert len(tool_starts) == 0, \ + f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" print("PASS: test_auto_heal_disabled") @@ -709,39 +646,23 @@ def test_metrics_accumulation_across_tool_iterations(): # First iteration: tool call tool_chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": { - "name": "web_search", - "arguments": '{"query": "test"}', - }, - } - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), ] tool_usage = {"prompt_tokens": 10, "completion_tokens": 5} tool_timings = {"predicted_ms": 50, "predicted_n": 5} - tool_sse = _build_sse_stream( - tool_chunks, final_usage = tool_usage, final_timings = tool_timings - ) + tool_sse = _build_sse_stream(tool_chunks, final_usage=tool_usage, final_timings=tool_timings) # Second iteration: plain text response (synthesis) synth_chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"content": "Based on my search, the answer is X."}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] synth_usage = {"prompt_tokens": 20, "completion_tokens": 8} synth_timings = {"predicted_ms": 100, "predicted_n": 8} - synth_sse = _build_sse_stream( - synth_chunks, final_usage = synth_usage, final_timings = synth_timings - ) + synth_sse = _build_sse_stream(synth_chunks, final_usage=synth_usage, final_timings=synth_timings) # We need to return different SSE streams for each iteration call_count = [0] @@ -750,52 +671,36 @@ def test_metrics_accumulation_across_tool_iterations(): fake_responses = [FakeResponse(tool_sse), FakeResponse(synth_sse)] @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): + def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): idx = min(call_count[0], len(fake_responses) - 1) call_count[0] += 1 yield fake_responses[idx] - def fake_execute_tool( - tool_name, arguments, cancel_event = None, timeout = None, session_id = None - ): + def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): return "Search result: success" backend._stream_with_retry = fake_stream_with_retry events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): + with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): for event in backend.generate_chat_completion_with_tools( - messages = [{"role": "user", "content": "Search for test"}], - tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - } - ], + messages=[{"role": "user", "content": "Search for test"}], + tools=[{"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}], ): events.append(event) meta_events = [e for e in events if e["type"] == "metadata"] - assert ( - len(meta_events) == 1 - ), f"Expected exactly 1 metadata event, got: {meta_events}" + assert len(meta_events) == 1, f"Expected exactly 1 metadata event, got: {meta_events}" meta = meta_events[0] # completion_tokens should be accumulated: 5 (tool iter) + 8 (synthesis) = 13 - assert ( - meta["usage"]["completion_tokens"] == 13 - ), f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" + assert meta["usage"]["completion_tokens"] == 13, \ + f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" # predicted_ms and predicted_n should also accumulate - assert ( - meta["timings"]["predicted_n"] == 13 - ), f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" + assert meta["timings"]["predicted_n"] == 13, \ + f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" print("PASS: test_metrics_accumulation_across_tool_iterations") @@ -809,10 +714,11 @@ if __name__ == "__main__": test_xml_tool_call_at_start, test_xml_function_tag_at_start, test_whitespace_before_tool_xml, - test_content_then_tool_xml_safety_net, + test_content_then_tool_xml_no_retroactive_execution, test_multiple_structured_tool_calls, test_reasoning_tokens_stream_immediately, test_reasoning_then_tool_call, + test_reasoning_only_no_content, test_empty_response, test_buffer_prefix_timeout, test_buffer_resolves_to_streaming_on_non_xml_first_char, @@ -834,7 +740,6 @@ if __name__ == "__main__": failed += 1 errors.append((test_fn.__name__, str(e))) import traceback - print(f"FAIL: {test_fn.__name__}: {e}") traceback.print_exc() print() From 24f92fac86e2c03f36bfabbf5ef0a7f607283953 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 27 Mar 2026 08:22:54 +0000 Subject: [PATCH 11/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 22 +- tests/test_streaming_tool_detection.py | 324 +++++++++++++++------ 2 files changed, 237 insertions(+), 109 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 6ce7a9b51c..fdd5420851 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2118,28 +2118,24 @@ class LlamaCppBackend: _fc = _fu.get("completion_tokens", 0) _fp = _fu.get("prompt_tokens", 0) _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings or _accumulated_completion_tokens: - _mt = ( - dict(_iter_timings) if _iter_timings else {} - ) - if ( - _accumulated_predicted_ms - or _accumulated_predicted_n - ): + if ( + _iter_usage + or _iter_timings + or _accumulated_completion_tokens + ): + _mt = dict(_iter_timings) if _iter_timings else {} + if _accumulated_predicted_ms or _accumulated_predicted_n: _mt["predicted_ms"] = ( _mt.get("predicted_ms", 0) + _accumulated_predicted_ms ) _tn = ( - _mt.get("predicted_n", 0) - + _accumulated_predicted_n + _mt.get("predicted_n", 0) + _accumulated_predicted_n ) _mt["predicted_n"] = _tn _tms = _mt["predicted_ms"] if _tms > 0: - _mt["predicted_per_second"] = ( - _tn / (_tms / 1000.0) - ) + _mt["predicted_per_second"] = _tn / (_tms / 1000.0) yield { "type": "metadata", "usage": { diff --git a/tests/test_streaming_tool_detection.py b/tests/test_streaming_tool_detection.py index 53b8adb87c..25e25fba2e 100644 --- a/tests/test_streaming_tool_detection.py +++ b/tests/test_streaming_tool_detection.py @@ -16,6 +16,7 @@ import sys, os # ── helpers ────────────────────────────────────────────────────────────── + def _sse_line(data: dict) -> str: """One SSE data line (no trailing blank line -- we add those in the stream).""" return f"data: {json.dumps(data)}" @@ -25,7 +26,7 @@ def _sse_done() -> str: return "data: [DONE]" -def _make_chunk(delta: dict, finish_reason=None, usage=None, timings=None): +def _make_chunk(delta: dict, finish_reason = None, usage = None, timings = None): """Build a chat-completions streaming chunk.""" choice = {"index": 0, "delta": delta} if finish_reason: @@ -38,7 +39,7 @@ def _make_chunk(delta: dict, finish_reason=None, usage=None, timings=None): return chunk -def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) -> str: +def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = None) -> str: """ Build a complete SSE text stream from a list of chunk dicts. Includes the role chunk, content/tool chunks, and [DONE]. @@ -46,7 +47,7 @@ def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) lines = [] for c in chunks: lines.append(_sse_line(c)) - lines.append("") # blank line separator + lines.append("") # blank line separator # Final usage chunk (if provided) if final_usage or final_timings: meta = {} @@ -64,6 +65,7 @@ def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) class FakeResponse: """Mimics httpx.Response for streaming.""" + def __init__(self, text: str, status_code: int = 200): self._text = text self.status_code = status_code @@ -82,6 +84,7 @@ class FakeResponse: class FakeClient: """Mimics httpx.Client context manager.""" + def __init__(self, response: FakeResponse): self._response = response @@ -92,12 +95,13 @@ class FakeClient: pass @contextlib.contextmanager - def stream(self, method, url, json=None, timeout=None, headers=None): + def stream(self, method, url, json = None, timeout = None, headers = None): yield self._response # ── Build a minimal LlamaCppBackend for testing ───────────────────────── + def _make_backend(): """Create a minimal mock backend with just enough to run the method.""" # We need the real class but only care about generate_chat_completion_with_tools @@ -111,9 +115,9 @@ def _make_backend(): from studio.backend.core.inference.llama_cpp import LlamaCppBackend backend = object.__new__(LlamaCppBackend) - backend._process = True # is_loaded checks _process is not None - backend._healthy = True # is_loaded checks _healthy - backend._port = 9999 # base_url property reads _port + backend._process = True # is_loaded checks _process is not None + backend._healthy = True # is_loaded checks _healthy + backend._port = 9999 # base_url property reads _port backend._api_key = None backend._supports_reasoning = False return backend @@ -124,13 +128,13 @@ def _synthesis_sse(): chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"content": "Done."}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] usage = {"prompt_tokens": 20, "completion_tokens": 1} - return _build_sse_stream(chunks, final_usage=usage) + return _build_sse_stream(chunks, final_usage = usage) -def _collect_events(backend, sse_text, tools=None, **kwargs): +def _collect_events(backend, sse_text, tools = None, **kwargs): """ Run generate_chat_completion_with_tools with a fake SSE stream and collect all yielded events. @@ -139,14 +143,24 @@ def _collect_events(backend, sse_text, tools=None, **kwargs): return a plain text synthesis response so the agentic loop terminates. """ if tools is None: - tools = [{"type": "function", "function": {"name": "web_search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}] + tools = [ + { + "type": "function", + "function": { + "name": "web_search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } + ] call_count = [0] synth_sse = _synthesis_sse() @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): + def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): idx = call_count[0] call_count[0] += 1 # First call: use the provided SSE. Subsequent: plain text synthesis. @@ -154,23 +168,26 @@ def _collect_events(backend, sse_text, tools=None, **kwargs): yield FakeResponse(text) # Patch execute_tool to return a dummy result - def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): + def fake_execute_tool( + tool_name, arguments, cancel_event = None, timeout = None, session_id = None + ): return f"Tool {tool_name} result: OK" original_stream = backend._stream_with_retry backend._stream_with_retry = fake_stream_with_retry events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): + with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): try: for event in backend.generate_chat_completion_with_tools( - messages=[{"role": "user", "content": "Hello"}], - tools=tools, + messages = [{"role": "user", "content": "Hello"}], + tools = tools, **kwargs, ): events.append(event) except Exception as e: import traceback + traceback.print_exc() events.append({"type": "error", "error": str(e)}) @@ -180,6 +197,7 @@ def _collect_events(backend, sse_text, tools=None, **kwargs): # ── The actual tests ───────────────────────────────────────────────────── + def test_no_tool_call_plain_text(): """90% case: model responds with plain text, no tool call. Should stream content immediately without delay.""" @@ -189,11 +207,11 @@ def test_no_tool_call_plain_text(): _make_chunk({"role": "assistant"}), _make_chunk({"content": "Hello"}), _make_chunk({"content": " there"}), - _make_chunk({"content": "!"}, finish_reason="stop"), + _make_chunk({"content": "!"}, finish_reason = "stop"), ] usage = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13} timings = {"predicted_ms": 100, "predicted_n": 3, "predicted_per_second": 30.0} - sse = _build_sse_stream(chunks, final_usage=usage, final_timings=timings) + sse = _build_sse_stream(chunks, final_usage = usage, final_timings = timings) events = _collect_events(backend, sse) @@ -223,21 +241,34 @@ def test_structured_tool_calls(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", - "function": {"name": "web_search", "arguments": ""}}]}), - _make_chunk({"tool_calls": [{"index": 0, - "function": {"arguments": '{"query":'}}]}), - _make_chunk({"tool_calls": [{"index": 0, - "function": {"arguments": ' "test"}'}}]}), - _make_chunk({}, finish_reason="tool_calls"), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "function": {"name": "web_search", "arguments": ""}, + } + ] + } + ), + _make_chunk( + {"tool_calls": [{"index": 0, "function": {"arguments": '{"query":'}}]} + ), + _make_chunk( + {"tool_calls": [{"index": 0, "function": {"arguments": ' "test"}'}}]} + ), + _make_chunk({}, finish_reason = "tool_calls"), ] usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - sse = _build_sse_stream(chunks, final_usage=usage) + sse = _build_sse_stream(chunks, final_usage = usage) events = _collect_events(backend, sse) # Should have status update for tool execution - status_events = [e for e in events if e["type"] == "status" and "Searching" in e.get("text", "")] + status_events = [ + e for e in events if e["type"] == "status" and "Searching" in e.get("text", "") + ] assert len(status_events) >= 1, f"Expected search status, got: {events}" # Should have tool_start event @@ -265,10 +296,14 @@ def test_xml_tool_call_at_start(): chunks = [_make_chunk({"role": "assistant"})] for char in content: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) - usage = {"prompt_tokens": 10, "completion_tokens": len(content), "total_tokens": 10 + len(content)} - sse = _build_sse_stream(chunks, final_usage=usage) + usage = { + "prompt_tokens": 10, + "completion_tokens": len(content), + "total_tokens": 10 + len(content), + } + sse = _build_sse_stream(chunks, final_usage = usage) events = _collect_events(backend, sse) @@ -288,15 +323,15 @@ def test_xml_function_tag_at_start(): Buffer should detect {tc_json}" for char in rest: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 1, f"Expected 1 tool_start after whitespace, got: {events}" + assert ( + len(tool_starts) == 1 + ), f"Expected 1 tool_start after whitespace, got: {events}" print("PASS: test_whitespace_before_tool_xml") @@ -346,16 +383,16 @@ def test_content_then_tool_xml_no_retroactive_execution(): chunks = [_make_chunk({"role": "assistant"})] chunks.append(_make_chunk({"content": "Let me search for that. "})) chunks.append(_make_chunk({"content": f"{tc_json}"})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) # Tool should NOT be executed (visible content was already emitted) tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 0, ( - f"Should NOT retroactively execute tools after visible content: {tool_starts}" - ) + assert ( + len(tool_starts) == 0 + ), f"Should NOT retroactively execute tools after visible content: {tool_starts}" # Content should be present (tool XML stripped by _strip_tool_markup) content_events = [e for e in events if e["type"] == "content"] @@ -370,29 +407,65 @@ def test_multiple_structured_tool_calls(): backend = _make_backend() tools = [ - {"type": "function", "function": {"name": "web_search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}, - {"type": "function", "function": {"name": "python", - "parameters": {"type": "object", "properties": {"code": {"type": "string"}}}}}, + { + "type": "function", + "function": { + "name": "web_search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + }, + { + "type": "function", + "function": { + "name": "python", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + }, + }, + }, ] chunks = [ _make_chunk({"role": "assistant"}), # Two tool calls streamed with different indices - _make_chunk({"tool_calls": [ - {"index": 0, "id": "call_0", "function": {"name": "web_search", "arguments": ""}}, - {"index": 1, "id": "call_1", "function": {"name": "python", "arguments": ""}}, - ]}), - _make_chunk({"tool_calls": [ - {"index": 0, "function": {"arguments": '{"query": "test"}'}}, - ]}), - _make_chunk({"tool_calls": [ - {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, - ]}), - _make_chunk({}, finish_reason="tool_calls"), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "function": {"name": "web_search", "arguments": ""}, + }, + { + "index": 1, + "id": "call_1", + "function": {"name": "python", "arguments": ""}, + }, + ] + } + ), + _make_chunk( + { + "tool_calls": [ + {"index": 0, "function": {"arguments": '{"query": "test"}'}}, + ] + } + ), + _make_chunk( + { + "tool_calls": [ + {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, + ] + } + ), + _make_chunk({}, finish_reason = "tool_calls"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, tools=tools) + events = _collect_events(backend, sse, tools = tools) tool_starts = [e for e in events if e["type"] == "tool_start"] assert len(tool_starts) == 2, f"Expected 2 tool_start events, got: {tool_starts}" @@ -415,10 +488,10 @@ def test_reasoning_tokens_stream_immediately(): _make_chunk({"reasoning_content": "Let me think..."}), _make_chunk({"reasoning_content": " about this."}), _make_chunk({"content": "The answer is 42."}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking=True) + events = _collect_events(backend, sse, enable_thinking = True) content_events = [e for e in events if e["type"] == "content"] assert len(content_events) >= 1, f"Expected content events, got: {content_events}" @@ -444,12 +517,24 @@ def test_reasoning_then_tool_call(): chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"reasoning_content": "I need to search for this."}), - _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", - "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), - _make_chunk({}, finish_reason="tool_calls"), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "function": { + "name": "web_search", + "arguments": '{"query": "test"}', + }, + } + ] + } + ), + _make_chunk({}, finish_reason = "tool_calls"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking=True) + events = _collect_events(backend, sse, enable_thinking = True) # Reasoning should NOT be yielded during tool detection # (prevents prev_text corruption in consumer). Instead it's @@ -474,10 +559,10 @@ def test_reasoning_only_no_content(): chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"reasoning_content": "The answer is simply 42."}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking=True) + events = _collect_events(backend, sse, enable_thinking = True) content_events = [e for e in events if e["type"] == "content"] assert len(content_events) >= 1, f"Should yield reasoning as content: {events}" @@ -485,7 +570,9 @@ def test_reasoning_only_no_content(): final = content_events[-1]["text"] assert "42" in final, f"Should contain reasoning text: {final}" # Should NOT have wrapper (reasoning-only fallback) - assert "" not in final, f"Reasoning-only should not have wrapper: {final}" + assert ( + "" not in final + ), f"Reasoning-only should not have wrapper: {final}" print("PASS: test_reasoning_only_no_content") @@ -496,7 +583,7 @@ def test_empty_response(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -518,7 +605,7 @@ def test_buffer_prefix_timeout(): # Stream char by char for char in content: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -568,7 +655,7 @@ def test_draining_false_positive(): _make_chunk({"role": "assistant"}), _make_chunk({"content": "Use a screwdriver"}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -594,21 +681,32 @@ def test_structured_tool_args_json_parsing(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({"tool_calls": [{"index": 0, "id": "call_abc", - "function": {"name": "web_search", "arguments": ""}}]}), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_abc", + "function": {"name": "web_search", "arguments": ""}, + } + ] + } + ), ] for part in arg_parts: - chunks.append(_make_chunk({"tool_calls": [{"index": 0, - "function": {"arguments": part}}]})) - chunks.append(_make_chunk({}, finish_reason="tool_calls")) + chunks.append( + _make_chunk({"tool_calls": [{"index": 0, "function": {"arguments": part}}]}) + ) + chunks.append(_make_chunk({}, finish_reason = "tool_calls")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) tool_starts = [e for e in events if e["type"] == "tool_start"] assert len(tool_starts) == 1 - assert tool_starts[0]["arguments"] == {"query": "what is python?"}, \ - f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" + assert tool_starts[0]["arguments"] == { + "query": "what is python?" + }, f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" assert tool_starts[0]["tool_call_id"] == "call_abc" print("PASS: test_structured_tool_args_json_parsing") @@ -625,15 +723,16 @@ def test_auto_heal_disabled(): chunks = [_make_chunk({"role": "assistant"})] # Send as one big content chunk chunks.append(_make_chunk({"content": content})) - chunks.append(_make_chunk({}, finish_reason="stop")) + chunks.append(_make_chunk({}, finish_reason = "stop")) sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, auto_heal_tool_calls=False) + events = _collect_events(backend, sse, auto_heal_tool_calls = False) # With auto_heal disabled, the XML should NOT be parsed as a tool call tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 0, \ - f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" + assert ( + len(tool_starts) == 0 + ), f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" print("PASS: test_auto_heal_disabled") @@ -646,23 +745,39 @@ def test_metrics_accumulation_across_tool_iterations(): # First iteration: tool call tool_chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", - "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), - _make_chunk({}, finish_reason="tool_calls"), + _make_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "function": { + "name": "web_search", + "arguments": '{"query": "test"}', + }, + } + ] + } + ), + _make_chunk({}, finish_reason = "tool_calls"), ] tool_usage = {"prompt_tokens": 10, "completion_tokens": 5} tool_timings = {"predicted_ms": 50, "predicted_n": 5} - tool_sse = _build_sse_stream(tool_chunks, final_usage=tool_usage, final_timings=tool_timings) + tool_sse = _build_sse_stream( + tool_chunks, final_usage = tool_usage, final_timings = tool_timings + ) # Second iteration: plain text response (synthesis) synth_chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"content": "Based on my search, the answer is X."}), - _make_chunk({}, finish_reason="stop"), + _make_chunk({}, finish_reason = "stop"), ] synth_usage = {"prompt_tokens": 20, "completion_tokens": 8} synth_timings = {"predicted_ms": 100, "predicted_n": 8} - synth_sse = _build_sse_stream(synth_chunks, final_usage=synth_usage, final_timings=synth_timings) + synth_sse = _build_sse_stream( + synth_chunks, final_usage = synth_usage, final_timings = synth_timings + ) # We need to return different SSE streams for each iteration call_count = [0] @@ -671,36 +786,52 @@ def test_metrics_accumulation_across_tool_iterations(): fake_responses = [FakeResponse(tool_sse), FakeResponse(synth_sse)] @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): + def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): idx = min(call_count[0], len(fake_responses) - 1) call_count[0] += 1 yield fake_responses[idx] - def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): + def fake_execute_tool( + tool_name, arguments, cancel_event = None, timeout = None, session_id = None + ): return "Search result: success" backend._stream_with_retry = fake_stream_with_retry events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): + with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): for event in backend.generate_chat_completion_with_tools( - messages=[{"role": "user", "content": "Search for test"}], - tools=[{"type": "function", "function": {"name": "web_search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}], + messages = [{"role": "user", "content": "Search for test"}], + tools = [ + { + "type": "function", + "function": { + "name": "web_search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } + ], ): events.append(event) meta_events = [e for e in events if e["type"] == "metadata"] - assert len(meta_events) == 1, f"Expected exactly 1 metadata event, got: {meta_events}" + assert ( + len(meta_events) == 1 + ), f"Expected exactly 1 metadata event, got: {meta_events}" meta = meta_events[0] # completion_tokens should be accumulated: 5 (tool iter) + 8 (synthesis) = 13 - assert meta["usage"]["completion_tokens"] == 13, \ - f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" + assert ( + meta["usage"]["completion_tokens"] == 13 + ), f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" # predicted_ms and predicted_n should also accumulate - assert meta["timings"]["predicted_n"] == 13, \ - f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" + assert ( + meta["timings"]["predicted_n"] == 13 + ), f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" print("PASS: test_metrics_accumulation_across_tool_iterations") @@ -740,6 +871,7 @@ if __name__ == "__main__": failed += 1 errors.append((test_fn.__name__, str(e))) import traceback + print(f"FAIL: {test_fn.__name__}: {e}") traceback.print_exc() print() From 2507d269704326407bedd4bb43ea4509490b221e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 08:26:10 +0000 Subject: [PATCH 12/18] 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 before logging. --- studio/backend/core/inference/llama_cpp.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index fdd5420851..62fbdeb1c6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -950,7 +950,12 @@ class LlamaCppBackend: else: self._api_key = None - logger.info(f"Starting llama-server: {' '.join(cmd)}") + _log_cmd = list(cmd) + if "--api-key" in _log_cmd: + _ki = _log_cmd.index("--api-key") + 1 + if _ki < len(_log_cmd): + _log_cmd[_ki] = "" + logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") # Set library paths so llama-server can find its shared libs and CUDA DLLs import os From 5392d736d227ac89b73c53741b788171e670dbe3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 08:28:18 +0000 Subject: [PATCH 13/18] Remove test file temporarily --- tests/test_streaming_tool_detection.py | 885 ------------------------- 1 file changed, 885 deletions(-) delete mode 100644 tests/test_streaming_tool_detection.py diff --git a/tests/test_streaming_tool_detection.py b/tests/test_streaming_tool_detection.py deleted file mode 100644 index 25e25fba2e..0000000000 --- a/tests/test_streaming_tool_detection.py +++ /dev/null @@ -1,885 +0,0 @@ -""" -Exhaustive tests for the speculative-buffer streaming tool detection in -generate_chat_completion_with_tools(). - -We mock the HTTP layer so llama-server is not required. Each test constructs -the exact SSE byte stream that llama-server would emit, feeds it through the -real method, and asserts on the yielded events. -""" - -import json -import threading -import types -import contextlib -from unittest.mock import MagicMock, patch, PropertyMock -import sys, os - -# ── helpers ────────────────────────────────────────────────────────────── - - -def _sse_line(data: dict) -> str: - """One SSE data line (no trailing blank line -- we add those in the stream).""" - return f"data: {json.dumps(data)}" - - -def _sse_done() -> str: - return "data: [DONE]" - - -def _make_chunk(delta: dict, finish_reason = None, usage = None, timings = None): - """Build a chat-completions streaming chunk.""" - choice = {"index": 0, "delta": delta} - if finish_reason: - choice["finish_reason"] = finish_reason - chunk = {"choices": [choice]} - if usage: - chunk["usage"] = usage - if timings: - chunk["timings"] = timings - return chunk - - -def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = None) -> str: - """ - Build a complete SSE text stream from a list of chunk dicts. - Includes the role chunk, content/tool chunks, and [DONE]. - """ - lines = [] - for c in chunks: - lines.append(_sse_line(c)) - lines.append("") # blank line separator - # Final usage chunk (if provided) - if final_usage or final_timings: - meta = {} - if final_usage: - meta["usage"] = final_usage - if final_timings: - meta["timings"] = final_timings - meta["choices"] = [] - lines.append(_sse_line(meta)) - lines.append("") - lines.append(_sse_done()) - lines.append("") - return "\n".join(lines) - - -class FakeResponse: - """Mimics httpx.Response for streaming.""" - - def __init__(self, text: str, status_code: int = 200): - self._text = text - self.status_code = status_code - self._closed = False - - def iter_text(self): - # Yield the whole thing in one shot (simplest case) - yield self._text - - def read(self): - return self._text.encode() - - def close(self): - self._closed = True - - -class FakeClient: - """Mimics httpx.Client context manager.""" - - def __init__(self, response: FakeResponse): - self._response = response - - def __enter__(self): - return self - - def __exit__(self, *args): - pass - - @contextlib.contextmanager - def stream(self, method, url, json = None, timeout = None, headers = None): - yield self._response - - -# ── Build a minimal LlamaCppBackend for testing ───────────────────────── - - -def _make_backend(): - """Create a minimal mock backend with just enough to run the method.""" - # We need the real class but only care about generate_chat_completion_with_tools - # Import the real module - _repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) - if _repo_root not in sys.path: - sys.path.insert(0, _repo_root) - - # Instead of importing the full module (which has other deps), we'll - # build a lightweight object that has the method and its dependencies. - from studio.backend.core.inference.llama_cpp import LlamaCppBackend - - backend = object.__new__(LlamaCppBackend) - backend._process = True # is_loaded checks _process is not None - backend._healthy = True # is_loaded checks _healthy - backend._port = 9999 # base_url property reads _port - backend._api_key = None - backend._supports_reasoning = False - return backend - - -def _synthesis_sse(): - """Build a simple text SSE response for post-tool synthesis.""" - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk({"content": "Done."}), - _make_chunk({}, finish_reason = "stop"), - ] - usage = {"prompt_tokens": 20, "completion_tokens": 1} - return _build_sse_stream(chunks, final_usage = usage) - - -def _collect_events(backend, sse_text, tools = None, **kwargs): - """ - Run generate_chat_completion_with_tools with a fake SSE stream - and collect all yielded events. - - After the first iteration (tool detection), subsequent iterations - return a plain text synthesis response so the agentic loop terminates. - """ - if tools is None: - tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - } - ] - - call_count = [0] - synth_sse = _synthesis_sse() - - @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): - idx = call_count[0] - call_count[0] += 1 - # First call: use the provided SSE. Subsequent: plain text synthesis. - text = sse_text if idx == 0 else synth_sse - yield FakeResponse(text) - - # Patch execute_tool to return a dummy result - def fake_execute_tool( - tool_name, arguments, cancel_event = None, timeout = None, session_id = None - ): - return f"Tool {tool_name} result: OK" - - original_stream = backend._stream_with_retry - backend._stream_with_retry = fake_stream_with_retry - - events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): - try: - for event in backend.generate_chat_completion_with_tools( - messages = [{"role": "user", "content": "Hello"}], - tools = tools, - **kwargs, - ): - events.append(event) - except Exception as e: - import traceback - - traceback.print_exc() - events.append({"type": "error", "error": str(e)}) - - backend._stream_with_retry = original_stream - return events - - -# ── The actual tests ───────────────────────────────────────────────────── - - -def test_no_tool_call_plain_text(): - """90% case: model responds with plain text, no tool call. - Should stream content immediately without delay.""" - backend = _make_backend() - - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk({"content": "Hello"}), - _make_chunk({"content": " there"}), - _make_chunk({"content": "!"}, finish_reason = "stop"), - ] - usage = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13} - timings = {"predicted_ms": 100, "predicted_n": 3, "predicted_per_second": 30.0} - sse = _build_sse_stream(chunks, final_usage = usage, final_timings = timings) - - events = _collect_events(backend, sse) - - # Should have content events with cumulative text - content_events = [e for e in events if e["type"] == "content"] - assert len(content_events) >= 1, f"Expected content events, got: {events}" - - # Final content should contain the full text - final_content = content_events[-1]["text"] - assert "Hello there!" in final_content, f"Missing text in: {final_content}" - - # Should have metadata - meta_events = [e for e in events if e["type"] == "metadata"] - assert len(meta_events) == 1, f"Expected 1 metadata event, got: {meta_events}" - - # Should have status clear - status_events = [e for e in events if e["type"] == "status"] - assert any(e["text"] == "" for e in status_events), "Missing status clear" - - print("PASS: test_no_tool_call_plain_text") - - -def test_structured_tool_calls(): - """Model emits structured delta.tool_calls (the standard path). - Should detect instantly and execute.""" - backend = _make_backend() - - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": {"name": "web_search", "arguments": ""}, - } - ] - } - ), - _make_chunk( - {"tool_calls": [{"index": 0, "function": {"arguments": '{"query":'}}]} - ), - _make_chunk( - {"tool_calls": [{"index": 0, "function": {"arguments": ' "test"}'}}]} - ), - _make_chunk({}, finish_reason = "tool_calls"), - ] - usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - sse = _build_sse_stream(chunks, final_usage = usage) - - events = _collect_events(backend, sse) - - # Should have status update for tool execution - status_events = [ - e for e in events if e["type"] == "status" and "Searching" in e.get("text", "") - ] - assert len(status_events) >= 1, f"Expected search status, got: {events}" - - # Should have tool_start event - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 1, f"Expected 1 tool_start, got: {tool_starts}" - assert tool_starts[0]["tool_name"] == "web_search" - assert tool_starts[0]["arguments"] == {"query": "test"} - - # Should have tool_end event - tool_ends = [e for e in events if e["type"] == "tool_end"] - assert len(tool_ends) == 1, f"Expected 1 tool_end, got: {tool_ends}" - - print("PASS: test_structured_tool_calls") - - -def test_xml_tool_call_at_start(): - """Model emits JSON instead of structured tool_calls. - Buffer should detect prefix and drain.""" - backend = _make_backend() - - tc_json = json.dumps({"name": "web_search", "arguments": {"query": "hello"}}) - content = f"{tc_json}" - - # Stream the XML content token by token to simulate real streaming - chunks = [_make_chunk({"role": "assistant"})] - for char in content: - chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) - - usage = { - "prompt_tokens": 10, - "completion_tokens": len(content), - "total_tokens": 10 + len(content), - } - sse = _build_sse_stream(chunks, final_usage = usage) - - events = _collect_events(backend, sse) - - # Should detect tool call and execute it - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 1, f"Expected 1 tool_start, got: {events}" - assert tool_starts[0]["tool_name"] == "web_search" - - tool_ends = [e for e in events if e["type"] == "tool_end"] - assert len(tool_ends) == 1, f"Expected 1 tool_end" - - print("PASS: test_xml_tool_call_at_start") - - -def test_xml_function_tag_at_start(): - """Model emits tag. - Buffer should detect . Buffer should strip - leading whitespace before prefix check.""" - backend = _make_backend() - - tc_json = json.dumps({"name": "web_search", "arguments": {"query": "test"}}) - content = f" \n {tc_json}" - - chunks = [_make_chunk({"role": "assistant"})] - # Send whitespace as one chunk, then the rest - chunks.append(_make_chunk({"content": " \n "})) - rest = f"{tc_json}" - for char in rest: - chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) - - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse) - - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert ( - len(tool_starts) == 1 - ), f"Expected 1 tool_start after whitespace, got: {events}" - - print("PASS: test_whitespace_before_tool_xml") - - -def test_content_then_tool_xml_no_retroactive_execution(): - """Rare case: model emits normal content first, then tool XML later. - Once visible content has been emitted to the user, we must NOT - retroactively switch to tool execution -- that would violate the - streaming contract and corrupt the route-layer cumulative delta - tracker. The tool XML is stripped by _strip_tool_markup, and the - user sees the cleaned content as a normal response.""" - backend = _make_backend() - - tc_json = json.dumps({"name": "web_search", "arguments": {"query": "q"}}) - # Start with normal text (triggers STREAMING), then tool XML - chunks = [_make_chunk({"role": "assistant"})] - chunks.append(_make_chunk({"content": "Let me search for that. "})) - chunks.append(_make_chunk({"content": f"{tc_json}"})) - chunks.append(_make_chunk({}, finish_reason = "stop")) - - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse) - - # Tool should NOT be executed (visible content was already emitted) - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert ( - len(tool_starts) == 0 - ), f"Should NOT retroactively execute tools after visible content: {tool_starts}" - - # Content should be present (tool XML stripped by _strip_tool_markup) - content_events = [e for e in events if e["type"] == "content"] - assert len(content_events) >= 1, f"Should have content: {events}" - assert "Let me search" in content_events[0]["text"] - - print("PASS: test_content_then_tool_xml_no_retroactive_execution") - - -def test_multiple_structured_tool_calls(): - """Model calls two tools in one response (parallel tool calls).""" - backend = _make_backend() - - tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - }, - { - "type": "function", - "function": { - "name": "python", - "parameters": { - "type": "object", - "properties": {"code": {"type": "string"}}, - }, - }, - }, - ] - - chunks = [ - _make_chunk({"role": "assistant"}), - # Two tool calls streamed with different indices - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": {"name": "web_search", "arguments": ""}, - }, - { - "index": 1, - "id": "call_1", - "function": {"name": "python", "arguments": ""}, - }, - ] - } - ), - _make_chunk( - { - "tool_calls": [ - {"index": 0, "function": {"arguments": '{"query": "test"}'}}, - ] - } - ), - _make_chunk( - { - "tool_calls": [ - {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), - ] - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, tools = tools) - - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 2, f"Expected 2 tool_start events, got: {tool_starts}" - - names = {ts["tool_name"] for ts in tool_starts} - assert names == {"web_search", "python"}, f"Wrong tool names: {names}" - - print("PASS: test_multiple_structured_tool_calls") - - -def test_reasoning_tokens_stream_immediately(): - """Thinking model: reasoning_content is accumulated during BUFFERING - and flushed together with content when transitioning to STREAMING. - The final output includes ... wrapping.""" - backend = _make_backend() - backend._supports_reasoning = True - - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk({"reasoning_content": "Let me think..."}), - _make_chunk({"reasoning_content": " about this."}), - _make_chunk({"content": "The answer is 42."}), - _make_chunk({}, finish_reason = "stop"), - ] - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking = True) - - content_events = [e for e in events if e["type"] == "content"] - assert len(content_events) >= 1, f"Expected content events, got: {content_events}" - - # Content should contain both tags and the answer - final = content_events[-1]["text"] - assert "" in final, f"Should have tag: {final}" - assert "Let me think" in final, f"Should have reasoning: {final}" - assert "42" in final, f"Should have answer: {final}" - assert "" in final, f"Should have closing : {final}" - - print("PASS: test_reasoning_tokens_stream_immediately") - - -def test_reasoning_then_tool_call(): - """Thinking model that reasons then calls a tool. - Reasoning is silently accumulated during tool detection (matching - old non-streaming behavior) so the consumer's prev_text is not - corrupted for subsequent iterations. Tool is still detected.""" - backend = _make_backend() - backend._supports_reasoning = True - - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk({"reasoning_content": "I need to search for this."}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": { - "name": "web_search", - "arguments": '{"query": "test"}', - }, - } - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), - ] - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking = True) - - # Reasoning should NOT be yielded during tool detection - # (prevents prev_text corruption in consumer). Instead it's - # accumulated silently, matching old non-streaming behavior. - # After tool execution, the synthesis pass handles display. - - # Tool should be executed - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 1, f"Expected tool_start: {events}" - assert tool_starts[0]["tool_name"] == "web_search" - - print("PASS: test_reasoning_then_tool_call") - - -def test_reasoning_only_no_content(): - """Thinking model produces only reasoning_content with no content tokens. - Should yield reasoning as plain text (no wrapper), matching - the final streaming pass behavior for models like Qwen3 always-think.""" - backend = _make_backend() - backend._supports_reasoning = True - - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk({"reasoning_content": "The answer is simply 42."}), - _make_chunk({}, finish_reason = "stop"), - ] - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking = True) - - content_events = [e for e in events if e["type"] == "content"] - assert len(content_events) >= 1, f"Should yield reasoning as content: {events}" - - final = content_events[-1]["text"] - assert "42" in final, f"Should contain reasoning text: {final}" - # Should NOT have wrapper (reasoning-only fallback) - assert ( - "" not in final - ), f"Reasoning-only should not have wrapper: {final}" - - print("PASS: test_reasoning_only_no_content") - - -def test_empty_response(): - """Model returns empty stream (just role + [DONE]). Should not crash.""" - backend = _make_backend() - - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk({}, finish_reason = "stop"), - ] - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse) - - # Should not crash, just return with no content - error_events = [e for e in events if e.get("type") == "error"] - assert len(error_events) == 0, f"Should not error: {error_events}" - - print("PASS: test_empty_response") - - -def test_buffer_prefix_timeout(): - """Content starts with '<' but is not a tool call (e.g., '

Hello

'). - Buffer should hold briefly then flush when no prefix match at 32 chars.""" - backend = _make_backend() - - content = "

This is a paragraph of HTML content that is not a tool call

" - chunks = [_make_chunk({"role": "assistant"})] - # Stream char by char - for char in content: - chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) - - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse) - - content_events = [e for e in events if e["type"] == "content"] - assert len(content_events) >= 1, f"Should have content events: {events}" - - # No tool calls should be detected - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 0, f"Should not detect tools in HTML: {tool_starts}" - - # Final content should contain the HTML - final = content_events[-1]["text"] - assert "

" in final, f"HTML content should pass through: {final}" - - print("PASS: test_buffer_prefix_timeout") - - -def test_buffer_resolves_to_streaming_on_non_xml_first_char(): - """First content char is not '<' and not whitespace. - Should immediately transition to STREAMING.""" - backend = _make_backend() - - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk({"content": "H"}), # 'H' is not '<', instant STREAMING - _make_chunk({"content": "ello"}), - ] - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse) - - content_events = [e for e in events if e["type"] == "content"] - # First content event should appear immediately with just "H" - assert len(content_events) >= 1 - assert "H" in content_events[0]["text"] - - print("PASS: test_buffer_resolves_to_streaming_on_non_xml_first_char") - - -def test_draining_false_positive(): - """Buffer detects 'Use a screwdriver"}), - _make_chunk({}, finish_reason = "stop"), - ] - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse) - - # "" so it enters BUFFERING. - # Then "_tip>" does NOT match "" since the buffer becomes - # "..." which doesn't start with "" or "32 chars the buffer should flush. - # No tool should be executed. - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 0, f"Should not detect tool in : {tool_starts}" - - print("PASS: test_draining_false_positive") - - -def test_structured_tool_args_json_parsing(): - """Verify that arguments streamed across multiple chunks get reassembled - and parsed correctly as JSON.""" - backend = _make_backend() - - # Arguments split across 4 chunks - arg_parts = ['{"qu', 'ery":', ' "wha', 't is python?"}'] - - chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_abc", - "function": {"name": "web_search", "arguments": ""}, - } - ] - } - ), - ] - for part in arg_parts: - chunks.append( - _make_chunk({"tool_calls": [{"index": 0, "function": {"arguments": part}}]}) - ) - chunks.append(_make_chunk({}, finish_reason = "tool_calls")) - - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse) - - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) == 1 - assert tool_starts[0]["arguments"] == { - "query": "what is python?" - }, f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" - assert tool_starts[0]["tool_call_id"] == "call_abc" - - print("PASS: test_structured_tool_args_json_parsing") - - -def test_auto_heal_disabled(): - """When auto_heal_tool_calls=False, XML tool calls in content should NOT - be parsed -- only structured tool_calls are honored.""" - backend = _make_backend() - - tc_json = json.dumps({"name": "web_search", "arguments": {"query": "test"}}) - content = f"{tc_json}" - - chunks = [_make_chunk({"role": "assistant"})] - # Send as one big content chunk - chunks.append(_make_chunk({"content": content})) - chunks.append(_make_chunk({}, finish_reason = "stop")) - - sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, auto_heal_tool_calls = False) - - # With auto_heal disabled, the XML should NOT be parsed as a tool call - tool_starts = [e for e in events if e["type"] == "tool_start"] - assert ( - len(tool_starts) == 0 - ), f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" - - print("PASS: test_auto_heal_disabled") - - -def test_metrics_accumulation_across_tool_iterations(): - """When tools are called, metrics from the tool iteration should be - accumulated and included in the final metadata.""" - backend = _make_backend() - - # First iteration: tool call - tool_chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": { - "name": "web_search", - "arguments": '{"query": "test"}', - }, - } - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), - ] - tool_usage = {"prompt_tokens": 10, "completion_tokens": 5} - tool_timings = {"predicted_ms": 50, "predicted_n": 5} - tool_sse = _build_sse_stream( - tool_chunks, final_usage = tool_usage, final_timings = tool_timings - ) - - # Second iteration: plain text response (synthesis) - synth_chunks = [ - _make_chunk({"role": "assistant"}), - _make_chunk({"content": "Based on my search, the answer is X."}), - _make_chunk({}, finish_reason = "stop"), - ] - synth_usage = {"prompt_tokens": 20, "completion_tokens": 8} - synth_timings = {"predicted_ms": 100, "predicted_n": 8} - synth_sse = _build_sse_stream( - synth_chunks, final_usage = synth_usage, final_timings = synth_timings - ) - - # We need to return different SSE streams for each iteration - call_count = [0] - original_sse = [tool_sse, synth_sse] - - fake_responses = [FakeResponse(tool_sse), FakeResponse(synth_sse)] - - @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): - idx = min(call_count[0], len(fake_responses) - 1) - call_count[0] += 1 - yield fake_responses[idx] - - def fake_execute_tool( - tool_name, arguments, cancel_event = None, timeout = None, session_id = None - ): - return "Search result: success" - - backend._stream_with_retry = fake_stream_with_retry - - events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): - for event in backend.generate_chat_completion_with_tools( - messages = [{"role": "user", "content": "Search for test"}], - tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - } - ], - ): - events.append(event) - - meta_events = [e for e in events if e["type"] == "metadata"] - assert ( - len(meta_events) == 1 - ), f"Expected exactly 1 metadata event, got: {meta_events}" - - meta = meta_events[0] - # completion_tokens should be accumulated: 5 (tool iter) + 8 (synthesis) = 13 - assert ( - meta["usage"]["completion_tokens"] == 13 - ), f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" - - # predicted_ms and predicted_n should also accumulate - assert ( - meta["timings"]["predicted_n"] == 13 - ), f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" - - print("PASS: test_metrics_accumulation_across_tool_iterations") - - -# ── Run all tests ──────────────────────────────────────────────────────── - -if __name__ == "__main__": - tests = [ - test_no_tool_call_plain_text, - test_structured_tool_calls, - test_xml_tool_call_at_start, - test_xml_function_tag_at_start, - test_whitespace_before_tool_xml, - test_content_then_tool_xml_no_retroactive_execution, - test_multiple_structured_tool_calls, - test_reasoning_tokens_stream_immediately, - test_reasoning_then_tool_call, - test_reasoning_only_no_content, - test_empty_response, - test_buffer_prefix_timeout, - test_buffer_resolves_to_streaming_on_non_xml_first_char, - test_draining_false_positive, - test_structured_tool_args_json_parsing, - test_auto_heal_disabled, - test_metrics_accumulation_across_tool_iterations, - ] - - passed = 0 - failed = 0 - errors = [] - - for test_fn in tests: - try: - test_fn() - passed += 1 - except Exception as e: - failed += 1 - errors.append((test_fn.__name__, str(e))) - import traceback - - print(f"FAIL: {test_fn.__name__}: {e}") - traceback.print_exc() - print() - - print(f"\n{'='*60}") - print(f"Results: {passed} passed, {failed} failed, {len(tests)} total") - if errors: - print(f"\nFailed tests:") - for name, err in errors: - print(f" - {name}: {err}") - print(f"{'='*60}") From f86d1143522cd7049109069e95c56644db1e9167 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 08:55:58 +0000 Subject: [PATCH 14/18] 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] --- studio/backend/core/inference/llama_cpp.py | 13 +++---- studio/backend/routes/inference.py | 40 ++++++++-------------- 2 files changed, 20 insertions(+), 33 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 62fbdeb1c6..fb4404ed42 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2033,15 +2033,12 @@ class LlamaCppBackend: # ── STREAMING path: no tool call ── if detect_state == _S_STREAMING: # Safety net: check for XML tool signals in content. - # Only if we have NOT already emitted visible text -- - # retroactively switching to tool mode after the user - # has seen content violates the streaming contract and - # corrupts the route-layer cumulative delta tracker. + # The route layer resets prev_text on tool_start, so + # post-tool synthesis streams correctly even if + # content was already emitted before the tool XML. _safety_tc = None - if ( - auto_heal_tool_calls - and not _last_emitted - and any(s in content_accum for s in _TOOL_XML_SIGNALS) + if auto_heal_tool_calls and any( + s in content_accum for s in _TOOL_XML_SIGNALS ): _safety_tc = self._parse_tool_calls_from_text( content_accum, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index cde29246b9..f57342b59c 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -930,13 +930,10 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - _dc_counter = 0 for chunk_text in audio_input_generate(): - _dc_counter += 1 - if _dc_counter % 20 == 0: - if await request.is_disconnected(): - cancel_event.set() - return + if await request.is_disconnected(): + cancel_event.set() + return if chunk_text: chunk = ChatCompletionChunk( id = completion_id, @@ -1110,13 +1107,10 @@ async def openai_chat_completions( prev_text = "" _stream_usage = None _stream_timings = None - _dc_counter = 0 while True: - _dc_counter += 1 - if _dc_counter % 20 == 0: - if await request.is_disconnected(): - cancel_event.set() - return + if await request.is_disconnected(): + cancel_event.set() + return event = await asyncio.to_thread(next, gen, _tool_sentinel) if event is _tool_sentinel: @@ -1134,6 +1128,8 @@ async def openai_chat_completions( continue if event["type"] in ("tool_start", "tool_end"): + if event["type"] == "tool_start": + prev_text = "" yield f"data: {json.dumps(event)}\n\n" continue @@ -1262,13 +1258,10 @@ async def openai_chat_completions( prev_text = "" _stream_usage = None _stream_timings = None - _dc_counter = 0 while True: - _dc_counter += 1 - if _dc_counter % 20 == 0: - if await request.is_disconnected(): - cancel_event.set() - return + if await request.is_disconnected(): + cancel_event.set() + return cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) if cumulative is _gguf_sentinel: break @@ -1474,7 +1467,6 @@ async def openai_chat_completions( _DONE = object() # sentinel for generator exhaustion loop = asyncio.get_event_loop() gen = generate() - _dc_counter = 0 while True: # next(gen, _DONE) returns _DONE instead of raising # StopIteration — StopIteration cannot propagate @@ -1482,12 +1474,10 @@ async def openai_chat_completions( cumulative = await loop.run_in_executor(None, next, gen, _DONE) if cumulative is _DONE: break - _dc_counter += 1 - if _dc_counter % 20 == 0: - if await request.is_disconnected(): - cancel_event.set() - backend.reset_generation_state() - return + if await request.is_disconnected(): + cancel_event.set() + backend.reset_generation_state() + return new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: From 9bef60db23f8b12d65c5f139169f7e66ac1add61 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 08:59:01 +0000 Subject: [PATCH 15/18] Use None instead of {} for empty auth headers in TTS methods --- studio/backend/core/inference/llama_cpp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index fb4404ed42..071e2ba195 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2416,7 +2416,7 @@ class LlamaCppBackend: return None try: _auth_headers = ( - {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None ) with httpx.Client(timeout = 10, headers = _auth_headers) as client: @@ -2537,7 +2537,7 @@ class LlamaCppBackend: payload["n_probs"] = 1 _auth_headers = ( - {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None ) with httpx.Client( timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers From 10555db27b96b95563741116d6918da3f3c7c8b5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 09:14:18 +0000 Subject: [PATCH 16/18] Include accumulated metrics in STREAMING metadata check --- studio/backend/core/inference/llama_cpp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 071e2ba195..103a1234e3 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2050,7 +2050,7 @@ class LlamaCppBackend: _fc = _fu.get("completion_tokens", 0) _fp = _fu.get("prompt_tokens", 0) _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings: + if _iter_usage or _iter_timings or _accumulated_completion_tokens: _mt = dict(_iter_timings) if _iter_timings else {} if _accumulated_predicted_ms or _accumulated_predicted_n: _mt["predicted_ms"] = ( From 91982268dc3beeb9a1c48264d1e458008f75cb7a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 27 Mar 2026 09:15:55 +0000 Subject: [PATCH 17/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 103a1234e3..00d5db018b 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2050,7 +2050,11 @@ class LlamaCppBackend: _fc = _fu.get("completion_tokens", 0) _fp = _fu.get("prompt_tokens", 0) _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings or _accumulated_completion_tokens: + if ( + _iter_usage + or _iter_timings + or _accumulated_completion_tokens + ): _mt = dict(_iter_timings) if _iter_timings else {} if _accumulated_predicted_ms or _accumulated_predicted_n: _mt["predicted_ms"] = ( From 4967f52b9a0048cf900025a9c1b6766bcb8a42b3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 10:30:37 +0000 Subject: [PATCH 18/18] Guard against late tool_calls after visible content, filter incomplete fragments 1. If visible content was already emitted (_last_emitted is non-empty) when delta.tool_calls arrives, ignore the tool_calls instead of reclassifying the turn as a tool call. llama-server never interleaves content and tool_calls (they are mutually exclusive), but this guard is defensive for other OpenAI-compatible backends. [9/10 reviewers] 2. Filter out incomplete structured tool_calls fragments before execution. Entries with empty function.name (from truncation by max_tokens, disconnect, or interruption) are skipped instead of being passed to execute_tool(). [2/10 reviewers] --- studio/backend/core/inference/llama_cpp.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 00d5db018b..4fa5d20968 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1867,6 +1867,11 @@ class LlamaCppBackend: # ── Structured tool_calls ── tc_deltas = delta.get("tool_calls") if tc_deltas: + # Once visible content has been + # emitted, do not reclassify this + # turn as a tool call. + if _last_emitted: + continue has_structured_tc = True detect_state = _S_DRAINING for tc_d in tc_deltas: @@ -2094,7 +2099,14 @@ class LlamaCppBackend: tool_calls = None content_text = content_accum if has_structured_tc: - tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] + # Filter out incomplete fragments (e.g. from + # truncation by max_tokens or disconnect). + tool_calls = [ + tool_calls_acc[i] + for i in sorted(tool_calls_acc) + if (tool_calls_acc[i].get("function", {}) + .get("name", "").strip()) + ] or None if ( not tool_calls and auto_heal_tool_calls