chore: move tool end detection to parser lib

Signed-off-by: ayushag <ayushag@nvidia.com>

Author: ayushag-nv

lib/llm/src/protocols/openai/chat_completions/jail.rs

@@ -3,12 +3,14 @@
 
 use async_stream::stream;
 use dynamo_async_openai::types::{
-    ChatChoiceStream, ChatCompletionMessageToolCallChunk, ChatCompletionStreamResponseDelta,
-    FinishReason, FunctionCallStream, Role,
+    ChatChoiceLogprobs, ChatChoiceStream, ChatCompletionMessageToolCallChunk,
+    ChatCompletionStreamResponseDelta, FinishReason, FunctionCallStream, Role,
 };
 
 use dynamo_parsers::tool_calling::parsers::get_tool_parser_map;
-use dynamo_parsers::tool_calling::{detect_tool_call_start, try_tool_call_parse_aggregate};
+use dynamo_parsers::tool_calling::{
+    detect_tool_call_start, find_tool_call_end_position, try_tool_call_parse_aggregate,
+};
 use dynamo_runtime::protocols::annotated::Annotated;
 use futures::{Stream, StreamExt};
 
@@ -74,7 +76,7 @@ struct ChoiceJailState {
 
 fn create_choice_stream(
     content: &str,
-    role: Role,
+    role: Option<Role>,
     index: u32,
     finish_reason: Option<FinishReason>,
     logprobs: Option<ChatChoiceLogprobs>,
@@ -83,7 +85,7 @@ fn create_choice_stream(
     ChatChoiceStream {
         index,
         delta: ChatCompletionStreamResponseDelta {
-            role: Some(role),
+            role,
             content: Some(content.to_string()),
             tool_calls: None,
             function_call: None,
@@ -327,7 +329,7 @@ impl ChoiceJailState {
             #[allow(deprecated)]
             let dummy_choice = create_choice_stream(
                 &self.accumulated_content,
-                Role::Assistant,
+                Some(Role::Assistant),
                 self.index,
                 None,
                 None,
@@ -648,7 +650,7 @@ impl JailedStream {
             if let Some(parser) = &self.tool_call_parser {
                 if let Ok((_, _)) = try_tool_call_parse_aggregate(accumulated_content, Some(parser))
                 {
-                    let split_pos = self.find_tool_call_end_position(accumulated_content, parser);
+                    let split_pos = find_tool_call_end_position(accumulated_content, Some(parser));
                     (true, split_pos)
                 } else {
                     (false, accumulated_content.len())
@@ -732,63 +734,6 @@ impl JailedStream {
         }
         false
     }
-
-    /// Find the exact position where the tool call ends for splitting content
-    /// This handles the early exit case where we have trailing content after the tool call
-    fn find_tool_call_end_position(&self, content: &str, parser: &str) -> usize {
-        match parser {
-            "hermes" => {
-                // For Hermes, look for </tool_call> marker
-                if let Some(pos) = content.find("</tool_call>") {
-                    pos + "</tool_call>".len()
-                } else {
-                    content.len()
-                }
-            }
-            "nemotron_deci" => {
-                // For Nemotron, look for </TOOLCALL> marker
-                if let Some(pos) = content.find("</TOOLCALL>") {
-                    pos + "</TOOLCALL>".len()
-                } else {
-                    content.len()
-                }
-            }
-            "mistral" => {
-                // For Mistral, look for [/TOOL_CALLS] marker or end of JSON array
-                if let Some(pos) = content.find("[/TOOL_CALLS]") {
-                    pos + "[/TOOL_CALLS]".len()
-                } else if let Some(pos) = content.rfind(']') {
-                    // Find the last ] which should be the end of the tool calls array
-                    pos + 1
-                } else {
-                    content.len()
-                }
-            }
-            "phi4" => {
-                // For Phi4, look for <|tool_call|> end marker
-                if let Some(pos) = content.rfind("<|tool_call|>") {
-                    // Look for the next occurrence after this position
-                    if let Some(end_pos) = content[pos..].find(">") {
-                        pos + end_pos + 1
-                    } else {
-                        content.len()
-                    }
-                } else {
-                    content.len()
-                }
-            }
-            "llama3_json" => {
-                // For Llama3 JSON, there's no explicit end marker
-                // The end is determined by complete JSON parsing
-                // Return full content length to avoid early splitting
-                content.len()
-            }
-            _ => {
-                // Unknown parser, default to full content
-                content.len()
-            }
-        }
-    }
 }
 
 /// Builder for configuring a JailedStream

lib/parsers/src/tool_calling/harmony/mod.rs

@@ -7,3 +7,7 @@ pub use super::{config, response};
 pub use harmony_parser::{
     detect_tool_call_start_harmony, parse_tool_calls_harmony, parse_tool_calls_harmony_complete,
 };
+
+pub fn find_tool_call_end_position_harmony(chunk: &str) -> usize {
+    chunk.len()
+}

lib/parsers/src/tool_calling/json/mod.rs

@@ -41,3 +41,31 @@ pub fn detect_tool_call_start_json(chunk: &str, config: &JsonParserConfig) -> bo
         JsonParserType::DeepseekV31 => detect_tool_call_start_deepseek_v3_1(chunk, config),
     }
 }
+
+pub fn find_tool_call_end_position_json(
+    chunk: &str,
+    parser: &str,
+    config: &JsonParserConfig,
+) -> usize {
+    match parser {
+        "hermes" | "nemotron_deci" => {
+            if let Some(end_token) = config.tool_call_end_tokens.first() {
+                if let Some(pos) = chunk.find(end_token) {
+                    pos + end_token.len()
+                } else {
+                    chunk.len()
+                }
+            } else {
+                chunk.len()
+            }
+        }
+        "mistral" | "phi4" => {
+            if let Some(pos) = chunk.rfind(']') {
+                pos + 1
+            } else {
+                chunk.len()
+            }
+        }
+        _ => chunk.len(),
+    }
+}

lib/parsers/src/tool_calling/mod.rs

@@ -13,7 +13,10 @@ pub mod tools;
 pub use config::{JsonParserConfig, ToolCallConfig, ToolCallParserType};
 pub use harmony::{parse_tool_calls_harmony, parse_tool_calls_harmony_complete};
 pub use json::try_tool_call_parse_json;
-pub use parsers::{detect_and_parse_tool_call, detect_tool_call_start, try_tool_call_parse};
+pub use parsers::{
+    detect_and_parse_tool_call, detect_tool_call_start, find_tool_call_end_position,
+    try_tool_call_parse,
+};
 pub use pythonic::try_tool_call_parse_pythonic;
 pub use response::{CalledFunction, ToolCallResponse, ToolCallType};
 pub use tools::{try_tool_call_parse_aggregate, try_tool_call_parse_stream};

lib/parsers/src/tool_calling/parsers.rs

@@ -2,9 +2,17 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::config::{ToolCallConfig, ToolCallParserType};
-use super::harmony::{detect_tool_call_start_harmony, parse_tool_calls_harmony_complete};
-use super::json::{detect_tool_call_start_json, try_tool_call_parse_json};
-use super::pythonic::{detect_tool_call_start_pythonic, try_tool_call_parse_pythonic};
+use super::harmony::{
+    detect_tool_call_start_harmony, find_tool_call_end_position_harmony,
+    parse_tool_calls_harmony_complete,
+};
+use super::json::{
+    detect_tool_call_start_json, find_tool_call_end_position_json, try_tool_call_parse_json,
+};
+use super::pythonic::{
+    detect_tool_call_start_pythonic, find_tool_call_end_position_pythonic,
+    try_tool_call_parse_pythonic,
+};
 use super::response::ToolCallResponse;
 use std::collections::HashMap;
 use std::sync::OnceLock;
@@ -116,6 +124,35 @@ pub fn detect_tool_call_start(chunk: &str, parser_str: Option<&str>) -> anyhow::
     }
 }
 
+pub fn find_tool_call_end_position(chunk: &str, parser_str: Option<&str>) -> usize {
+    let parser_map = get_tool_parser_map();
+    let parser_key = match parser_str {
+        Some(s) if !s.is_empty() => s,
+        _ => "default",
+    };
+
+    match parser_map.get(parser_key) {
+        Some(config) => match config.format {
+            ToolCallParserType::Json => {
+                find_tool_call_end_position_json(chunk, parser_key, &config.json)
+            }
+            ToolCallParserType::Harmony => find_tool_call_end_position_harmony(chunk),
+            ToolCallParserType::Pythonic => find_tool_call_end_position_pythonic(chunk),
+            ToolCallParserType::Typescript => {
+                // Typescript parser not implemented
+                chunk.len()
+            }
+            ToolCallParserType::Xml => {
+                // Xml parser not implemented
+                chunk.len()
+            }
+        },
+        None => {
+            // Unknown parser, return full content length
+            chunk.len()
+        }
+    }
+}
 // Tests
 // cargo test postprocessor::tool_calling::parsers
 #[cfg(test)]

lib/parsers/src/tool_calling/pythonic/mod.rs

@@ -5,3 +5,7 @@ pub mod pythonic_parser;
 
 pub use super::{config, response};
 pub use pythonic_parser::{detect_tool_call_start_pythonic, try_tool_call_parse_pythonic};
+
+pub fn find_tool_call_end_position_pythonic(chunk: &str) -> usize {
+    chunk.len()
+}