Tool-call healing (default on) and opt-in nudging for the client-tool passthrough (#6801)
* inference: add passthrough tool-call healing core (heal_gate, heal_openai_message, StreamToolCallHealer, nudge helpers)
Small GGUF models often emit tool calls as text (<tool_call>{...}</tool_call>,
Gemma <|tool_call>, <function=> XML) instead of structured tool_calls. Studio's
enable-tools loop already heals these, but the client-tool passthrough
(unsloth run --disable-tools, unsloth start agents) relays them verbatim, so
the agent sees prose and the turn dies.
This module is the shared response-side repair layer the passthrough routes
will call: promote parsed text-form calls to structured calls, but only for
function names the client actually declared; coerce arguments through the same
canonical-key healing as the tool loop; never touch the upstream request body
(llama-server KV/slot reuse stays byte-identical). StreamToolCallHealer is the
streaming buffer-and-repair state machine: prose forwards immediately, only a
partial-signal tail or a suspected tool block is held, false alarms flush
verbatim, and a 64 KiB bound caps memory. nudge_should_retry/nudge_messages
support an opt-in single-retry nudge for non-streaming routes (wired later).
Kill-switch: UNSLOTH_DISABLE_TOOL_CALL_HEALING=1. Reuses
core/tool_healing.parse_tool_calls_from_text, strip_tool_call_markup, and
tool_loop_controller.coerce_tool_arguments unchanged.
* inference: heal text-form tool calls on the OpenAI and Responses passthrough
Wire the passthrough healing core into /v1/chat/completions and /v1/responses,
default ON whenever the request declares client tools:
Non-streaming: heal_openai_message runs inside the existing response-mutation
loop; a promoted call flips finish_reason to tool_calls and nulls the content,
and the verbatim-bytes fast path still applies when nothing was healed.
/v1/responses non-streaming inherits this through openai_chat_completions.
Streaming: a StreamToolCallHealer per stream. Ordinary prose relays
byte-for-byte (a fast path keeps upstream bytes when the healer passes a chunk
through whole); once a tool signal appears, content is held, and at the
finish/[DONE] boundary either synthetic delta.tool_calls chunks replace the
markup (finish_reason rewritten to tool_calls, including the synthetic-finish
path) or a false alarm flushes the held text verbatim. Structured upstream
deltas put the healer to sleep after flushing anything held, so grammar-mode
responses stay byte-identical. The Responses stream feeds healed calls through
the same per-call state machinery as structured deltas (indexes live in a
disjoint range so a healed call can never merge into a structured call's
state), and the visible/reasoning split runs first so reasoning text is never
promoted. parallel_tool_calls=false caps healed calls on every path.
The upstream request body is never touched and healing issues no extra
generation, so llama-server slot/KV-cache reuse is unchanged. Opt-out per
request with auto_heal_tool_calls=false (Responses reads it from the
extra-body); requests without tools relay verbatim.
* inference: heal text-form tool calls on the Anthropic /v1/messages passthrough
Streaming: AnthropicPassthroughEmitter.enable_healing(allowed_tools) routes
content deltas through the shared StreamToolCallHealer. A promoted call closes
any open text block (only the safe prose prefix ever streamed into it), opens a
synthetic tool_use block with a fresh toolu_* id, carries one input_json_delta,
and closes; finish() then forces stop_reason to tool_use unless a truncation
(max_tokens) wins. Structured upstream deltas flush anything held and put the
healer to sleep, so grammar-mode responses are untouched, as is every stream
where enable_healing is never called (Studio's own loop, no-tools requests).
disable_parallel_tool_use caps healed calls too.
Non-streaming: the OpenAI message dict is healed BEFORE block building, so the
existing tool_use promotion loop and stop_reason line treat promoted calls
exactly like native ones (finish_reason length still maps to max_tokens). The
legacy tool-XML strip still runs on remaining text, so opted-out requests keep
today's cleanup behavior byte-for-byte.
auto_heal_tool_calls is now a typed field on AnthropicMessagesRequest
(default True, mirroring Chat Completions) and threads into both passthrough
calls. Healing never touches the upstream request body.
* inference: opt-in single-retry tool-call nudge on the non-streaming passthrough
When the model clearly tried to call a tool (a tool signal in the text) but
healing produced nothing usable, re-ask once: the retry body is the original
body plus an assistant turn (the model's own failed text) and a short user
nudge naming the declared tools. The prompt prefix stays byte-identical, so
llama-server reuses the slot's KV cache and only the two-message suffix is
prefilled. The retry replaces the original response only when it actually
yields a promotable or structured call; on any error or still-garbage output
the original response is returned unchanged. Exactly one retry, non-streaming
OpenAI and Anthropic passthroughs only (a stream has already emitted bytes).
OPT-IN per user decision: nudge_tool_calls=true per request (typed on both
ChatCompletionRequest and AnthropicMessagesRequest, lifted from the Responses
extra-body), or UNSLOTH_TOOL_CALL_NUDGE=1 to flip the process default.
auto_heal_tool_calls=false disables healing AND the nudge.
Also align the non-streaming heal on allow_incomplete=True: the response is
final, so a trailing unclosed tool block is a model failure worth repairing,
matching the enable-tools loop's drain semantics.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* inference: never assume the upstream response shape in the nudge helpers
llama-server error bodies can carry message: null (or no choices at all), and
_last_assistant_text / response_has_promotable_calls / nudge_should_retry
called .get() on the message without a dict check, so a malformed upstream
response raised an AttributeError the surrounding except tuples did not catch,
failing the request instead of degrading to 'nothing to heal'. Route the shape
probing through one _first_choice_message helper that returns None for any
non-dict message, and add a parametrized test over the malformed shapes.
* inference: constrain healing by tool_choice, preserve length finish_reason, keep healed event order in Responses streams
Three review findings on the passthrough healer:
- heal_gate now honors the request's tool_choice: "none" disables healing
outright and a forced function narrows the promotion allowlist to that
one function, so healing can never contradict the request's tool-choice
constraint. Wired through the OpenAI chat (stream and non-stream),
Responses, and Anthropic (converted shape) passthroughs.
- The OpenAI non-streaming heal only upgrades finish_reason "stop" to
"tool_calls"; a truncated generation keeps "length" (the healed call
stays attached) matching the streaming and Anthropic paths.
- The Responses stream emits healer events in order instead of collapsing
all text ahead of the healed calls, so text after a healed call no longer
jumps ahead of the function_call item and output indexes are claimed in
the order the model produced them.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* inference: all-or-nothing promotion when a response mixes declared and undeclared text-form calls
Promoting a subset used to strip ALL tool markup from the content, which
silently deleted the text of any call naming an undeclared tool. The heal
now declines entirely when any parsed call is unpromotable, so the whole
message relays verbatim (pre-PR behavior) and no bytes are ever lost. In
streaming, a declared call that completed before an undeclared one arrived
is already emitted; the late undeclared markup still flushes as raw text.
The nudge helpers mirror the same contract via a shared predicate.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* tests: wrap long lines in the Responses healing tests to the project style
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* inference: span-exact healing, disjoint healed stream indexes, per-call Responses message items, allowlisted nudge acceptance
Four review findings on the passthrough healer:
- parse_tool_calls_from_text gains an optional with_spans return so healing
removes EXACTLY the promoted calls' markup. This supersedes the previous
all-or-nothing rule: declared calls promote and every unpromoted byte
(undeclared calls, unparseable closed blocks, suppressed alternate
formats such as a <function=...> block after a JSON call) relays as text.
The stream healer also processes one block per pass, so text between two
healed calls keeps its document position instead of trailing them.
- The OpenAI chat stream shifts native tool-call delta indexes past any
already-emitted healed calls; clients merge deltas by index, so a healed
call and a later native call can no longer merge into one.
- A healed call in the Responses stream closes the open message item and
trailing text opens a fresh one with a later output index, matching the
native stream shape; response.completed snapshots every message item
with its own text.
- The nudge retry only replaces the original response when the retry's
structured call names a DECLARED tool; a hallucinated undeclared call is
not an improvement.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: stop the heal path folding trailing prose into a closed function call
parse_tool_calls_from_text(allow_incomplete=True) cut a <function=...> body only
at an end-anchored </function>, so a fully closed call followed by trailing prose
(<function=..>..</parameter></function> words) folded </parameter></function> and
the prose into the tool argument and deleted the prose from visible content. The
strict path (allow_incomplete=False) already cut at the real </function> via rfind.
Do the same in both modes: trim the body at the real </function> when present and
end the removal span there, falling back to the end-anchored strip and body_end
only when the call is genuinely truncated. Add a regression test.
* inference: one shared single-call budget for healed and native calls
Codex round 5: the parallel-call caps counted healed and native calls
separately, so a healed text-form call followed by a native structured
delta double-emitted on all three streaming surfaces when the client
disabled parallel calls.
- OpenAI SSE: once a healed call went out with parallel_tool_calls
false, native tool_call deltas are dropped instead of index-shifted.
- Anthropic emitter: native deltas skip block allocation when the
healed-plus-native count already filled the single slot, and healed
emission counts open native states too.
- Responses stream: native deltas that survived the chunk-level cap are
skipped once a healed call claimed the slot.
Also adds a span assertion for the closed-</function> trailing-prose
parse fixed in the previous commit.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: relay undeclared text-form calls as text on Anthropic non-streaming
heal_openai_message promotes only declared text-form tool calls and
span-trims just their markup, deliberately leaving every unpromoted byte
(undeclared text-form calls included) in the content to relay as text.
The Anthropic non-streaming builder then ran a blanket _TOOL_XML_RE strip
over that content unconditionally, deleting the undeclared block before
building the text part, so Anthropic clients silently lost a call the
OpenAI non-streaming path preserves. The strip was harmless when healing
was all-or-nothing but became data loss once healing turned span-exact.
Gate the legacy strip on whether healing promoted a call, matching the
OpenAI passthrough and the intent already stated in the comment above.
Add a route-level regression test for the mixed declared+undeclared case.
* inference: require fully declared nudge retries; keep unpromoted Anthropic text
Codex round 6, two findings:
- response_has_promotable_calls accepted a nudge retry when any one
structured call named a declared tool, so a mixed retry (hallucinated
undeclared call plus a declared one) replaced the original and the
caller forwarded the undeclared call, or with parallel_tool_calls
false could keep only it. All structured retry calls must be declared.
- The Anthropic non-streaming builder still ran the legacy _TOOL_XML_RE
strip after span-exact healing, deleting undeclared or malformed call
text that healing deliberately preserved. The legacy strip now runs
only when healing is off (no declared tools, or opted out), matching
the OpenAI passthrough.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* inference: keep unpromoted Anthropic text whenever healing is active
The previous commit skipped the legacy strip only when a call was
actually promoted, so an undeclared-only (or malformed-only) response
was still silently emptied: exactly the dead-turn shape this path
exists to fix, and inconsistent with the OpenAI passthrough, which
relays those bytes verbatim. Gate the strip on healing being active
instead; opt-out and no-tools requests keep the legacy strip.
* Fix schema-aware tool healing for PR #6801
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix passthrough healing ordering for PR #6801
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix stream finish ordering for PR #6801
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
b8400f40df
commit
308ea5a93c
8 changed files with 2911 additions and 194 deletions
|
|
@ -494,6 +494,29 @@ class AnthropicPassthroughEmitter:
|
|||
self._usage: dict = {}
|
||||
self._stop_reason: str = "end_turn"
|
||||
self._stop_sequence: Optional[str] = None
|
||||
# Optional text-form tool-call healing (client-tool passthrough only).
|
||||
self._healer = None
|
||||
self._healed_tool_use = False
|
||||
self._healed_call_count = 0
|
||||
self._heal_disable_parallel = False
|
||||
|
||||
def enable_healing(
|
||||
self,
|
||||
allowed_tools: set,
|
||||
tools: Optional[list] = None,
|
||||
*,
|
||||
disable_parallel_tool_use: bool = False,
|
||||
) -> None:
|
||||
"""Promote text-form tool calls in streamed content to tool_use blocks.
|
||||
|
||||
Only calls naming a tool in ``allowed_tools`` (the client's declared
|
||||
tools) are promoted; everything else streams as text exactly as before.
|
||||
Never enabled for Studio's own tool loop.
|
||||
"""
|
||||
from core.inference.passthrough_healing import StreamToolCallHealer
|
||||
|
||||
self._healer = StreamToolCallHealer(allowed_tools, tools)
|
||||
self._heal_disable_parallel = disable_parallel_tool_use
|
||||
|
||||
def start(
|
||||
self,
|
||||
|
|
@ -542,29 +565,42 @@ class AnthropicPassthroughEmitter:
|
|||
delta = choice.get("delta") or {}
|
||||
finish_reason = choice.get("finish_reason")
|
||||
|
||||
# ── Structured tool calls take precedence over healing ──
|
||||
# Grammar mode worked: flush anything the healer held (it preceded the
|
||||
# call in the model's output) and relay verbatim from here on.
|
||||
if delta.get("tool_calls") and self._healer is not None and not self._healer.dormant:
|
||||
for kind, value in self._healer.structured_tool_call_seen():
|
||||
if kind == "text" and value:
|
||||
events.extend(self._emit_text_delta(value))
|
||||
|
||||
# ── Text content ──
|
||||
content = delta.get("content")
|
||||
if content:
|
||||
if self._current_block_type != "text":
|
||||
if self._current_block_type is not None:
|
||||
events.append(self._close_current_block())
|
||||
events.extend(self._open_text_block())
|
||||
events.append(
|
||||
build_anthropic_sse_event(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": self.block_index,
|
||||
"delta": {"type": "text_delta", "text": content},
|
||||
},
|
||||
)
|
||||
)
|
||||
if content and self._healer is not None and not self._healer.dormant:
|
||||
# Route text through the healer: held/promoted portions become
|
||||
# synthetic tool_use blocks, the rest streams as text unchanged.
|
||||
for kind, value in self._healer.feed(content):
|
||||
if kind == "text":
|
||||
events.extend(self._emit_text_delta(value))
|
||||
else:
|
||||
events.extend(self._emit_healed_tool_use(value))
|
||||
elif content:
|
||||
events.extend(self._emit_text_delta(content))
|
||||
|
||||
# ── Tool calls (streaming deltas) ──
|
||||
tool_calls = delta.get("tool_calls") or []
|
||||
for tc in tool_calls:
|
||||
tc_idx = tc.get("index", 0)
|
||||
fn = tc.get("function") or {}
|
||||
if (
|
||||
self._heal_disable_parallel
|
||||
and tc_idx not in self._tool_call_states
|
||||
and (self._healed_call_count + len(self._tool_call_states)) >= 1
|
||||
):
|
||||
# disable_parallel_tool_use: a healed call already consumed the
|
||||
# single allowed slot. The caller's chunk-level cap only sees
|
||||
# native indexes, so drop this native call (and its later
|
||||
# argument deltas, which never allocate a state either).
|
||||
continue
|
||||
if tc_idx not in self._tool_call_states:
|
||||
# New tool call — close prior block, open tool_use block
|
||||
if self._current_block_type is not None:
|
||||
|
|
@ -618,6 +654,17 @@ class AnthropicPassthroughEmitter:
|
|||
|
||||
def finish(self) -> list[str]:
|
||||
events: list[str] = []
|
||||
if self._healer is not None:
|
||||
# Last-chance heal of any held residue (e.g. an unclosed tool block).
|
||||
for kind, value in self._healer.finalize():
|
||||
if kind == "text" and value:
|
||||
events.extend(self._emit_text_delta(value))
|
||||
elif kind == "tool_call":
|
||||
events.extend(self._emit_healed_tool_use(value))
|
||||
if self._healed_tool_use and self._stop_reason != "max_tokens":
|
||||
# A promoted call must stop for tool use; a truncation still wins
|
||||
# (its arguments may be incomplete).
|
||||
self._stop_reason = "tool_use"
|
||||
if self._current_block_type is not None:
|
||||
events.append(self._close_current_block())
|
||||
events.append(
|
||||
|
|
@ -641,6 +688,76 @@ class AnthropicPassthroughEmitter:
|
|||
)
|
||||
return events
|
||||
|
||||
def _emit_text_delta(self, content: str) -> list[str]:
|
||||
events: list[str] = []
|
||||
if self._current_block_type != "text":
|
||||
if self._current_block_type is not None:
|
||||
events.append(self._close_current_block())
|
||||
events.extend(self._open_text_block())
|
||||
events.append(
|
||||
build_anthropic_sse_event(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": self.block_index,
|
||||
"delta": {"type": "text_delta", "text": content},
|
||||
},
|
||||
)
|
||||
)
|
||||
return events
|
||||
|
||||
def _emit_healed_tool_use(self, call: dict) -> list[str]:
|
||||
# A healed call arrives complete, so its tool_use block opens, carries
|
||||
# one input_json_delta, and closes immediately; an open text block is
|
||||
# closed first (only the safe prefix ever streamed into it).
|
||||
if (
|
||||
self._heal_disable_parallel
|
||||
and (self._healed_call_count + len(self._tool_call_states)) >= 1
|
||||
):
|
||||
# Healed and native calls share the single allowed slot.
|
||||
return []
|
||||
events: list[str] = []
|
||||
if self._current_block_type is not None:
|
||||
events.append(self._close_current_block())
|
||||
function = call.get("function") or {}
|
||||
tool_id = anthropic_tool_use_id("")
|
||||
self.block_index += 1
|
||||
self._current_block_type = "tool_use"
|
||||
events.append(
|
||||
build_anthropic_sse_event(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": self.block_index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": tool_id,
|
||||
"name": function.get("name", ""),
|
||||
"input": {},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
arguments = function.get("arguments") or ""
|
||||
if arguments:
|
||||
events.append(
|
||||
build_anthropic_sse_event(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": self.block_index,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": arguments,
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
events.append(self._close_current_block())
|
||||
self._healed_tool_use = True
|
||||
self._healed_call_count += 1
|
||||
return events
|
||||
|
||||
def _open_text_block(self) -> list[str]:
|
||||
self.block_index += 1
|
||||
self._current_block_type = "text"
|
||||
|
|
|
|||
535
studio/backend/core/inference/passthrough_healing.py
Normal file
535
studio/backend/core/inference/passthrough_healing.py
Normal file
|
|
@ -0,0 +1,535 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tool-call healing for the client-tool passthrough.
|
||||
|
||||
With server-side tools disabled (``unsloth run --disable-tools``, every
|
||||
``unsloth start`` coding agent), requests carrying the client's own ``tools``
|
||||
bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small
|
||||
GGUF models often emit their tool calls as TEXT (``<tool_call>{...}</tool_call>``,
|
||||
Gemma ``<|tool_call>...``, ``<function=...>`` XML) instead of structured
|
||||
``tool_calls`` -- on the passthrough that text reaches the agent as prose and
|
||||
the turn dies. This module promotes such text back into structured calls on the
|
||||
RESPONSE side only: the upstream request body is never touched, no extra
|
||||
generation is issued, so llama-server slot/KV-cache reuse is byte-identical.
|
||||
|
||||
Healing only ever fires when the request declared client tools, and only
|
||||
promotes calls whose function name exactly matches a declared tool. Promotion
|
||||
removes EXACTLY the promoted calls' markup spans (the parser reports them):
|
||||
undeclared calls, unparseable blocks, and suppressed alternate formats keep
|
||||
every byte and relay as text, so healing can never silently delete model
|
||||
output. Responses without a tool signal, requests without tools, and Studio's
|
||||
own enable-tools loop are untouched. Per-request opt-out:
|
||||
``auto_heal_tool_calls: false``. Process kill-switch:
|
||||
``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Optional
|
||||
|
||||
from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal
|
||||
from core.inference.tool_loop_controller import coerce_tool_arguments
|
||||
from core.tool_healing import parse_tool_calls_from_text
|
||||
|
||||
# Read once at import (same convention as the other UNSLOTH_* switches).
|
||||
_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1"
|
||||
# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process
|
||||
# default with UNSLOTH_TOOL_CALL_NUDGE=1 (e.g. an `unsloth run` operator).
|
||||
_NUDGE_DEFAULT = os.environ.get("UNSLOTH_TOOL_CALL_NUDGE", "0") == "1"
|
||||
|
||||
|
||||
def nudge_enabled(request_flag: Optional[bool]) -> bool:
|
||||
return _NUDGE_DEFAULT if request_flag is None else bool(request_flag)
|
||||
|
||||
|
||||
_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS)
|
||||
# A suspected-but-unclosed tool block larger than this is declared a false
|
||||
# alarm and flushed, bounding memory on a model rambling XML-lookalike text.
|
||||
_MAX_HOLD_CHARS = 64 * 1024
|
||||
|
||||
|
||||
def heal_gate(
|
||||
auto_heal: Optional[bool],
|
||||
tools: Optional[list],
|
||||
tool_choice: Any = None,
|
||||
) -> Optional[set]:
|
||||
"""Return the declared client-tool name set when healing applies, else None.
|
||||
|
||||
``tools`` is the OpenAI-shaped list forwarded to llama-server
|
||||
(``[{"type": "function", "function": {"name": ...}}, ...]``). The name set
|
||||
doubles as the promotion allowlist so healed calls can never invent a tool
|
||||
the client did not declare.
|
||||
|
||||
``tool_choice`` (OpenAI shape) constrains the allowlist so healing never
|
||||
contradicts the request: ``"none"`` forbids tool calls outright (text-form
|
||||
markup stays text), and a forced ``{"type": "function", "function":
|
||||
{"name": N}}`` narrows promotion to that one function. ``"auto"`` /
|
||||
``"required"`` / absent keep the full declared set.
|
||||
"""
|
||||
if _HEALING_DISABLED or auto_heal is False:
|
||||
return None
|
||||
if tool_choice == "none":
|
||||
return None
|
||||
names = set()
|
||||
for tool in tools or []:
|
||||
if not isinstance(tool, dict):
|
||||
continue
|
||||
function = tool.get("function")
|
||||
if isinstance(function, dict) and isinstance(function.get("name"), str):
|
||||
names.add(function["name"])
|
||||
if isinstance(tool_choice, dict):
|
||||
function = tool_choice.get("function")
|
||||
forced = function.get("name") if isinstance(function, dict) else None
|
||||
if isinstance(forced, str):
|
||||
names &= {forced}
|
||||
return names or None
|
||||
|
||||
|
||||
def _tool_schemas_by_name(tools: Optional[list]) -> dict[str, Any]:
|
||||
schemas: dict[str, Any] = {}
|
||||
for tool in tools or []:
|
||||
if not isinstance(tool, dict):
|
||||
continue
|
||||
function = tool.get("function")
|
||||
if not isinstance(function, dict):
|
||||
continue
|
||||
name = function.get("name")
|
||||
if isinstance(name, str):
|
||||
schemas[name] = function.get("parameters")
|
||||
return schemas
|
||||
|
||||
|
||||
def _string_arg_key_from_schema(schema: Any) -> Optional[str]:
|
||||
if not isinstance(schema, dict):
|
||||
return None
|
||||
properties = schema.get("properties")
|
||||
required = schema.get("required")
|
||||
if not isinstance(properties, dict) or not isinstance(required, list):
|
||||
return None
|
||||
required_names = [name for name in required if isinstance(name, str)]
|
||||
if len(required_names) != 1:
|
||||
return None
|
||||
key = required_names[0]
|
||||
|
||||
if key not in properties:
|
||||
return None
|
||||
prop_schema = properties.get(key)
|
||||
if isinstance(prop_schema, dict):
|
||||
prop_type = prop_schema.get("type")
|
||||
if isinstance(prop_type, list):
|
||||
if "string" not in prop_type:
|
||||
return None
|
||||
elif prop_type is not None and prop_type != "string":
|
||||
return None
|
||||
return key
|
||||
|
||||
|
||||
def _coerce_promoted_arguments(
|
||||
raw_args: Any, tool_name: str, tool_schemas: Optional[dict]
|
||||
) -> Optional[dict]:
|
||||
if isinstance(raw_args, Mapping):
|
||||
return dict(raw_args)
|
||||
if isinstance(raw_args, str):
|
||||
try:
|
||||
parsed = json.loads(raw_args)
|
||||
if isinstance(parsed, Mapping):
|
||||
return dict(parsed)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
if tool_schemas is not None:
|
||||
key = _string_arg_key_from_schema(tool_schemas.get(tool_name))
|
||||
return {key: raw_args} if key else None
|
||||
coerced = coerce_tool_arguments(raw_args, heal = True, tool_name = tool_name)
|
||||
return coerced.arguments
|
||||
|
||||
|
||||
def _promote(
|
||||
calls: list,
|
||||
allowed_tools: set,
|
||||
id_offset: int = 0,
|
||||
tool_schemas: Optional[dict] = None,
|
||||
) -> list:
|
||||
"""Filter parsed calls to declared tools and normalize their arguments.
|
||||
|
||||
Bare string arguments on the client-tool passthrough use the declared
|
||||
schema's single required string property. If the schema is ambiguous, the
|
||||
call stays text instead of inventing a generic key.
|
||||
"""
|
||||
promoted = []
|
||||
for call in calls:
|
||||
function = call.get("function") if isinstance(call, dict) else None
|
||||
name = function.get("name") if isinstance(function, dict) else None
|
||||
if name not in allowed_tools:
|
||||
continue
|
||||
arguments = _coerce_promoted_arguments(function.get("arguments"), name, tool_schemas)
|
||||
if arguments is None:
|
||||
continue
|
||||
promoted.append(
|
||||
{
|
||||
"id": f"call_{id_offset + len(promoted)}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"arguments": json.dumps(arguments, ensure_ascii = False),
|
||||
},
|
||||
}
|
||||
)
|
||||
return promoted
|
||||
|
||||
|
||||
def _remove_spans(text: str, spans: list) -> str:
|
||||
"""Text with the given non-overlapping, sorted (start, end) ranges removed."""
|
||||
pieces = []
|
||||
pos = 0
|
||||
for start, end in spans:
|
||||
pieces.append(text[pos:start])
|
||||
pos = end
|
||||
pieces.append(text[pos:])
|
||||
return "".join(pieces)
|
||||
|
||||
|
||||
def heal_openai_message_events(
|
||||
msg: dict,
|
||||
allowed_tools: set,
|
||||
tools: Optional[list] = None,
|
||||
) -> Optional[list]:
|
||||
if not isinstance(msg, dict) or msg.get("tool_calls"):
|
||||
return None
|
||||
content = msg.get("content")
|
||||
if not isinstance(content, str) or not has_tool_signal(content):
|
||||
return None
|
||||
parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True)
|
||||
tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
|
||||
events: list = []
|
||||
pos = 0
|
||||
call_count = 0
|
||||
for call, (start, end) in zip(parsed, spans):
|
||||
promoted = _promote([call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas)
|
||||
if promoted:
|
||||
if content[pos:start]:
|
||||
events.append(("text", content[pos:start]))
|
||||
events.append(("tool_call", promoted[0]))
|
||||
call_count += 1
|
||||
else:
|
||||
events.append(("text", content[pos:end]))
|
||||
pos = end
|
||||
if not call_count:
|
||||
return None
|
||||
if content[pos:]:
|
||||
events.append(("text", content[pos:]))
|
||||
return events
|
||||
|
||||
|
||||
def heal_openai_message(
|
||||
msg: dict,
|
||||
allowed_tools: set,
|
||||
tools: Optional[list] = None,
|
||||
) -> bool:
|
||||
"""Promote text-form tool calls in a non-streaming OpenAI message. In place.
|
||||
|
||||
No-op (returns False) unless the message has NO structured ``tool_calls``
|
||||
(grammar mode already worked when it does) and its content carries a tool
|
||||
signal that parses into at least one declared call. Only the promoted
|
||||
calls' markup spans are removed from the content; undeclared calls and
|
||||
anything the parser did not consume stay in the text byte-intact.
|
||||
"""
|
||||
events = heal_openai_message_events(msg, allowed_tools, tools)
|
||||
if not events:
|
||||
return False
|
||||
calls = [value for kind, value in events if kind == "tool_call"]
|
||||
content = "".join(value for kind, value in events if kind == "text").strip()
|
||||
msg["tool_calls"] = calls
|
||||
# OpenAI requires content = null on a pure tool-call turn.
|
||||
msg["content"] = content or None
|
||||
return True
|
||||
|
||||
|
||||
def _earliest_signal(buffer: str) -> int:
|
||||
best = -1
|
||||
for signal in TOOL_XML_SIGNALS:
|
||||
index = buffer.find(signal)
|
||||
if index >= 0 and (best < 0 or index < best):
|
||||
best = index
|
||||
return best
|
||||
|
||||
|
||||
def _closed_signal_span(buffer: str) -> Optional[tuple[int, int]]:
|
||||
spans = []
|
||||
for open_tag, close_tag in (
|
||||
("<tool_call>", "</tool_call>"),
|
||||
("<|tool_call>", "<tool_call|>"),
|
||||
("<function=", "</function>"),
|
||||
):
|
||||
start = buffer.find(open_tag)
|
||||
if start < 0:
|
||||
continue
|
||||
end = buffer.find(close_tag, start)
|
||||
if end >= 0:
|
||||
spans.append((start, end + len(close_tag)))
|
||||
return min(spans, key = lambda span: span[0]) if spans else None
|
||||
|
||||
|
||||
def _partial_signal_suffix(buffer: str) -> int:
|
||||
"""Length of the longest buffer suffix that is a proper prefix of a signal."""
|
||||
for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1):
|
||||
tail = buffer[-length:]
|
||||
if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS):
|
||||
return length
|
||||
return 0
|
||||
|
||||
|
||||
class StreamToolCallHealer:
|
||||
"""Buffer-and-repair state machine for streamed passthrough content.
|
||||
|
||||
``feed(text)`` / ``finalize()`` yield ``("text", str)`` events for content
|
||||
to relay and ``("tool_call", dict)`` events carrying an OpenAI-shaped call
|
||||
(string ``function.arguments``). Normal prose is forwarded immediately; only
|
||||
a trailing partial-signal window (< max signal length) or a suspected tool
|
||||
block is ever withheld, so streaming latency stays bounded. A false alarm
|
||||
(the buffer can no longer become a parseable declared call) flushes the held
|
||||
text verbatim.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
allowed_tools: set,
|
||||
tools: Optional[list] = None,
|
||||
) -> None:
|
||||
self._allowed = set(allowed_tools)
|
||||
|
||||
self._tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
|
||||
self._buffer = ""
|
||||
self._holding = False
|
||||
self._id_offset = 0
|
||||
# Structured delta.tool_calls seen upstream: grammar mode already
|
||||
# worked, so healing goes dormant and text relays verbatim.
|
||||
self.dormant = False
|
||||
|
||||
@property
|
||||
def healed(self) -> bool:
|
||||
return self._id_offset > 0
|
||||
|
||||
def structured_tool_call_seen(self) -> list:
|
||||
"""Go dormant; flush anything held so no text is swallowed."""
|
||||
self.dormant = True
|
||||
held, self._buffer, self._holding = self._buffer, "", False
|
||||
return [("text", held)] if held else []
|
||||
|
||||
def feed(self, text: str) -> list:
|
||||
if self.dormant:
|
||||
return [("text", text)] if text else []
|
||||
self._buffer += text
|
||||
return self._drain()
|
||||
|
||||
def _drain(self) -> list:
|
||||
events: list = []
|
||||
while True:
|
||||
if not self._holding:
|
||||
start = _earliest_signal(self._buffer)
|
||||
if start >= 0:
|
||||
if start:
|
||||
events.append(("text", self._buffer[:start]))
|
||||
self._buffer = self._buffer[start:]
|
||||
self._holding = True
|
||||
else:
|
||||
keep = _partial_signal_suffix(self._buffer)
|
||||
emit = self._buffer[: len(self._buffer) - keep]
|
||||
if emit:
|
||||
events.append(("text", emit))
|
||||
self._buffer = self._buffer[len(self._buffer) - keep :]
|
||||
return events
|
||||
# HOLD: handle the FIRST complete block per pass so events keep
|
||||
# document order (a later declared call must not overtake an
|
||||
# earlier undeclared one flushing as text).
|
||||
parsed, spans = parse_tool_calls_from_text(
|
||||
self._buffer,
|
||||
id_offset = self._id_offset,
|
||||
allow_incomplete = False,
|
||||
with_spans = True,
|
||||
)
|
||||
if not parsed:
|
||||
closed_span = _closed_signal_span(self._buffer)
|
||||
if closed_span:
|
||||
_start, end = closed_span
|
||||
events.append(("text", self._buffer[:end]))
|
||||
self._buffer = self._buffer[end:]
|
||||
self._holding = False
|
||||
continue
|
||||
if len(self._buffer) > _MAX_HOLD_CHARS:
|
||||
events.append(("text", self._buffer))
|
||||
self._buffer = ""
|
||||
self._holding = False
|
||||
continue
|
||||
return events
|
||||
start, end = spans[0]
|
||||
promoted = _promote(
|
||||
[parsed[0]],
|
||||
self._allowed,
|
||||
id_offset = self._id_offset,
|
||||
tool_schemas = self._tool_schemas,
|
||||
)
|
||||
if promoted:
|
||||
if start:
|
||||
events.append(("text", self._buffer[:start]))
|
||||
events.append(("tool_call", promoted[0]))
|
||||
self._id_offset += 1
|
||||
# Drop exactly the promoted markup span; everything else
|
||||
# (leading text, later blocks) stays and is rescanned.
|
||||
self._buffer = self._buffer[end:]
|
||||
else:
|
||||
# Undeclared or unusable name: its markup is DATA, flush it
|
||||
# (and anything before it) verbatim, then rescan the rest.
|
||||
events.append(("text", self._buffer[:end]))
|
||||
self._buffer = self._buffer[end:]
|
||||
self._holding = False
|
||||
|
||||
def finalize(self) -> list:
|
||||
"""End of stream: last-chance heal of the residue, else flush it.
|
||||
|
||||
Events keep document order; only the promoted calls' markup spans are
|
||||
dropped, every other residue byte flushes as text.
|
||||
"""
|
||||
if not self._buffer:
|
||||
return []
|
||||
residue, self._buffer = self._buffer, ""
|
||||
holding, self._holding = self._holding, False
|
||||
if self.dormant or not holding:
|
||||
return [("text", residue)]
|
||||
parsed, spans = parse_tool_calls_from_text(
|
||||
residue,
|
||||
id_offset = self._id_offset,
|
||||
allow_incomplete = True,
|
||||
with_spans = True,
|
||||
)
|
||||
events: list = []
|
||||
pos = 0
|
||||
any_promoted = False
|
||||
for call, (start, end) in zip(parsed, spans):
|
||||
promoted = _promote(
|
||||
[call],
|
||||
self._allowed,
|
||||
id_offset = self._id_offset,
|
||||
tool_schemas = self._tool_schemas,
|
||||
)
|
||||
if promoted:
|
||||
if residue[pos:start]:
|
||||
events.append(("text", residue[pos:start]))
|
||||
events.append(("tool_call", promoted[0]))
|
||||
self._id_offset += 1
|
||||
any_promoted = True
|
||||
else:
|
||||
events.append(("text", residue[pos:end]))
|
||||
pos = end
|
||||
if not any_promoted:
|
||||
return [("text", residue)]
|
||||
tail = residue[pos:].strip()
|
||||
if tail:
|
||||
events.append(("text", tail))
|
||||
return events
|
||||
|
||||
|
||||
def _first_choice_message(data: Any) -> Optional[dict]:
|
||||
"""First-choice message dict of a non-streaming chat response, else None.
|
||||
|
||||
Upstream error bodies can carry ``"message": null`` (or no choices at all),
|
||||
so never assume the shape: a non-dict message means "nothing to heal".
|
||||
"""
|
||||
try:
|
||||
message = data["choices"][0]["message"]
|
||||
except (KeyError, IndexError, TypeError):
|
||||
return None
|
||||
return message if isinstance(message, dict) else None
|
||||
|
||||
|
||||
def _last_assistant_text(data: Any) -> str:
|
||||
"""First-choice assistant content of a non-streaming chat response, or ''."""
|
||||
message = _first_choice_message(data)
|
||||
content = message.get("content") if message else None
|
||||
return content if isinstance(content, str) else ""
|
||||
|
||||
|
||||
def _heal_would_promote(
|
||||
text: str,
|
||||
allowed_tools: set,
|
||||
tools: Optional[list] = None,
|
||||
) -> bool:
|
||||
"""Whether ``heal_openai_message`` would promote at least one call."""
|
||||
parsed = parse_tool_calls_from_text(text, allow_incomplete = True)
|
||||
tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
|
||||
return bool(_promote(parsed, allowed_tools, tool_schemas = tool_schemas))
|
||||
|
||||
|
||||
def response_has_promotable_calls(
|
||||
data: Any,
|
||||
allowed_tools: set,
|
||||
tools: Optional[list] = None,
|
||||
) -> bool:
|
||||
"""True when a non-streaming chat response carries a usable tool call
|
||||
(structured naming a DECLARED tool, or text-form that healing would
|
||||
promote). Used to decide whether a nudge retry actually improved on the
|
||||
original response; a hallucinated undeclared call is not an improvement."""
|
||||
message = _first_choice_message(data)
|
||||
if not message:
|
||||
return False
|
||||
tool_calls = message.get("tool_calls")
|
||||
if tool_calls:
|
||||
# ALL structured calls must be declared: the caller forwards the whole
|
||||
# list (and a parallel cap could keep only the FIRST one), so a mixed
|
||||
# response with a single hallucinated name could still hand the client
|
||||
# an undeclared tool.
|
||||
return all(
|
||||
isinstance(tc, dict)
|
||||
and isinstance(tc.get("function"), dict)
|
||||
and tc["function"].get("name") in allowed_tools
|
||||
for tc in tool_calls
|
||||
)
|
||||
text = message.get("content")
|
||||
if not isinstance(text, str):
|
||||
return False
|
||||
return _heal_would_promote(text, allowed_tools, tools)
|
||||
|
||||
|
||||
def nudge_should_retry(
|
||||
data: Any,
|
||||
allowed_tools: Optional[set],
|
||||
tools: Optional[list] = None,
|
||||
) -> bool:
|
||||
"""True when the first response tried to call a tool but nothing healed.
|
||||
|
||||
Trigger only on: healing enabled (allowed_tools set), zero structured
|
||||
calls, a tool signal present in the text, and zero promotable calls -- the
|
||||
exact failure a single re-ask can fix. Clean prose never retries.
|
||||
"""
|
||||
if not allowed_tools:
|
||||
return False
|
||||
message = _first_choice_message(data)
|
||||
if not message or message.get("tool_calls"):
|
||||
return False
|
||||
text = message.get("content")
|
||||
if not isinstance(text, str) or not has_tool_signal(text):
|
||||
return False
|
||||
return not _heal_would_promote(text, allowed_tools, tools)
|
||||
|
||||
|
||||
def nudge_messages(data: Any, allowed_tools: set) -> list:
|
||||
"""The two-message suffix appended for the single nudge retry.
|
||||
|
||||
The retry body is the original body plus this suffix, so the prompt prefix
|
||||
is byte-identical and llama-server's slot/prefix cache is reused (same
|
||||
shape as the enable-tools loop's reprompt).
|
||||
"""
|
||||
tool_hint = " or ".join(f"`{name}`" for name in sorted(allowed_tools)) or "an available tool"
|
||||
return [
|
||||
{"role": "assistant", "content": _last_assistant_text(data)},
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"You have access to the declared tools. If a tool is needed to "
|
||||
f"complete the action you described, call {tool_hint} now using the "
|
||||
"native tool-call format with valid JSON arguments, not prose. If no "
|
||||
"tool is needed, provide the final answer directly."
|
||||
),
|
||||
},
|
||||
]
|
||||
|
|
@ -301,28 +301,28 @@ def parse_tool_calls_from_text(
|
|||
*,
|
||||
id_offset: int = 0,
|
||||
allow_incomplete: bool = True,
|
||||
) -> list[dict]:
|
||||
with_spans: bool = False,
|
||||
):
|
||||
"""Parse OpenAI-format tool calls from model text.
|
||||
|
||||
Handles formats like:
|
||||
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
|
||||
<|tool_call>call:web_search{query:"..."}<tool_call|>
|
||||
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
|
||||
|
||||
With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]``
|
||||
is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup
|
||||
in ``content`` (including its close tag when present), so a caller can
|
||||
remove exactly the parsed markup and keep every other byte intact.
|
||||
"""
|
||||
tool_calls: list[dict] = []
|
||||
# Collect JSON- and Gemma-format candidates with their byte spans, then
|
||||
# accept them in document order. Both order and spans matter:
|
||||
# * tools execute in returned order, so a call appearing earlier in the
|
||||
# text must be emitted first even across the two formats;
|
||||
# * a tool-call marker INSIDE another call's argument string is data, not a
|
||||
# call, so a candidate starting within an already accepted span is
|
||||
# skipped (covers a JSON marker nested in a Gemma arg and a Gemma marker
|
||||
# nested in a JSON arg alike, regardless of which format is outer).
|
||||
call_spans: list[tuple] = []
|
||||
# Collect every supported call format with spans, then emit in document
|
||||
# order. A marker inside another call's argument string is data, not a
|
||||
# separate executable call.
|
||||
parsed_items = [] # (start, span_end, name, arguments)
|
||||
candidates = [] # (start, brace_end, kind, match)
|
||||
for m in _TC_JSON_START_RE.finditer(content):
|
||||
# A marker that begins inside an open <function=...><parameter=...> value
|
||||
# is that parameter's data, not its own call; skip it (same guard the
|
||||
# XML-style parser below applies to nested <function= markers).
|
||||
if _inside_open_parameter(content, m.start()):
|
||||
continue
|
||||
end = _balanced_brace_end(content, m.end() - 1)
|
||||
|
|
@ -336,14 +336,9 @@ def parse_tool_calls_from_text(
|
|||
candidates.append((m.start(), end, "gemma", m))
|
||||
candidates.sort(key = lambda c: c[0])
|
||||
|
||||
spans = [(s, e) for s, e, _kind, _m in candidates]
|
||||
candidate_spans = [(s, e) for s, e, _kind, _m in candidates]
|
||||
for idx, (start, end, kind, m) in enumerate(candidates):
|
||||
# Skip a candidate nested inside another candidate's brace span: it is
|
||||
# the enclosing call's argument data, not its own call. Checked against
|
||||
# every candidate span (not only the ones that parsed successfully), so a
|
||||
# marker inside an outer call that later fails to normalize is still
|
||||
# never promoted to its own executable tool call.
|
||||
if any(s <= start and end <= e for j, (s, e) in enumerate(spans) if j != idx):
|
||||
if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx):
|
||||
continue
|
||||
if not allow_incomplete:
|
||||
tail = content[end + 1 :].lstrip()
|
||||
|
|
@ -362,6 +357,85 @@ def parse_tool_calls_from_text(
|
|||
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end]))
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
span_end = end + 1
|
||||
close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
|
||||
ws = len(content[span_end:]) - len(content[span_end:].lstrip())
|
||||
close_m = close_re.match(content, span_end + ws)
|
||||
if close_m:
|
||||
span_end = close_m.end()
|
||||
parsed_items.append((start, span_end, name, arguments))
|
||||
|
||||
func_starts = [
|
||||
fm
|
||||
for fm in _TC_FUNC_START_RE.finditer(content)
|
||||
if not _inside_open_parameter(content, fm.start())
|
||||
and not any(s <= fm.start() <= e for s, e in candidate_spans)
|
||||
]
|
||||
for idx, fm in enumerate(func_starts):
|
||||
func_name = fm.group(1)
|
||||
body_start = fm.end()
|
||||
next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content)
|
||||
end_tag = _TC_END_TAG_RE.search(content[body_start:])
|
||||
if end_tag:
|
||||
body_end = body_start + end_tag.start()
|
||||
else:
|
||||
body_end = len(content)
|
||||
body_end = min(body_end, next_func)
|
||||
body = content[body_start:body_end]
|
||||
close_idx = body.rfind(_FUNC_CLOSE_TAG)
|
||||
if close_idx >= 0:
|
||||
span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG)
|
||||
body = body[:close_idx]
|
||||
elif not allow_incomplete:
|
||||
continue
|
||||
else:
|
||||
body = _TC_FUNC_CLOSE_RE.sub("", body)
|
||||
span_end = body_end
|
||||
|
||||
arguments: dict = {}
|
||||
param_starts = list(_TC_PARAM_START_RE.finditer(body))
|
||||
if len(param_starts) == 1:
|
||||
pm = param_starts[0]
|
||||
val = body[pm.end() :]
|
||||
if not allow_incomplete:
|
||||
stripped_val = val.rstrip()
|
||||
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
|
||||
continue
|
||||
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
|
||||
else:
|
||||
val = _TC_PARAM_CLOSE_RE.sub("", val)
|
||||
arguments[pm.group(1)] = val.strip()
|
||||
else:
|
||||
valid_params = True
|
||||
for pidx, pm in enumerate(param_starts):
|
||||
param_name = pm.group(1)
|
||||
val_start = pm.end()
|
||||
next_param = (
|
||||
param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body)
|
||||
)
|
||||
val = body[val_start:next_param]
|
||||
if not allow_incomplete:
|
||||
stripped_val = val.rstrip()
|
||||
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
|
||||
valid_params = False
|
||||
break
|
||||
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
|
||||
else:
|
||||
val = _TC_PARAM_CLOSE_RE.sub("", val)
|
||||
arguments[param_name] = val.strip()
|
||||
if not valid_params:
|
||||
continue
|
||||
|
||||
span_start = fm.start()
|
||||
wrap_open = re.search(r"<tool_call>\s*$", content[:span_start])
|
||||
wrap_close = re.match(r"\s*</tool_call>", content[span_end:])
|
||||
if wrap_open and wrap_close:
|
||||
span_start = wrap_open.start()
|
||||
span_end += wrap_close.end()
|
||||
parsed_items.append((span_start, span_end, func_name, json.dumps(arguments)))
|
||||
|
||||
parsed_items.sort(key = lambda item: item[0])
|
||||
for start, span_end, name, arguments in parsed_items:
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": f"call_{id_offset + len(tool_calls)}",
|
||||
|
|
@ -369,77 +443,9 @@ def parse_tool_calls_from_text(
|
|||
"function": {"name": name, "arguments": arguments},
|
||||
}
|
||||
)
|
||||
|
||||
if not tool_calls:
|
||||
func_starts = [
|
||||
fm
|
||||
for fm in _TC_FUNC_START_RE.finditer(content)
|
||||
if not _inside_open_parameter(content, fm.start())
|
||||
]
|
||||
for idx, fm in enumerate(func_starts):
|
||||
func_name = fm.group(1)
|
||||
body_start = fm.end()
|
||||
next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content)
|
||||
end_tag = _TC_END_TAG_RE.search(content[body_start:])
|
||||
if end_tag:
|
||||
body_end = body_start + end_tag.start()
|
||||
else:
|
||||
body_end = len(content)
|
||||
body_end = min(body_end, next_func)
|
||||
body = content[body_start:body_end]
|
||||
if not allow_incomplete:
|
||||
close_idx = body.rfind(_FUNC_CLOSE_TAG)
|
||||
if close_idx < 0:
|
||||
continue
|
||||
body = body[:close_idx]
|
||||
else:
|
||||
body = _TC_FUNC_CLOSE_RE.sub("", body)
|
||||
|
||||
arguments: dict = {}
|
||||
param_starts = list(_TC_PARAM_START_RE.finditer(body))
|
||||
if len(param_starts) == 1:
|
||||
pm = param_starts[0]
|
||||
val = body[pm.end() :]
|
||||
if not allow_incomplete:
|
||||
stripped_val = val.rstrip()
|
||||
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
|
||||
continue
|
||||
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
|
||||
else:
|
||||
val = _TC_PARAM_CLOSE_RE.sub("", val)
|
||||
arguments[pm.group(1)] = val.strip()
|
||||
else:
|
||||
valid_params = True
|
||||
for pidx, pm in enumerate(param_starts):
|
||||
param_name = pm.group(1)
|
||||
val_start = pm.end()
|
||||
next_param = (
|
||||
param_starts[pidx + 1].start()
|
||||
if pidx + 1 < len(param_starts)
|
||||
else len(body)
|
||||
)
|
||||
val = body[val_start:next_param]
|
||||
if not allow_incomplete:
|
||||
stripped_val = val.rstrip()
|
||||
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
|
||||
valid_params = False
|
||||
break
|
||||
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
|
||||
else:
|
||||
val = _TC_PARAM_CLOSE_RE.sub("", val)
|
||||
arguments[param_name] = val.strip()
|
||||
if not valid_params:
|
||||
continue
|
||||
|
||||
tc = {
|
||||
"id": f"call_{id_offset + len(tool_calls)}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": func_name,
|
||||
"arguments": json.dumps(arguments),
|
||||
},
|
||||
}
|
||||
tool_calls.append(tc)
|
||||
call_spans.append((start, span_end))
|
||||
if with_spans:
|
||||
return tool_calls, call_spans
|
||||
return tool_calls
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -780,6 +780,16 @@ class ChatCompletionRequest(BaseModel):
|
|||
True,
|
||||
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
|
||||
)
|
||||
nudge_tool_calls: Optional[bool] = Field(
|
||||
None,
|
||||
description = (
|
||||
"[x-unsloth] Opt-in, non-streaming client-tool passthrough only: when the "
|
||||
"model emitted a tool signal that healing could not repair, retry ONCE with "
|
||||
"a short nudge appended (the retry shares the full prompt prefix, so the "
|
||||
"server's KV cache is reused). Default off; UNSLOTH_TOOL_CALL_NUDGE=1 flips "
|
||||
"the process default."
|
||||
),
|
||||
)
|
||||
context_overflow: Optional[Literal["error", "truncate_middle"]] = Field(
|
||||
None,
|
||||
description = (
|
||||
|
|
@ -1612,6 +1622,14 @@ class AnthropicMessagesRequest(BaseModel):
|
|||
False,
|
||||
description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.",
|
||||
)
|
||||
auto_heal_tool_calls: Optional[bool] = Field(
|
||||
True,
|
||||
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).",
|
||||
)
|
||||
nudge_tool_calls: Optional[bool] = Field(
|
||||
None,
|
||||
description = "[x-unsloth] Opt-in, non-streaming only: retry once with a nudge when the model emitted a tool signal healing could not repair (mirrors the Chat Completions field).",
|
||||
)
|
||||
model_config = {"extra": "allow"}
|
||||
|
||||
@model_validator(mode = "before")
|
||||
|
|
|
|||
|
|
@ -1135,6 +1135,16 @@ from core.inference.key_exchange import decrypt_api_key
|
|||
from core.inference.model_ids import public_model_id
|
||||
from core.inference.api_monitor import api_monitor
|
||||
from core.inference.llama_http import nonstreaming_client
|
||||
from core.inference.passthrough_healing import (
|
||||
StreamToolCallHealer,
|
||||
heal_gate,
|
||||
heal_openai_message,
|
||||
heal_openai_message_events,
|
||||
nudge_enabled,
|
||||
nudge_messages,
|
||||
nudge_should_retry,
|
||||
response_has_promotable_calls,
|
||||
)
|
||||
from core.inference.providers import get_base_url
|
||||
from core.inference.external_provider import ExternalProviderClient
|
||||
from core.inference.chat_templates import resolve_effective_chat_template_override
|
||||
|
|
@ -8065,6 +8075,13 @@ def _build_chat_request(
|
|||
if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw:
|
||||
chat_kwargs["enable_thinking"] = bool(_tpl_kw["enable_thinking"])
|
||||
explicit_enable_thinking = True
|
||||
# auto_heal_tool_calls / nudge_tool_calls are not typed on
|
||||
# ResponsesRequest; lift them from the extra-body so passthrough
|
||||
# healing (and the opt-in nudge) honor them on both paths.
|
||||
if isinstance(_extra.get("auto_heal_tool_calls"), bool):
|
||||
chat_kwargs["auto_heal_tool_calls"] = _extra["auto_heal_tool_calls"]
|
||||
if isinstance(_extra.get("nudge_tool_calls"), bool):
|
||||
chat_kwargs["nudge_tool_calls"] = _extra["nudge_tool_calls"]
|
||||
|
||||
if isinstance(payload.reasoning, dict):
|
||||
effort = payload.reasoning.get("effort")
|
||||
|
|
@ -8299,16 +8316,112 @@ async def _responses_stream(
|
|||
parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend)
|
||||
)
|
||||
reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False}
|
||||
message_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False}
|
||||
message_state: dict[str, Any] = {
|
||||
"output_index": None,
|
||||
"item_id": None,
|
||||
"opened": False,
|
||||
"text": "",
|
||||
}
|
||||
# Message items already closed mid-stream (a healed tool call splits
|
||||
# the assistant text into separate message items, as native Responses
|
||||
# streams do). Kept for the final response.completed snapshot.
|
||||
closed_message_states: list[dict] = []
|
||||
# Per-tool-call state keyed by Chat Completions `tool_calls[].index`,
|
||||
# stable across chunks for the same call. Values:
|
||||
# {output_index, item_id, call_id, name, arguments, opened}
|
||||
tool_call_state: dict[int, dict] = {}
|
||||
next_output_index = 0
|
||||
# Text-form tool calls promoted back to structured calls (declared
|
||||
# client tools only); dormant once grammar-mode structured deltas appear.
|
||||
_allowed_tools = heal_gate(
|
||||
getattr(chat_req, "auto_heal_tool_calls", None),
|
||||
body.get("tools"),
|
||||
body.get("tool_choice"),
|
||||
)
|
||||
healer = StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None
|
||||
healed_tc_index = 0
|
||||
|
||||
def _healed_tc(call: dict):
|
||||
# Chat-delta shape for a healed call. Indexes live in a disjoint
|
||||
# range so a healed call can never merge into a structured call's
|
||||
# state slot; parallel_tool_calls=false caps healed calls too (the
|
||||
# upstream cap ran before injection).
|
||||
nonlocal healed_tc_index
|
||||
if payload.parallel_tool_calls is False and healed_tc_index >= 1:
|
||||
return None
|
||||
tc = {
|
||||
"index": 1_000_000 + healed_tc_index,
|
||||
"id": call["id"],
|
||||
"type": "function",
|
||||
"function": call["function"],
|
||||
}
|
||||
healed_tc_index += 1
|
||||
return tc
|
||||
|
||||
def _sse(event_name: str, payload: dict) -> str:
|
||||
return f"event: {event_name}\ndata: {json.dumps(payload)}\n\n"
|
||||
|
||||
def _tool_call_delta_events(tc: dict) -> list:
|
||||
# One Chat Completions tool_calls delta -> Responses SSE events,
|
||||
# allocating/merging per-call state (shared by the structured loop
|
||||
# and the healer's promoted calls).
|
||||
events = []
|
||||
idx = tc.get("index", 0)
|
||||
st = tool_call_state.get(idx)
|
||||
fn = tc.get("function") or {}
|
||||
if st is None:
|
||||
# First chunk for this tool call -- allocate an
|
||||
# output_index and emit output_item.added.
|
||||
st = {
|
||||
"output_index": _claim_output_index(),
|
||||
"item_id": f"fc_{uuid.uuid4().hex[:12]}",
|
||||
"call_id": tc.get("id") or "",
|
||||
"name": fn.get("name") or "",
|
||||
"arguments": "",
|
||||
"opened": False,
|
||||
}
|
||||
tool_call_state[idx] = st
|
||||
else:
|
||||
# Later chunks sometimes carry id/name only once; merge
|
||||
# when present.
|
||||
if tc.get("id") and not st["call_id"]:
|
||||
st["call_id"] = tc["id"]
|
||||
if fn.get("name") and not st["name"]:
|
||||
st["name"] = fn["name"]
|
||||
|
||||
if not st["opened"] and st["call_id"] and st["name"]:
|
||||
item_added = {
|
||||
"type": "response.output_item.added",
|
||||
"output_index": st["output_index"],
|
||||
"item": {
|
||||
"type": "function_call",
|
||||
"id": st["item_id"],
|
||||
"status": "in_progress",
|
||||
"call_id": st["call_id"],
|
||||
"name": st["name"],
|
||||
"arguments": "",
|
||||
},
|
||||
}
|
||||
events.append(_sse("response.output_item.added", item_added))
|
||||
st["opened"] = True
|
||||
|
||||
arg_delta = fn.get("arguments") or ""
|
||||
if arg_delta and st["opened"]:
|
||||
st["arguments"] += arg_delta
|
||||
args_delta_event = {
|
||||
"type": "response.function_call_arguments.delta",
|
||||
"item_id": st["item_id"],
|
||||
"output_index": st["output_index"],
|
||||
"delta": arg_delta,
|
||||
}
|
||||
events.append(_sse("response.function_call_arguments.delta", args_delta_event))
|
||||
elif arg_delta:
|
||||
# Buffer args until we can open the item (some models
|
||||
# send id/name in the same chunk as the first arg delta;
|
||||
# if not, stash).
|
||||
st["arguments"] += arg_delta
|
||||
return events
|
||||
|
||||
def _claim_output_index() -> int:
|
||||
nonlocal next_output_index
|
||||
output_index = next_output_index
|
||||
|
|
@ -8393,6 +8506,98 @@ async def _responses_stream(
|
|||
),
|
||||
]
|
||||
|
||||
def _close_message_item() -> list[str]:
|
||||
"""Close the open message item so later text opens a fresh one.
|
||||
|
||||
Emits the same done-event triplet the end-of-stream close loop
|
||||
would, records the item for the final snapshot, and resets the
|
||||
state in place. No-op when no message item is open.
|
||||
"""
|
||||
if not message_state["opened"]:
|
||||
return []
|
||||
text = message_state["text"]
|
||||
events = [
|
||||
_sse(
|
||||
"response.output_text.done",
|
||||
{
|
||||
"type": "response.output_text.done",
|
||||
"item_id": message_state["item_id"],
|
||||
"output_index": message_state["output_index"],
|
||||
"content_index": 0,
|
||||
"text": text,
|
||||
},
|
||||
),
|
||||
_sse(
|
||||
"response.content_part.done",
|
||||
{
|
||||
"type": "response.content_part.done",
|
||||
"item_id": message_state["item_id"],
|
||||
"output_index": message_state["output_index"],
|
||||
"content_index": 0,
|
||||
"part": {"type": "output_text", "text": text, "annotations": []},
|
||||
},
|
||||
),
|
||||
_sse(
|
||||
"response.output_item.done",
|
||||
{
|
||||
"type": "response.output_item.done",
|
||||
"output_index": message_state["output_index"],
|
||||
"item": {
|
||||
"type": "message",
|
||||
"id": message_state["item_id"],
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": text, "annotations": []}],
|
||||
},
|
||||
},
|
||||
),
|
||||
]
|
||||
closed_message_states.append(dict(message_state))
|
||||
message_state.update(
|
||||
{"output_index": None, "item_id": None, "opened": False, "text": ""}
|
||||
)
|
||||
return events
|
||||
|
||||
def _healed_event_sse(events) -> list[str]:
|
||||
"""Serialize healer events preserving their order.
|
||||
|
||||
Text around a healed call must keep its position relative to the
|
||||
function_call item (output indexes are claimed in emission order),
|
||||
so never split an event list into all-text-then-all-calls. A healed
|
||||
call also CLOSES any open message item, so trailing text opens a
|
||||
fresh message with a later output index, exactly like a native
|
||||
Responses stream that interleaves messages and calls.
|
||||
"""
|
||||
nonlocal full_text
|
||||
out: list[str] = []
|
||||
for kind, value in events:
|
||||
if kind == "text":
|
||||
if not value:
|
||||
continue
|
||||
out.extend(_ensure_message_open())
|
||||
full_text += value
|
||||
message_state["text"] += value
|
||||
api_monitor.append_reply(monitor_id, value)
|
||||
out.append(
|
||||
_sse(
|
||||
"response.output_text.delta",
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": message_state["item_id"],
|
||||
"output_index": message_state["output_index"],
|
||||
"content_index": 0,
|
||||
"delta": value,
|
||||
},
|
||||
)
|
||||
)
|
||||
else:
|
||||
tc = _healed_tc(value)
|
||||
if tc is None:
|
||||
continue
|
||||
out.extend(_close_message_item())
|
||||
out.extend(_tool_call_delta_events(tc))
|
||||
return out
|
||||
|
||||
def _snapshot_output() -> list[dict]:
|
||||
"""Snapshot of all completed output items for response.completed."""
|
||||
indexed_items: list[tuple[int, dict]] = []
|
||||
|
|
@ -8409,19 +8614,23 @@ async def _responses_stream(
|
|||
},
|
||||
)
|
||||
)
|
||||
if message_state["opened"]:
|
||||
# Closed copies keep opened=True (snapshotted before reset); the
|
||||
# live state contributes only when a message is currently open.
|
||||
for msg_st in [*closed_message_states, message_state]:
|
||||
if not msg_st["opened"]:
|
||||
continue
|
||||
indexed_items.append(
|
||||
(
|
||||
message_state["output_index"],
|
||||
msg_st["output_index"],
|
||||
{
|
||||
"type": "message",
|
||||
"id": message_state["item_id"],
|
||||
"id": msg_st["item_id"],
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "output_text",
|
||||
"text": full_text,
|
||||
"text": msg_st["text"],
|
||||
"annotations": [],
|
||||
}
|
||||
],
|
||||
|
|
@ -8605,10 +8814,30 @@ async def _responses_stream(
|
|||
"delta": reasoning_delta,
|
||||
},
|
||||
)
|
||||
# Heal text-form tool calls in the visible stream (never in
|
||||
# reasoning text): promoted calls join the structured tc loop
|
||||
# below through the same state machinery, and healer events are
|
||||
# emitted IN ORDER so text after a healed call never jumps ahead
|
||||
# of the function_call item. Once a structured delta arrives,
|
||||
# grammar mode worked and the healer goes dormant.
|
||||
if healer is not None and not healer.dormant:
|
||||
healed_events = []
|
||||
if delta.get("tool_calls"):
|
||||
# Held text preceded the structured call; the call's own
|
||||
# deltas follow in the structured loop below.
|
||||
healed_events = healer.structured_tool_call_seen()
|
||||
if visible_delta:
|
||||
healed_events.append(("text", visible_delta))
|
||||
elif visible_delta:
|
||||
healed_events = healer.feed(visible_delta)
|
||||
visible_delta = ""
|
||||
for event in _healed_event_sse(healed_events):
|
||||
yield event
|
||||
if visible_delta:
|
||||
for event in _ensure_message_open():
|
||||
yield event
|
||||
full_text += visible_delta
|
||||
message_state["text"] += visible_delta
|
||||
api_monitor.append_reply(monitor_id, visible_delta)
|
||||
yield _sse(
|
||||
"response.output_text.delta",
|
||||
|
|
@ -8622,60 +8851,19 @@ async def _responses_stream(
|
|||
)
|
||||
|
||||
for tc in delta.get("tool_calls") or []:
|
||||
idx = tc.get("index", 0)
|
||||
st = tool_call_state.get(idx)
|
||||
fn = tc.get("function") or {}
|
||||
if st is None:
|
||||
# First chunk for this tool call -- allocate an
|
||||
# output_index and emit output_item.added.
|
||||
st = {
|
||||
"output_index": _claim_output_index(),
|
||||
"item_id": f"fc_{uuid.uuid4().hex[:12]}",
|
||||
"call_id": tc.get("id") or "",
|
||||
"name": fn.get("name") or "",
|
||||
"arguments": "",
|
||||
"opened": False,
|
||||
}
|
||||
tool_call_state[idx] = st
|
||||
else:
|
||||
# Later chunks sometimes carry id/name only once; merge
|
||||
# when present.
|
||||
if tc.get("id") and not st["call_id"]:
|
||||
st["call_id"] = tc["id"]
|
||||
if fn.get("name") and not st["name"]:
|
||||
st["name"] = fn["name"]
|
||||
|
||||
if not st["opened"] and st["call_id"] and st["name"]:
|
||||
item_added = {
|
||||
"type": "response.output_item.added",
|
||||
"output_index": st["output_index"],
|
||||
"item": {
|
||||
"type": "function_call",
|
||||
"id": st["item_id"],
|
||||
"status": "in_progress",
|
||||
"call_id": st["call_id"],
|
||||
"name": st["name"],
|
||||
"arguments": "",
|
||||
},
|
||||
}
|
||||
yield _sse("response.output_item.added", item_added)
|
||||
st["opened"] = True
|
||||
|
||||
arg_delta = fn.get("arguments") or ""
|
||||
if arg_delta and st["opened"]:
|
||||
st["arguments"] += arg_delta
|
||||
args_delta_event = {
|
||||
"type": "response.function_call_arguments.delta",
|
||||
"item_id": st["item_id"],
|
||||
"output_index": st["output_index"],
|
||||
"delta": arg_delta,
|
||||
}
|
||||
yield _sse("response.function_call_arguments.delta", args_delta_event)
|
||||
elif arg_delta:
|
||||
# Buffer args until we can open the item (some models
|
||||
# send id/name in the same chunk as the first arg delta;
|
||||
# if not, stash).
|
||||
st["arguments"] += arg_delta
|
||||
if (
|
||||
payload.parallel_tool_calls is False
|
||||
and healed_tc_index >= 1
|
||||
and tc.get("index", 0) not in tool_call_state
|
||||
):
|
||||
# A healed call already consumed the single allowed slot;
|
||||
# _drop_parallel_tool_call_deltas only sees native indexes,
|
||||
# so a native index-0 call would still open a second
|
||||
# function_call item. Skip it (and its later argument
|
||||
# deltas, which never allocate a state either).
|
||||
continue
|
||||
for event in _tool_call_delta_events(tc):
|
||||
yield event
|
||||
|
||||
_apply_usage(chunk_data.get("usage"))
|
||||
except asyncio.CancelledError:
|
||||
|
|
@ -8731,10 +8919,19 @@ async def _responses_stream(
|
|||
"delta": final_reasoning,
|
||||
},
|
||||
)
|
||||
# Last-chance heal of any held residue (e.g. a tool block the model
|
||||
# never closed) before the trailing visible text is flushed; events
|
||||
# keep healer order so trailing text stays behind a healed call.
|
||||
if healer is not None:
|
||||
events = (healer.feed(final_visible) if final_visible else []) + healer.finalize()
|
||||
final_visible = ""
|
||||
for event in _healed_event_sse(events):
|
||||
yield event
|
||||
if final_visible:
|
||||
for event in _ensure_message_open():
|
||||
yield event
|
||||
full_text += final_visible
|
||||
message_state["text"] += final_visible
|
||||
api_monitor.append_reply(monitor_id, final_visible)
|
||||
yield _sse(
|
||||
"response.output_text.delta",
|
||||
|
|
@ -8793,6 +8990,10 @@ async def _responses_stream(
|
|||
continue
|
||||
|
||||
if kind == "message":
|
||||
# Per-item text: message items closed mid-stream (healed-call
|
||||
# rotation) already emitted their done events, so this state
|
||||
# carries only its own text, not the whole stream's.
|
||||
_msg_text = st["text"]
|
||||
yield _sse(
|
||||
"response.output_text.done",
|
||||
{
|
||||
|
|
@ -8800,7 +9001,7 @@ async def _responses_stream(
|
|||
"item_id": st["item_id"],
|
||||
"output_index": st["output_index"],
|
||||
"content_index": 0,
|
||||
"text": full_text,
|
||||
"text": _msg_text,
|
||||
},
|
||||
)
|
||||
yield _sse(
|
||||
|
|
@ -8810,7 +9011,7 @@ async def _responses_stream(
|
|||
"item_id": st["item_id"],
|
||||
"output_index": st["output_index"],
|
||||
"content_index": 0,
|
||||
"part": {"type": "output_text", "text": full_text, "annotations": []},
|
||||
"part": {"type": "output_text", "text": _msg_text, "annotations": []},
|
||||
},
|
||||
)
|
||||
yield _sse(
|
||||
|
|
@ -8824,7 +9025,7 @@ async def _responses_stream(
|
|||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "output_text", "text": full_text, "annotations": []}
|
||||
{"type": "output_text", "text": _msg_text, "annotations": []}
|
||||
],
|
||||
},
|
||||
},
|
||||
|
|
@ -9430,6 +9631,7 @@ async def anthropic_messages(
|
|||
session_id = payload.session_id,
|
||||
cancel_id = payload.cancel_id,
|
||||
disable_parallel_tool_use = _disable_parallel,
|
||||
auto_heal_tool_calls = payload.auto_heal_tool_calls,
|
||||
)
|
||||
)
|
||||
return await _monitored_anthropic(
|
||||
|
|
@ -9449,6 +9651,8 @@ async def anthropic_messages(
|
|||
presence_penalty = presence_penalty,
|
||||
tool_choice = openai_tool_choice,
|
||||
disable_parallel_tool_use = _disable_parallel,
|
||||
auto_heal_tool_calls = payload.auto_heal_tool_calls,
|
||||
nudge_tool_calls = payload.nudge_tool_calls,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -10019,6 +10223,7 @@ async def _anthropic_passthrough_stream(
|
|||
session_id = None,
|
||||
cancel_id = None,
|
||||
disable_parallel_tool_use = False,
|
||||
auto_heal_tool_calls = None,
|
||||
):
|
||||
"""Streaming client-side pass-through: forward tools to llama-server and
|
||||
translate its stream to Anthropic SSE without executing anything."""
|
||||
|
|
@ -10055,6 +10260,16 @@ async def _anthropic_passthrough_stream(
|
|||
|
||||
async def _stream():
|
||||
emitter = AnthropicPassthroughEmitter()
|
||||
# Promote text-form tool calls (declared client tools only) into
|
||||
# tool_use blocks; verbatim behavior when healing is off or no tools.
|
||||
# tool_choice arrives here already converted to the OpenAI shape.
|
||||
_allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice)
|
||||
if _allowed_tools:
|
||||
emitter.enable_healing(
|
||||
_allowed_tools,
|
||||
openai_tools,
|
||||
disable_parallel_tool_use = disable_parallel_tool_use,
|
||||
)
|
||||
for line in emitter.start(message_id, model_name, input_tokens = input_tokens):
|
||||
yield line
|
||||
|
||||
|
|
@ -10191,6 +10406,8 @@ async def _anthropic_passthrough_non_streaming(
|
|||
presence_penalty = None,
|
||||
tool_choice = "auto",
|
||||
disable_parallel_tool_use = False,
|
||||
auto_heal_tool_calls = None,
|
||||
nudge_tool_calls = None,
|
||||
):
|
||||
"""Non-streaming client-side pass-through."""
|
||||
target_url = f"{llama_backend.base_url}/v1/chat/completions"
|
||||
|
|
@ -10223,34 +10440,98 @@ async def _anthropic_passthrough_non_streaming(
|
|||
)
|
||||
|
||||
data = resp.json()
|
||||
# tool_choice arrives here already converted to the OpenAI shape.
|
||||
_allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice)
|
||||
|
||||
# Opt-in single-retry nudge (mirrors the OpenAI passthrough): the model
|
||||
# tried to call a tool but nothing usable came out; re-ask once with the
|
||||
# prompt prefix intact so llama-server's KV cache is reused.
|
||||
if (
|
||||
_allowed_tools
|
||||
and nudge_enabled(nudge_tool_calls)
|
||||
and nudge_should_retry(data, _allowed_tools, openai_tools)
|
||||
):
|
||||
retry_body = {
|
||||
**body,
|
||||
"messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)],
|
||||
}
|
||||
try:
|
||||
retry_resp = await nonstreaming_client().post(
|
||||
target_url,
|
||||
json = retry_body,
|
||||
timeout = _llama_non_streaming_generation_timeout(),
|
||||
)
|
||||
if retry_resp.status_code == 200:
|
||||
retry_data = retry_resp.json()
|
||||
if response_has_promotable_calls(retry_data, _allowed_tools, openai_tools):
|
||||
data = retry_data
|
||||
except (httpx.RequestError, ValueError) as exc:
|
||||
logger.warning("tool-call nudge retry failed; keeping original: %s", exc)
|
||||
|
||||
choice = (data.get("choices") or [{}])[0]
|
||||
message = choice.get("message") or {}
|
||||
finish_reason = choice.get("finish_reason")
|
||||
|
||||
content_blocks = []
|
||||
text = message.get("content") or ""
|
||||
if text:
|
||||
text = _TOOL_XML_RE.sub("", text).strip()
|
||||
if text:
|
||||
content_blocks.append(AnthropicResponseTextBlock(text = text))
|
||||
healing_active = bool(_allowed_tools)
|
||||
healed_events = (
|
||||
heal_openai_message_events(message, _allowed_tools, openai_tools)
|
||||
if healing_active
|
||||
else None
|
||||
)
|
||||
|
||||
tool_calls = message.get("tool_calls") or []
|
||||
# disable_parallel_tool_use: keep only the first tool_use block.
|
||||
if disable_parallel_tool_use and len(tool_calls) > 1:
|
||||
tool_calls = tool_calls[:1]
|
||||
for tc in tool_calls:
|
||||
fn = tc.get("function") or {}
|
||||
try:
|
||||
args = json.loads(fn.get("arguments", "{}"))
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
content_blocks.append(
|
||||
AnthropicResponseToolUseBlock(
|
||||
id = anthropic_tool_use_id(tc.get("id")),
|
||||
name = fn.get("name", ""),
|
||||
input = args,
|
||||
content_blocks = []
|
||||
tool_calls = []
|
||||
if healed_events:
|
||||
emitted_tool_uses = 0
|
||||
for kind, value in healed_events:
|
||||
if kind == "text":
|
||||
text = str(value).strip()
|
||||
if text:
|
||||
content_blocks.append(AnthropicResponseTextBlock(text = text))
|
||||
continue
|
||||
if disable_parallel_tool_use and emitted_tool_uses >= 1:
|
||||
continue
|
||||
fn = value.get("function") or {}
|
||||
try:
|
||||
args = json.loads(fn.get("arguments", "{}"))
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
tool_calls.append(value)
|
||||
emitted_tool_uses += 1
|
||||
content_blocks.append(
|
||||
AnthropicResponseToolUseBlock(
|
||||
id = anthropic_tool_use_id(value.get("id")),
|
||||
name = fn.get("name", ""),
|
||||
input = args,
|
||||
)
|
||||
)
|
||||
else:
|
||||
text = message.get("content") or ""
|
||||
if text:
|
||||
# Keep unpromoted bytes when healing is active; legacy stripping is
|
||||
# only for opted-out or no-client-tool requests.
|
||||
if not healing_active:
|
||||
text = _TOOL_XML_RE.sub("", text)
|
||||
text = text.strip()
|
||||
if text:
|
||||
content_blocks.append(AnthropicResponseTextBlock(text = text))
|
||||
|
||||
tool_calls = message.get("tool_calls") or []
|
||||
if disable_parallel_tool_use and len(tool_calls) > 1:
|
||||
tool_calls = tool_calls[:1]
|
||||
for tc in tool_calls:
|
||||
fn = tc.get("function") or {}
|
||||
try:
|
||||
args = json.loads(fn.get("arguments", "{}"))
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
content_blocks.append(
|
||||
AnthropicResponseToolUseBlock(
|
||||
id = anthropic_tool_use_id(tc.get("id")),
|
||||
name = fn.get("name", ""),
|
||||
input = args,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls))
|
||||
|
||||
|
|
@ -10598,6 +10879,13 @@ async def _openai_passthrough_stream(
|
|||
body = _build_openai_passthrough_body(
|
||||
payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend
|
||||
)
|
||||
# Text-form tool calls from small models get promoted to structured calls on
|
||||
# the way back (declared client tools only); requests without tools or with
|
||||
# auto_heal_tool_calls=false keep the verbatim relay. tool_choice constrains
|
||||
# the allowlist ("none" disables, a forced function narrows to it).
|
||||
_allowed_tools = heal_gate(
|
||||
payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice")
|
||||
)
|
||||
|
||||
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
|
||||
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
|
||||
|
|
@ -10700,9 +10988,14 @@ async def _openai_passthrough_stream(
|
|||
last_chunk_id = completion_id
|
||||
last_chunk_model = model_name
|
||||
last_chunk_created = int(time.time())
|
||||
healer = (
|
||||
StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None
|
||||
)
|
||||
healed_call_index = 0
|
||||
|
||||
def _synthetic_finish_line() -> str:
|
||||
finish_reason = "tool_calls" if saw_tool_call_delta else "stop"
|
||||
healed = healer is not None and healer.healed
|
||||
finish_reason = "tool_calls" if (saw_tool_call_delta or healed) else "stop"
|
||||
chunk = ChatCompletionChunk(
|
||||
id = last_chunk_id,
|
||||
created = last_chunk_created,
|
||||
|
|
@ -10716,6 +11009,108 @@ async def _openai_passthrough_stream(
|
|||
)
|
||||
return f"data: {chunk.model_dump_json(exclude_none = True)}"
|
||||
|
||||
def _healer_sse_lines(events) -> list:
|
||||
# Serialize healer events as chunks matching the upstream stream's
|
||||
# id/model/created so clients see one coherent completion.
|
||||
nonlocal healed_call_index
|
||||
lines = []
|
||||
for kind, value in events:
|
||||
if kind == "text":
|
||||
if not value:
|
||||
continue
|
||||
delta = {"content": value}
|
||||
else:
|
||||
# parallel_tool_calls=false caps healed calls too (the SSE
|
||||
# line cap only sees structured upstream deltas).
|
||||
if payload.parallel_tool_calls is False and healed_call_index >= 1:
|
||||
continue
|
||||
delta = {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": healed_call_index,
|
||||
"id": value["id"],
|
||||
"type": "function",
|
||||
"function": value["function"],
|
||||
}
|
||||
]
|
||||
}
|
||||
healed_call_index += 1
|
||||
chunk = {
|
||||
"id": last_chunk_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": last_chunk_created,
|
||||
"model": last_chunk_model,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": None}],
|
||||
}
|
||||
lines.append("data: " + json.dumps(chunk, ensure_ascii = False))
|
||||
return lines
|
||||
|
||||
def _heal_transform(chunk_data: dict, raw_line: str) -> list:
|
||||
"""SSE lines to emit in place of one upstream line (healing on)."""
|
||||
choices = chunk_data.get("choices")
|
||||
if not (isinstance(choices, list) and choices and isinstance(choices[0], dict)):
|
||||
return [raw_line]
|
||||
choice = choices[0]
|
||||
delta = choice.get("delta")
|
||||
delta = delta if isinstance(delta, dict) else {}
|
||||
if delta.get("tool_calls"):
|
||||
# Structured call streamed: grammar mode worked. Flush any held
|
||||
# text (it preceded the call) and relay verbatim from here on.
|
||||
lines = _healer_sse_lines(healer.structured_tool_call_seen())
|
||||
if healed_call_index:
|
||||
if payload.parallel_tool_calls is False:
|
||||
# A healed call already consumed the single allowed
|
||||
# slot; the upstream SSE cap keeps native index 0, so
|
||||
# drop the native call here or the client gets two.
|
||||
del delta["tool_calls"]
|
||||
if delta or choice.get("finish_reason") or chunk_data.get("usage"):
|
||||
lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False))
|
||||
return lines
|
||||
# A healed call already went out on index 0..n-1; OpenAI
|
||||
# clients merge tool-call deltas by index, so shift the
|
||||
# native calls into the next indexes or they would merge
|
||||
# into the healed call.
|
||||
for tc in delta["tool_calls"]:
|
||||
if isinstance(tc, dict) and isinstance(tc.get("index"), int):
|
||||
tc["index"] += healed_call_index
|
||||
return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)]
|
||||
return lines + [raw_line]
|
||||
content = delta.get("content")
|
||||
finish = choice.get("finish_reason")
|
||||
if not isinstance(content, str) or not content:
|
||||
if not finish:
|
||||
return [raw_line]
|
||||
# Finish chunk: last-chance heal of the residue, and rewrite a
|
||||
# "stop" into "tool_calls" when text-form calls were promoted.
|
||||
lines = _healer_sse_lines(healer.finalize())
|
||||
if healer.healed and finish == "stop":
|
||||
choice["finish_reason"] = "tool_calls"
|
||||
return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)]
|
||||
return lines + [raw_line]
|
||||
events = healer.feed(content)
|
||||
if finish:
|
||||
events += healer.finalize()
|
||||
if not finish and events == [("text", content)]:
|
||||
# Nothing held or promoted: the healer passed the chunk
|
||||
# through whole, so keep the verbatim upstream bytes.
|
||||
return [raw_line]
|
||||
del delta["content"]
|
||||
prefix_lines = []
|
||||
if delta:
|
||||
prefix_chunk = {k: v for k, v in chunk_data.items() if k != "usage"}
|
||||
prefix_choice = dict(choice)
|
||||
prefix_choice["delta"] = dict(delta)
|
||||
prefix_choice["finish_reason"] = None
|
||||
prefix_chunk["choices"] = [prefix_choice]
|
||||
prefix_lines.append("data: " + json.dumps(prefix_chunk, ensure_ascii = False))
|
||||
delta.clear()
|
||||
lines = prefix_lines + _healer_sse_lines(events)
|
||||
if delta or finish or chunk_data.get("usage"):
|
||||
if healer.healed and finish == "stop":
|
||||
choice["finish_reason"] = "tool_calls"
|
||||
lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False))
|
||||
return lines
|
||||
|
||||
try:
|
||||
lines_iter = resp.aiter_lines()
|
||||
async for raw_line in _aiter_llama_stream_items(
|
||||
|
|
@ -10732,6 +11127,14 @@ async def _openai_passthrough_stream(
|
|||
data_text = raw_line[6:].strip()
|
||||
if data_text == "[DONE]":
|
||||
saw_done = True
|
||||
# Upstream ended without a finish chunk: heal the residue
|
||||
# first so the synthetic finish sees healer.healed.
|
||||
if healer is not None and not saw_stream_error:
|
||||
for held_line in _healer_sse_lines(healer.finalize()):
|
||||
_monitor_openai_sse_line(
|
||||
monitor_id, held_line, llama_backend.context_length
|
||||
)
|
||||
yield held_line + "\n\n"
|
||||
if (
|
||||
not saw_finish_reason
|
||||
and not saw_stream_error
|
||||
|
|
@ -10787,13 +11190,18 @@ async def _openai_passthrough_stream(
|
|||
# emit a successful finish_reason after a failed stream.
|
||||
if _monitor_openai_error_message(chunk_data):
|
||||
saw_stream_error = True
|
||||
monitor_event = _monitor_openai_sse_line(
|
||||
monitor_id,
|
||||
raw_line,
|
||||
llama_backend.context_length,
|
||||
)
|
||||
if monitor_event == "error":
|
||||
saw_stream_error = True
|
||||
# With healing active, a content-bearing line may be replaced by
|
||||
# held/promoted chunks; otherwise the single upstream line
|
||||
# relays verbatim (monitored exactly as emitted either way).
|
||||
if (
|
||||
healer is not None
|
||||
and not healer.dormant
|
||||
and isinstance(chunk_data, dict)
|
||||
and not saw_stream_error
|
||||
):
|
||||
out_lines = _heal_transform(chunk_data, raw_line)
|
||||
else:
|
||||
out_lines = [raw_line]
|
||||
# If a trailing usage-only chunk (include_usage) arrives before
|
||||
# any finish chunk, emit the synthetic finish first so the order
|
||||
# stays finish -> usage -> [DONE], matching the other streams.
|
||||
|
|
@ -10807,23 +11215,46 @@ async def _openai_passthrough_stream(
|
|||
and not saw_stream_error
|
||||
and not cancel_event.is_set()
|
||||
):
|
||||
if healer is not None:
|
||||
# Residue must precede the finish it may upgrade.
|
||||
held = _healer_sse_lines(healer.finalize())
|
||||
for held_line in held:
|
||||
_monitor_openai_sse_line(
|
||||
monitor_id, held_line, llama_backend.context_length
|
||||
)
|
||||
yield held_line + "\n\n"
|
||||
finish_line = _synthetic_finish_line()
|
||||
_monitor_openai_sse_line(
|
||||
monitor_id, finish_line, llama_backend.context_length
|
||||
)
|
||||
yield finish_line + "\n\n"
|
||||
saw_finish_reason = True
|
||||
# Relay verbatim to preserve llama-server's native id,
|
||||
# finish_reason, delta.tool_calls, and usage chunks.
|
||||
yield raw_line + "\n\n"
|
||||
if monitor_event == "done":
|
||||
monitor_done = True
|
||||
for out_line in out_lines:
|
||||
monitor_event = _monitor_openai_sse_line(
|
||||
monitor_id,
|
||||
out_line,
|
||||
llama_backend.context_length,
|
||||
)
|
||||
if monitor_event == "error":
|
||||
saw_stream_error = True
|
||||
# Relay to preserve llama-server's native id,
|
||||
# finish_reason, delta.tool_calls, and usage chunks.
|
||||
yield out_line + "\n\n"
|
||||
if monitor_event == "done":
|
||||
monitor_done = True
|
||||
if monitor_done:
|
||||
break
|
||||
if not saw_done and not saw_stream_error and not cancel_event.is_set():
|
||||
# Synthesize a finish chunk only if one was not already
|
||||
# emitted (e.g. before a trailing usage-only chunk), but
|
||||
# always close with [DONE] whenever the upstream omitted it,
|
||||
# so the stream ends on the [DONE] sentinel either way.
|
||||
if healer is not None:
|
||||
for held_line in _healer_sse_lines(healer.finalize()):
|
||||
_monitor_openai_sse_line(
|
||||
monitor_id, held_line, llama_backend.context_length
|
||||
)
|
||||
yield held_line + "\n\n"
|
||||
if not saw_finish_reason:
|
||||
finish_line = _synthetic_finish_line()
|
||||
_monitor_openai_sse_line(
|
||||
|
|
@ -10962,6 +11393,9 @@ async def _openai_passthrough_non_streaming(
|
|||
_guided_fence = bool((payload.model_extra or {}).get("_unsloth_guided_fence"))
|
||||
_do_fence = _guided_fence and _extract_response_format(payload) is not None
|
||||
_cap_parallel = payload.parallel_tool_calls is False
|
||||
_allowed_tools = heal_gate(
|
||||
payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice")
|
||||
)
|
||||
|
||||
try:
|
||||
data = resp.json()
|
||||
|
|
@ -10974,6 +11408,33 @@ async def _openai_passthrough_non_streaming(
|
|||
api_monitor.finish(monitor_id)
|
||||
return Response(content = resp.content, media_type = "application/json")
|
||||
|
||||
# Opt-in single-retry nudge: the model clearly tried to call a tool (signal
|
||||
# present) but nothing parseable/declared came out, so re-ask once with the
|
||||
# original prompt prefix intact (llama-server reuses the slot's KV cache)
|
||||
# plus a two-message nudge suffix. The retry replaces the original response
|
||||
# only when it actually yields a usable call.
|
||||
if (
|
||||
_allowed_tools
|
||||
and nudge_enabled(payload.nudge_tool_calls)
|
||||
and nudge_should_retry(data, _allowed_tools, body.get("tools"))
|
||||
):
|
||||
retry_body = {
|
||||
**body,
|
||||
"messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)],
|
||||
}
|
||||
try:
|
||||
retry_resp = await nonstreaming_client().post(
|
||||
target_url,
|
||||
json = retry_body,
|
||||
timeout = _llama_non_streaming_generation_timeout(),
|
||||
)
|
||||
if retry_resp.status_code == 200:
|
||||
retry_data = retry_resp.json()
|
||||
if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")):
|
||||
resp, data = retry_resp, retry_data
|
||||
except (httpx.RequestError, ValueError) as exc:
|
||||
logger.warning("tool-call nudge retry failed; keeping original: %s", exc)
|
||||
|
||||
changed = False
|
||||
for choice in data.get("choices", []):
|
||||
if not isinstance(choice, dict):
|
||||
|
|
@ -10982,6 +11443,17 @@ async def _openai_passthrough_non_streaming(
|
|||
if not isinstance(msg, dict):
|
||||
continue
|
||||
|
||||
# Small models emit tool calls as text instead of structured tool_calls;
|
||||
# promote them (declared client tools only) so the agent sees a real call.
|
||||
# Truncation wins over the upgrade (same rule as the streaming and
|
||||
# Anthropic paths): a call cut off at max_tokens keeps
|
||||
# finish_reason="length" so the client knows the arguments may be
|
||||
# incomplete, while the healed call itself stays attached.
|
||||
if _allowed_tools and heal_openai_message(msg, _allowed_tools, body.get("tools")):
|
||||
if choice.get("finish_reason") == "stop":
|
||||
choice["finish_reason"] = "tool_calls"
|
||||
changed = True
|
||||
|
||||
# OpenAI requires content=null on a pure tool-call turn; llama-server
|
||||
# emits content="".
|
||||
if msg.get("tool_calls") and msg.get("content") == "":
|
||||
|
|
|
|||
1358
studio/backend/tests/test_passthrough_healing.py
Normal file
1358
studio/backend/tests/test_passthrough_healing.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -1986,3 +1986,177 @@ class TestTranslatedMessagesValidate:
|
|||
msgs = _normalise_responses_input(payload)
|
||||
for m in msgs:
|
||||
ChatMessage(**m.model_dump(exclude_none = True))
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Streaming passthrough healing — text-form calls promoted in order
|
||||
# =====================================================================
|
||||
|
||||
|
||||
class TestResponsesStreamHealing:
|
||||
"""Route-level healing on the /v1/responses stream: text-form tool calls
|
||||
are promoted through the same per-call item state machinery as structured
|
||||
deltas, and healer events keep their order (text around a healed call must
|
||||
not move relative to the function_call item)."""
|
||||
|
||||
_XML = '<tool_call>{"name":"lookup","arguments":{"q":"x"}}</tool_call>'
|
||||
_TOOL = {"type": "function", "name": "lookup", "parameters": {"type": "object"}}
|
||||
|
||||
@staticmethod
|
||||
def _ordered_events(lines):
|
||||
events = []
|
||||
for line in lines:
|
||||
if not line.startswith("event: "):
|
||||
continue
|
||||
name, _, rest = line.partition("\n")
|
||||
payload = json.loads(rest.split("data: ", 1)[1].strip())
|
||||
events.append((name[len("event: ") :], payload))
|
||||
return events
|
||||
|
||||
def _run_stream(self, monkeypatch, content, **payload_kwargs):
|
||||
TestResponsesStreamAdapter._install_stream_mock(
|
||||
monkeypatch, [{"choices": [{"delta": {"content": content}}]}]
|
||||
)
|
||||
payload = ResponsesRequest(input = "hi", stream = True, tools = [self._TOOL], **payload_kwargs)
|
||||
messages = [ChatMessage(role = "user", content = "hi")]
|
||||
|
||||
async def run():
|
||||
response = await _responses_stream(
|
||||
payload, messages, TestResponsesStreamAdapter._Request()
|
||||
)
|
||||
return await TestResponsesStreamAdapter._collect(response)
|
||||
|
||||
return self._ordered_events(asyncio.run(run()))
|
||||
|
||||
def test_text_around_healed_call_keeps_order(self, monkeypatch):
|
||||
events = self._run_stream(monkeypatch, f"before {self._XML} after.")
|
||||
pos_before = pos_item = pos_after = None
|
||||
for i, (name, payload) in enumerate(events):
|
||||
if name == "response.output_text.delta":
|
||||
if "before" in payload["delta"] and pos_before is None:
|
||||
pos_before = i
|
||||
if "after" in payload["delta"]:
|
||||
pos_after = i
|
||||
if (
|
||||
name == "response.output_item.added"
|
||||
and payload["item"]["type"] == "function_call"
|
||||
and pos_item is None
|
||||
):
|
||||
pos_item = i
|
||||
assert payload["item"]["name"] == "lookup"
|
||||
assert pos_before is not None and pos_item is not None and pos_after is not None
|
||||
assert pos_before < pos_item < pos_after
|
||||
|
||||
def test_call_before_trailing_text_claims_lower_output_index(self, monkeypatch):
|
||||
events = self._run_stream(monkeypatch, f"{self._XML} done.")
|
||||
item_added = [
|
||||
(name, payload) for name, payload in events if name == "response.output_item.added"
|
||||
]
|
||||
# The call came first in the model output, so its item is added first
|
||||
# and claims the lower output_index; the trailing text's message item
|
||||
# follows.
|
||||
assert [payload["item"]["type"] for _, payload in item_added] == [
|
||||
"function_call",
|
||||
"message",
|
||||
]
|
||||
call_idx = item_added[0][1]["output_index"]
|
||||
msg_idx = item_added[1][1]["output_index"]
|
||||
assert call_idx < msg_idx
|
||||
text = "".join(
|
||||
payload["delta"] for name, payload in events if name == "response.output_text.delta"
|
||||
)
|
||||
assert "done." in text
|
||||
assert "<tool_call>" not in text
|
||||
|
||||
def test_tool_choice_none_streams_raw_text(self, monkeypatch):
|
||||
events = self._run_stream(monkeypatch, self._XML, tool_choice = "none")
|
||||
assert not any(
|
||||
payload["item"]["type"] == "function_call"
|
||||
for name, payload in events
|
||||
if name == "response.output_item.added"
|
||||
)
|
||||
text = "".join(
|
||||
payload["delta"] for name, payload in events if name == "response.output_text.delta"
|
||||
)
|
||||
assert text == self._XML
|
||||
|
||||
def test_healed_call_splits_message_items(self, monkeypatch):
|
||||
# Text on both sides of a healed call becomes TWO message items: the
|
||||
# healed function_call closes the first, trailing text opens a fresh
|
||||
# one with a later output index (native Responses stream shape).
|
||||
events = self._run_stream(monkeypatch, f"before {self._XML} after.")
|
||||
added = [
|
||||
(payload["output_index"], payload["item"]["type"], payload["item"].get("id"))
|
||||
for name, payload in events
|
||||
if name == "response.output_item.added"
|
||||
]
|
||||
assert [item_type for _, item_type, _ in added] == [
|
||||
"message",
|
||||
"function_call",
|
||||
"message",
|
||||
]
|
||||
assert [idx for idx, _, _ in added] == sorted(idx for idx, _, _ in added)
|
||||
assert added[0][2] != added[2][2] # distinct message item ids
|
||||
# Text deltas attribute to their OWN message item.
|
||||
deltas = [
|
||||
(payload["item_id"], payload["delta"])
|
||||
for name, payload in events
|
||||
if name == "response.output_text.delta"
|
||||
]
|
||||
assert [d for i, d in deltas if i == added[0][2]] == ["before "]
|
||||
assert [d for i, d in deltas if i == added[2][2]] == [" after."]
|
||||
# The completed snapshot lists all three items with per-item text.
|
||||
completed = [payload for name, payload in events if name == "response.completed"]
|
||||
output = completed[0]["response"]["output"]
|
||||
assert [item["type"] for item in output] == ["message", "function_call", "message"]
|
||||
assert output[0]["content"][0]["text"] == "before "
|
||||
assert output[2]["content"][0]["text"] == " after."
|
||||
|
||||
def test_parallel_cap_drops_native_after_healed(self, monkeypatch):
|
||||
# parallel_tool_calls=false: a healed call consumed the single allowed
|
||||
# slot; a later native structured call (index 0, so it survives
|
||||
# _drop_parallel_tool_call_deltas) must not open a second
|
||||
# function_call item.
|
||||
TestResponsesStreamAdapter._install_stream_mock(
|
||||
monkeypatch,
|
||||
[
|
||||
{"choices": [{"delta": {"content": self._XML}}]},
|
||||
{
|
||||
"choices": [
|
||||
{
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_up",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
],
|
||||
)
|
||||
payload = ResponsesRequest(
|
||||
input = "hi",
|
||||
stream = True,
|
||||
tools = [self._TOOL],
|
||||
parallel_tool_calls = False,
|
||||
)
|
||||
messages = [ChatMessage(role = "user", content = "hi")]
|
||||
|
||||
async def run():
|
||||
response = await _responses_stream(
|
||||
payload, messages, TestResponsesStreamAdapter._Request()
|
||||
)
|
||||
return await TestResponsesStreamAdapter._collect(response)
|
||||
|
||||
events = self._ordered_events(asyncio.run(run()))
|
||||
calls = [
|
||||
payload
|
||||
for name, payload in events
|
||||
if name == "response.output_item.added" and payload["item"]["type"] == "function_call"
|
||||
]
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["item"]["name"] == "lookup"
|
||||
|
|
|
|||
|
|
@ -71,6 +71,28 @@ class TestFunctionStyleTrailingText:
|
|||
call = _only(text)
|
||||
assert call == {"name": "python", "arguments": {"code": 'print("</function>")'}}
|
||||
|
||||
def test_closed_function_with_trailing_prose_heal_path(self):
|
||||
# Regression: the heal / finalize path (allow_incomplete=True) used to fold
|
||||
# </parameter></function> and the trailing prose into the argument and drop
|
||||
# the prose from visible content. It must now match the strict path -- keep a
|
||||
# clean argument and leave the trailing prose outside the call span.
|
||||
text = "<function=web_search><parameter=query>cats</parameter></function> trailing words"
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = True)
|
||||
assert len(calls) == 1
|
||||
fn = calls[0]["function"]
|
||||
assert fn["name"] == "web_search"
|
||||
assert json.loads(fn["arguments"]) == {"query": "cats"}
|
||||
# The trailing prose sits outside the removed span, so it stays visible.
|
||||
from core.tool_healing import (
|
||||
parse_tool_calls_from_text as _parse_with_spans,
|
||||
)
|
||||
|
||||
_calls, spans = _parse_with_spans(text, allow_incomplete = True, with_spans = True)
|
||||
out = text
|
||||
for s, e in sorted(spans, reverse = True):
|
||||
out = out[:s] + out[e:]
|
||||
assert out == " trailing words"
|
||||
|
||||
def test_incomplete_function_without_close_is_still_rejected(self):
|
||||
text = "<function=web_search><parameter=query>weather london"
|
||||
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
|
||||
|
|
@ -160,3 +182,18 @@ class TestHealingPathUnaffected:
|
|||
calls = parse_tool_calls_from_text(text, allow_incomplete = True)
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["function"]["name"] == "web_search"
|
||||
|
||||
def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self):
|
||||
# allow_incomplete exists for truncated output; a call that DID close
|
||||
# must parse identically to strict mode, leaving prose after
|
||||
# </function> out of the last parameter and out of the removal span.
|
||||
from core.tool_healing import parse_tool_calls_from_text as parse_with_spans
|
||||
|
||||
text = "<function=web_search><parameter=query>cats</parameter></function> trailing"
|
||||
calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True)
|
||||
(call,) = calls
|
||||
assert json.loads(call["function"]["arguments"]) == {"query": "cats"}
|
||||
(span,) = spans
|
||||
assert text[span[0] : span[1]] == (
|
||||
"<function=web_search><parameter=query>cats</parameter></function>"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue