@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 158 lines 5.46 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
4
5 """Gemma4 tool call parsing utilities for offline inference.
6
7 Standalone functions that parse decoded model text to extract tool calls
8 from Gemma4 models. These are pure-Python utilities with zero heavy
9 dependencies — they work on raw decoded strings from any inference
10 backend (vLLM, HuggingFace, TGI, etc.).
11
12 For the OpenAI-compatible API server tool parser (streaming +
13 non-streaming), see ``vllm.tool_parsers.gemma4_tool_parser``.
14 For thinking/reasoning output parsing, see
15 ``vllm.reasoning.gemma4_utils``.
16
17 Usage with vLLM offline inference::
18
19 from vllm import LLM, SamplingParams
20 from vllm.tool_parsers.gemma4_utils import (
21 parse_tool_calls,
22 has_tool_response_tag,
23 )
24
25 llm = LLM(model="google/gemma-4-it")
26 outputs = llm.generate(prompt, SamplingParams(...))
27 text = tokenizer.decode(outputs[0].outputs[0].token_ids, skip_special_tokens=False)
28
29 # Extract tool calls
30 tool_calls = parse_tool_calls(text)
31 for tc in tool_calls:
32 print(f"{tc['name']}({tc['arguments']})")
33
34 Ported from ``transformers.models.gemma4.utils_gemma4`` so that vLLM users
35 do not need a transformers dependency for output parsing.
36 """
37
38 import regex as re
39
40 _TOOL_RESPONSE_START_TAG = "<|tool_response>"
41
42
43 def _parse_tool_arguments(args_str: str) -> dict[str, str]:
44 """Parse tool call arguments from the Gemma4 compact format.
45
46 Delegates to the native ``<|"|>``-aware parser from
47 ``vllm.parser.gemma4``, which handles internal quotes, nested
48 objects, arrays, and all Gemma4 value types correctly.
49
50 Args:
51 args_str: Raw argument string from inside ``call:name{...}``.
52
53 Returns:
54 Dictionary of argument name → string value.
55
56 """
57 if not args_str or not args_str.strip():
58 return {}
59
60 from vllm.parser.gemma4 import _parse_gemma4_args
61
62 parsed = _parse_gemma4_args(args_str)
63 return {k: str(v) if not isinstance(v, str) else v for k, v in parsed.items()}
64
65
66 def parse_tool_calls(text: str, *, strict: bool = False) -> list[dict]:
67 """Parse tool calls from decoded Gemma4 model output.
68
69 Uses a tiered parsing strategy to handle known output variations in
70 Gemma4 models, which may emit
71 non-standard tool call formats.
72
73 Parsing tiers:
74 1. **Standard**: ``<|tool_call>call:name{args}<tool_call|>``
75 (special token IDs 48/49 in decoded text)
76 2. **Fallback** (when ``strict=False``): bare ``call:name{args}``
77 patterns, including ``<call>name{args}`` (fragmented tokens from
78 multimodal inputs)
79
80 Args:
81 text: Decoded model output text (from ``tokenizer.decode(...,
82 skip_special_tokens=False)``).
83 strict: If ``True``, only match the standard ``<|tool_call>`` format.
84 If ``False`` (default), also try fallback patterns for
85 known Gemma4 output variations.
86
87 Returns:
88 A list of dicts, each with keys:
89 - ``"name"``: The tool function name (e.g. ``"get_weather"``).
90 - ``"arguments"``: A dict of argument name → value.
91
92 Example::
93
94 >>> from vllm.tool_parsers.gemma4_utils import parse_tool_calls
95 >>> output = tokenizer.decode(outputs[0], skip_special_tokens=False)
96 >>> tool_calls = parse_tool_calls(output)
97 >>> for tc in tool_calls:
98 ... print(f"Call: {tc['name']}({tc['arguments']})")
99
100 """
101 results = []
102
103 # Tier 1: Standard format with special tokens.
104 # <|tool_call>call:name{args}<tool_call|>
105 # Note: Some Gemma4 models emit <turn|> instead of <tool_call|>.
106 standard_pattern = r"<\|tool_call\>call:(\w+)\{(.*?)\}(?:<tool_call\|>|<turn\|>)"
107 for match in re.finditer(standard_pattern, text, re.DOTALL):
108 name, args_str = match.group(1), match.group(2)
109 results.append(
110 {
111 "name": name,
112 "arguments": _parse_tool_arguments(args_str),
113 }
114 )
115
116 if results or strict:
117 return results
118
119 # Tier 2: Fallback for known Gemma4 output variations.
120 # Matches: <call>name{args}, call:name{args}, or bare call:name{args}<eos>
121 fallback_pattern = r"(?:<call>|(?:^|\s)call:)(\w+)\{(.*?)\}"
122 for match in re.finditer(fallback_pattern, text, re.DOTALL):
123 name, args_str = match.group(1), match.group(2)
124 results.append(
125 {
126 "name": name,
127 "arguments": _parse_tool_arguments(args_str),
128 }
129 )
130
131 return results
132
133
134 def has_tool_response_tag(text: str) -> bool:
135 """Check if model output properly ends with a tool response tag.
136
137 Some Gemma4 models sometimes emit ``<eos>`` instead of
138 ``<|tool_response>`` after a tool call. This helper detects
139 whether the model used the proper termination, so callers can
140 decide whether to inject ``<|tool_response>`` into the next prompt.
141
142 Args:
143 text: Decoded model output text.
144
145 Returns:
146 ``True`` if the output ends with ``<|tool_response>``
147 (proper behavior), ``False`` otherwise.
148
149 Example::
150
151 >>> from vllm.tool_parsers.gemma4_utils import has_tool_response_tag
152 >>> if not has_tool_response_tag(model_output):
153 ... # Model used <eos> instead — inject <|tool_response> manually
154 ... next_prompt = "<|tool_response>" + tool_result
155
156 """
157 stripped = text.rstrip()
158 return stripped.endswith(_TOOL_RESPONSE_START_TAG)