Connorli/fix func call parsing

bindings/python/src/smg/serve.py

@@ -11,7 +11,9 @@
 import logging
 import os
 import random
+import shutil
 import signal
+import tempfile
 import socket
 import subprocess
 import sys
@@ -550,6 +552,7 @@ def __init__(self, backend: str, args: argparse.Namespace, backend_args: list[st
         self.launcher: WorkerLauncher = BACKEND_LAUNCHERS[backend]()
         self.workers: list[tuple[subprocess.Popen, int]] = []
         self._shutting_down = False
+        self._prometheus_dir: str | None = None
 
     # -- public API ---------------------------------------------------------
 
@@ -571,6 +574,15 @@ def run(self) -> None:
     def _launch_workers(self) -> None:
         ports = _find_available_ports(self.args.worker_base_port, self.args.data_parallel_size)
         host = self.args.worker_host
+
+        if getattr(self.args, "connection_mode", "grpc") == "grpc":
+            self._prometheus_dir = tempfile.mkdtemp(prefix="smg_prometheus_")
+            os.environ["PROMETHEUS_MULTIPROC_DIR"] = self._prometheus_dir
+            logger.info(
+                "Set PROMETHEUS_MULTIPROC_DIR=%s for gRPC metrics collection",
+                self._prometheus_dir,
+            )
+
         for dp_rank, port in enumerate(ports):
             env = self.launcher.gpu_env(self.args, dp_rank)
             proc = self.launcher.launch(self.args, self.backend_args, host, port, env)
@@ -641,6 +653,23 @@ def _cleanup_workers(self) -> None:
                 except (ProcessLookupError, OSError):
                     pass
 
+        self._cleanup_prometheus_dir()
+
+    def _cleanup_prometheus_dir(self) -> None:
+        """Remove the temporary prometheus multiprocess directory and its .db files."""
+        if self._prometheus_dir is None:
+            return
+        try:
+            shutil.rmtree(self._prometheus_dir)
+            logger.info("Cleaned up PROMETHEUS_MULTIPROC_DIR=%s", self._prometheus_dir)
+        except OSError as e:
+            logger.warning(
+                "Failed to clean up PROMETHEUS_MULTIPROC_DIR=%s: %s",
+                self._prometheus_dir,
+                e,
+            )
+        self._prometheus_dir = None
+
 
 # ---------------------------------------------------------------------------
 # Entry point

crates/reasoning_parser/src/parsers/base.rs

@@ -63,6 +63,12 @@ impl ReasoningParser for BaseReasoningParser {
             .to_string();
 
         if !processed_text.contains(&self.config.think_end_token) {
+            // Don't consume tool call markers as reasoning content
+            if let Some(tool_pos) = processed_text.find("<|tool_calls_section_begin|>") {
+                let reasoning_text = processed_text[..tool_pos].trim().to_string();
+                let normal_text = processed_text[tool_pos..].to_string();
+                return Ok(ParserResult::new(normal_text, reasoning_text));
+            }
             // Assume reasoning was truncated before end token
             return Ok(ParserResult::reasoning(processed_text));
         }
@@ -133,6 +139,14 @@ impl ReasoningParser for BaseReasoningParser {
 
         // Continue with reasoning content
         if self.in_reasoning && self.config.stream_reasoning {
+            // Some models skip </think> and go straight to tool calls
+            if let Some(tool_pos) = current_text.find("<|tool_calls_section_begin|>") {
+                let reasoning_text = current_text[..tool_pos].trim().to_string();
+                let normal_text = current_text[tool_pos..].to_string();
+                self.buffer.clear();
+                self.in_reasoning = false;
+                return Ok(ParserResult::new(normal_text, reasoning_text));
+            }
             // Stream the content immediately
             let reasoning_text = current_text;
             self.buffer.clear();

crates/tool_parser/src/parsers/kimik2.rs

@@ -57,16 +57,15 @@ impl KimiK2Parser {
         reason = "regex patterns are compile-time string literals"
     )]
     pub fn new() -> Self {
-        // Pattern for complete tool calls
-        let tool_call_pattern = r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*?\})\s*<\|tool_call_end\|>";
+        // Supports alternative delimiters: <|func_start|>/<|func_end|>; (?s) for multi-line JSON
+        let tool_call_pattern = r"(?s)<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*(?:<\|tool_call_argument_begin\|>\s*|<\|func_start\|>\s*)?(?P<function_arguments>\{.*?\})\s*(?:<\|tool_call_end\|>|<\|func_end\|>)";
         let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
 
-        // Pattern for streaming (partial) tool calls
-        let stream_pattern = r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*)";
+        let stream_pattern = r"(?s)<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*(?:<\|tool_call_argument_begin\|>\s*|<\|func_start\|>\s*)?(?P<function_arguments>\{.*)";
         let stream_tool_call_extractor = Regex::new(stream_pattern).expect("Valid regex pattern");
 
         // Pattern for removing completed tool calls
-        let end_pattern = r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>";
+        let end_pattern = r"<\|tool_call_begin\|>.*?(?:<\|tool_call_end\|>|<\|func_end\|>)";
         let tool_call_end_pattern = Regex::new(end_pattern).expect("Valid regex pattern");
 
         // Robust parser for ids like "functions.search:0" or fallback "search:0"
@@ -181,7 +180,7 @@ impl ToolParser for KimiK2Parser {
             // No tool markers detected - return all buffered content as normal text
             let mut normal_text = std::mem::take(&mut self.buffer);
             // Remove end tokens if present
-            for e_token in ["<|tool_calls_section_end|>", "<|tool_call_end|>"] {
+            for e_token in ["<|tool_calls_section_end|>", "<|tool_call_end|>", "<|func_end|>"] {
                 normal_text = normal_text.replace(e_token, "");
             }
             return Ok(StreamingParseResult {
@@ -242,13 +241,14 @@ impl ToolParser for KimiK2Parser {
                             function_args
                         };
 
-                        // Split by end token before sending (like Python does)
-                        let parsed_args_diff =
-                            if let Some(pos) = argument_diff.find("<|tool_call_end|>") {
-                                &argument_diff[..pos]
-                            } else {
-                                argument_diff
-                            };
+                        // Split by end token before sending
+                        let end_pos = argument_diff.find("<|tool_call_end|>")
+                            .or_else(|| argument_diff.find("<|func_end|>"));
+                        let parsed_args_diff = if let Some(pos) = end_pos {
+                            &argument_diff[..pos]
+                        } else {
+                            argument_diff
+                        };
 
                         if !parsed_args_diff.is_empty() {
                             calls.push(ToolCallItem {
@@ -265,8 +265,9 @@ impl ToolParser for KimiK2Parser {
                         }
 
                         // Check completeness - split by end token first
-                        let parsed_args = if let Some(pos) = function_args.find("<|tool_call_end|>")
-                        {
+                        let end_pos2 = function_args.find("<|tool_call_end|>")
+                            .or_else(|| function_args.find("<|func_end|>"));
+                        let parsed_args = if let Some(pos) = end_pos2 {
                             &function_args[..pos]
                         } else {
                             function_args

model_gateway/src/core/worker_manager.rs

@@ -17,7 +17,7 @@ use tokio::{
     sync::{watch, Mutex},
     task::JoinHandle,
 };
-use tracing::{debug, info};
+use tracing::{debug, info, warn};
 
 use crate::{
     core::{
@@ -85,6 +85,36 @@ impl IntoResponse for EngineMetricsResult {
     }
 }
 
+/// Collect gRPC worker metrics by aggregating `PROMETHEUS_MULTIPROC_DIR` via a python3 subprocess.
+async fn collect_prometheus_multiproc_metrics() -> Result<String, String> {
+    let dir = std::env::var("PROMETHEUS_MULTIPROC_DIR").map_err(|_| {
+        "PROMETHEUS_MULTIPROC_DIR not set; cannot collect metrics from gRPC workers".to_string()
+    })?;
+
+    let output = tokio::process::Command::new("python3")
+        .args([
+            "-c",
+            "import sys\n\
+             from prometheus_client import CollectorRegistry, generate_latest\n\
+             from prometheus_client.multiprocess import MultiProcessCollector\n\
+             registry = CollectorRegistry()\n\
+             MultiProcessCollector(registry)\n\
+             sys.stdout.buffer.write(generate_latest(registry))\n",
+        ])
+        .env("PROMETHEUS_MULTIPROC_DIR", &dir)
+        .output()
+        .await
+        .map_err(|e| format!("failed to run python3 prometheus collector: {e}"))?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(format!("python3 prometheus collector failed: {stderr}"));
+    }
+
+    String::from_utf8(output.stdout)
+        .map_err(|e| format!("prometheus collector output is not valid UTF-8: {e}"))
+}
+
 pub struct WorkerManager;
 
 impl WorkerManager {
@@ -273,22 +303,49 @@ impl WorkerManager {
             return EngineMetricsResult::Err("No available workers".to_string());
         }
 
-        let responses = fan_out(&workers, client, "metrics", reqwest::Method::GET).await;
-
         let mut metric_packs = Vec::new();
-        for resp in responses {
-            if let Ok(r) = resp.result {
-                if r.status().is_success() {
-                    if let Ok(text) = r.text().await {
-                        metric_packs.push(MetricPack {
-                            labels: vec![("worker_addr".into(), resp.url)],
-                            metrics_text: text,
-                        });
+
+        let http_workers: Vec<_> = workers
+            .iter()
+            .filter(|w| matches!(w.connection_mode(), ConnectionMode::Http))
+            .cloned()
+            .collect();
+        let has_grpc = workers
+            .iter()
+            .any(|w| matches!(w.connection_mode(), ConnectionMode::Grpc));
+
+        if !http_workers.is_empty() {
+            let responses = fan_out(&http_workers, client, "metrics", reqwest::Method::GET).await;
+            for resp in responses {
+                if let Ok(r) = resp.result {
+                    if r.status().is_success() {
+                        if let Ok(text) = r.text().await {
+                            metric_packs.push(MetricPack {
+                                labels: vec![("worker_addr".into(), resp.url)],
+                                metrics_text: text,
+                            });
+                        }
                     }
                 }
             }
         }
 
+        if has_grpc {
+            match collect_prometheus_multiproc_metrics().await {
+                Ok(text) => {
+                    metric_packs.push(MetricPack {
+                        labels: vec![],
+                        metrics_text: text,
+                    });
+                }
+                Err(e) => {
+                    warn!(
+                        "Failed to collect gRPC worker metrics from PROMETHEUS_MULTIPROC_DIR: {e}"
+                    );
+                }
+            }
+        }
+
         if metric_packs.is_empty() {
             return EngineMetricsResult::Err("All backend requests failed".to_string());
         }

model_gateway/src/routers/grpc/regular/processor.rs

@@ -148,7 +148,16 @@ impl ResponseProcessor {
                 _ => false,
             };
 
-            if used_json_schema {
+            if self.configured_tool_parser.is_some() && tool_parser_available {
+                // Explicitly configured parser takes priority (models may emit native tokens regardless of tool_choice)
+                (tool_calls, processed_text) = self
+                    .parse_tool_calls(
+                        &processed_text,
+                        &original_request.model,
+                        history_tool_calls_count,
+                    )
+                    .await;
+            } else if used_json_schema {
                 (tool_calls, processed_text) = utils::parse_json_schema_response(
                     &processed_text,
                     original_request.tool_choice.as_ref(),
@@ -183,7 +192,58 @@ impl ResponseProcessor {
             utils::convert_proto_to_openai_logprobs(proto_logprobs, tokenizer)
         });
 
-        // Step 5: Build ChatCompletionMessage (proper response message type)
+        // Strip leaked chatml tokens only when a model-specific parser is configured
+        if self.configured_tool_parser.is_some() || self.configured_reasoning_parser.is_some() {
+            for token in [
+                "<|im_end|>",
+                "<|im_start|>",
+                "<|im_user|>",
+                "<|im_assistant|>",
+                "<|im_system|>",
+                "<|im_middle|>",
+            ] {
+                processed_text = processed_text.replace(token, "");
+            }
+            processed_text = processed_text.trim().to_string();
+        }
+
+        let is_json_response = matches!(
+            &original_request.response_format,
+            Some(openai_protocol::common::ResponseFormat::JsonObject)
+                | Some(openai_protocol::common::ResponseFormat::JsonSchema { .. })
+        );
+        if is_json_response {
+            if processed_text.starts_with("```json") || processed_text.starts_with("```JSON") {
+                if let Some(start) = processed_text.find('\n') {
+                    let inner = &processed_text[start + 1..];
+                    if let Some(end) = inner.rfind("```") {
+                        processed_text = inner[..end].trim().to_string();
+                    }
+                }
+            }
+
+            if processed_text.starts_with('{') {
+                let mut depth = 0i32;
+                let mut in_string = false;
+                let mut escape = false;
+                let mut json_end = None;
+                for (i, ch) in processed_text.char_indices() {
+                    if escape { escape = false; continue; }
+                    if ch == '\\' && in_string { escape = true; continue; }
+                    if ch == '"' { in_string = !in_string; continue; }
+                    if in_string { continue; }
+                    if ch == '{' { depth += 1; } else if ch == '}' {
+                        depth -= 1;
+                        if depth == 0 { json_end = Some(i + 1); break; }
+                    }
+                }
+                if let Some(end) = json_end {
+                    processed_text = processed_text[..end].to_string();
+                }
+            }
+        }
+
+        // Build ChatCompletionMessage (proper response message type)
         let chat_message = ChatCompletionMessage {
             role: "assistant".to_string(),
             content: if processed_text.is_empty() {
@@ -621,7 +681,17 @@ impl ResponseProcessor {
                 Some(messages::ToolChoice::Tool { .. } | messages::ToolChoice::Any { .. })
             );
 
-            if used_json_schema {
+            if self.configured_tool_parser.is_some() && tool_parser_available {
+                (tool_calls, processed_text) = self
+                    .parse_tool_calls(
+                        &processed_text,
+                        &messages_request.model,
+                        utils::message_utils::get_history_tool_calls_count_messages(
+                            &messages_request,
+                        ),
+                    )
+                    .await;
+            } else if used_json_schema {
                 // Bridge Messages ToolChoice to Chat ToolChoice for reuse
                 let chat_tool_choice = messages_request
                     .tool_choice
@@ -647,7 +717,22 @@ impl ResponseProcessor {
             }
         }
 
-        // Step 3: Build content blocks
+        // Strip leaked chatml tokens only when a model-specific parser is configured
+        if self.configured_tool_parser.is_some() || self.configured_reasoning_parser.is_some() {
+            for token in [
+                "<|im_end|>",
+                "<|im_start|>",
+                "<|im_user|>",
+                "<|im_assistant|>",
+                "<|im_system|>",
+                "<|im_middle|>",
+            ] {
+                processed_text = processed_text.replace(token, "");
+            }
+            processed_text = processed_text.trim().to_string();
+        }
+
+        // Build content blocks
         let mut content_blocks: Vec<messages::ContentBlock> = Vec::new();
 
         // Thinking block first (if present)