[model-gateway] change sgl-router to sgl-model-gateway (#14312)

2025-12-05 12:04:48 -08:00
parent 1ea6b740a7
commit 49dfa1d891
431 changed files with 86 additions and 93 deletions
@@ -0,0 +1,414 @@
+use std::sync::OnceLock;
+
+/// Pythonic format parser for tool calls
+///
+/// Handles Python function call syntax within square brackets:
+/// ```text
+/// [tool1(arg1=val1, arg2=val2), tool2(arg1=val3)]
+/// ```
+///
+/// This format is used by Llama models and uses Python literals
+/// rather than JSON for arguments.
+/// Reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct?chat_template=default
+use async_trait::async_trait;
+use num_traits::ToPrimitive;
+use regex::Regex;
+use rustpython_parser::{
+    ast::{Constant, Expr, Mod, UnaryOp},
+    parse, Mode,
+};
+use serde_json::{Map, Number, Value};
+
+use crate::{
+    protocols::common::Tool,
+    tool_parser::{
+        errors::{ParserError, ParserResult},
+        parsers::helpers,
+        traits::ToolParser,
+        types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
+    },
+};
+
+static PYTHONIC_BLOCK_REGEX: OnceLock<Regex> = OnceLock::new();
+
+/// Lazily compiled regex that locates pythonic tool call blocks.
+fn pythonic_block_regex() -> &'static Regex {
+    PYTHONIC_BLOCK_REGEX.get_or_init(|| {
+        // Matches one or more function calls inside a list. The `(?s)` flag allows
+        // newlines inside argument lists while keeping the pattern anchored to
+        // identifiers followed by parentheses, preventing plain lists like
+        // `[1, 2, 3]` from matching.
+        Regex::new(r"(?s)\[\s*[A-Za-z_]\w*\s*\(.*?\)\s*(?:,\s*[A-Za-z_]\w*\s*\(.*?\)\s*)*\]")
+            .expect("pythonic tool call regex must compile")
+    })
+}
+
+/// Parser for Pythonic tool call format
+pub struct PythonicParser {
+    /// Buffer for accumulating chunks
+    buffer: String,
+}
+
+impl Default for PythonicParser {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl PythonicParser {
+    /// Create a new Pythonic parser
+    pub fn new() -> Self {
+        Self {
+            buffer: String::new(),
+        }
+    }
+
+    /// Extract the first pythonic tool call block and return it along with the
+    /// surrounding "normal" content.
+    fn extract_tool_calls(&self, text: &str) -> Option<(String, String)> {
+        pythonic_block_regex().find(text).map(|mat| {
+            let block = mat.as_str().to_string();
+            let normal = format!("{}{}", &text[..mat.start()], &text[mat.end()..]);
+            (block, normal)
+        })
+    }
+
+    /// Strip special tokens that Llama models might output
+    fn strip_special_tokens(text: &str) -> String {
+        text.replace("<|python_start|>", "")
+            .replace("<|python_end|>", "")
+    }
+
+    fn parse_tool_call_block(&self, block: &str) -> ParserResult<Vec<ToolCall>> {
+        let expr = parse_python_expression(block)?;
+        match expr {
+            Expr::List(list_expr) => list_expr
+                .elts
+                .into_iter()
+                .enumerate()
+                .map(|(idx, call_expr)| build_tool_call(call_expr, idx))
+                .collect(),
+            _ => Err(ParserError::ParsingFailed(
+                "Expected a list of function calls in pythonic tool call".to_string(),
+            )),
+        }
+    }
+}
+
+#[async_trait]
+impl ToolParser for PythonicParser {
+    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
+        let cleaned = Self::strip_special_tokens(text);
+
+        if let Some((tool_calls_text, normal_text)) = self.extract_tool_calls(&cleaned) {
+            match self.parse_tool_call_block(&tool_calls_text) {
+                Ok(calls) => {
+                    if calls.is_empty() {
+                        // No tools successfully parsed despite having markers
+                        Ok((text.to_string(), vec![]))
+                    } else {
+                        Ok((normal_text, calls))
+                    }
+                }
+                Err(e) => {
+                    // Log warning and return entire text as fallback
+                    tracing::warn!("Failed to parse pythonic tool calls: {}", e);
+                    Ok((text.to_string(), vec![]))
+                }
+            }
+        } else {
+            Ok((text.to_string(), vec![]))
+        }
+    }
+
+    async fn parse_incremental(
+        &mut self,
+        chunk: &str,
+        tools: &[Tool],
+    ) -> ParserResult<StreamingParseResult> {
+        self.buffer.push_str(chunk);
+
+        let cleaned = Self::strip_special_tokens(&self.buffer);
+
+        // Look for opening bracket
+        if let Some(start) = cleaned.find('[') {
+            let normal_text = if start > 0 {
+                cleaned[..start].to_string()
+            } else {
+                String::new()
+            };
+
+            // Look for matching closing bracket
+            if let Some(end) = find_matching_bracket(&cleaned, start) {
+                // Found complete tool call - extract it and parse using parse_complete
+                let call_text = &cleaned[start..=end];
+
+                match self.parse_complete(call_text).await {
+                    Ok((_, calls)) => {
+                        // Update buffer with remaining text after tool call
+                        let remaining_text = &cleaned[end + 1..];
+                        self.buffer = remaining_text.to_string();
+
+                        // Validate tool names and convert ToolCall to ToolCallItem
+                        let tool_indices = helpers::get_tool_indices(tools);
+                        let items: Vec<ToolCallItem> = calls
+                            .into_iter()
+                            .enumerate()
+                            .filter_map(|(idx, tool)| {
+                                if !tool_indices.contains_key(&tool.function.name) {
+                                    tracing::warn!(
+                                        "Invalid tool name '{}' - skipping",
+                                        tool.function.name
+                                    );
+                                    return None;
+                                }
+
+                                Some(ToolCallItem {
+                                    tool_index: idx,
+                                    name: Some(tool.function.name),
+                                    parameters: tool.function.arguments,
+                                })
+                            })
+                            .collect();
+
+                        return Ok(StreamingParseResult {
+                            normal_text,
+                            calls: items,
+                        });
+                    }
+                    Err(e) => {
+                        tracing::warn!("Failed to parse pythonic tool call: {}", e);
+                        // Clear buffer on error
+                        self.buffer.clear();
+                        return Ok(StreamingParseResult::default());
+                    }
+                }
+            } else {
+                // We have an opening bracket but no closing bracket yet
+                // Put back everything from the bracket onwards
+                self.buffer = cleaned[start..].to_string();
+
+                if !normal_text.is_empty() {
+                    return Ok(StreamingParseResult {
+                        normal_text,
+                        calls: vec![],
+                    });
+                }
+
+                // Still accumulating a potential tool call
+                return Ok(StreamingParseResult::default());
+            }
+        }
+
+        // No tool call bracket found
+        self.buffer.clear();
+        Ok(StreamingParseResult {
+            normal_text: cleaned,
+            calls: vec![],
+        })
+    }
+
+    fn has_tool_markers(&self, text: &str) -> bool {
+        let cleaned = Self::strip_special_tokens(text);
+        if pythonic_block_regex().is_match(&cleaned) {
+            return true;
+        }
+
+        false
+    }
+}
+
+/// Find the matching closing bracket for the opening bracket at start position.
+/// Properly handles nested brackets.
+fn find_matching_bracket(buffer: &str, start: usize) -> Option<usize> {
+    let mut bracket_count = 0;
+    let chars: Vec<char> = buffer.chars().collect();
+
+    for (i, &ch) in chars.iter().enumerate().skip(start) {
+        if ch == '[' {
+            bracket_count += 1;
+        } else if ch == ']' {
+            bracket_count -= 1;
+            if bracket_count == 0 {
+                return Some(i);
+            }
+        }
+    }
+    None // No matching bracket found
+}
+
+fn parse_python_expression(source: &str) -> ParserResult<Expr> {
+    let module = parse(source, Mode::Expression, "<pythonic_tool_call>")
+        .map_err(|err| ParserError::ParsingFailed(err.to_string()))?;
+
+    match module {
+        Mod::Expression(expr_mod) => Ok(*expr_mod.body),
+        _ => Err(ParserError::ParsingFailed(
+            "Expected a Python expression".to_string(),
+        )),
+    }
+}
+
+fn build_tool_call(expr: Expr, _index: usize) -> ParserResult<ToolCall> {
+    match expr {
+        Expr::Call(call_expr) => {
+            if !call_expr.args.is_empty() {
+                return Err(ParserError::ParsingFailed(
+                    "Positional arguments are not supported in pythonic tool calls".to_string(),
+                ));
+            }
+
+            let function_name = match *call_expr.func {
+                Expr::Name(name_expr) => name_expr.id.to_string(),
+                _ => {
+                    return Err(ParserError::ParsingFailed(
+                        "Unsupported function reference in pythonic tool call".to_string(),
+                    ))
+                }
+            };
+
+            let mut arguments_map = Map::with_capacity(call_expr.keywords.len());
+            for keyword in call_expr.keywords {
+                let arg_name = keyword.arg.ok_or_else(|| {
+                    ParserError::ParsingFailed(
+                        "pythonic tool calls do not support **kwargs".to_string(),
+                    )
+                })?;
+                let value_json = expression_to_json(&keyword.value)?;
+                arguments_map.insert(arg_name.to_string(), value_json);
+            }
+
+            let arguments_json = Value::Object(arguments_map);
+            let arguments_string = serde_json::to_string(&arguments_json)?;
+
+            Ok(ToolCall {
+                function: FunctionCall {
+                    name: function_name,
+                    arguments: arguments_string,
+                },
+            })
+        }
+        _ => Err(ParserError::ParsingFailed(
+            "Expected function calls inside pythonic tool call list".to_string(),
+        )),
+    }
+}
+
+fn expression_to_json(expr: &Expr) -> ParserResult<Value> {
+    match expr {
+        Expr::Constant(expr_constant) => constant_to_json(&expr_constant.value),
+        Expr::List(list_expr) => collect_sequence(&list_expr.elts).map(Value::Array),
+        Expr::Tuple(tuple_expr) => collect_sequence(&tuple_expr.elts).map(Value::Array),
+        Expr::Dict(dict_expr) => {
+            collect_dict(&dict_expr.keys, &dict_expr.values).map(Value::Object)
+        }
+        Expr::UnaryOp(unary_expr) => match unary_expr.op {
+            UnaryOp::USub => match unary_expr.operand.as_ref() {
+                Expr::Constant(const_expr) => negate_constant(&const_expr.value),
+                _ => Err(ParserError::ParsingFailed(
+                    "Unsupported unary operand in pythonic tool call".to_string(),
+                )),
+            },
+            UnaryOp::UAdd => expression_to_json(unary_expr.operand.as_ref()),
+            _ => Err(ParserError::ParsingFailed(format!(
+                "Unsupported unary operator in pythonic tool call: {:?}",
+                unary_expr.op
+            ))),
+        },
+        Expr::Name(name_expr) => Ok(Value::String(name_expr.id.to_string())),
+        _ => Err(ParserError::ParsingFailed(format!(
+            "Unsupported expression in pythonic tool call: {:?}",
+            expr
+        ))),
+    }
+}
+
+fn constant_to_json(constant: &Constant) -> ParserResult<Value> {
+    match constant {
+        Constant::None => Ok(Value::Null),
+        Constant::Bool(b) => Ok(Value::Bool(*b)),
+        Constant::Int(value) => Ok(integer_constant_to_value(value, false)),
+        Constant::Float(f) => Number::from_f64(*f).map(Value::Number).ok_or_else(|| {
+            ParserError::ParsingFailed("Invalid float literal in pythonic tool call".to_string())
+        }),
+        Constant::Str(s) => Ok(Value::String(s.clone())),
+        Constant::Bytes(bytes) => Ok(Value::String(String::from_utf8_lossy(bytes).into_owned())),
+        Constant::Tuple(values) => constant_tuple_to_array(values).map(Value::Array),
+        Constant::Ellipsis | Constant::Complex { .. } => Err(ParserError::ParsingFailed(
+            "Unsupported literal in pythonic tool call".to_string(),
+        )),
+    }
+}
+
+fn negate_constant(constant: &Constant) -> ParserResult<Value> {
+    match constant {
+        Constant::Int(value) => Ok(integer_constant_to_value(value, true)),
+        Constant::Float(f) => Number::from_f64(-f).map(Value::Number).ok_or_else(|| {
+            ParserError::ParsingFailed("Invalid float literal in pythonic tool call".to_string())
+        }),
+        _ => Err(ParserError::ParsingFailed(
+            "Unsupported unary operand in pythonic tool call".to_string(),
+        )),
+    }
+}
+
+fn value_to_key_string(value: Value) -> ParserResult<String> {
+    match value {
+        Value::String(s) => Ok(s),
+        Value::Number(num) => Ok(num.to_string()),
+        Value::Bool(b) => Ok(b.to_string()),
+        Value::Null => Ok("null".to_string()),
+        other => Err(ParserError::ParsingFailed(format!(
+            "Unsupported key type in pythonic tool call: {:?}",
+            other
+        ))),
+    }
+}
+
+fn collect_sequence(elements: &[Expr]) -> ParserResult<Vec<Value>> {
+    elements.iter().map(expression_to_json).collect()
+}
+
+fn collect_dict(keys: &[Option<Expr>], values: &[Expr]) -> ParserResult<Map<String, Value>> {
+    let mut map = Map::with_capacity(keys.len());
+    for (key_expr, value_expr) in keys.iter().zip(values.iter()) {
+        let key_expr = key_expr.as_ref().ok_or_else(|| {
+            ParserError::ParsingFailed("pythonic tool calls do not support **kwargs".to_string())
+        })?;
+        let key_value = expression_to_json(key_expr)?;
+        let key = value_to_key_string(key_value)?;
+        let value_json = expression_to_json(value_expr)?;
+        map.insert(key, value_json);
+    }
+    Ok(map)
+}
+
+fn constant_tuple_to_array(values: &[Constant]) -> ParserResult<Vec<Value>> {
+    values.iter().map(constant_to_json).collect()
+}
+
+fn integer_constant_to_value<T>(value: &T, negate: bool) -> Value
+where
+    T: ToPrimitive + std::fmt::Display,
+{
+    if let Some(mut i) = value.to_i64() {
+        if negate {
+            i = -i;
+        }
+        return Value::Number(Number::from(i));
+    }
+
+    if negate {
+        if let Some(u) = value.to_u64() {
+            if u <= i64::MAX as u64 {
+                return Value::Number(Number::from(-(u as i64)));
+            }
+            return Value::String(format!("-{}", value));
+        }
+        Value::String(format!("-{}", value))
+    } else if let Some(u) = value.to_u64() {
+        Value::Number(Number::from(u))
+    } else {
+        Value::String(value.to_string())
+    }
+}