diff --git a/src/transformers/cli/serving/chat_completion.py b/src/transformers/cli/serving/chat_completion.py index e39765020bdb..161a25a02f41 100644 --- a/src/transformers/cli/serving/chat_completion.py +++ b/src/transformers/cli/serving/chat_completion.py @@ -40,9 +40,9 @@ BaseGenerateManager, BaseHandler, Modality, - ToolCallParser, _StreamError, - detect_tool_format, + get_tool_call_config, + parse_tool_calls, ) @@ -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: @@ -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( @@ -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 ----- @@ -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) @@ -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" @@ -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( @@ -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}", 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: diff --git a/src/transformers/cli/serving/response.py b/src/transformers/cli/serving/response.py index 2ab2d7291d04..4d29dfd1d6a2 100644 --- a/src/transformers/cli/serving/response.py +++ b/src/transformers/cli/serving/response.py @@ -56,9 +56,9 @@ BaseGenerateManager, BaseHandler, Modality, - ToolCallParser, _StreamError, - detect_tool_format, + get_tool_call_config, + parse_tool_calls, ) @@ -118,7 +118,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse # Two-step input conversion (chat completions skips step 1 since messages are already standard): # 1. Normalize Responses API input (string/list/dict + instructions) → standard messages list # 2. Transform message content for the HF processor (VLM image handling, text joining, etc.) - messages = self._input_to_messages(body) + messages = self._normalize_input(body) processor_inputs = self.get_processor_inputs_from_messages(messages, modality) has_video = any( @@ -131,10 +131,12 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse chat_template_kwargs = {} if has_video: chat_template_kwargs["num_frames"] = 32 + # updates the flat tool structure to the one expected by the `apply_chat_template` method. + tools = self._normalize_tools(body.get("tools")) inputs = processor.apply_chat_template( processor_inputs, add_generation_prompt=True, - tools=body.get("tools"), + tools=tools, return_tensors=None if use_cb else "pt", return_dict=True, tokenize=True, @@ -148,7 +150,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse # TODO: remove when CB supports per-request generation config if use_cb: gen_manager.init_cb(model, gen_config) - 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", True) if streaming: @@ -161,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, ) else: return await self._non_streaming( @@ -173,58 +175,111 @@ 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, ) # ----- input conversion ----- @staticmethod - def _input_to_messages(body: dict) -> list[dict]: - """Convert the Responses API ``input`` field to a list of chat messages. + def _normalize_tools(tools: list[dict] | None) -> list[dict] | None: + """Normalize Responses API tool definitions for ``apply_chat_template``. - The Responses API ``input`` field accepts several formats. This method normalizes - all of them into a standard list of messages with ``role`` and ``content`` keys. - - Supported input formats: - 1. **String**: ``input="Hello"`` → ``[{"role": "user", "content": "Hello"}]`` - 2. **Flat content list** (Responses API native, no ``role`` key): - ``input=[{"type": "input_text", "text": "..."}, {"type": "input_image", ...}]`` - → wrapped as a single user message. - 3. **Messages list** (multi-turn, with ``role`` keys): - ``input=[{"role": "user", "content": [...]}, {"role": "assistant", ...}]`` - → passed through as-is. - - If ``instructions`` is provided, it is prepended as a system message (or replaces - an existing one). - - Args: - body (`dict`): The raw request body containing ``input`` and optionally ``instructions``. + The Responses API uses a flat format: ``{"type": "function", "name": ..., "parameters": ...}`` + while ``apply_chat_template`` expects a nested format: + ``{"type": "function", "function": {"name": ..., "parameters": ...}}``. + Already-nested tools are passed through unchanged. + """ + if not tools: + return tools + return [ + {"type": "function", "function": {k: v for k, v in t.items() if k != "type"}} if "function" not in t else t + for t in tools + ] - Returns: - `list[dict]`: Standard chat messages with ``role`` and ``content`` keys. + @staticmethod + def _normalize_input(body: dict) -> list[dict]: + """Normalize the Responses API ``input`` field into chat messages. + + The Responses API accepts multiple input formats. This method converts them + into a structure close to what ``apply_chat_template`` expects (messages with + ``role``, ``content``, ``tool_calls``, ``tool_call_id``). Further processing + is done by ``get_processor_inputs_from_messages``. + + NOTE: if this conversion logic grows too complex, consider having separate + ``get_processor_inputs_from_messages`` implementations for chat completions + and the Responses API instead of funneling both through the same path. + + Formats handled: + - **String** → single user message. + - **Flat content list** (``input_text``, ``input_image``, no ``role``) → user message. + - **Multi-turn list** — messages and tool call items (``function_call``, + ``function_call_output``) from a previous response, converted via + :meth:`_normalize_response_items`. + + If ``instructions`` is present, it is prepended as a system message. """ inp = body["input"] instructions = body.get("instructions") if isinstance(inp, str): - messages = [{"role": "system", "content": instructions}] if instructions else [] - messages.append({"role": "user", "content": inp}) + messages = [{"role": "user", "content": inp}] elif isinstance(inp, list): - # Flat content list (no "role" key) — wrap as a single user message - if inp and "type" in inp[0] and "role" not in inp[0]: - messages = [{"role": "system", "content": instructions}] if instructions else [] - messages.append({"role": "user", "content": inp}) - elif instructions: - if inp[0]["role"] != "system": - messages = [{"role": "system", "content": instructions}, *inp] - else: - messages = list(inp) - messages[0]["content"] = instructions + if inp and "role" not in inp[0]: + # Flat content list (single-turn, e.g. input_text/input_image) + messages = [{"role": "user", "content": inp}] else: - messages = inp + messages = ResponseHandler._normalize_response_items(inp) else: raise HTTPException(status_code=422, detail="'input' must be a string or list") + # Prepend instructions as a system message + if instructions: + if messages and messages[0]["role"] == "system": + messages[0]["content"] = instructions + else: + messages.insert(0, {"role": "system", "content": instructions}) + + return messages + + @staticmethod + def _normalize_response_items(items: list[dict]) -> list[dict]: + """Convert a list of Responses API items into chat messages. + + Input items may be a mix of: + - Messages (``EasyInputMessageParam`` with ``role``, or ``type: "message"``). + - ``function_call`` — merged as ``tool_calls`` onto the preceding assistant message. + - ``function_call_output`` — converted to ``role: "tool"`` messages. + """ + messages = [] + + for item in items: + item_type = item.get("type") + + if "role" in item: + messages.append({"role": item["role"], "content": item.get("content", "")}) + + elif item_type == "function_call": + tc = { + "id": item["call_id"], + "function": {"name": item["name"], "arguments": item["arguments"]}, + } + if messages and messages[-1]["role"] == "assistant": + messages[-1].setdefault("tool_calls", []).append(tc) + else: + messages.append({"role": "assistant", "tool_calls": [tc]}) + + elif item_type == "function_call_output": + messages.append( + { + "role": "tool", + "tool_call_id": item["call_id"], + "content": item["output"], + } + ) + + else: + raise HTTPException(status_code=422, detail=f"Unsupported input item type: {item_type!r}") + return messages # ----- streaming ----- @@ -239,14 +294,20 @@ def _streaming( inputs: dict, gen_config: "GenerationConfig", gen_manager: BaseGenerateManager, - tool_format: dict | None = None, + tool_config: dict | None = None, ) -> StreamingResponse: """Generate a streaming Responses API reply (SSE) using 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 seq = 0 output_index = 0 @@ -362,59 +423,6 @@ async def event_stream() -> AsyncGenerator[str, None]: yield "".join(sse_parts) return - # Tool call parsing - if parser is not None and (result := parser.feed(text)) is not None: - if result is not ToolCallParser.CONSUMED: - tc_id = f"{request_id}_tool_call" - name = result["name"] - arguments = result["arguments"] - tc_item = ResponseFunctionToolCall( - id=tc_id, - call_id=tc_id, - type="function_call", - name=name, - arguments=arguments, - status="completed", - ) - tool_calls.append(tc_item) - output_index += 1 - sse_parts.append( - self.chunk_to_sse( - ResponseOutputItemAddedEvent( - type="response.output_item.added", - sequence_number=seq, - output_index=output_index, - item=tc_item, - ) - ) - ) - seq += 1 - sse_parts.append( - self.chunk_to_sse( - ResponseFunctionCallArgumentsDoneEvent( - type="response.function_call_arguments.done", - sequence_number=seq, - item_id=tc_id, - output_index=output_index, - arguments=arguments, - name=name, - ) - ) - ) - seq += 1 - sse_parts.append( - self.chunk_to_sse( - ResponseOutputItemDoneEvent( - type="response.output_item.done", - sequence_number=seq, - output_index=output_index, - item=tc_item, - ) - ) - ) - seq += 1 - continue - full_text += text sse_parts.append( self.chunk_to_sse( @@ -434,7 +442,54 @@ async def event_stream() -> AsyncGenerator[str, None]: if sse_parts: yield "".join(sse_parts) - # 5. Close text output + # 5. Tool calls are parsed after generation completes (not during streaming), + # because the full token sequence is needed for reliable parsing. + if tool_config: + parsed = parse_tool_calls(processor, streamer.generated_token_ids, tool_config["schema"]) + if parsed: + for i, tc in enumerate(parsed): + tc_id = f"{request_id}_tool_call_{i}" + tc_item = ResponseFunctionToolCall( + id=tc_id, + call_id=tc_id, + type="function_call", + name=tc["name"], + arguments=tc["arguments"], + status="completed", + ) + tool_calls.append(tc_item) + output_index += 1 + yield self.chunk_to_sse( + ResponseOutputItemAddedEvent( + type="response.output_item.added", + sequence_number=seq, + output_index=output_index, + item=tc_item, + ) + ) + seq += 1 + yield self.chunk_to_sse( + ResponseFunctionCallArgumentsDoneEvent( + type="response.function_call_arguments.done", + sequence_number=seq, + item_id=tc_id, + output_index=output_index, + arguments=tc["arguments"], + name=tc["name"], + ) + ) + seq += 1 + yield self.chunk_to_sse( + ResponseOutputItemDoneEvent( + type="response.output_item.done", + sequence_number=seq, + output_index=output_index, + item=tc_item, + ) + ) + seq += 1 + + # 6. Close text output output_text_part = ResponseOutputText(type="output_text", text=full_text, annotations=[]) yield self.chunk_to_sse( ResponseTextDoneEvent( @@ -478,7 +533,7 @@ async def event_stream() -> AsyncGenerator[str, None]: ) seq += 1 - # 6. Completed + # 7. Completed all_output = [msg_item] + list(tool_calls) usage = compute_usage(input_len, streamer.total_tokens) yield self.chunk_to_sse( @@ -509,7 +564,7 @@ async def _non_streaming( inputs: dict, gen_config: "GenerationConfig", gen_manager: BaseGenerateManager, - tool_format: dict | None = None, + tool_config: dict | None = None, ) -> JSONResponse: """Generate a non-streaming Responses API reply (single JSON).""" full_text, input_len, generated_ids = await gen_manager.generate_non_streaming( @@ -527,12 +582,11 @@ async def _non_streaming( ) ] - # Parse tool calls from the generated text - if tool_format is not None: - parsed_calls = ToolCallParser.parse(full_text, tool_format) - if parsed_calls is not None: - for i, tc in enumerate(parsed_calls): - tc_id = f"{request_id}_tool_call" + if tool_config is not None: + parsed = parse_tool_calls(processor, generated_ids, tool_config["schema"]) + if parsed: + for i, tc in enumerate(parsed): + tc_id = f"{request_id}_tool_call_{i}" output_items.append( ResponseFunctionToolCall( id=tc_id, diff --git a/src/transformers/cli/serving/utils.py b/src/transformers/cli/serving/utils.py index 3ef55d5f0041..3a21cc7093f6 100644 --- a/src/transformers/cli/serving/utils.py +++ b/src/transformers/cli/serving/utils.py @@ -73,132 +73,87 @@ class _GenerationCancelled(Exception): """Raised inside ``DirectStreamer.put()`` to abort ``model.generate()``.""" -# Model-specific tokens that mark the start/end of a tool call block. -# TODO: extract these from the chat template at runtime instead of hardcoding. -# Qwen/Hermes use /, Mistral uses [TOOL_CALLS], etc. -# The markers are defined in each model's Jinja chat template. -_TOOL_CALL_TOKENS = { +# Fallback tool call configs for models that don't declare stc_token/etc_token/response_schema +# on their tokenizer. +# Keys are matched via substring against model_type (e.g. "qwen" matches "qwen2", "qwen3_vl", etc.). +# If a model family changes its tool call format, split into separate keys (e.g. "qwen2", "qwen3"). +_TOOL_CALL_FALLBACKS = { "qwen": { - "start": "", - "end": "", + "stc": "", + "etc": "", + "schema": { + "x-regex-iterator": r"(.*?)", + "type": "array", + "items": {"type": "object", "x-parser": "json"}, + }, }, } -def detect_tool_format(model: "PreTrainedModel") -> dict | None: - """Return the tool call token format for a model, if supported. +def get_tool_call_config(processor, model: "PreTrainedModel") -> dict | None: + """Return tool call config for the model, or ``None`` if tool calls are not supported. - Args: - model (`PreTrainedModel`): The loaded model. - - Returns: - `dict | None`: A dict ``{"start": str, "end": str}`` with the model's tool call - delimiters, or ``None`` if the model family is not recognized. + Returns a dict with: + - ``schema`` (`dict`): Schema to pass to ``tokenizer.parse_response(block, schema)``. + - ``stc_id`` (`int`): Token ID of the start-of-tool-call delimiter. + - ``etc_id`` (`int`): Token ID of the end-of-tool-call delimiter. """ - architecture = model.config.architectures[0].lower() - for family in _TOOL_CALL_TOKENS: - if family in architecture: - return _TOOL_CALL_TOKENS[family] - return None - - -class ToolCallParser: - """Parses tool calls from model output. + tokenizer = getattr(processor, "tokenizer", processor) + stc = getattr(tokenizer, "stc_token", None) + etc = getattr(tokenizer, "etc_token", None) + response_schema = getattr(tokenizer, "response_schema", None) + + # Models with full tokenizer config (e.g. Gemma 4) + if stc and etc and response_schema: + schema = response_schema["properties"]["tool_calls"] + else: + # Fallback: known model families without full tokenizer config + fallback = next((v for k, v in _TOOL_CALL_FALLBACKS.items() if k in model.config.model_type), None) + if fallback is None: + return None + stc, etc, schema = fallback["stc"], fallback["etc"], fallback["schema"] - The model emits tool calls as structured text between start/end tokens - (e.g. ``{"name": "fn", "arguments": {...}}``). + stc_id = tokenizer.convert_tokens_to_ids(stc) + etc_id = tokenizer.convert_tokens_to_ids(etc) + return {"schema": schema, "stc_id": stc_id, "etc_id": etc_id} - **Streaming** (``feed``): buffers tokens between start/end markers, parses - the complete block when the end marker is seen, returns a ``ChoiceDeltaToolCall``. - **Non-streaming** (``parse``): extracts all tool call blocks from complete text. +def _normalize_tool_call(tool_call: dict) -> dict: + """Normalize a parsed tool call to ``{"name": str, "arguments": str}``. - Usage:: + Different models return different structures from ``parse_response``: + - Gemma: ``{"function": {"name": ..., "arguments": {...}}}`` (nested, arguments as dict) + - Qwen: ``{"name": ..., "arguments": {...}}`` (flat, arguments as dict) - parser = ToolCallParser(tool_format={"start": ..., "end": ...}) - for text_chunk in streamer: - result = parser.feed(text_chunk) - if result is None: - # Normal text — emit as content - elif result is ToolCallParser.CONSUMED: - # Buffering — skip - else: - # result is a ChoiceDeltaToolCall — emit it + The OpenAI API expects ``arguments`` as a JSON **string**, so we ``json.dumps`` it. """ + function = tool_call.get("function", tool_call) + arguments = function.get("arguments", {}) + return { + "name": function["name"], + "arguments": json.dumps(arguments) if not isinstance(arguments, str) else arguments, + } - def __init__(self, tool_format: dict): - self._tokens = tool_format - self._inside = False - self._buffer = "" - - # Sentinel: token was consumed by the parser but produced no output. - CONSUMED = object() - def feed(self, text: str) -> object | dict | None: - """Feed a text chunk (streaming). +def parse_tool_calls(processor, generated_ids, schema: dict) -> list[dict] | None: + """Parse tool calls from generated token IDs using ``tokenizer.parse_response``. - Returns: - - ``None`` — normal text, not a tool token. Emit as content. - - ``CONSUMED`` — token consumed internally (buffering/markers). Skip. - - A ``ChoiceDeltaToolCall`` — emit as a tool call delta. - """ - if text.strip() == self._tokens["start"]: - self._inside = True - self._buffer = "" - return self.CONSUMED - - if text.strip() == self._tokens["end"]: - self._inside = False - block = self._buffer.strip() - self._buffer = "" - return self._parse_block(block) or self.CONSUMED - - if self._inside: - self._buffer += text - return self.CONSUMED + Args: + processor: The processor or tokenizer. + generated_ids: Token IDs from generation. Passed directly to ``parse_response`` + which decodes them internally, preserving special tokens that + ``skip_special_tokens=True`` would strip (e.g. Gemma's ``<|tool_call>``). + schema: The tool call schema (from ``response_schema`` or ``_TOOL_CALL_FALLBACKS``). + Returns a list of ``{"name": str, "arguments": str}`` dicts, or ``None`` if none found. + """ + parsed = processor.parse_response(generated_ids, schema) + if not parsed: return None - - @staticmethod - def _extract_name_and_args(block: str) -> tuple[str, str] | None: - """Extract (name, arguments_json) from a tool call block, or None if invalid.""" - if not block: - return None - parsed = json.loads(block) - name = parsed.get("name") - if name is None: - return None - arguments = parsed.get("arguments", {}) - return name, json.dumps(arguments) - - @staticmethod - def parse(text: str, tool_format: dict) -> list[dict] | None: - """Parse tool calls from complete text. - - Returns a list of ``{"name": str, "arguments": str}`` dicts, or ``None`` if none found. - """ - start, end = tool_format["start"], tool_format["end"] - tool_calls = [] - pos = 0 - while True: - s = text.find(start, pos) - if s < 0: - break - e = text.find(end, s + len(start)) - if e < 0: - break - result = ToolCallParser._extract_name_and_args(text[s + len(start) : e].strip()) - if result is not None: - tool_calls.append({"name": result[0], "arguments": result[1]}) - pos = e + len(end) - return tool_calls if tool_calls else None - - def _parse_block(self, block: str) -> dict | None: - """Parse a buffered tool call block. Returns ``{"name": str, "arguments": str}`` or None.""" - result = self._extract_name_and_args(block) - if result is None: - return None - return {"name": result[0], "arguments": result[1]} + if not isinstance(parsed, list): + parsed = [parsed] + tool_calls = [_normalize_tool_call(tool_call) for tool_call in parsed] + return tool_calls if tool_calls else None class DownloadAggregator: @@ -330,6 +285,7 @@ def __init__( loop: asyncio.AbstractEventLoop, queue: asyncio.Queue, skip_special_tokens: bool = True, + tool_config: dict | None = None, ): """ Args: @@ -338,6 +294,9 @@ def __init__( queue (`asyncio.Queue`): The queue that receives decoded text chunks. skip_special_tokens (`bool`, *optional*, defaults to `True`): Whether to strip special tokens during decoding. + tool_config (`dict`, *optional*): Tool call config from ``get_tool_call_config``. + When set, tokens between stc/etc delimiters (inclusive) are suppressed + from the queue so tool call markup is never streamed to the client. """ from tokenizers.decoders import DecodeStream @@ -345,9 +304,13 @@ def __init__( self._loop = loop self._queue = queue self._decode_stream = DecodeStream([], skip_special_tokens) + self._stc_id = tool_config["stc_id"] if tool_config else None + self._etc_id = tool_config["etc_id"] if tool_config else None + self._inside_tool_call = False self._first = True self._cancelled = threading.Event() self.total_tokens = 0 + self.generated_token_ids: list[int] = [] def put(self, value: "torch.Tensor") -> None: """Called by ``model.generate()`` after each decode step with new token(s).""" @@ -359,8 +322,15 @@ def put(self, value: "torch.Tensor") -> None: return for token_id in value.tolist(): self.total_tokens += 1 + self.generated_token_ids.append(token_id) + + if token_id == self._stc_id: + self._inside_tool_call = True + elif token_id == self._etc_id: + self._inside_tool_call = False + text = self._decode_stream.step(self._tokenizer, token_id) - if text is not None: + if text is not None and not self._inside_tool_call and token_id != self._etc_id: self._loop.call_soon_threadsafe(self._queue.put_nowait, text) def end(self) -> None: @@ -388,6 +358,7 @@ def __init__( tokenizer: "tokenizers.Tokenizer", loop: asyncio.AbstractEventLoop, queue: asyncio.Queue, + tool_config: dict | None = None, ): """ Args: @@ -396,6 +367,7 @@ def __init__( tokenizer: The Rust tokenizer (``tokenizer._tokenizer``). loop (`asyncio.AbstractEventLoop`): The event loop to push decoded text to. queue (`asyncio.Queue`): The queue that receives decoded text chunks. + tool_config (`dict`, *optional*): Tool call config (see ``DirectStreamer``). """ from tokenizers.decoders import DecodeStream @@ -405,8 +377,12 @@ def __init__( self._queue = queue self._tokenizer = tokenizer self._decode_stream = DecodeStream([], True) + self._stc_id = tool_config["stc_id"] if tool_config else None + self._etc_id = tool_config["etc_id"] if tool_config else None + self._inside_tool_call = False self._prev_len = 0 self.total_tokens = 0 + self.generated_token_ids: list[int] = [] def put(self, output: "GenerationOutput") -> None: """Decode new tokens from a CB ``GenerationOutput`` and push text to the queue.""" @@ -414,8 +390,15 @@ def put(self, output: "GenerationOutput") -> None: self._prev_len = len(output.generated_tokens) for token_id in new_tokens: self.total_tokens += 1 + self.generated_token_ids.append(token_id) + + if token_id == self._stc_id: + self._inside_tool_call = True + elif token_id == self._etc_id: + self._inside_tool_call = False + text = self._decode_stream.step(self._tokenizer, token_id) - if text is not None: + if text is not None and not self._inside_tool_call and token_id != self._etc_id: self._queue.put_nowait(text) def end(self) -> None: @@ -502,6 +485,7 @@ def generate_streaming( inputs: dict, gen_config: "GenerationConfig", request_id: str, + tool_config: dict | None = None, ) -> tuple[asyncio.Queue, "DirectStreamer | CBStreamer"]: """Start streaming generation. @@ -511,6 +495,8 @@ def generate_streaming( inputs (`dict`): Tokenized inputs (tensors for sequential, lists for CB). gen_config (`GenerationConfig`): Generation parameters. request_id (`str`): Unique request identifier. + tool_config (`dict`, *optional*): Tool call config from ``get_tool_call_config``. + When set, tool call tokens (between stc/etc) are suppressed from output. Returns: `tuple[asyncio.Queue, DirectStreamer | CBStreamer]`: A ``(queue, streamer)`` pair @@ -558,13 +544,14 @@ def generate_streaming( inputs: dict, gen_config: "GenerationConfig", request_id: str, + tool_config: dict | None = None, ) -> tuple[asyncio.Queue, DirectStreamer]: """Start streaming generation via ``model.generate()`` on the inference thread.""" loop = asyncio.get_running_loop() queue: asyncio.Queue = asyncio.Queue() # ProcessorMixin exposes the fast tokenizer as .tokenizer; PreTrainedTokenizerFast is already one. rust_tokenizer = getattr(processor, "tokenizer", processor)._tokenizer # type: ignore[union-attr] - streamer = DirectStreamer(rust_tokenizer, loop, queue, skip_special_tokens=True) + streamer = DirectStreamer(rust_tokenizer, loop, queue, tool_config=tool_config) gen_kwargs = {**inputs, "streamer": streamer, "generation_config": gen_config, "tokenizer": processor} if hasattr(model, "has_talker"): gen_kwargs["generation_mode"] = "text" @@ -655,6 +642,7 @@ def generate_streaming( inputs: dict, gen_config: "GenerationConfig", request_id: str, + tool_config: dict | None = None, ) -> tuple[asyncio.Queue, CBStreamer]: """Start streaming CB generation. Registers a per-request output handler.""" cb = self._cb @@ -674,7 +662,7 @@ def generate_streaming( ) # ProcessorMixin exposes the fast tokenizer as .tokenizer; PreTrainedTokenizerFast is already one. rust_tokenizer = getattr(processor, "tokenizer", processor)._tokenizer # type: ignore[union-attr] - streamer = CBStreamer(self._cb, request_id, rust_tokenizer, loop, text_queue) + streamer = CBStreamer(self._cb, request_id, rust_tokenizer, loop, text_queue, tool_config=tool_config) # Register a direct callback: the dispatcher calls this on the event loop with each GenerationOutput. # This decodes tokens and pushes text straight to the SSE text_queue @@ -951,14 +939,16 @@ def get_processor_inputs_from_messages(messages: list[dict], modality: Modality) if "tool_call_id" in message: parsed["tool_call_id"] = message["tool_call_id"] - raw_content = message.get("content", []) + # When tool_calls are present, ignore content — it's either empty or contains + # raw tool call markup that would confuse the chat template if rendered. + raw_content = [] if "tool_calls" in message else (message.get("content") or []) if isinstance(raw_content, str): raw_content = [{"type": "text", "text": raw_content}] for content in raw_content: content_type = content["type"] # Text: chat completions ("text") and Responses API ("input_text") - if content_type in ("text", "input_text"): + if content_type in ("text", "input_text", "output_text"): parsed["content"].append({"type": "text", "text": content["text"]}) # Image: chat completions ("image_url") and Responses API ("input_image") elif content_type in ("image_url", "input_image") and modality in (Modality.VLM, Modality.MULTIMODAL): diff --git a/tests/cli/test_serve.py b/tests/cli/test_serve.py index c7f9b96f4790..f9e61681c056 100644 --- a/tests/cli/test_serve.py +++ b/tests/cli/test_serve.py @@ -32,11 +32,12 @@ from transformers.cli.serving.server import build_server from transformers.cli.serving.transcription import TranscriptionHandler from transformers.cli.serving.utils import ( + _TOOL_CALL_FALLBACKS, BaseHandler, GenerationState, Modality, - ToolCallParser, - detect_tool_format, + get_tool_call_config, + parse_tool_calls, ) from transformers.testing_utils import ( require_librosa, @@ -47,6 +48,7 @@ require_vision, slow, ) +from transformers.utils.chat_parsing_utils import recursive_parse from transformers.utils.import_utils import is_serve_available @@ -504,130 +506,7 @@ def test_chunk_to_sse_wraps_plain_string(self): self.assertEqual(result, "data: hello\n\n") -QWEN_TOOL_FORMAT = {"start": "", "end": ""} - - @require_serve -class TestToolParser(unittest.TestCase): - def test_detect_tool_format_qwen(self): - model = MagicMock() - model.config.architectures = ["Qwen2ForCausalLM"] - fmt = detect_tool_format(model) - self.assertEqual(fmt, QWEN_TOOL_FORMAT) - - def test_detect_tool_format_unsupported(self): - model = MagicMock() - model.config.architectures = ["LlamaForCausalLM"] - self.assertIsNone(detect_tool_format(model)) - - def test_parser_start_token(self): - parser = ToolCallParser(QWEN_TOOL_FORMAT) - result = parser.feed("") - self.assertIs(result, ToolCallParser.CONSUMED) - - def test_parser_end_token(self): - parser = ToolCallParser(QWEN_TOOL_FORMAT) - parser.feed("") - result = parser.feed("") - self.assertIs(result, ToolCallParser.CONSUMED) - - def test_parser_buffers_until_end(self): - parser = ToolCallParser(QWEN_TOOL_FORMAT) - parser.feed("") - # Intermediate tokens are buffered - result = parser.feed('{"name": "my_tool", "arguments": {"x": 1}}') - self.assertIs(result, ToolCallParser.CONSUMED) - # Tool call is emitted on end token - result = parser.feed("") - self.assertIsNot(result, ToolCallParser.CONSUMED) - self.assertEqual(result["name"], "my_tool") - - def test_parser_normal_text_returns_none(self): - parser = ToolCallParser(QWEN_TOOL_FORMAT) - result = parser.feed("Hello world") - self.assertIsNone(result) - - def test_parser_full_flow(self): - """Simulate a complete tool call token sequence.""" - - parser = ToolCallParser(QWEN_TOOL_FORMAT) - tool_calls = [] - - for token in [ - "", - '{"name": "get_weather",', - ' "arguments": {', - '"city": "Paris"', - "}}", - "\n", - "", - ]: - result = parser.feed(token) - if result is not None and result is not ToolCallParser.CONSUMED: - tool_calls.append(result) - - # Single tool call emitted on with both name and arguments - self.assertEqual(len(tool_calls), 1) - self.assertEqual(tool_calls[0]["name"], "get_weather") - self.assertIn("Paris", tool_calls[0]["arguments"]) - - def test_parse_tool_calls_from_text(self): - """Non-streaming tool call parsing from complete text.""" - - text = '\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n' - calls = ToolCallParser.parse(text, QWEN_TOOL_FORMAT) - self.assertIsNotNone(calls) - self.assertEqual(len(calls), 1) - self.assertEqual(calls[0]["name"], "get_weather") - self.assertIn("Paris", calls[0]["arguments"]) - - def test_parse_tool_calls_no_tool_call(self): - """Non-streaming: normal text returns None.""" - - calls = ToolCallParser.parse("Hello, how can I help?", QWEN_TOOL_FORMAT) - self.assertIsNone(calls) - - def test_parse_multiple_tool_calls(self): - """Non-streaming: multiple tool calls in one response.""" - - text = ( - '\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n\n' - '\n{"name": "get_weather", "arguments": {"city": "London"}}\n' - ) - calls = ToolCallParser.parse(text, QWEN_TOOL_FORMAT) - self.assertIsNotNone(calls) - self.assertEqual(len(calls), 2) - self.assertEqual(calls[0]["name"], "get_weather") - self.assertIn("Paris", calls[0]["arguments"]) - self.assertEqual(calls[1]["name"], "get_weather") - self.assertIn("London", calls[1]["arguments"]) - - def test_feed_multiple_tool_calls(self): - """Streaming: multiple tool calls emitted sequentially.""" - - parser = ToolCallParser(QWEN_TOOL_FORMAT) - tool_calls = [] - - tokens = [ - "", - '{"name": "get_weather", "arguments": {"city": "Paris"}}', - "", - "", - '{"name": "get_weather", "arguments": {"city": "London"}}', - "", - ] - for token in tokens: - result = parser.feed(token) - if result is not None and result is not ToolCallParser.CONSUMED: - tool_calls.append(result) - - self.assertEqual(len(tool_calls), 2) - self.assertEqual(tool_calls[0]["name"], "get_weather") - self.assertIn("Paris", tool_calls[0]["arguments"]) - self.assertEqual(tool_calls[1]["name"], "get_weather") - self.assertIn("London", tool_calls[1]["arguments"]) - - @require_serve class TestAppRoutes(unittest.TestCase): @classmethod @@ -784,110 +663,6 @@ def test_streaming_usage(self): self.assertGreater(last.usage.completion_tokens, 0) self.assertEqual(last.usage.total_tokens, last.usage.prompt_tokens + last.usage.completion_tokens) - def test_tool_call(self): - """Tool calls should be parsed and emitted as ChoiceDeltaToolCall objects.""" - # Qwen2.5-0.5B-Instruct supports tools (Qwen family) - tool_def = { - "function": { - "name": "get_weather", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - "description": "Get the weather for a city.", - }, - "type": "function", - } - chunks = list( - self.client.chat.completions.create( - model=self.MODEL, - messages=[{"role": "user", "content": "What is the weather in Paris?"}], - stream=True, - max_tokens=50, - temperature=0.0, - tools=[tool_def], - ) - ) - - # First chunk should have role="assistant" - self.assertEqual(chunks[0].choices[0].delta.role, "assistant") - - # Model should make a tool call for this prompt - tool_chunks = [c for c in chunks if c.choices[0].delta.tool_calls] - self.assertGreater(len(tool_chunks), 0, "Model did not produce a tool call") - - # First tool call delta should have the function name - first_tool = tool_chunks[0].choices[0].delta.tool_calls[0] - self.assertEqual(first_tool.function.name, "get_weather") - - # finish_reason should be "tool_calls" - last = chunks[-1] - self.assertEqual(last.choices[0].finish_reason, "tool_calls") - - # Arguments should be valid JSON with no trailing brace - args_json = first_tool.function.arguments - import json as json_mod - - parsed_args = json_mod.loads(args_json) - self.assertIsInstance(parsed_args, dict) - - def test_tool_call_non_streaming(self): - """Non-streaming tool calls should return tool_calls in the message.""" - tool_def = { - "function": { - "name": "get_weather", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, - "description": "Get the weather for a city.", - }, - "type": "function", - } - resp = self.client.chat.completions.create( - model=self.MODEL, - messages=[{"role": "user", "content": "What is the weather in Paris?"}], - stream=False, - max_tokens=50, - temperature=0.0, - tools=[tool_def], - ) - self.assertEqual(resp.choices[0].finish_reason, "tool_calls") - self.assertIsNotNone(resp.choices[0].message.tool_calls) - tc = resp.choices[0].message.tool_calls[0] - self.assertEqual(tc.function.name, "get_weather") - - import json as json_mod - - parsed_args = json_mod.loads(tc.function.arguments) - self.assertIsInstance(parsed_args, dict) - - def test_tool_call_multi(self): - """Model should be able to call multiple tools when asked.""" - tool_def = { - "function": { - "name": "get_weather", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, - "description": "Get the weather for a city.", - }, - "type": "function", - } - # Ask for two cities to encourage multiple tool calls - chunks = list( - self.client.chat.completions.create( - model=self.MODEL, - messages=[{"role": "user", "content": "What is the weather in Paris and London?"}], - stream=True, - max_tokens=100, - temperature=0.0, - tools=[tool_def], - ) - ) - tool_chunks = [c for c in chunks if c.choices[0].delta.tool_calls] - # Should have two tool calls — one for Paris, one for London - self.assertEqual(len(tool_chunks), 2, f"Expected 2 tool calls, got {len(tool_chunks)}") - cities = {tc.choices[0].delta.tool_calls[0].function.name for tc in tool_chunks} - self.assertEqual(cities, {"get_weather"}) - last = chunks[-1] - self.assertEqual(last.choices[0].finish_reason, "tool_calls") - def test_concurrent_non_streaming(self): """Two concurrent non-streaming requests should both complete without interference.""" import concurrent.futures @@ -978,19 +753,19 @@ def _make_handler(self): def test_string_input(self): handler = self._make_handler() - msgs = handler._input_to_messages({"input": "Hello"}) + msgs = handler._normalize_input({"input": "Hello"}) self.assertEqual(msgs, [{"role": "user", "content": "Hello"}]) def test_string_input_with_instructions(self): handler = self._make_handler() - msgs = handler._input_to_messages({"input": "Hello", "instructions": "Be brief"}) + msgs = handler._normalize_input({"input": "Hello", "instructions": "Be brief"}) self.assertEqual(len(msgs), 2) self.assertEqual(msgs[0], {"role": "system", "content": "Be brief"}) self.assertEqual(msgs[1], {"role": "user", "content": "Hello"}) def test_list_input(self): handler = self._make_handler() - msgs = handler._input_to_messages( + msgs = handler._normalize_input( {"input": [{"role": "user", "content": "A"}, {"role": "assistant", "content": "B"}]} ) self.assertEqual(len(msgs), 2) @@ -998,14 +773,14 @@ def test_list_input(self): def test_list_input_with_instructions_prepends_system(self): handler = self._make_handler() - msgs = handler._input_to_messages({"input": [{"role": "user", "content": "Hi"}], "instructions": "Be helpful"}) + msgs = handler._normalize_input({"input": [{"role": "user", "content": "Hi"}], "instructions": "Be helpful"}) self.assertEqual(len(msgs), 2) self.assertEqual(msgs[0]["role"], "system") self.assertEqual(msgs[0]["content"], "Be helpful") def test_list_input_with_instructions_replaces_existing_system(self): handler = self._make_handler() - msgs = handler._input_to_messages( + msgs = handler._normalize_input( {"input": [{"role": "system", "content": "Old"}, {"role": "user", "content": "Hi"}], "instructions": "New"} ) self.assertEqual(len(msgs), 2) @@ -1018,7 +793,7 @@ def test_flat_content_list(self): {"type": "input_text", "text": "Hello"}, {"type": "input_image", "image_url": "https://example.com/img.jpg"}, ] - msgs = handler._input_to_messages({"input": flat_input}) + msgs = handler._normalize_input({"input": flat_input}) self.assertEqual(len(msgs), 1) self.assertEqual(msgs[0]["role"], "user") self.assertEqual(msgs[0]["content"], flat_input) @@ -1027,7 +802,7 @@ def test_flat_content_list_with_instructions(self): """Flat content list with instructions prepends a system message.""" handler = self._make_handler() flat_input = [{"type": "input_text", "text": "Hello"}] - msgs = handler._input_to_messages({"input": flat_input, "instructions": "Be brief"}) + msgs = handler._normalize_input({"input": flat_input, "instructions": "Be brief"}) self.assertEqual(len(msgs), 2) self.assertEqual(msgs[0], {"role": "system", "content": "Be brief"}) self.assertEqual(msgs[1]["role"], "user") @@ -1221,96 +996,6 @@ def test_streaming_usage(self): self.assertGreater(usage.output_tokens, 0) self.assertEqual(usage.total_tokens, usage.input_tokens + usage.output_tokens) - def test_tool_call_streaming(self): - """Streaming responses with tools should emit function_call events.""" - tool_def = { - "function": { - "name": "get_weather", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, - "description": "Get the weather for a city.", - }, - "type": "function", - } - events = list( - self.client.responses.create( - model=self.MODEL, - input="What is the weather in Paris?", - stream=True, - max_output_tokens=50, - tools=[tool_def], - ) - ) - types = [e.type for e in events] - self.assertIn("response.created", types) - self.assertIn("response.completed", types) - - # Should have function call events - self.assertIn("response.output_item.added", types) - self.assertIn("response.function_call_arguments.done", types) - - # Check the arguments done event - args_done = [e for e in events if e.type == "response.function_call_arguments.done"] - self.assertGreater(len(args_done), 0) - self.assertEqual(args_done[0].name, "get_weather") - - import json as json_mod - - parsed = json_mod.loads(args_done[0].arguments) - self.assertIsInstance(parsed, dict) - - def test_tool_call_non_streaming(self): - """Non-streaming responses with tools should include function_call output items.""" - tool_def = { - "function": { - "name": "get_weather", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, - "description": "Get the weather for a city.", - }, - "type": "function", - } - resp = self.client.responses.create( - model=self.MODEL, - input="What is the weather in Paris?", - stream=False, - max_output_tokens=50, - tools=[tool_def], - ) - self.assertEqual(resp.status, "completed") - - # Should have at least message + function_call in output - self.assertGreater(len(resp.output), 1) - fc_items = [o for o in resp.output if o.type == "function_call"] - self.assertGreater(len(fc_items), 0) - self.assertEqual(fc_items[0].name, "get_weather") - - import json as json_mod - - parsed = json_mod.loads(fc_items[0].arguments) - self.assertIsInstance(parsed, dict) - - def test_tool_call_multi(self): - """Model should produce multiple tool calls when asked about two cities.""" - tool_def = { - "function": { - "name": "get_weather", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, - "description": "Get the weather for a city.", - }, - "type": "function", - } - events = list( - self.client.responses.create( - model=self.MODEL, - input="What is the weather in Paris and London?", - stream=True, - max_output_tokens=100, - tools=[tool_def], - ) - ) - args_done = [e for e in events if e.type == "response.function_call_arguments.done"] - self.assertEqual(len(args_done), 2, f"Expected 2 tool calls, got {len(args_done)}") - self.assertEqual(events[-1].type, "response.completed") - def test_multi_turn(self): """Multi-turn conversation via list input.""" resp = self.client.responses.create( @@ -1872,6 +1557,377 @@ def test_responses_with_video_streaming(self): self._assert_video_description(text) +class TestToolCallUnit(unittest.TestCase): + """Unit tests for tool call parsing utilities (no server needed).""" + + def test_get_tool_call_config_fallback(self): + """Fallback config is returned for known model families (Qwen).""" + model = MagicMock() + model.config.model_type = "qwen2" + processor = MagicMock(spec=["convert_tokens_to_ids"]) + processor.convert_tokens_to_ids.return_value = 151657 + config = get_tool_call_config(processor, model) + self.assertIsNotNone(config) + self.assertEqual(config["stc_id"], 151657) + self.assertEqual(config["etc_id"], 151657) + + def test_get_tool_call_config_unsupported(self): + """None is returned for models without tool call support.""" + model = MagicMock() + model.config.model_type = "llama" + processor = MagicMock(spec=[]) + self.assertIsNone(get_tool_call_config(processor, model)) + + def test_parse_tool_calls_from_text(self): + text = '\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n' + processor = MagicMock() + processor.parse_response = lambda t, s: recursive_parse(t, s) + calls = parse_tool_calls(processor, text, _TOOL_CALL_FALLBACKS["qwen"]["schema"]) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["name"], "get_weather") + + def test_parse_multiple_tool_calls_from_text(self): + text = ( + '\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n\n' + '\n{"name": "get_weather", "arguments": {"city": "London"}}\n' + ) + processor = MagicMock() + processor.parse_response = lambda t, s: recursive_parse(t, s) + calls = parse_tool_calls(processor, text, _TOOL_CALL_FALLBACKS["qwen"]["schema"]) + self.assertEqual(len(calls), 2) + + +class _TestToolCallBase: + """Base class for tool call integration tests. Subclasses set MODEL and inherit all tests.""" + + MODEL: str + + @classmethod + def setUpClass(cls): + cls.serve, port = _start_serve() + cls.base_url = f"http://localhost:{port}" + cls.client = OpenAI(base_url=f"{cls.base_url}/v1", api_key="unused") + + @classmethod + def tearDownClass(cls): + cls.serve.kill_server() + + def _get_tool_def(self): + return { + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + "description": "Get the weather for a city.", + }, + "type": "function", + } + + def test_chat_non_streaming(self): + resp = self.client.chat.completions.create( + model=self.MODEL, + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + stream=False, + max_tokens=50, + temperature=0.0, + tools=[self._get_tool_def()], + ) + self.assertEqual(resp.choices[0].finish_reason, "tool_calls") + self.assertIsNotNone(resp.choices[0].message.tool_calls) + tc = resp.choices[0].message.tool_calls[0] + self.assertEqual(tc.function.name, "get_weather") + parsed_args = json.loads(tc.function.arguments) + self.assertIsInstance(parsed_args, dict) + + def test_chat_streaming(self): + chunks = list( + self.client.chat.completions.create( + model=self.MODEL, + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + stream=True, + max_tokens=50, + temperature=0.0, + tools=[self._get_tool_def()], + ) + ) + tool_chunks = [c for c in chunks if c.choices[0].delta.tool_calls] + self.assertGreater(len(tool_chunks), 0, "Model did not produce a tool call") + first_tool = tool_chunks[0].choices[0].delta.tool_calls[0] + self.assertEqual(first_tool.function.name, "get_weather") + self.assertEqual(chunks[-1].choices[0].finish_reason, "tool_calls") + parsed_args = json.loads(first_tool.function.arguments) + self.assertIsInstance(parsed_args, dict) + + def test_chat_multiple_tool_calls_non_streaming(self): + resp = self.client.chat.completions.create( + model=self.MODEL, + messages=[{"role": "user", "content": "What is the weather in Paris and London?"}], + stream=False, + max_tokens=100, + temperature=0.0, + tools=[self._get_tool_def()], + ) + self.assertEqual(resp.choices[0].finish_reason, "tool_calls") + self.assertEqual(len(resp.choices[0].message.tool_calls), 2) + + def test_chat_multiple_tool_calls_streaming(self): + chunks = list( + self.client.chat.completions.create( + model=self.MODEL, + messages=[{"role": "user", "content": "What is the weather in Paris and London?"}], + stream=True, + max_tokens=100, + temperature=0.0, + tools=[self._get_tool_def()], + ) + ) + tool_chunks = [c for c in chunks if c.choices[0].delta.tool_calls] + self.assertEqual(len(tool_chunks), 2, f"Expected 2 tool calls, got {len(tool_chunks)}") + self.assertEqual(chunks[-1].choices[0].finish_reason, "tool_calls") + + def test_chat_multi_turn_non_streaming(self): + tool_def = self._get_tool_def() + resp1 = self.client.chat.completions.create( + model=self.MODEL, + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + stream=False, + max_tokens=50, + temperature=0.0, + tools=[tool_def], + ) + self.assertEqual(resp1.choices[0].finish_reason, "tool_calls") + tc = resp1.choices[0].message.tool_calls[0] + + resp2 = self.client.chat.completions.create( + model=self.MODEL, + messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + resp1.choices[0].message, + {"role": "tool", "tool_call_id": tc.id, "content": '{"temperature": 22, "condition": "sunny"}'}, + ], + stream=False, + max_tokens=100, + temperature=0.0, + tools=[tool_def], + ) + self.assertIn(resp2.choices[0].finish_reason, ("stop", "length")) + content = resp2.choices[0].message.content + self.assertIsNotNone(content) + self.assertTrue( + "22" in content.lower() or "sunny" in content.lower(), + f"Expected model to reference tool result, got: {content}", + ) + + def test_chat_multi_turn_streaming(self): + tool_def = self._get_tool_def() + + # Turn 1: streaming — accumulate tool call from deltas + chunks = list( + self.client.chat.completions.create( + model=self.MODEL, + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + stream=True, + max_tokens=50, + temperature=0.0, + tools=[tool_def], + ) + ) + self.assertEqual(chunks[-1].choices[0].finish_reason, "tool_calls") + tool_chunks = [c for c in chunks if c.choices[0].delta.tool_calls] + self.assertGreater(len(tool_chunks), 0) + tc = tool_chunks[0].choices[0].delta.tool_calls[0] + + # Reconstruct assistant message from deltas + content = "".join(c.choices[0].delta.content for c in chunks if c.choices[0].delta.content) + assistant_msg = { + "role": "assistant", + "content": content, + "tool_calls": [{"id": tc.id, "type": "function", "function": tc.function.model_dump()}], + } + + # Turn 2: streaming — send back tool result + chunks2 = list( + self.client.chat.completions.create( + model=self.MODEL, + messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + assistant_msg, + {"role": "tool", "tool_call_id": tc.id, "content": '{"temperature": 22, "condition": "sunny"}'}, + ], + stream=True, + max_tokens=100, + temperature=0.0, + tools=[tool_def], + ) + ) + content = "".join(c.choices[0].delta.content for c in chunks2 if c.choices[0].delta.content) + self.assertTrue( + "22" in content.lower() or "sunny" in content.lower(), + f"Expected model to reference tool result, got: {content}", + ) + + def test_responses_non_streaming(self): + resp = self.client.responses.create( + model=self.MODEL, + input="What is the weather in Paris?", + stream=False, + max_output_tokens=50, + tools=[self._get_tool_def()], + ) + self.assertEqual(resp.status, "completed") + fc_items = [o for o in resp.output if o.type == "function_call"] + self.assertGreater(len(fc_items), 0) + self.assertEqual(fc_items[0].name, "get_weather") + parsed = json.loads(fc_items[0].arguments) + self.assertIsInstance(parsed, dict) + + def test_responses_streaming(self): + events = list( + self.client.responses.create( + model=self.MODEL, + input="What is the weather in Paris?", + stream=True, + max_output_tokens=50, + tools=[self._get_tool_def()], + ) + ) + types = [e.type for e in events] + self.assertIn("response.created", types) + self.assertIn("response.completed", types) + self.assertIn("response.function_call_arguments.done", types) + + args_done = [e for e in events if e.type == "response.function_call_arguments.done"] + self.assertGreater(len(args_done), 0) + self.assertEqual(args_done[0].name, "get_weather") + parsed = json.loads(args_done[0].arguments) + self.assertIsInstance(parsed, dict) + + def test_responses_multiple_tool_calls_non_streaming(self): + resp = self.client.responses.create( + model=self.MODEL, + input="What is the weather in Paris and London?", + stream=False, + max_output_tokens=100, + tools=[self._get_tool_def()], + ) + self.assertEqual(resp.status, "completed") + fc_items = [o for o in resp.output if o.type == "function_call"] + self.assertEqual(len(fc_items), 2, f"Expected 2 tool calls, got {len(fc_items)}") + + def test_responses_multiple_tool_calls_streaming(self): + events = list( + self.client.responses.create( + model=self.MODEL, + input="What is the weather in Paris and London?", + stream=True, + max_output_tokens=100, + tools=[self._get_tool_def()], + ) + ) + args_done = [e for e in events if e.type == "response.function_call_arguments.done"] + self.assertEqual(len(args_done), 2, f"Expected 2 tool calls, got {len(args_done)}") + self.assertEqual(events[-1].type, "response.completed") + + def test_responses_multi_turn_non_streaming(self): + tool_def = self._get_tool_def() + resp1 = self.client.responses.create( + model=self.MODEL, + input="What is the weather in Paris?", + stream=False, + max_output_tokens=50, + tools=[tool_def], + ) + self.assertEqual(resp1.status, "completed") + fc_items = [o for o in resp1.output if o.type == "function_call"] + self.assertGreater(len(fc_items), 0) + + input_list = [{"role": "user", "content": "What is the weather in Paris?"}] + input_list += resp1.output + input_list.append( + { + "type": "function_call_output", + "call_id": fc_items[0].call_id, + "output": '{"temperature": 22, "condition": "sunny"}', + } + ) + resp2 = self.client.responses.create( + model=self.MODEL, + input=input_list, + stream=False, + max_output_tokens=100, + tools=[tool_def], + ) + self.assertEqual(resp2.status, "completed") + msg_items = [o for o in resp2.output if o.type == "message"] + self.assertGreater(len(msg_items), 0) + content = msg_items[0].content[0].text + self.assertTrue( + "22" in content.lower() or "sunny" in content.lower(), + f"Expected model to reference tool result, got: {content}", + ) + + def test_responses_multi_turn_streaming(self): + tool_def = self._get_tool_def() + + # Turn 1: streaming — get completed response with tool calls + events = list( + self.client.responses.create( + model=self.MODEL, + input="What is the weather in Paris?", + stream=True, + max_output_tokens=50, + tools=[tool_def], + ) + ) + completed = [e for e in events if e.type == "response.completed"] + self.assertEqual(len(completed), 1) + resp1_output = completed[0].response.output + fc_items = [o for o in resp1_output if o.type == "function_call"] + self.assertGreater(len(fc_items), 0) + + # Turn 2: streaming — send back tool result + input_list = [{"role": "user", "content": "What is the weather in Paris?"}] + input_list += resp1_output + input_list.append( + { + "type": "function_call_output", + "call_id": fc_items[0].call_id, + "output": '{"temperature": 22, "condition": "sunny"}', + } + ) + events2 = list( + self.client.responses.create( + model=self.MODEL, + input=input_list, + stream=True, + max_output_tokens=100, + tools=[tool_def], + ) + ) + content = "".join(e.delta for e in events2 if e.type == "response.output_text.delta") + self.assertTrue( + "22" in content.lower() or "sunny" in content.lower(), + f"Expected model to reference tool result, got: {content}", + ) + + +@slow +@require_serve +@require_torch_accelerator +class TestToolCallQwen(_TestToolCallBase, unittest.TestCase): + """Tool call tests with Qwen (fallback config, no response_schema).""" + + MODEL = "Qwen/Qwen2.5-0.5B-Instruct" + + +@slow +@require_serve +@require_torch_accelerator +class TestToolCallGemma(_TestToolCallBase, unittest.TestCase): + """Tool call tests with Gemma 4 (response_schema + stc/etc special tokens).""" + + MODEL = "google/gemma-4-E2B-it" + + @slow @require_librosa @require_multipart