fix: forward tool schemas to response parsers (#50385)

* fix: forward tool schemas to response parsers

* test: cover tool schema propagation across pipelines and serving

* refactor: reuse normalized tools and simplify tests

---------

Co-authored-by: BlazeSR10hawk <175792037+BlazeSR10hawk@users.noreply.github.com>

Author: BlazeSR10hawk

src/transformers/cli/serving/chat_completion.py

@@ -163,6 +163,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse
                 inputs,
                 gen_config,
                 gen_manager=gen_manager,
+                tools=body.get("tools"),
             )
         else:
             return await self._non_streaming(
@@ -173,6 +174,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse
                 inputs,
                 gen_config,
                 gen_manager=gen_manager,
+                tools=body.get("tools"),
             )
 
     # ----- streaming -----
@@ -186,9 +188,10 @@ def _streaming(
         inputs: dict,
         gen_config: "GenerationConfig",
         gen_manager: BaseGenerateManager,
+        tools=None,
     ) -> StreamingResponse:
         """Stream tokens as SSE via DirectStreamer."""
-        response_parser = build_response_parser(processor, model, inputs["input_ids"])
+        response_parser = build_response_parser(processor, model, inputs["input_ids"], tools=tools)
         queue, streamer = gen_manager.generate_streaming(
             model,
             processor,
@@ -288,6 +291,7 @@ async def _non_streaming(
         inputs: dict,
         gen_config: "GenerationConfig",
         gen_manager: BaseGenerateManager,
+        tools=None,
     ) -> JSONResponse:
         """Run generation and return a JSONResponse."""
         content, input_len, generated_ids = await gen_manager.generate_non_streaming(
@@ -303,7 +307,7 @@ async def _non_streaming(
         )
 
         content, reasoning_content, parsed_tool_calls = parse_assistant_message(
-            processor, model, generated_ids, input_ids=inputs["input_ids"], cleaned_content=content
+            processor, model, generated_ids, input_ids=inputs["input_ids"], cleaned_content=content, tools=tools
         )
         tool_calls = None
         if parsed_tool_calls:

src/transformers/cli/serving/response.py

@@ -430,6 +430,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse
                 inputs,
                 gen_config,
                 gen_manager=gen_manager,
+                tools=tools,
             )
         else:
             return await self._non_streaming(
@@ -441,6 +442,7 @@ async def handle_request(self, body: dict, request_id: str) -> StreamingResponse
                 inputs,
                 gen_config,
                 gen_manager=gen_manager,
+                tools=tools,
             )
 
     # ----- input conversion -----
@@ -569,9 +571,10 @@ def _streaming(
         inputs: dict,
         gen_config: "GenerationConfig",
         gen_manager: BaseGenerateManager,
+        tools=None,
     ) -> StreamingResponse:
         """Generate a streaming Responses API reply (SSE) using DirectStreamer."""
-        response_parser = build_response_parser(processor, model, inputs["input_ids"])
+        response_parser = build_response_parser(processor, model, inputs["input_ids"], tools=tools)
         queue, streamer = gen_manager.generate_streaming(
             model,
             processor,
@@ -675,14 +678,20 @@ async def _non_streaming(
         inputs: dict,
         gen_config: "GenerationConfig",
         gen_manager: BaseGenerateManager,
+        tools=None,
     ) -> JSONResponse:
         """Generate a non-streaming Responses API reply (single JSON)."""
         content, input_len, generated_ids = await gen_manager.generate_non_streaming(
             model, processor, inputs, gen_config, request_id=request_id
         )
 
         content, reasoning_content, parsed_tool_calls = parse_assistant_message(
-            processor, model, generated_ids, input_ids=inputs["input_ids"], cleaned_content=content
+            processor,
+            model,
+            generated_ids,
+            input_ids=inputs["input_ids"],
+            cleaned_content=content,
+            tools=tools,
         )
 
         output_items = []

src/transformers/cli/serving/utils.py

@@ -237,6 +237,7 @@ def build_response_parser(
     processor,
     model: "PreTrainedModel",
     input_ids,
+    tools=None,
 ) -> "ResponseParser | None":
     """Build a streaming :class:`ResponseParser` for the current generation.
 
@@ -253,7 +254,7 @@ def build_response_parser(
     if template is None:
         return None
     tokenizer = getattr(processor, "tokenizer", processor)
-    return ResponseParser(template, prefix=_decode_prefix(tokenizer, input_ids))
+    return ResponseParser(template, prefix=_decode_prefix(tokenizer, input_ids), tools=tools)
 
 
 def _normalize_tool_call(value: dict) -> ToolCall:
@@ -293,6 +294,7 @@ def parse_assistant_message(
     generated_ids,
     input_ids,
     cleaned_content: str = "",
+    tools=None,
 ) -> tuple[str, str | None, list[ToolCall] | None]:
     """Parse a full assistant message into ``(content, reasoning, tool_calls)``.
 
@@ -308,7 +310,9 @@ def parse_assistant_message(
     if template is None:
         return cleaned_content, None, None
     tokenizer = getattr(processor, "tokenizer", processor)
-    parsed = tokenizer.parse_response(generated_ids, template, prefix=_decode_prefix(tokenizer, input_ids))
+    parsed = tokenizer.parse_response(
+        generated_ids, template, prefix=_decode_prefix(tokenizer, input_ids), tools=tools
+    )
     # Absent ``content`` key means the model emitted only tool calls / thinking;
     # surface that as an empty string rather than swapping in the raw decode.
     content = parsed.get("content", "")

src/transformers/pipelines/any_to_any.py

@@ -233,6 +233,8 @@ def _sanitize_parameters(
         if skip_special_tokens is not None:
             postprocess_params["skip_special_tokens"] = skip_special_tokens
         postprocess_params["generation_mode"] = generation_mode
+        if "tools" in kwargs:
+            postprocess_params["tools"] = kwargs["tools"]
         return preprocess_params, forward_kwargs, postprocess_params
 
     @overload
@@ -439,6 +441,7 @@ def postprocess(
         return_type=None,
         continue_final_message=None,
         skip_special_tokens=None,
+        tools=None,
         **postprocess_kwargs,
     ):
         input_texts = model_outputs["prompt_text"]
@@ -514,7 +517,9 @@ def postprocess(
                             # New-style templates need to see the prompt as `prefix`, because chat
                             # templates often pre-write part of the assistant message (e.g. an
                             # opening <think> tag), which affects parsing.
-                            assistant_message = self.tokenizer.parse_response(generated_text, prefix=decoded_input)
+                            assistant_message = self.tokenizer.parse_response(
+                                generated_text, prefix=decoded_input, tools=tools
+                            )
                         else:
                             assistant_message = {"role": "assistant", "content": generated_text}
                         generated_text = list(prompt_text.messages) + [assistant_message]

src/transformers/pipelines/image_text_to_text.py

@@ -194,6 +194,8 @@ def _sanitize_parameters(
         if skip_special_tokens is not None:
             postprocess_params["skip_special_tokens"] = skip_special_tokens
 
+        if "tools" in kwargs:
+            postprocess_params["tools"] = kwargs["tools"]
         return preprocess_params, forward_kwargs, postprocess_params
 
     @overload
@@ -403,6 +405,7 @@ def postprocess(
         return_type=ReturnType.FULL_TEXT,
         continue_final_message=None,
         skip_special_tokens=None,
+        tools=None,
         **postprocess_kwargs,
     ):
         input_texts = model_outputs["prompt_text"]
@@ -472,7 +475,9 @@ def postprocess(
                             # New-style templates need to see the prompt as `prefix`, because chat
                             # templates often pre-write part of the assistant message (e.g. an
                             # opening <think> tag), which affects parsing.
-                            assistant_message = self.tokenizer.parse_response(generated_text, prefix=decoded_input)
+                            assistant_message = self.tokenizer.parse_response(
+                                generated_text, prefix=decoded_input, tools=tools
+                            )
                         else:
                             assistant_message = {"role": "assistant", "content": generated_text}
                         generated_text = list(prompt_text.messages) + [assistant_message]

src/transformers/pipelines/text_generation.py

@@ -204,6 +204,8 @@ def _sanitize_parameters(
 
         # postprocess kwargs
         postprocess_params = {}
+        if tools is not None:
+            postprocess_params["tools"] = tools
         if return_full_text is not None and return_type is None:
             if return_text is not None:
                 raise ValueError("`return_text` is mutually exclusive with `return_full_text`")
@@ -427,6 +429,7 @@ def postprocess(
         clean_up_tokenization_spaces=None,
         continue_final_message=None,
         skip_special_tokens=None,
+        tools=None,
     ):
         generated_sequence = model_outputs["generated_sequence"][0]
         input_ids = model_outputs["input_ids"]
@@ -487,7 +490,9 @@ def postprocess(
                                     skip_special_tokens=skip_special_tokens,
                                     clean_up_tokenization_spaces=clean_up_tokenization_spaces,
                                 )
-                                assistant_message = self.tokenizer.parse_response(all_text, prefix=prompt_prefix)
+                                assistant_message = self.tokenizer.parse_response(
+                                    all_text, prefix=prompt_prefix, tools=tools
+                                )
                             else:
                                 # If there's no template, then we have to assume it's all content
                                 assistant_message = {"role": "assistant", "content": all_text}

tests/cli/test_serve.py

@@ -39,6 +39,7 @@
     GenerationState,
     Modality,
     ToolCall,
+    build_response_parser,
     get_response_template,
     parse_assistant_message,
     response_events_to_chunks,
@@ -1904,6 +1905,41 @@ def test_parse_assistant_message_tool_call(self):
         called_template = tokenizer.parse_response.call_args.args[1]
         self.assertIs(called_template, template)
 
+    def test_tool_schema_forwarded_to_serving_parsers(self):
+        tokenizer = MagicMock()
+        tokenizer.response_template = {
+            "defaults": {"role": "assistant"},
+            "fields": {},
+        }
+        tokenizer.decode.return_value = ""
+        tokenizer.parse_response.return_value = {"role": "assistant"}
+
+        processor = MagicMock()
+        model = MagicMock()
+        processor.tokenizer = tokenizer
+
+        tools = [
+            {
+                "type": "function",
+                "function": {"name": "set_alarm"},
+            }
+        ]
+
+        for schema in (None, tools):
+            with self.subTest(tools=schema):
+                parse_assistant_message(processor, model, [2], [1], tools=schema)
+                self.assertEqual(
+                    tokenizer.parse_response.call_args.kwargs["tools"],
+                    schema,
+                )
+
+                with patch("transformers.cli.serving.utils.ResponseParser") as parser:
+                    build_response_parser(processor, model, [1], tools=schema)
+                    self.assertEqual(
+                        parser.call_args.kwargs["tools"],
+                        schema,
+                    )
+
     def test_parse_assistant_message_multiple_tool_calls(self):
         tokenizer = MagicMock()
         tokenizer.response_template = None

tests/pipelines/test_pipelines_text_generation.py

@@ -302,6 +302,67 @@ def test_small_chat_model_with_response_template_prefix(self):
         self.assertIsInstance(parsed_message["reasoning_content"], str)
         self.assertGreater(len(parsed_message["reasoning_content"]), 0)
 
+    @require_torch
+    def test_small_chat_model_forwards_tool_schema(self):
+        text_generator = pipeline(
+            "text-generation",
+            model="hf-internal-testing/tiny-gpt2-with-chatml-template",
+        )
+        text_generator.tokenizer.chat_template = (
+            "{% for m in messages %}"
+            "{{ '<|im_start|>' + m['role'] + '\n' + m['content'] + '<|im_end|>\n' }}"
+            "{% endfor %}"
+            "{% if add_generation_prompt %}"
+            "{{ '<|im_start|>assistant\n<tool_call>\n<function=set_alarm>\n'"
+            " + '<parameter=hour>7</parameter>\n'"
+            " + '<parameter=enabled>false</parameter>\n'"
+            " + '</function>\n</tool_call>' }}"
+            "{% endif %}"
+        )
+        text_generator.tokenizer.response_template = {
+            "defaults": {"role": "assistant"},
+            "start_anchor": "<|im_start|>assistant\n",
+            "fields": {
+                "tool_calls": {
+                    "open_pattern": r"<tool_call>\s*<function=(?P<name>\w+)>",
+                    "close": "</tool_call>",
+                    "repeats": True,
+                    "content": "xml-inline",
+                    "content_args": {"tag_pattern": r"<parameter=(?P<key>\w+)>\s*(?P<value>.*?)\s*</parameter>"},
+                    "transform": {
+                        "type": "function",
+                        "function": {
+                            "name": "{name}",
+                            "arguments": "{content}",
+                        },
+                    },
+                }
+            },
+        }
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "set_alarm",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "hour": {"type": "integer"},
+                            "enabled": {"type": "boolean"},
+                        },
+                    },
+                },
+            }
+        ]
+        result = text_generator(
+            [{"role": "user", "content": "Set alarm"}],
+            tools=tools,
+            do_sample=False,
+            max_new_tokens=1,
+        )
+        arguments = result[0]["generated_text"][-1]["tool_calls"][0]["function"]["arguments"]
+        self.assertEqual(arguments, {"hour": 7, "enabled": False})
+
     @require_torch
     def test_return_full_text_false_with_chat_template(self):
         """Regression test for #45854: return_full_text=False must not include prompt when using chat template."""