ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Harden Gemma array parsing, XML-parameter guard, and stream teardown

Browse source 
Address five review findings on the Gemma tool-call and OpenAI passthrough
streaming paths:

- parse_tool_calls_from_text collected JSON and Gemma markers without the
  _inside_open_parameter guard, so a marker embedded in an existing
  <function=...><parameter=...> value was promoted to a separate tool call.
  Candidates that start inside an open XML parameter are now skipped, matching
  the guard the XML-style parser already applies.

- _quote_gemma_array_elements preserved array elements starting with { or [
  verbatim, so an array of objects (items:[{path:a}]) or a nested array failed
  json.loads and the whole call was dropped. Object and nested-array elements
  are now normalised recursively.

- _openai_passthrough_stream synthesized a finish chunk before a trailing
  usage-only chunk and set saw_finish_reason, which made the EOF guard skip the
  [DONE] sentinel. The EOF path now emits [DONE] whenever the upstream omitted
  it, even after a finish chunk was already synthesized.

- /generate/stream drove generation through asyncio.to_thread with no
  disconnect watcher, so a client disconnect during a long generation went
  unnoticed until the next send. It now runs _await_disconnect_then_cancel
  against the request, matching the other local streaming endpoints.

- _SameTaskStreamingResponse closed the body iterator with aclose() on a
  send-side disconnect, raising GeneratorExit so the generators' cancellation
  handlers (which finish the api_monitor entry) never ran. It now throws
  CancelledError, falling back to aclose() when athrow is unavailable.

Extends tests/test_gemma_tool_parse_edge_cases.py with array-of-objects,
nested-array, and marker-inside-XML-parameter cases.
This commit is contained in:
danielhanchen 2026-06-22 15:14:42 +00:00
commit 1e58c3707d
3 changed files with 116 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

33
studio/backend/core/tool_healing.py
View file

@ -136,16 +136,32 @@ def _split_top_level_commas(src: str) -> list:
def _quote_gemma_array_elements(body: str) -> str:
"""Quote bare (unquoted) string elements in a Gemma array value. Gemma may
emit ``labels:[bug,ui]`` without per-element quotes; left as-is json.loads
fails and the whole call is dropped. Quoted strings (already normalised from
``<|"|>``), numbers, and JSON literals are preserved."""
"""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."""
out: list[str] = []
for element in _split_top_level_commas(body):
stripped = element.strip()
if not stripped or stripped[0] in '"{[':
if not stripped or stripped[0] == '"':
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]) + "]")
else:
out.append(element)
continue
try:
json.loads(stripped)
out.append(element)
@ -301,10 +317,17 @@ def parse_tool_calls_from_text(
# nested in a JSON arg alike, regardless of which format is outer).
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)
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))

64
studio/backend/routes/inference.py
View file

@ -812,9 +812,22 @@ class _SameTaskStreamingResponse(StreamingResponse):
try:
await self.stream_response(send)
except OSError:
aclose = getattr(self.body_iterator, "aclose", None)
if aclose is not None:
await aclose()
# Client disconnected mid-send. Throw CancelledError into the body
# generator instead of aclose() (which raises GeneratorExit): the
# generators run their `except asyncio.CancelledError` handler, which
# finishes the api_monitor entry as "cancelled", whereas GeneratorExit
# skips it and only runs `finally`, leaving the monitor entry active.
# Fall back to aclose() for iterators without athrow.
athrow = getattr(self.body_iterator, "athrow", None)
if athrow is not None:
try:
await athrow(asyncio.CancelledError())
except (asyncio.CancelledError, StopAsyncIteration, RuntimeError):
pass
else:
aclose = getattr(self.body_iterator, "aclose", None)
if aclose is not None:
await aclose()
raise ClientDisconnect()
if self.background is not None:
await self.background()
@ -3332,7 +3345,9 @@ async def get_api_monitor_entry(entry_id: str, current_subject: str = Depends(ge
@router.post("/generate/stream")
async def generate_stream(
request: GenerateRequest, current_subject: str = Depends(get_current_subject)
request: GenerateRequest,
fastapi_request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
Generate a chat response with Server-Sent Events (SSE) streaming.
@ -3382,6 +3397,13 @@ async def generate_stream(
async def stream():
gen = None
completed = False
# Cancel the generation when the client disconnects. The generator only
# awaits asyncio.to_thread(next, gen, ...), so without a concurrent
# watcher a disconnect during a long prefill/generation would go
# unnoticed until the next send and the backend would keep generating.
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(fastapi_request, cancel_event)
)
try:
gen = backend.generate_chat_response(
messages = request.messages,
@ -3396,12 +3418,15 @@ async def generate_stream(
)
_DONE = object()
while True:
if cancel_event.is_set():
break
chunk = await asyncio.to_thread(next, gen, _DONE)
if chunk is _DONE:
completed = True
break
yield f"data: {json.dumps({'content': chunk})}\n\n"
completed = True
yield "data: [DONE]\n\n"
if completed:
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
@ -3413,6 +3438,7 @@ async def generate_stream(
logger.error(f"Error during generation: {e}", exc_info = True)
yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n"
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
if not completed and not cancel_event.is_set():
cancel_event.set()
backend.reset_generation_state()
@ -9624,19 +9650,19 @@ async def _openai_passthrough_stream(
if monitor_event == "done":
monitor_done = True
break
if (
not saw_done
and not saw_finish_reason
and not saw_stream_error
and not cancel_event.is_set()
):
finish_line = _synthetic_finish_line()
_monitor_openai_sse_line(
monitor_id,
finish_line,
llama_backend.context_length,
)
yield finish_line + "\n\n"
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 not saw_finish_reason:
finish_line = _synthetic_finish_line()
_monitor_openai_sse_line(
monitor_id,
finish_line,
llama_backend.context_length,
)
yield finish_line + "\n\n"
done_line = "data: [DONE]"
_monitor_openai_sse_line(
monitor_id,

43
studio/backend/tests/test_gemma_tool_parse_edge_cases.py
View file

@ -100,3 +100,46 @@ def test_array_keeps_numbers_and_quoted_elements():
'<|tool_call>call:f{nums:[1,2],tags:[<|"|>a,b<|"|>,c]}<tool_call|>'
)
assert _args(calls[0]) == {"nums": [1, 2], "tags": ["a,b", "c"]}
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|>"
)
assert len(calls) == 1, calls
assert _args(calls[0]) == {
"items": [{"path": "a", "mode": "r"}, {"path": "b", "mode": "w"}]
}
def test_nested_array_elements_are_normalised():
calls = parse_tool_calls_from_text(
"<|tool_call>call:grid{cells:[[a,b],[c,d]]}<tool_call|>"
)
assert _args(calls[0]) == {"cells": [["a", "b"], ["c", "d"]]}
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|>"
"</parameter></function></tool_call>"
)
calls = parse_tool_calls_from_text(content)
assert [c["function"]["name"] for c in calls] == ["python"], calls
assert "terminal" in _args(calls[0])["code"]
def test_json_marker_inside_xml_parameter_is_not_a_second_call():
content = (
"<tool_call><function=python><parameter=code>"
'run(<tool_call>{"name":"terminal","arguments":{"command":"ls"}}</tool_call>)'
"</parameter></function></tool_call>"
)
calls = parse_tool_calls_from_text(content)
assert [c["function"]["name"] for c in calls] == ["python"], calls