* Quote-aware Gemma strip, symmetric unstarted cleanup, ReDoS anchor
Address review findings on the tool-strip and streaming paths:
- strip_tool_call_markup stripped Gemma-native spans with a plain regex that
stops at the first <tool_call|>, so a literal close marker inside a
<|"|>-quoted argument truncated the span and leaked its suffix into visible
text. A brace/quote-aware _strip_gemma_native_spans now removes complete
spans (keeping an incomplete one unless final), matching the parser's own
balance logic.
- The Gemma close pattern this PR added (<\|tool_call>.*?<tool_call\|>) had no
\Z fallback, so a run of unclosed markers backtracked from every open
position (quadratic, and the streaming stripper re-scans per token). It is
now anchored to (?:<tool_call|>|\Z) like routes/inference.py's _TOOL_XML_RE,
linear with identical output on well-formed input.
- _SameTaskStreamingResponse added unstarted_cleanup for the OpenAI passthrough,
but the local GGUF/safetensors streams that enter _TrackedCancel before
returning only unregister in the generator finally, which never runs if the
client disconnects before the body iterator starts, leaking cancel-registry
entries. Each such stream now passes unstarted_cleanup to exit its tracker.
- __call__ reads _unstarted_cleanup via getattr so a response built through
__new__ (the cancel-timing test) without __init__ does not raise
AttributeError; the test also sets the attribute explicitly.
- Document that the verbatim /v1/chat/completions passthrough delegates
<think>/<|tool_call> splitting to llama-server (--jinja, --reasoning-format
auto) and is intentionally not re-parsed locally, noting the llama.cpp
dependency.
Adds a regression test for the close-marker-inside-quoted-argument strip.
* Tighten comments on the tool-strip and streaming paths
Compress the verbose comment blocks added with the Gemma tool-call / streaming
work to crisp one or two liners, drop restatements of obvious code, and shorten
docstrings, keeping the load-bearing rationale (ReDoS anchor, quote-aware strip,
unstarted-cleanup, llama.cpp passthrough dependency). Code is unchanged
(verified comment-only via AST/ast signature, docstrings stripped).
* Harden Gemma parse/strip: span-aware XML fallback and quote-aware streaming
- Security: the XML fallback in parse_tool_calls_from_text scanned the whole
content for <function=...> markers and only skipped those inside an open XML
parameter, not those inside a collected JSON/Gemma candidate span. A balanced
but unparsable Gemma call whose argument data contained XML tool markup
(<|tool_call>call:outer{code:<function=terminal>...}<tool_call|>) therefore
fell through to the fallback and returned an executable terminal call. The
fallback now also excludes <function=> markers inside any candidate span,
including ones that failed to parse.
- strip_tool_call_markup no longer skips the generic Gemma regex after running
the quote-aware _strip_gemma_native_spans, so a closed Gemma span the helper
cannot match (malformed, e.g. <|tool_call>{"name":"x"}<tool_call|>) is still
stripped instead of leaking its opener and payload into visible text.
- _strip_gemma_native_spans stops at the first unbalanced start instead of
re-scanning every later start to EOF, keeping it linear on a run of unclosed
markers rather than quadratic.
- The GGUF and safetensors streaming strippers run _strip_gemma_native_spans
before the regex patterns, so a well-formed streamed call whose quoted
argument contains a literal close marker no longer leaks its suffix into
incremental display.
Adds regression tests for the nested-XML escape and the malformed-span strip.
* Avoid remainder copy in _strip_gemma_native_spans
Match the Gemma close marker with re pos directly on the buffer instead
of slicing tail = text[brace_end + 1:] on every span. The streaming
strippers re-scan a growing cumulative buffer per token, so the per-span
remainder copy was quadratic. Behavior is unchanged.
* Exclude unclosed Gemma/JSON starts from the XML tool-call fallback
The nested-XML guard only skipped <function=> markers inside recorded
candidate spans, but a span is recorded only when the braces balance. An
unbalanced call such as <|tool_call>call:outer{code:<function=terminal>...
recorded no span, so the fallback still promoted the inner <function=> to
an executable terminal call. Treat unclosed JSON/Gemma starts as exclusion
spans through EOF before scanning. Standalone <function=> calls with no
preceding unclosed start still parse. Regression tests added.
* Skip doomed tool-strip passes to avoid quadratic rescans
The lazy closed-pair strip patterns (<tool_call>.*?</tool_call>,
<function=...>.*?</function>) rescan to EOF from every opener when their
close token is absent, which is O(n^2) and re-runs per streamed token. Add
strip_tool_patterns, which skips a pass whose close token is not present in
the text; output is identical to the per-pattern loop (verified by fuzz),
and a degenerate run drops from ~minutes to milliseconds. Used by
strip_tool_call_markup and the GGUF/safetensors streaming strippers.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Use full tool-call envelopes to close nested-XML escape variants
Key the parser and stripper off the full <|tool_call>...<tool_call|> /
<tool_call>...</tool_call> envelope (start to close marker, searched after
the braces; EOF if unclosed) instead of just the braces:
- XML between the closing brace and the close marker
(call:outer{broken:{x}}<function=terminal>...<tool_call|>) is now inside
the envelope, so the fallback no longer promotes it to a tool call.
- A balanced inner call inside an unclosed outer
(call:outer{code:<|tool_call>call:terminal{...}<tool_call|>) is skipped
via the envelope nested check, not just the XML fallback.
- strip_tool_call_markup searches for the close marker after the braces, so
junk before <tool_call|> is stripped through the close and text after it is
preserved instead of truncated to EOF; a no-close run stops early (linear).
Regression tests added; standalone XML and well-formed calls unaffected.
* Fix non-final Gemma strip and missing-close recovery for PR #6611
Split the nested-skip from the XML fallback exclusion: nesting is decided by
each marker's brace region, so a balanced call after one with a missing close
marker is recovered instead of being swallowed to EOF. Only the XML fallback
keeps the search-to-close envelope, so trailing nested markup still cannot
escape as an executable call.
Use a closed-only Gemma pattern in the non-final strip list so an incomplete
block is preserved (matching the JSON and function paths); the final list keeps
the close-or-EOF Gemma pattern in its original position, so streaming display
output is byte-for-byte unchanged.
Add regression tests for both cases.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Block gap-nested tool markers and fix XML strip order for PR #6611
Decide candidate nesting by a per-marker coverage region paired with a
per-format stack (a close after the braces pops the nearest still-open marker
of that format). A closed outer call now covers up to its own close marker, so a
JSON or Gemma tool marker smuggled between the outer braces and that close is
treated as data instead of being executed. An outer that balances but has no
close of its own covers only its brace region, so a later sibling after an
omitted close marker is still recovered (adjacent calls use an exclusive end
bound so the next call is not misread as nested).
Strip every closed pair (JSON, Gemma, function) before any to-EOF sweep, so a
closed function call whose parameter text contains a bare Gemma opener is
removed as a unit and the to-EOF sweep can no longer drop the visible text after
the close.
Add regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Strip closed tool blocks before the Gemma final sweep for PR #6611
The final display strip ran the quote-aware Gemma helper before the closed
JSON/function patterns. A closed <tool_call>...</tool_call> or
<function=...>...</function> block whose argument data held a call-form Gemma
opener (e.g. a "<|tool_call>call:t{" string) was read as an incomplete Gemma
span and truncated to EOF, dropping the block's close and any visible text after
it.
Strip closed JSON/function blocks first, so such a block is removed as a unit
before the helper runs. Centralize the final strip order in a shared
strip_tool_markup_final so strip_tool_call_markup and both streaming display
wrappers (safetensors, llama_cpp) stay in sync, and apply the same closed-block
pre-pass to the non-final path.
Add regression tests for the JSON and function variants.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Recover XML/JSON siblings after a close-less tool marker for PR #6611
Two fixes so the XML fallback and marker coverage recover a later valid call
after an earlier marker omits its close, matching the candidate loop:
Reuse the candidate marker-coverage in the XML fallback instead of a separate
search-to-close-or-EOF envelope. A balanced but close-less marker now covers
only its brace region there too, so a following <function=...> sibling is
recovered rather than filtered as nested data; an unbalanced marker still covers
to EOF and a closed one still covers through its close, so nested XML stays
blocked.
Ignore a close token that falls inside another call's balanced braces when
pairing closes in _marker_coverage. Such a token is that call's quoted argument
data, so it no longer pops an earlier close-less marker and extends its coverage
over a later valid sibling.
Add regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Make the closed-block strip pre-pass Gemma-span-aware
The final display strip ran the closed JSON/function regex pre-pass before
removing Gemma-native spans, so a literal <function=...> quoted inside a Gemma
argument plus any later </function> (a real call's close or even prose) was
deleted across the Gemma boundary. That mangled the Gemma close marker, the
quote-aware helper then saw an unclosed opener, and the whole visible tail
after the call was truncated.
The pre-pass now skips matches that start inside a complete Gemma span (that
text is the span's argument data) and resumes scanning at the end of the
covering span, so a real function-XML call after the Gemma call is still
stripped. The original ordering rationale is preserved: a Gemma opener inside
a JSON or function argument still cannot truncate that block, covered by
regression tests for both directions.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments in the Gemma streaming and strip pipeline to essentials
* Tighten comments in the Gemma strip and streaming disconnect paths
* Fold marker-collection comment to two lines
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
e9ea45b6a5
commit
eb1ef44255
6 changed files with 562 additions and 217 deletions
|
|
@ -530,13 +530,23 @@ def parse_tool_calls_from_text(
|
|||
|
||||
# Formats tool_healing does not cover: ``<function name="...">`` (MiniCPM-5 / MiniMax-M2),
|
||||
# Llama-3 and Mistral. Run only after tool_healing found nothing, so a strict-rejected
|
||||
# call is never re-healed here.
|
||||
# call is never re-healed here. Blank any JSON/Gemma marker coverage first: markup inside
|
||||
# a marker's span (even one that failed to parse) is that call's data, not a sibling, so
|
||||
# a nested ``<function=...>`` / ``<|python_tag|>`` / ``[TOOL_CALLS]`` must not be promoted.
|
||||
fallback_content = content
|
||||
coverage = _tool_healing.marker_coverage(content)
|
||||
if coverage:
|
||||
chars = list(content)
|
||||
for cov_start, cov_end in coverage:
|
||||
for i in range(cov_start, min(cov_end, len(chars))):
|
||||
chars[i] = " "
|
||||
fallback_content = "".join(chars)
|
||||
for parser in (
|
||||
_parse_function_xml, # <function name="..."> attribute form
|
||||
_parse_llama3_python_tag, # Llama-3 <|python_tag|>
|
||||
_parse_mistral_tool_calls, # Mistral [TOOL_CALLS]
|
||||
):
|
||||
calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete)
|
||||
calls = parser(fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete)
|
||||
if calls:
|
||||
return calls
|
||||
|
||||
|
|
|
|||
|
|
@ -10,20 +10,45 @@ orchestrator, structlog, httpx, or the rest of the studio backend.
|
|||
import json
|
||||
import re
|
||||
|
||||
# Pre-compiled patterns for tool XML stripping. The hyphen in the name
|
||||
# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues,
|
||||
# issue-number) parse alongside the built-ins.
|
||||
# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs
|
||||
# strip first so a closed call goes as a unit before any to-EOF sweep reaches
|
||||
# nested markup; only the final list adds the .*$ EOF sweeps.
|
||||
_TC_JSON_CLOSED_PAT = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL)
|
||||
_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL)
|
||||
_TC_FUNC_CLOSED_PAT = re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL)
|
||||
_TC_GEMMA_END_PAT = re.compile(r"<tool_call\|>")
|
||||
_TOOL_CLOSED_PATS = [
|
||||
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
|
||||
re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL),
|
||||
re.compile(r"<tool_call\|>"),
|
||||
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
|
||||
_TC_JSON_CLOSED_PAT,
|
||||
_TC_GEMMA_CLOSED_PAT,
|
||||
_TC_FUNC_CLOSED_PAT,
|
||||
_TC_GEMMA_END_PAT,
|
||||
]
|
||||
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
|
||||
re.compile(r"<tool_call>.*$", re.DOTALL),
|
||||
re.compile(r"<\|tool_call>.*$", re.DOTALL),
|
||||
re.compile(r"<tool_call>.*$", re.DOTALL),
|
||||
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
|
||||
]
|
||||
# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in
|
||||
# their argument data cannot make the helper truncate the block and its tail.
|
||||
_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT]
|
||||
# A lazy closed-pair pattern whose close token is absent rescans to EOF from
|
||||
# every opener (quadratic, re-run per streamed token); skip that doomed pass.
|
||||
_PAT_REQUIRED_TOKEN = {
|
||||
_TC_JSON_CLOSED_PAT: "</tool_call>",
|
||||
_TC_GEMMA_CLOSED_PAT: "<tool_call|>",
|
||||
_TC_FUNC_CLOSED_PAT: "</function>",
|
||||
}
|
||||
|
||||
|
||||
def strip_tool_patterns(text: str, patterns) -> str:
|
||||
"""Apply ``patterns`` in order, skipping closed-pair passes with no close token."""
|
||||
for pat in patterns:
|
||||
token = _PAT_REQUIRED_TOKEN.get(pat)
|
||||
if token is not None and token not in text:
|
||||
continue
|
||||
text = pat.sub("", text)
|
||||
return text
|
||||
|
||||
|
||||
# Pre-compiled patterns for tool-call XML parsing.
|
||||
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
|
||||
|
|
@ -40,13 +65,9 @@ _TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
|
|||
_GEMMA_QUOTE = '<|"|>'
|
||||
_PARAM_CLOSE_TAG = "</parameter>"
|
||||
_FUNC_CLOSE_TAG = "</function>"
|
||||
# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next
|
||||
# `key:` pair. A comma NOT followed by a key token is part of the value (e.g.
|
||||
# `location:New York, NY`), so it must not terminate the value. The key token
|
||||
# must be identifier-shaped (start with a letter or underscore); a comma
|
||||
# followed by digits-then-colon is value text such as a timestamp or ratio
|
||||
# (`meet at 10:00, 11:00 tomorrow`), not a new key.
|
||||
# Dots match the key-quoting scanner: a dotted key after a bare value must end the value at the comma.
|
||||
# A bare (unquoted) Gemma value ends at `}` or at a comma beginning the next
|
||||
# identifier-shaped `key:` pair; a comma before a non-key (`New York, NY`,
|
||||
# `10:00, 11:00`) stays in the value. Dots let a dotted key end the value.
|
||||
_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:")
|
||||
|
||||
|
||||
|
|
@ -143,14 +164,8 @@ def _split_top_level_commas(src: str) -> list:
|
|||
|
||||
|
||||
def _quote_gemma_array_elements(body: str) -> str:
|
||||
"""Normalise the elements of a Gemma array value so json.loads succeeds.
|
||||
|
||||
Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of
|
||||
objects (``items:[{path:a}]``) whose keys/values also lack quotes; left
|
||||
as-is json.loads fails and the whole call is dropped. Bare string elements
|
||||
are quoted, object and nested-array elements are normalised recursively, and
|
||||
quoted strings (already normalised from ``<|"|>``), numbers, and JSON
|
||||
literals are preserved."""
|
||||
"""Normalise a Gemma array value (``labels:[bug,ui]``) so json.loads succeeds:
|
||||
quote bare strings, recurse into objects/arrays, keep quoted/JSON literals."""
|
||||
out: list[str] = []
|
||||
for element in _split_top_level_commas(body):
|
||||
stripped = element.strip()
|
||||
|
|
@ -158,11 +173,9 @@ def _quote_gemma_array_elements(body: str) -> str:
|
|||
out.append(element)
|
||||
continue
|
||||
if stripped[0] == "{":
|
||||
# Object element: quote its keys/bare values like a top-level object.
|
||||
out.append(_quote_gemma_object_keys(stripped))
|
||||
continue
|
||||
if stripped[0] == "[":
|
||||
# Nested array: normalise its elements too.
|
||||
inner_end = _balanced_bracket_end(stripped, 0)
|
||||
if inner_end == len(stripped) - 1:
|
||||
out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]")
|
||||
|
|
@ -241,15 +254,12 @@ def _quote_gemma_object_keys(src: str) -> str:
|
|||
parts.append(src[i:colon_pos])
|
||||
parts.append(":")
|
||||
i = colon_pos + 1
|
||||
# Gemma may emit bare string values ({unit:celsius}); quote them so
|
||||
# json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is.
|
||||
# Quote bare string values ({unit:celsius}); JSON stays as-is.
|
||||
ws = i
|
||||
while i < len(src) and src[i].isspace():
|
||||
i += 1
|
||||
parts.append(src[ws:i])
|
||||
if i < len(src) and src[i] == "[":
|
||||
# Array value: quote bare string elements (e.g. labels:[bug,ui])
|
||||
# so json.loads succeeds instead of dropping the call.
|
||||
arr_end = _balanced_bracket_end(src, i)
|
||||
if arr_end < 0:
|
||||
parts.append(src[i:])
|
||||
|
|
@ -259,9 +269,7 @@ def _quote_gemma_object_keys(src: str) -> str:
|
|||
i = arr_end + 1
|
||||
elif i < len(src) and src[i] not in '"{':
|
||||
v_start = i
|
||||
# Consume the bare value up to `}` or a comma that starts the
|
||||
# next key:value pair; a comma inside the value (e.g.
|
||||
# `New York, NY`) does not terminate it.
|
||||
# Bare value: up to `}` or a comma that starts the next key:pair.
|
||||
while i < len(src):
|
||||
if src[i] == "}":
|
||||
break
|
||||
|
|
@ -329,6 +337,68 @@ def _trim_param_value(val: str) -> str:
|
|||
return val
|
||||
|
||||
|
||||
def _marker_coverage(content: str, markers) -> list[tuple[int, int]]:
|
||||
"""Coverage ``[start, end]`` per marker, used to skip markers that are another
|
||||
call's data. Closes pair to markers via a per-format stack so an inner close
|
||||
is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with
|
||||
a paired close cover through it (markers before the close are data); balanced
|
||||
without one cover only the braces, so a later sibling is still recovered."""
|
||||
n = len(content)
|
||||
brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0]
|
||||
events = [] # (position, order) with order 0 = braces-done, 1 = close marker
|
||||
for idx, (_start, brace_end, _kind, _m) in enumerate(markers):
|
||||
if brace_end >= 0:
|
||||
events.append((brace_end, 0, _kind, idx))
|
||||
for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)):
|
||||
for cm in close_re.finditer(content):
|
||||
# A close inside another call's balanced braces is quoted data; it
|
||||
# must not pop an earlier close-less marker and swallow a sibling.
|
||||
if any(s < cm.start() < be for s, be in brace_regions):
|
||||
continue
|
||||
events.append((cm.start(), 1, kind, cm.end()))
|
||||
events.sort(key = lambda e: (e[0], e[1]))
|
||||
waiting = {"json": [], "gemma": []}
|
||||
close_end_for: dict[int, int] = {}
|
||||
for _pos, order, kind, payload in events:
|
||||
if order == 0:
|
||||
waiting[kind].append(payload) # marker index, now awaiting its close
|
||||
elif waiting[kind]:
|
||||
close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here
|
||||
coverage = []
|
||||
for idx, (start, brace_end, _kind, _m) in enumerate(markers):
|
||||
if brace_end < 0:
|
||||
coverage.append((start, n))
|
||||
elif idx in close_end_for:
|
||||
coverage.append((start, close_end_for[idx]))
|
||||
else:
|
||||
coverage.append((start, brace_end))
|
||||
return coverage
|
||||
|
||||
|
||||
def _build_markers(content: str):
|
||||
"""JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document
|
||||
order; ``brace_end < 0`` marks an unbalanced (to-EOF) open."""
|
||||
markers = []
|
||||
for start_re, gemma, kind in (
|
||||
(_TC_JSON_START_RE, False, "json"),
|
||||
(_TC_GEMMA_START_RE, True, "gemma"),
|
||||
):
|
||||
for m in start_re.finditer(content):
|
||||
if _inside_open_parameter(content, m.start()):
|
||||
continue
|
||||
brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma)
|
||||
markers.append((m.start(), brace_end, kind, m))
|
||||
markers.sort(key = lambda c: c[0])
|
||||
return markers
|
||||
|
||||
|
||||
def marker_coverage(content: str) -> list[tuple[int, int]]:
|
||||
"""Coverage spans of JSON/Gemma tool markers so other parsers can treat markup
|
||||
inside a marker's coverage (even a marker that failed to parse) as that call's
|
||||
data rather than a sibling call."""
|
||||
return _marker_coverage(content, _build_markers(content))
|
||||
|
||||
|
||||
def parse_tool_calls_from_text(
|
||||
content: str,
|
||||
*,
|
||||
|
|
@ -350,37 +420,26 @@ def parse_tool_calls_from_text(
|
|||
"""
|
||||
tool_calls: list[dict] = []
|
||||
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):
|
||||
if _inside_open_parameter(content, m.start()):
|
||||
continue
|
||||
end = _balanced_brace_end(content, m.end() - 1)
|
||||
if end >= 0:
|
||||
candidates.append((m.start(), end, "json", m))
|
||||
for m in _TC_GEMMA_START_RE.finditer(content):
|
||||
if _inside_open_parameter(content, m.start()):
|
||||
continue
|
||||
end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True)
|
||||
if end >= 0:
|
||||
candidates.append((m.start(), end, "gemma", m))
|
||||
candidates.sort(key = lambda c: c[0])
|
||||
|
||||
candidate_spans = [(s, e) for s, e, _kind, _m in candidates]
|
||||
for idx, (start, end, kind, m) in enumerate(candidates):
|
||||
if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx):
|
||||
# Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside
|
||||
# another call's coverage, or an open <parameter=> value, is data not executed.
|
||||
markers = _build_markers(content)
|
||||
coverage = _marker_coverage(content, markers)
|
||||
parsed_items = [] # (start, span_end, name, arguments) in document order
|
||||
for idx, (start, brace_end, kind, m) in enumerate(markers):
|
||||
# A marker starting inside another's coverage is that call's data. The
|
||||
# end is exclusive so a marker at a close's end is an adjacent sibling.
|
||||
if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx):
|
||||
continue
|
||||
if brace_end < 0:
|
||||
continue # unclosed: not parseable; the fallback still excludes its XML
|
||||
if not allow_incomplete:
|
||||
tail = content[end + 1 :].lstrip()
|
||||
tail = content[brace_end + 1 :].lstrip()
|
||||
close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
|
||||
if close_re.match(tail) is None:
|
||||
continue
|
||||
try:
|
||||
if kind == "json":
|
||||
obj = json.loads(content[m.end() - 1 : end + 1])
|
||||
obj = json.loads(content[m.end() - 1 : brace_end + 1])
|
||||
name = obj.get("name", "")
|
||||
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes <tool_call>).
|
||||
arguments = obj.get("arguments")
|
||||
|
|
@ -390,10 +449,11 @@ def parse_tool_calls_from_text(
|
|||
arguments = json.dumps(arguments)
|
||||
else:
|
||||
name = m.group(1)
|
||||
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end]))
|
||||
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end]))
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
span_end = end + 1
|
||||
# Span reaches through the close tag when present, else just the braces.
|
||||
span_end = brace_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)
|
||||
|
|
@ -401,11 +461,15 @@ def parse_tool_calls_from_text(
|
|||
span_end = close_m.end()
|
||||
parsed_items.append((start, span_end, name, arguments))
|
||||
|
||||
# Function-XML calls promote in document order alongside marker calls (the
|
||||
# #6801 contract). A <function=> inside any marker's coverage is excluded --
|
||||
# even if that marker failed to parse -- so nested XML cannot escape; one
|
||||
# after a balanced close-less marker is a sibling, not swallowed to EOF.
|
||||
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)
|
||||
and not any(s <= fm.start() < e for s, e in coverage)
|
||||
]
|
||||
for idx, fm in enumerate(func_starts):
|
||||
func_name = fm.group(1)
|
||||
|
|
@ -481,90 +545,106 @@ def parse_tool_calls_from_text(
|
|||
)
|
||||
call_spans.append((start, span_end))
|
||||
|
||||
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]
|
||||
# Span for with_spans callers: through the </function> close if present, else body end.
|
||||
span_end = body_end
|
||||
if not allow_incomplete:
|
||||
close_idx = _func_close_index(content, body_start, body)
|
||||
if close_idx < 0:
|
||||
continue
|
||||
body = body[:close_idx]
|
||||
span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG)
|
||||
else:
|
||||
# Terminate at the real close so trailing prose doesn't leak in; no close -> whole body.
|
||||
close_idx = _func_close_index(content, body_start, body)
|
||||
if close_idx >= 0:
|
||||
body = body[:close_idx]
|
||||
span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG)
|
||||
|
||||
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)] = _trim_param_value(val)
|
||||
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] = _trim_param_value(val)
|
||||
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((fm.start(), span_end))
|
||||
|
||||
if with_spans:
|
||||
return tool_calls, call_spans
|
||||
return tool_calls
|
||||
|
||||
|
||||
def _strip_gemma_native_spans(text: str, *, final: bool) -> str:
|
||||
"""Remove complete Gemma-native spans, brace/quote-balanced so a literal
|
||||
``<tool_call|>`` in a quoted argument cannot truncate the span. An incomplete
|
||||
span is dropped to EOF when ``final``, else kept (still streaming)."""
|
||||
out: list[str] = []
|
||||
cursor = 0
|
||||
for match in _TC_GEMMA_START_RE.finditer(text):
|
||||
start = match.start()
|
||||
if start < cursor:
|
||||
continue
|
||||
brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True)
|
||||
if brace_end < 0:
|
||||
# Unbalanced: nothing completes from here on. Drop the rest if final,
|
||||
# else keep it; stop either way (rescanning would be quadratic).
|
||||
if final:
|
||||
out.append(text[cursor:start])
|
||||
cursor = len(text)
|
||||
break
|
||||
# Junk between } and <tool_call|> is malformed-call markup: strip through
|
||||
# the close, keep text after it. No close anywhere means stop (linear).
|
||||
close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1)
|
||||
if close is None:
|
||||
if final:
|
||||
out.append(text[cursor:start])
|
||||
cursor = len(text)
|
||||
break
|
||||
out.append(text[cursor:start])
|
||||
cursor = close.end()
|
||||
out.append(text[cursor:])
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _gemma_span_ranges(text: str) -> list:
|
||||
"""``(start, end)`` of each complete Gemma-native span; same walk as
|
||||
``_strip_gemma_native_spans`` without stripping."""
|
||||
ranges: list[tuple] = []
|
||||
cursor = 0
|
||||
for match in _TC_GEMMA_START_RE.finditer(text):
|
||||
start = match.start()
|
||||
if start < cursor:
|
||||
continue
|
||||
brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True)
|
||||
if brace_end < 0:
|
||||
break
|
||||
close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1)
|
||||
if close is None:
|
||||
break
|
||||
ranges.append((start, close.end()))
|
||||
cursor = close.end()
|
||||
return ranges
|
||||
|
||||
|
||||
def _strip_closed_blocks_outside_gemma(text: str) -> str:
|
||||
"""Closed JSON/function pre-pass that skips matches starting inside a complete
|
||||
Gemma span: deleting across the span boundary would mangle the Gemma close and
|
||||
truncate the tail. A skipped match resumes at the covering span's end, so a
|
||||
real function-XML call after the span is still stripped."""
|
||||
ranges = _gemma_span_ranges(text)
|
||||
if not ranges:
|
||||
return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS)
|
||||
for pat in _TOOL_CLOSED_BLOCK_PATS:
|
||||
token = _PAT_REQUIRED_TOKEN.get(pat)
|
||||
if token is not None and token not in text:
|
||||
continue
|
||||
out: list[str] = []
|
||||
pos = 0
|
||||
while True:
|
||||
m = pat.search(text, pos)
|
||||
if m is None:
|
||||
out.append(text[pos:])
|
||||
break
|
||||
covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None)
|
||||
if covering is not None:
|
||||
out.append(text[pos : covering[1]])
|
||||
pos = covering[1]
|
||||
continue
|
||||
out.append(text[pos : m.start()])
|
||||
pos = m.end()
|
||||
new_text = "".join(out)
|
||||
if new_text != text:
|
||||
text = new_text
|
||||
ranges = _gemma_span_ranges(text)
|
||||
return text
|
||||
|
||||
|
||||
def strip_tool_markup_final(text: str) -> str:
|
||||
"""Final display strip, shared with the streaming wrappers so all paths order
|
||||
the passes identically: Gemma-aware closed JSON/function blocks first, then
|
||||
well-formed Gemma spans (quote-aware), then the regex sweeps mop up malformed
|
||||
spans and drop any unclosed remainder to EOF. Whitespace is kept."""
|
||||
text = _strip_closed_blocks_outside_gemma(text)
|
||||
text = _strip_gemma_native_spans(text, final = True)
|
||||
return strip_tool_patterns(text, _TOOL_ALL_PATS)
|
||||
|
||||
|
||||
def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
|
||||
"""Strip tool-call XML markup from text.
|
||||
|
||||
|
|
@ -572,7 +652,9 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
|
|||
When ``final`` is True, trailing incomplete tool-call blocks are removed
|
||||
too, and the result is stripped of surrounding whitespace.
|
||||
"""
|
||||
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
|
||||
for pat in patterns:
|
||||
text = pat.sub("", text)
|
||||
return text.strip() if final else text
|
||||
if final:
|
||||
return strip_tool_markup_final(text).strip()
|
||||
# Non-final: same ordering as the final path, but incomplete blocks are kept.
|
||||
text = _strip_closed_blocks_outside_gemma(text)
|
||||
text = _strip_gemma_native_spans(text, final = False)
|
||||
return strip_tool_patterns(text, _TOOL_CLOSED_PATS)
|
||||
|
|
|
|||
|
|
@ -852,17 +852,14 @@ class _SameTaskStreamingResponse(StreamingResponse):
|
|||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
# Async callable invoked when the client disconnects before the body
|
||||
# iterator is ever advanced. A generator that never started cannot run
|
||||
# its own try/finally, so a stream that acquires resources before its
|
||||
# first yield (the passthrough opens an upstream httpx stream eagerly)
|
||||
# passes this to release them.
|
||||
# Released when the client disconnects before the body iterator starts:
|
||||
# its try/finally never runs, so a stream that opens resources before the
|
||||
# first yield (the passthrough's upstream httpx stream) passes this.
|
||||
self._unstarted_cleanup = unstarted_cleanup
|
||||
|
||||
async def __call__(self, scope, receive, send) -> None:
|
||||
# Track whether the body iterator was ever advanced: send() only emits a
|
||||
# body message after the generator yields its first chunk, so a failure
|
||||
# before then means it never entered its try/finally.
|
||||
# send() emits a body message only after the first chunk, so no body
|
||||
# message means the generator never entered its try/finally.
|
||||
body_started = False
|
||||
|
||||
async def _tracking_send(message) -> None:
|
||||
|
|
@ -873,15 +870,11 @@ class _SameTaskStreamingResponse(StreamingResponse):
|
|||
|
||||
try:
|
||||
await self.stream_response(_tracking_send)
|
||||
except OSError:
|
||||
# Client disconnected mid-send.
|
||||
except OSError: # client disconnected mid-send
|
||||
if body_started:
|
||||
# The generator produced at least one chunk and is suspended in
|
||||
# its try/finally. Throw CancelledError into it (not aclose's
|
||||
# GeneratorExit) so its `except asyncio.CancelledError` handler
|
||||
# runs and finishes any api_monitor entry; GeneratorExit would
|
||||
# skip it and only run `finally`. Fall back to aclose() without
|
||||
# athrow.
|
||||
# Generator is suspended in its try/finally: throw CancelledError
|
||||
# (not aclose's GeneratorExit) so its handler finishes the
|
||||
# api_monitor entry. Fall back to aclose() without athrow.
|
||||
athrow = getattr(self.body_iterator, "athrow", None)
|
||||
if athrow is not None:
|
||||
try:
|
||||
|
|
@ -893,16 +886,16 @@ class _SameTaskStreamingResponse(StreamingResponse):
|
|||
if aclose is not None:
|
||||
await aclose()
|
||||
else:
|
||||
# http.response.start failed before the body iterator advanced,
|
||||
# so its try/finally never armed and aclose()/athrow() are no-ops
|
||||
# on an unstarted generator. Release any resources acquired
|
||||
# before the first yield via the explicit cleanup hook.
|
||||
# Generator never started; aclose()/athrow() are no-ops on it, so
|
||||
# release eager resources via the hook. getattr guards a response
|
||||
# built through __new__ without __init__ (tests, pickling).
|
||||
aclose = getattr(self.body_iterator, "aclose", None)
|
||||
if aclose is not None:
|
||||
await aclose()
|
||||
if self._unstarted_cleanup is not None:
|
||||
cleanup = getattr(self, "_unstarted_cleanup", None)
|
||||
if cleanup is not None:
|
||||
try:
|
||||
await self._unstarted_cleanup()
|
||||
await cleanup()
|
||||
except Exception:
|
||||
pass
|
||||
raise ClientDisconnect()
|
||||
|
|
@ -910,6 +903,16 @@ class _SameTaskStreamingResponse(StreamingResponse):
|
|||
await self.background()
|
||||
|
||||
|
||||
def _tracked_cancel_unstarted_cleanup(tracker):
|
||||
"""unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when
|
||||
the generator's finally (which normally exits it) never runs."""
|
||||
|
||||
async def _cleanup() -> None:
|
||||
tracker.__exit__(None, None, None)
|
||||
|
||||
return _cleanup
|
||||
|
||||
|
||||
async def _aclose_stream_resources(
|
||||
*,
|
||||
watchers = (),
|
||||
|
|
@ -4069,12 +4072,9 @@ async def generate_stream(
|
|||
_DONE = object()
|
||||
while True:
|
||||
if cancel_event.is_set():
|
||||
# The disconnect watcher set cancel_event between chunks.
|
||||
# Reset the backend here: closing the Python generator does
|
||||
# not signal a subprocess backend, so without this it keeps
|
||||
# decoding after the client is gone. The finally's reset is
|
||||
# guarded on cancel_event being unset, so it will not run
|
||||
# again for this path.
|
||||
# Watcher set cancel_event between chunks. Reset here: closing
|
||||
# the generator does not signal a subprocess backend, so it would
|
||||
# keep decoding. The finally's reset is guarded, so no double-run.
|
||||
backend.reset_generation_state()
|
||||
break
|
||||
chunk = await asyncio.to_thread(next, gen, _DONE)
|
||||
|
|
@ -5684,6 +5684,7 @@ async def openai_chat_completions(
|
|||
|
||||
return _SameTaskStreamingResponse(
|
||||
audio_input_stream(),
|
||||
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
"Cache-Control": "no-cache",
|
||||
|
|
@ -6163,6 +6164,7 @@ async def openai_chat_completions(
|
|||
if payload.stream:
|
||||
return _SameTaskStreamingResponse(
|
||||
gguf_tool_stream(),
|
||||
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
"Cache-Control": "no-cache",
|
||||
|
|
@ -6419,6 +6421,7 @@ async def openai_chat_completions(
|
|||
|
||||
return _SameTaskStreamingResponse(
|
||||
gguf_stream_chunks(),
|
||||
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
"Cache-Control": "no-cache",
|
||||
|
|
@ -6852,6 +6855,7 @@ async def openai_chat_completions(
|
|||
if payload.stream:
|
||||
return _SameTaskStreamingResponse(
|
||||
sf_tool_stream(),
|
||||
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
"Cache-Control": "no-cache",
|
||||
|
|
@ -7067,6 +7071,7 @@ async def openai_chat_completions(
|
|||
|
||||
return _SameTaskStreamingResponse(
|
||||
stream_chunks(),
|
||||
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
"Cache-Control": "no-cache",
|
||||
|
|
@ -9977,11 +9982,8 @@ async def _anthropic_tool_stream(
|
|||
drop_until_tool_end = False
|
||||
|
||||
gen = run_gen()
|
||||
# Concurrent disconnect watcher: the loop only polls is_disconnected()
|
||||
# between events, so a client disconnect during a long prefill or
|
||||
# generation step would otherwise hold the decode slot until the next
|
||||
# event or a failed send. The watcher sets cancel_event so the backend
|
||||
# stops promptly.
|
||||
# Watcher to cancel on disconnect: the in-loop poll fires only between
|
||||
# events, so a mid-prefill disconnect would otherwise hold the decode slot.
|
||||
disconnect_watcher = asyncio.create_task(
|
||||
_await_disconnect_then_cancel(request, cancel_event)
|
||||
)
|
||||
|
|
@ -10073,11 +10075,8 @@ async def _anthropic_plain_stream(
|
|||
captured_finish_reason = None
|
||||
|
||||
gen = run_gen()
|
||||
# Concurrent disconnect watcher: the loop only polls is_disconnected()
|
||||
# between chunks, so a client disconnect during a long prefill or
|
||||
# generation step would otherwise hold the decode slot until the next
|
||||
# chunk or a failed send. The watcher sets cancel_event so the backend
|
||||
# stops promptly.
|
||||
# Watcher to cancel on disconnect: the in-loop poll fires only between
|
||||
# chunks, so a mid-prefill disconnect would otherwise hold the decode slot.
|
||||
disconnect_watcher = asyncio.create_task(
|
||||
_await_disconnect_then_cancel(request, cancel_event)
|
||||
)
|
||||
|
|
@ -11030,6 +11029,10 @@ async def _openai_passthrough_stream(
|
|||
response ``id``, ``finish_reason`` (including ``"tool_calls"``),
|
||||
``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so
|
||||
the client sees a standard OpenAI response.
|
||||
|
||||
Reasoning/tool-call splitting is delegated to llama-server (``--jinja
|
||||
--reasoning-format auto``), so ``delta.content`` carries no raw markup and is
|
||||
deliberately not re-parsed locally, unlike the ``/completion`` paths.
|
||||
"""
|
||||
target_url = f"{llama_backend.base_url}/v1/chat/completions"
|
||||
body = _build_openai_passthrough_body(
|
||||
|
|
@ -11446,11 +11449,9 @@ async def _openai_passthrough_stream(
|
|||
delta = choice.get("delta")
|
||||
if isinstance(delta, dict) and delta.get("tool_calls"):
|
||||
saw_tool_call_delta = True
|
||||
# Detect an upstream error chunk independently of API
|
||||
# monitoring: when monitor_id is None (skip_api_monitor),
|
||||
# _monitor_openai_sse_line returns before inspecting the
|
||||
# error, so without this the synthetic-finish guard would
|
||||
# emit a successful finish_reason after a failed stream.
|
||||
# Detect an error chunk independently of API monitoring
|
||||
# (skip_api_monitor returns early), else the synthetic
|
||||
# finish would fire after a failed stream.
|
||||
if _monitor_openai_error_message(chunk_data):
|
||||
saw_stream_error = True
|
||||
# With healing active, a content-bearing line may be replaced by
|
||||
|
|
|
|||
|
|
@ -1,15 +1,8 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Edge cases in Gemma-native tool-call parsing.
|
||||
|
||||
Covers two failure modes:
|
||||
1. A bare (unquoted) string argument that contains a comma, e.g.
|
||||
``location:New York, NY`` -- the comma must not be treated as the next
|
||||
key boundary, or the whole call is dropped.
|
||||
2. A tool-call marker that appears INSIDE another call's argument string is
|
||||
data, not a real call, so it must not be promoted to a second tool call.
|
||||
"""
|
||||
"""Gemma-native tool-call parsing edge cases: commas inside bare string values,
|
||||
and markers inside another call's argument data staying data."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -25,6 +18,7 @@ from core.inference.tool_call_parser import (
|
|||
_gemma_parse_value,
|
||||
parse_tool_calls_from_text,
|
||||
)
|
||||
from core.tool_healing import strip_tool_call_markup
|
||||
|
||||
|
||||
def _args(call: dict) -> dict:
|
||||
|
|
@ -43,8 +37,6 @@ def test_bare_string_argument_with_comma_is_kept():
|
|||
def test_normal_multi_key_arguments_still_split():
|
||||
calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}<tool_call|>')
|
||||
assert len(calls) == 1, calls
|
||||
# Numbers stay numeric, bare strings get quoted, an explicit quoted comma
|
||||
# stays inside its value.
|
||||
assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"}
|
||||
|
||||
|
||||
|
|
@ -60,8 +52,7 @@ def test_empty_bare_value_becomes_empty_string_not_dropped():
|
|||
|
||||
|
||||
def test_bare_value_with_timestamps_after_comma_is_kept():
|
||||
# A comma followed by digits-then-colon (a timestamp/ratio) is value text,
|
||||
# not a new key, so the whole query must be preserved as one argument.
|
||||
# A comma before digits-then-colon (timestamp/ratio) is value text, not a key.
|
||||
calls = parse_tool_calls_from_text(
|
||||
"<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}<tool_call|>"
|
||||
)
|
||||
|
|
@ -70,8 +61,6 @@ def test_bare_value_with_timestamps_after_comma_is_kept():
|
|||
|
||||
|
||||
def test_marker_inside_json_argument_is_not_a_second_call():
|
||||
# A python call whose `code` argument contains a Gemma marker string. The
|
||||
# marker is data and must not execute as a second `terminal` call.
|
||||
content = (
|
||||
'<tool_call>{"name":"python","arguments":{"code":'
|
||||
'"x = 1 # <|tool_call>call:terminal{command:ls}<tool_call|>"}}</tool_call>'
|
||||
|
|
@ -89,8 +78,6 @@ def test_two_separate_gemma_calls_both_parse():
|
|||
|
||||
|
||||
def test_mixed_format_calls_preserve_document_order():
|
||||
# A Gemma-native call precedes a JSON-format call in the text; tools execute
|
||||
# in returned order, so `create` must come before `read`.
|
||||
content = (
|
||||
"<|tool_call>call:create{path:a}<tool_call|> then "
|
||||
'<tool_call>{"name":"read","arguments":{"path":"a"}}</tool_call>'
|
||||
|
|
@ -100,8 +87,6 @@ def test_mixed_format_calls_preserve_document_order():
|
|||
|
||||
|
||||
def test_json_marker_inside_gemma_argument_is_not_a_second_call():
|
||||
# The reverse of the JSON-outer case: a JSON-style marker inside a Gemma
|
||||
# call's quoted argument is code text, not a second `terminal` call.
|
||||
content = (
|
||||
'<|tool_call>call:python{code:<|"|>'
|
||||
'print(<tool_call>{"name":"terminal","arguments":{"command":"ls"}}</tool_call>)'
|
||||
|
|
@ -112,18 +97,14 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call():
|
|||
|
||||
|
||||
def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call():
|
||||
# An UNQUOTED Gemma value containing a literal marker: the outer object fails
|
||||
# to normalize (the inner braces/marker break the JSON), but the inner marker
|
||||
# is nested in the outer candidate span, so it must not be promoted to a
|
||||
# standalone `terminal` call. The safe outcome is no executed tool call.
|
||||
# The outer object fails to normalize, but the nested marker is covered by
|
||||
# its span; safe outcome is no executed call at all.
|
||||
content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}<tool_call|>}<tool_call|>"
|
||||
calls = parse_tool_calls_from_text(content)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_bare_string_array_argument_is_quoted():
|
||||
# Gemma may emit an array of bare strings without per-element quotes; they
|
||||
# must be quoted so the call is not dropped.
|
||||
calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}<tool_call|>")
|
||||
assert len(calls) == 1, calls
|
||||
assert _args(calls[0]) == {"labels": ["bug", "ui"]}
|
||||
|
|
@ -137,8 +118,6 @@ def test_array_keeps_numbers_and_quoted_elements():
|
|||
|
||||
|
||||
def test_array_of_objects_is_normalised():
|
||||
# Arrays of objects are a common tool-schema shape; their (unquoted) keys and
|
||||
# bare values must be normalised too, not left verbatim, or the call drops.
|
||||
calls = parse_tool_calls_from_text(
|
||||
"<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}<tool_call|>"
|
||||
)
|
||||
|
|
@ -152,9 +131,6 @@ def test_nested_array_elements_are_normalised():
|
|||
|
||||
|
||||
def test_gemma_marker_inside_xml_parameter_is_not_a_second_call():
|
||||
# An XML-style <function=...> call whose <parameter=code> value contains a
|
||||
# Gemma marker: the marker is the parameter's data, not a separate terminal
|
||||
# call, so only the python call must be returned.
|
||||
content = (
|
||||
"<tool_call><function=python><parameter=code>"
|
||||
"x = 1 # <|tool_call>call:terminal{command:ls}<tool_call|>"
|
||||
|
|
@ -175,6 +151,165 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call():
|
|||
assert [c["function"]["name"] for c in calls] == ["python"], calls
|
||||
|
||||
|
||||
def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping():
|
||||
# Parse keeps the quoted close marker as data; strip removes the whole span.
|
||||
text = '<|tool_call>call:python{code:<|"|>print("<tool_call|>")<|"|>}<tool_call|>'
|
||||
calls = parse_tool_calls_from_text(text)
|
||||
assert len(calls) == 1, calls
|
||||
assert _args(calls[0]) == {"code": 'print("<tool_call|>")'}
|
||||
assert strip_tool_call_markup("before " + text + " after") == "before after"
|
||||
assert strip_tool_call_markup("before " + text + " after", final = True) == "before after"
|
||||
|
||||
|
||||
def test_nested_xml_in_malformed_gemma_call_does_not_execute():
|
||||
# The failed Gemma candidate's span still covers its nested <function=>.
|
||||
text = (
|
||||
"<|tool_call>call:outer{code:<function=terminal><parameter=command>id"
|
||||
"</parameter></function></tool_call>, broken:{x}}<tool_call|>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_unbalanced_gemma_call_with_xml_does_not_execute():
|
||||
# Unclosed braces cover to EOF, so the trailing <function=> is excluded.
|
||||
text = (
|
||||
"<|tool_call>call:outer{code:<function=terminal>"
|
||||
"<parameter=command>id</parameter></function>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_standalone_function_xml_still_parses():
|
||||
text = "<function=terminal><parameter=command>id</parameter></function>"
|
||||
calls = parse_tool_calls_from_text(text)
|
||||
assert [c["function"]["name"] for c in calls] == ["terminal"], calls
|
||||
|
||||
|
||||
def test_xml_between_braces_and_close_marker_does_not_execute():
|
||||
# Coverage runs to the close marker, so <function=> in the gap is data.
|
||||
text = (
|
||||
"<|tool_call>call:outer{broken:{x}}<function=terminal>"
|
||||
"<parameter=command>id</parameter></function><tool_call|>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_balanced_inner_call_inside_unclosed_outer_does_not_execute():
|
||||
text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}<tool_call|>"
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_strip_preserves_text_after_malformed_gemma_close():
|
||||
# Junk before the close is a malformed span: strip through it, keep the tail.
|
||||
text = "pre <|tool_call>call:t{a:1} note <tool_call|> post"
|
||||
assert strip_tool_call_markup(text) == "pre post"
|
||||
assert strip_tool_call_markup(text, final = True) == "pre post"
|
||||
|
||||
|
||||
def test_malformed_closed_gemma_span_is_stripped():
|
||||
assert (
|
||||
strip_tool_call_markup('before <|tool_call>{"name":"x"}<tool_call|> after')
|
||||
== "before after"
|
||||
)
|
||||
|
||||
|
||||
def test_valid_call_after_missing_close_is_recovered():
|
||||
# A close-less call covers only its braces, so the later call is recovered.
|
||||
text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}<tool_call|>"
|
||||
names_inc = [
|
||||
c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True)
|
||||
]
|
||||
assert "b" in names_inc, names_inc
|
||||
names_strict = [
|
||||
c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False)
|
||||
]
|
||||
assert names_strict == ["b"], names_strict
|
||||
|
||||
|
||||
def test_strip_non_final_keeps_incomplete_gemma_block():
|
||||
text = "before <|tool_call>call:t{"
|
||||
assert strip_tool_call_markup(text) == text
|
||||
assert strip_tool_call_markup(text, final = True) == "before"
|
||||
|
||||
|
||||
def test_json_call_between_gemma_braces_and_close_does_not_execute():
|
||||
# A JSON call between the outer's braces and its close is covered data.
|
||||
text = (
|
||||
"<|tool_call>call:outer{broken:{x}}"
|
||||
'<tool_call>{"name":"terminal","arguments":{"command":"id"}}</tool_call>'
|
||||
"<tool_call|>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_gemma_call_between_gemma_braces_and_close_does_not_execute():
|
||||
# Same escape with a Gemma-native inner marker.
|
||||
text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}<tool_call|><tool_call|>"
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener():
|
||||
# The to-EOF Gemma sweep must not eat visible text after </function>.
|
||||
text = (
|
||||
'before <function=python><parameter=code>print("<|tool_call>")</parameter></function> after'
|
||||
)
|
||||
assert strip_tool_call_markup(text, final = True) == "before after"
|
||||
assert strip_tool_call_markup(text) == "before after"
|
||||
|
||||
|
||||
def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener():
|
||||
# A call-form Gemma opener quoted in a closed block must not truncate it.
|
||||
xml = "<function=python><parameter=code><|tool_call>call:t{</parameter></function>"
|
||||
json_block = (
|
||||
'<tool_call>{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}</tool_call>'
|
||||
)
|
||||
for block in (xml, json_block):
|
||||
text = "before " + block + " after"
|
||||
assert strip_tool_call_markup(text, final = True) == "before after", block
|
||||
assert strip_tool_call_markup(text) == "before after", block
|
||||
|
||||
|
||||
def test_function_sibling_after_close_less_gemma_marker_is_recovered():
|
||||
# The close-less marker covers only its braces; the XML sibling is recovered.
|
||||
text = (
|
||||
"<|tool_call>call:bad{broken:{x}} "
|
||||
"<function=terminal><parameter=command>id</parameter></function>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert [c["function"]["name"] for c in calls] == ["terminal"], calls
|
||||
|
||||
|
||||
def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered():
|
||||
# A close token quoted in the later call must not extend the earlier
|
||||
# close-less marker's coverage over that call.
|
||||
gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|></tool_call><|"|>}<tool_call|>'
|
||||
names = [
|
||||
c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False)
|
||||
]
|
||||
assert names == ["b"], names
|
||||
json_text = (
|
||||
'<tool_call>{"name":"a","arguments":{}} '
|
||||
'<tool_call>{"name":"b","arguments":{"x":"</tool_call>"}}</tool_call>'
|
||||
)
|
||||
names_j = [
|
||||
c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False)
|
||||
]
|
||||
assert "b" in names_j, names_j
|
||||
|
||||
|
||||
def test_gemma_parse_value_always_advances_on_stray_delimiter():
|
||||
# A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the
|
||||
# parser, or a looping caller spins forever (DoS).
|
||||
|
|
|
|||
|
|
@ -1063,3 +1063,44 @@ class TestBareJsonStripRequiresTopLevelName:
|
|||
def test_real_call_still_strips_name_agnostic(self):
|
||||
from core.inference.tool_call_parser import strip_leading_bare_json_call
|
||||
assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') == ""
|
||||
|
||||
|
||||
class TestGemmaAwareClosedBlockPrePass:
|
||||
"""The closed JSON/function strip pre-pass must not delete across a complete
|
||||
Gemma span (a quoted <function=...> plus a later real </function>)."""
|
||||
|
||||
def test_literal_function_in_gemma_arg_with_later_real_call(self):
|
||||
from core.tool_healing import strip_tool_call_markup
|
||||
text = (
|
||||
'before <|tool_call>call:python{code:<|"|>print("<function=x>")<|"|>}'
|
||||
"<tool_call|> <function=terminal><parameter=cmd>ls</parameter>"
|
||||
"</function> after"
|
||||
)
|
||||
assert strip_tool_call_markup(text, final = True) == "before after"
|
||||
|
||||
def test_literal_function_in_gemma_arg_with_prose_closer(self):
|
||||
from core.tool_healing import strip_tool_call_markup
|
||||
|
||||
text = (
|
||||
'before <|tool_call>call:python{code:<|"|>print("<function=x>")<|"|>}'
|
||||
"<tool_call|> then use </function> to close. after"
|
||||
)
|
||||
out = strip_tool_call_markup(text, final = True)
|
||||
assert out.startswith("before")
|
||||
assert out.endswith("after")
|
||||
assert "call:python" not in out
|
||||
|
||||
def test_gemma_opener_inside_json_arg_still_strips_block(self):
|
||||
from core.tool_healing import strip_tool_call_markup
|
||||
text = (
|
||||
'<tool_call>{"name":"t","arguments":{"code":"<|tool_call>call:x{"}}</tool_call> after'
|
||||
)
|
||||
assert strip_tool_call_markup(text, final = True) == "after"
|
||||
|
||||
def test_gemma_opener_inside_function_param_still_strips_block(self):
|
||||
from core.tool_healing import strip_tool_call_markup
|
||||
text = (
|
||||
'<function=python><parameter=code>x = "<|tool_call>call:t{"</parameter>'
|
||||
"</function> after"
|
||||
)
|
||||
assert strip_tool_call_markup(text, final = True) == "after"
|
||||
|
|
|
|||
76
studio/backend/tests/test_tool_strip_guard.py
Normal file
76
studio/backend/tests/test_tool_strip_guard.py
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""strip_tool_patterns must match the plain per-pattern loop while skipping the
|
||||
quadratic no-match rescan of a closed-pair sweep whose close token is absent."""
|
||||
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
|
||||
if str(_BACKEND_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND_ROOT))
|
||||
|
||||
from core.tool_healing import (
|
||||
_TOOL_ALL_PATS,
|
||||
_TOOL_CLOSED_PATS,
|
||||
strip_tool_call_markup,
|
||||
strip_tool_patterns,
|
||||
)
|
||||
|
||||
|
||||
def _naive(text, patterns):
|
||||
for pat in patterns:
|
||||
text = pat.sub("", text)
|
||||
return text
|
||||
|
||||
|
||||
_TOKENS = [
|
||||
"<tool_call>",
|
||||
"</tool_call>",
|
||||
"<|tool_call>",
|
||||
"<tool_call|>",
|
||||
"<function=x>",
|
||||
"<function=mcp__s__a-b>",
|
||||
"</function>",
|
||||
"<parameter=p>",
|
||||
"</parameter>",
|
||||
"call:fn{",
|
||||
"}",
|
||||
"{",
|
||||
'<|"|>',
|
||||
"A",
|
||||
" ",
|
||||
"\n",
|
||||
"id",
|
||||
"x:1",
|
||||
"</tool",
|
||||
"call>",
|
||||
]
|
||||
|
||||
|
||||
def test_guard_matches_plain_loop_on_fuzz():
|
||||
rng = random.Random(1234)
|
||||
for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS):
|
||||
for _ in range(20000):
|
||||
s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10)))
|
||||
assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns)
|
||||
|
||||
|
||||
def test_strip_markup_representative_cases_unchanged():
|
||||
assert strip_tool_call_markup("a <tool_call>{}</tool_call> b") == "a b"
|
||||
assert strip_tool_call_markup("a <function=x><parameter=p>1</parameter></function> b") == "a b"
|
||||
# Non-final keeps an unclosed block; final strips it to EOF.
|
||||
assert strip_tool_call_markup("a <tool_call>{partial") == "a <tool_call>{partial"
|
||||
assert strip_tool_call_markup("a <tool_call>{partial", final = True) == "a"
|
||||
|
||||
|
||||
def test_no_quadratic_blowup_on_unclosed_markers():
|
||||
# Unguarded, this took minutes.
|
||||
big = "<tool_call>" * 20000 + "<function=x>" * 20000
|
||||
t0 = time.perf_counter()
|
||||
out = strip_tool_call_markup(big, final = True)
|
||||
assert time.perf_counter() - t0 < 2.0
|
||||
assert out == ""
|
||||
Loading…
Add table
Add a link
Reference in a new issue