1use async_trait::async_trait;
8use openai_protocol::common::Tool;
9use serde_json::{Map, Value};
10
11use crate::{
12 errors::ParserResult,
13 json_format,
14 parsers::helpers::coerce_by_schema_type,
15 traits::ToolParser,
16 types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
17};
18
19const START: &str = "<|tool_call>";
20const END: &str = "<tool_call|>";
21const STRING: &str = r#"<|"|>"#;
22const BOUNDARIES: [&str; 3] = [END, "<turn|>", "<|tool_response>"];
23
24#[derive(Default)]
25pub struct Gemma4Parser {
26 buffer: String,
27 next_tool_index: usize,
28 after_call: bool,
29}
30
31impl Gemma4Parser {
32 pub fn new() -> Self {
33 Self::default()
34 }
35
36 fn drain_ready(&mut self, tools: &[Tool]) -> StreamingParseResult {
37 let mut result = StreamingParseResult::default();
38 while !self.buffer.is_empty() {
39 if self.after_call {
40 let trimmed = self.buffer.trim_start();
41 if let Some(marker) = BOUNDARIES.iter().find(|m| trimmed.starts_with(**m)) {
42 let consumed = self.buffer.len() - trimmed.len() + marker.len();
43 self.buffer.drain(..consumed);
44 continue;
45 }
46 if trimmed.is_empty() || BOUNDARIES.iter().any(|m| m.starts_with(trimmed)) {
47 break;
48 }
49 self.after_call = false;
50 }
51
52 if self.buffer.starts_with(START) {
53 match decode_call(&self.buffer, tools) {
54 Some((name, args, consumed)) => {
55 result.calls.push(ToolCallItem {
56 tool_index: self.next_tool_index,
57 name: Some(name),
58 parameters: json_format::to_string(&args),
59 });
60 self.next_tool_index += 1;
61 self.buffer.drain(..consumed);
62 self.after_call = true;
63 }
64 None => {
65 let Some(end) = unquoted_marker(&self.buffer[START.len()..], END) else {
68 break;
69 };
70 let consumed = START.len() + end + END.len();
71 result.normal_text.push_str(&self.buffer[..consumed]);
72 self.buffer.drain(..consumed);
73 }
74 }
75 } else {
76 let consumed = self.buffer.find(START).unwrap_or_else(|| {
77 let partial =
78 super::helpers::ends_with_partial_token(&self.buffer, START).unwrap_or(0);
79 self.buffer.len() - partial
80 });
81 if consumed == 0 {
82 break;
83 }
84 result.normal_text.push_str(&self.buffer[..consumed]);
85 self.buffer.drain(..consumed);
86 }
87 }
88 result
89 }
90
91 fn parse_output(output: &str, tools: &[Tool]) -> (String, Vec<ToolCall>) {
92 let mut parser = Self {
93 buffer: output.to_owned(),
94 ..Self::default()
95 };
96 let result = parser.drain_ready(tools);
97 let mut text = result.normal_text;
98 text.push_str(&parser.take_unstreamed_normal_text());
99 let calls = result
100 .calls
101 .into_iter()
102 .map(|call| ToolCall {
103 function: FunctionCall {
104 name: call.name.unwrap_or_default(),
105 arguments: call.parameters,
106 },
107 })
108 .collect();
109 (text, calls)
110 }
111}
112
113#[async_trait]
114impl ToolParser for Gemma4Parser {
115 async fn parse_complete(&self, output: &str) -> ParserResult<(String, Vec<ToolCall>)> {
116 Ok(Self::parse_output(output, &[]))
117 }
118
119 async fn parse_complete_with_tools(
120 &self,
121 output: &str,
122 tools: &[Tool],
123 ) -> ParserResult<(String, Vec<ToolCall>)> {
124 Ok(Self::parse_output(output, tools))
125 }
126
127 async fn parse_incremental(
128 &mut self,
129 chunk: &str,
130 tools: &[Tool],
131 ) -> ParserResult<StreamingParseResult> {
132 self.buffer.push_str(chunk);
133 Ok(self.drain_ready(tools))
134 }
135
136 fn has_tool_markers(&self, text: &str) -> bool {
137 text.contains(START)
138 }
139
140 fn take_unstreamed_normal_text(&mut self) -> String {
141 std::mem::take(&mut self.buffer)
142 }
143
144 fn reset(&mut self) {
145 *self = Self::default();
146 }
147}
148
149fn decode_call(input: &str, tools: &[Tool]) -> Option<(String, Value, usize)> {
150 let body = input[START.len()..].trim_start();
151 let prefix = if body.starts_with("call:") {
152 "call:"
153 } else if body.starts_with(':') {
154 ":"
155 } else {
156 return None;
157 };
158 let rest = &body[prefix.len()..];
159 let open = rest.find(['{', '('])?;
160 let name = rest[..open].trim();
161 if name.is_empty() || name.chars().any(char::is_whitespace) || name.contains('<') {
162 return None;
163 }
164 let close = if rest.as_bytes()[open] == b'{' {
165 '}'
166 } else {
167 ')'
168 };
169 let args_start = input.len() - rest.len() + open + 1;
170 let mut reader = ArgumentReader {
171 input: &input[args_start..],
172 pos: 0,
173 };
174 let schema = tools
175 .iter()
176 .find(|tool| tool.function.name == name)
177 .map(|tool| tool.function.parameters_schema())
178 .unwrap_or(&Value::Null);
179 let args = reader.object(close, 0, schema)?;
180 Some((name.to_owned(), args, args_start + reader.pos))
181}
182
183fn unquoted_marker(input: &str, marker: &str) -> Option<usize> {
184 let mut quoted = false;
185 let mut pos = 0;
186 while pos < input.len() {
187 let tail = &input[pos..];
188 if tail.starts_with(STRING) {
189 quoted = !quoted;
190 pos += STRING.len();
191 } else if !quoted && tail.starts_with(marker) {
192 return Some(pos);
193 } else {
194 pos += tail.chars().next()?.len_utf8();
195 }
196 }
197 None
198}
199
200struct ArgumentReader<'a> {
202 input: &'a str,
203 pos: usize,
204}
205
206impl ArgumentReader<'_> {
207 fn remaining(&self) -> &str {
208 &self.input[self.pos..]
209 }
210
211 fn skip_separators(&mut self) {
212 let tail = self.remaining();
213 let trimmed = tail.trim_start_matches(|c: char| c.is_whitespace() || c == ',');
214 self.pos += tail.len() - trimmed.len();
215 }
216
217 fn string(&mut self) -> Option<String> {
218 self.pos += STRING.len();
219 let end = self.remaining().find(STRING)?;
220 let value = self.remaining()[..end].to_owned();
221 self.pos += end + STRING.len();
222 Some(value)
223 }
224
225 fn object(&mut self, close: char, depth: usize, schema: &Value) -> Option<Value> {
226 let mut map = Map::new();
227 loop {
228 self.skip_separators();
229 if self.remaining().starts_with(close) {
230 self.pos += close.len_utf8();
231 return Some(Value::Object(map));
232 }
233 let key = if self.remaining().starts_with(STRING) {
234 self.string()?
235 } else {
236 let end = self.remaining().find(':')?;
237 let key = self.remaining()[..end].trim().to_owned();
238 self.pos += end;
239 key
240 };
241 let whitespace = self.remaining().len() - self.remaining().trim_start().len();
242 self.pos += whitespace;
243 if key.is_empty() || key.contains(['{', '}', '[', ']', '(', ')', '<']) {
244 return None;
245 }
246 if !self.remaining().starts_with(':') {
247 return None;
248 }
249 self.pos += 1;
250 let property = schema
251 .get("properties")
252 .and_then(|properties| properties.get(&key))
253 .unwrap_or(&Value::Null);
254 map.insert(key, self.value(close, depth + 1, property)?);
255 }
256 }
257
258 fn value(&mut self, close: char, depth: usize, schema: &Value) -> Option<Value> {
259 if depth > 128 {
260 return None;
261 }
262 let whitespace = self.remaining().len() - self.remaining().trim_start().len();
263 self.pos += whitespace;
264 if self.remaining().starts_with(STRING) {
265 return self.string().map(Value::String);
266 }
267 let first = self.remaining().chars().next()?;
268 if first == '{' {
269 self.pos += 1;
270 return self.object('}', depth, schema);
271 }
272 if first == '[' {
273 self.pos += 1;
274 let mut values = Vec::new();
275 loop {
276 self.skip_separators();
277 if self.remaining().starts_with(']') {
278 self.pos += 1;
279 return Some(Value::Array(values));
280 }
281 values.push(self.value(
282 ']',
283 depth + 1,
284 schema.get("items").unwrap_or(&Value::Null),
285 )?);
286 }
287 }
288 let end = self.remaining().find([',', '}', ']', ')'])?;
289 let text = self.remaining()[..end].trim();
290 let delimiter = self.remaining().as_bytes()[end] as char;
291 if text.is_empty() || (delimiter != ',' && delimiter != close) || text.contains('<') {
292 return None;
293 }
294 let value = coerce_bare_argument(text, schema);
297 self.pos += end;
298 Some(value)
299 }
300}
301
302fn coerce_bare_argument(text: &str, schema: &Value) -> Value {
303 let branches = if schema.get("type").is_some() {
304 std::slice::from_ref(schema)
305 } else {
306 schema
307 .get("anyOf")
308 .and_then(Value::as_array)
309 .map(Vec::as_slice)
310 .unwrap_or_else(|| std::slice::from_ref(schema))
311 };
312 for branch in branches {
313 let Some(declared) = branch.get("type") else {
314 continue;
315 };
316 let types = declared
317 .as_array()
318 .map(Vec::as_slice)
319 .unwrap_or_else(|| std::slice::from_ref(declared));
320 for kind in types {
321 let converted = match kind.as_str() {
322 Some("string") => None,
323 Some("null") if text.eq_ignore_ascii_case("null") => Some(Value::Null),
324 Some("boolean") if text == "1" => Some(Value::Bool(true)),
325 Some("boolean") if text == "0" => Some(Value::Bool(false)),
326 Some(kind @ ("integer" | "number" | "boolean")) => {
327 coerce_by_schema_type(text, Some(kind))
328 }
329 _ => None,
330 };
331 if let Some(converted) = converted {
332 return converted;
333 }
334 }
335 }
336 Value::String(text.to_owned())
337}