Skip to content

[serve] Update tool call to switch to parse_response #45485

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SunMarc merged 4 commits into from
Apr 20, 2026
+688 −584
Merged

[serve] Update tool call to switch to parse_response #45485

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d562e04
fix tool call to support parse response
SunMarc Apr 16, 2026
ce495c2
simplify parser
SunMarc Apr 16, 2026
a384446
simplification !
SunMarc Apr 17, 2026
44c1771
Update src/transformers/cli/serving/response.py
SunMarc Apr 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 42 additions & 38 deletions src/transformers/cli/serving/chat_completion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,9 @@
BaseGenerateManager,
BaseHandler,
Modality,
ToolCallParser,
_StreamError,
detect_tool_format,
get_tool_call_config,
parse_tool_calls,
)


Expand Down Expand Up @@ -140,11 +140,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse
if use_cb:
gen_manager.init_cb(model, gen_config)

# Detect tool support for the loaded model
# TODO: after tool_call start token, use constrained generation to:
# 1. force generation to pick from the available tool names
# 2. force generation to produce valid JSON matching the tool's parameter schema
tool_format = detect_tool_format(model) if body.get("tools") else None
tool_config = get_tool_call_config(processor, model) if body.get("tools") else None

streaming = body.get("stream")
if streaming:
Expand All @@ -156,7 +152,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse
inputs,
gen_config,
gen_manager=gen_manager,
tool_format=tool_format,
tool_config=tool_config,
)
else:
return await self._non_streaming(
Expand All @@ -167,7 +163,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse
inputs,
gen_config,
gen_manager=gen_manager,
tool_format=tool_format,
tool_config=tool_config,
)

# ----- streaming -----
Expand All @@ -181,17 +177,22 @@ def _streaming(
inputs: dict,
gen_config: "GenerationConfig",
gen_manager: BaseGenerateManager,
tool_format: dict | None = None,
tool_config: dict | None = None,
) -> StreamingResponse:
"""Stream tokens as SSE via DirectStreamer."""
queue, streamer = gen_manager.generate_streaming(model, processor, inputs, gen_config, request_id=request_id)
queue, streamer = gen_manager.generate_streaming(
model,
processor,
inputs,
gen_config,
request_id=request_id,
tool_config=tool_config,
)
input_ids = inputs["input_ids"]
# CB returns plain lists, regular path returns tensors
input_len = len(input_ids) if isinstance(input_ids, list) else input_ids.shape[-1]
parser = ToolCallParser(tool_format) if tool_format else None

async def sse_gen() -> AsyncGenerator[str, None]:
has_tool_calls = False
try:
yield self._build_chunk_sse(request_id, role="assistant", model=model_id)

Expand All @@ -215,28 +216,32 @@ async def sse_gen() -> AsyncGenerator[str, None]:
yield "".join(sse_parts)
return

# Tool call parsing: None = normal text, CONSUMED = buffering, else = tool call dict
chunk_kwargs = {"content": text}
if parser is not None and (result := parser.feed(text)) is not None:
if result is ToolCallParser.CONSUMED:
continue
has_tool_calls = True
chunk_kwargs = {
"tool_calls": [
ChoiceDeltaToolCall(
index=0,
type="function",
id=f"{request_id}_tool_call",
function={"name": result["name"], "arguments": result["arguments"]},
)
]
}

sse_parts.append(self._build_chunk_sse(request_id, model=model_id, **chunk_kwargs))
sse_parts.append(self._build_chunk_sse(request_id, model=model_id, content=text))

if sse_parts:
yield "".join(sse_parts)

# Tool calls are parsed after generation completes (not during streaming),
# because the full token sequence is needed for reliable parsing.
has_tool_calls = False
if tool_config:
parsed = parse_tool_calls(processor, streamer.generated_token_ids, tool_config["schema"])
if parsed:
has_tool_calls = True
for i, tc in enumerate(parsed):
yield self._build_chunk_sse(
request_id,
model=model_id,
tool_calls=[
ChoiceDeltaToolCall(
index=i,
type="function",
id=f"{request_id}_tool_call_{i}",
function={"name": tc["name"], "arguments": tc["arguments"]},
)
],
)

hit_max = gen_config.max_new_tokens is not None and streamer.total_tokens >= gen_config.max_new_tokens
if has_tool_calls:
finish_reason = "tool_calls"
Expand Down Expand Up @@ -274,7 +279,7 @@ async def _non_streaming(
inputs: dict,
gen_config: "GenerationConfig",
gen_manager: BaseGenerateManager,
tool_format: dict | None = None,
tool_config: dict | None = None,
) -> JSONResponse:
"""Run generation and return a JSONResponse."""
content, input_len, generated_ids = await gen_manager.generate_non_streaming(
Expand All @@ -289,18 +294,17 @@ async def _non_streaming(
total_tokens=input_len + completion_tokens,
)

# Parse tool calls from the generated text
tool_calls = None
if tool_format is not None:
parsed = ToolCallParser.parse(content, tool_format)
if parsed is not None:
if tool_config is not None:
parsed = parse_tool_calls(processor, generated_ids, tool_config["schema"])
if parsed:
tool_calls = [
ChatCompletionMessageToolCall(
id=f"{request_id}_tool_call",
id=f"{request_id}_tool_call_{i}",
Comment thread
SunMarc marked this conversation as resolved.
type="function",
function={"name": tc["name"], "arguments": tc["arguments"]},
)
for tc in parsed
for i, tc in enumerate(parsed)
]

if tool_calls is not None:
Expand Down
Loading
Loading