Ciaran/gpt oss tool call finish reason (#1673)
## Motivation When gpt-oss truncates a response mid-tool-call (e.g. hitting max tokens), the finish_reason was swallowed because the parser was in tool-accumulation mode and never yielded it back to the caller. This caused the API to hang or fail to signal completion. ## Changes In parse_gpt_oss, when accumulating tool call arguments and a finish_reason arrives, yield the response (with empty text) so the finish reason propagates ## Why It Works The parser now checks for finish_reason on every chunk while inside a tool call. When the model stops generating (truncation), the finish signal is forwarded immediately rather than being silently consumed. ## Test Plan ### Automated Testing Added tests covering truncated tool calls (finish_reason="length") and plain text truncation
This commit is contained in:
@@ -116,6 +116,9 @@ def parse_gpt_oss(
|
||||
if current_tool_name is not None:
|
||||
if delta:
|
||||
tool_arg_parts.append(delta)
|
||||
if response.finish_reason is not None:
|
||||
yield response.model_copy(update={"text": "".join(tool_arg_parts)})
|
||||
tool_arg_parts = []
|
||||
continue
|
||||
|
||||
if ch == "analysis" and not thinking:
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from collections.abc import Generator
|
||||
|
||||
from exo.shared.types.api import FinishReason
|
||||
from exo.shared.types.worker.runner_response import (
|
||||
GenerationResponse,
|
||||
ToolCallResponse,
|
||||
@@ -83,6 +84,7 @@ THINKING_THEN_TOOL_TOKENS: list[tuple[int, str]] = [
|
||||
|
||||
def _make_gen_responses(
|
||||
tokens: list[tuple[int, str]],
|
||||
last_finish_reason: FinishReason = "stop",
|
||||
) -> list[GenerationResponse]:
|
||||
"""Build GenerationResponse list from (token_id, text) pairs."""
|
||||
responses: list[GenerationResponse] = []
|
||||
@@ -92,7 +94,7 @@ def _make_gen_responses(
|
||||
GenerationResponse(
|
||||
text=text,
|
||||
token=tid,
|
||||
finish_reason="stop" if is_last else None,
|
||||
finish_reason=last_finish_reason if is_last else None,
|
||||
usage=None,
|
||||
)
|
||||
)
|
||||
@@ -101,11 +103,12 @@ def _make_gen_responses(
|
||||
|
||||
def _collect(
|
||||
tokens: list[tuple[int, str]],
|
||||
last_finish_reason: FinishReason = "stop",
|
||||
) -> list[GenerationResponse | ToolCallResponse]:
|
||||
"""Feed tokens through parse_gpt_oss and collect all yielded responses."""
|
||||
|
||||
def _gen() -> Generator[GenerationResponse, None, None]:
|
||||
yield from _make_gen_responses(tokens)
|
||||
yield from _make_gen_responses(tokens, last_finish_reason)
|
||||
|
||||
return list(x for x in parse_gpt_oss(_gen()) if x is not None)
|
||||
|
||||
@@ -171,3 +174,80 @@ class TestParseGptOssThinkingThenToolCall:
|
||||
tc = _get_tool_call(results)
|
||||
assert tc.tool_calls[0].name == "get_current_weather"
|
||||
assert "Tokyo" in tc.tool_calls[0].arguments
|
||||
|
||||
|
||||
# fmt: off
|
||||
# Truncated tool call: recipient + channel + message + partial args, no <|call|>
|
||||
TRUNCATED_TOOL_CALL_TOKENS: list[tuple[int, str]] = [
|
||||
(316, " to"),
|
||||
(28, "="),
|
||||
(44580, "functions"),
|
||||
(775, ".get"),
|
||||
(23981, "_current"),
|
||||
(170154, "_weather"),
|
||||
(_CHANNEL, "<|channel|>"),
|
||||
(12606, "comment"),
|
||||
(815, "ary"),
|
||||
(5701, " json"),
|
||||
(_MESSAGE, "<|message|>"),
|
||||
(10848, '{"'),
|
||||
(7693, "location"),
|
||||
(1243, '":'),
|
||||
(392, ' "'),
|
||||
(173844, "Tokyo"),
|
||||
# No <|call|> — generation truncated here
|
||||
]
|
||||
|
||||
# Plain text tokens (no tool call)
|
||||
PLAIN_TEXT_TOKENS: list[tuple[int, str]] = [
|
||||
(_CHANNEL, "<|channel|>"),
|
||||
(35644, "analysis"),
|
||||
(_MESSAGE, "<|message|>"),
|
||||
(12845, "Let"),
|
||||
(668, " me"),
|
||||
(2411, " think"),
|
||||
(1078, " about"),
|
||||
(495, " this"),
|
||||
(13, "."),
|
||||
(_END, "<|end|>"),
|
||||
(_START, "<|start|>"),
|
||||
(_ASSISTANT, "assistant"),
|
||||
(_CHANNEL, "<|channel|>"),
|
||||
(12606, "comment"),
|
||||
(815, "ary"),
|
||||
(_MESSAGE, "<|message|>"),
|
||||
(9906, "Hello"),
|
||||
(14, ","),
|
||||
(2989, " world"),
|
||||
]
|
||||
# fmt: on
|
||||
|
||||
|
||||
class TestParseGptOssMaxTokensTruncation:
|
||||
"""Truncated tool calls must still yield finish_reason."""
|
||||
|
||||
def test_truncated_tool_call_yields_finish_reason(self):
|
||||
results = _collect(TRUNCATED_TOOL_CALL_TOKENS, last_finish_reason="length")
|
||||
gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
|
||||
finish_reasons = [
|
||||
r.finish_reason for r in gen_responses if r.finish_reason is not None
|
||||
]
|
||||
assert "length" in finish_reasons
|
||||
|
||||
def test_truncated_tool_call_emits_partial_args(self):
|
||||
results = _collect(TRUNCATED_TOOL_CALL_TOKENS, last_finish_reason="length")
|
||||
gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
|
||||
last = [r for r in gen_responses if r.finish_reason is not None][-1]
|
||||
assert len(last.text) > 0
|
||||
|
||||
def test_truncated_plain_text_still_works(self):
|
||||
results = _collect(PLAIN_TEXT_TOKENS, last_finish_reason="length")
|
||||
gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
|
||||
finish_reasons = [
|
||||
r.finish_reason for r in gen_responses if r.finish_reason is not None
|
||||
]
|
||||
assert "length" in finish_reasons
|
||||
# Verify non-empty text was yielded (delta text differs from raw token text
|
||||
# due to Harmony encoding, so we just check something was emitted)
|
||||
all_text = "".join(r.text for r in gen_responses)
|
||||
assert len(all_text) > 0
|
||||
|
||||
Reference in New Issue
Block a user