chore: JailStream Optimizations

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

@@ -8,27 +8,19 @@ use openai_harmony::{
     HarmonyEncoding, HarmonyEncodingName, StreamableParser, load_harmony_encoding,
 };
 use serde_json::Value;
+use std::sync::OnceLock;
 
-static GLOBAL_HARMONY_GPTOSS_ENCODING: tokio::sync::OnceCell<
-    Result<HarmonyEncoding, anyhow::Error>,
-> = tokio::sync::OnceCell::const_new();
+static GLOBAL_HARMONY_GPTOSS_ENCODING: OnceLock<Result<HarmonyEncoding, anyhow::Error>> =
+    OnceLock::new();
 
-pub async fn get_harmony_encoding() -> &'static Result<HarmonyEncoding, anyhow::Error> {
+pub fn get_harmony_encoding() -> &'static Result<HarmonyEncoding, anyhow::Error> {
     GLOBAL_HARMONY_GPTOSS_ENCODING
-        .get_or_init(|| async {
-            tokio::task::spawn_blocking(|| {
-                load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss)
-            })
-            .await
-            .map_err(anyhow::Error::msg)
-            .flatten()
-        })
-        .await
+        .get_or_init(|| load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss))
 }
 
 /// Parse tool calls from Harmony Format text
 /// <|channel|>analysis<|message|>Need to use function get_current_weather.<|end|><|start|>assistant<|channel|>commentary to=functions.get_current_weather <|constrain|>json<|message|>{"location":"San Francisco"}<|call|>
-pub async fn parse_tool_calls_harmony(
+pub fn parse_tool_calls_harmony(
     text: &str,
     config: &JsonParserConfig,
 ) -> anyhow::Result<(Vec<ToolCallResponse>, Option<String>)> {
@@ -52,7 +44,7 @@ pub async fn parse_tool_calls_harmony(
         trimmed.push_str(end_token);
     }
 
-    let enc = match get_harmony_encoding().await.as_ref() {
+    let enc = match get_harmony_encoding().as_ref() {
         Ok(e) => e,
         Err(e) => {
             tracing::debug!("Failed to load harmony encoding: {e}. Tool calls will not be parsed.");
@@ -181,11 +173,11 @@ pub async fn parse_tool_calls_harmony(
 /// # Returns
 /// * `Ok((tool_calls, normal_text))` - Tuple containing extracted tool calls and any normal text
 /// * `Err(e)` - If parsing fails due to encoding or tokenization errors
-pub async fn parse_tool_calls_harmony_complete(
+pub fn parse_tool_calls_harmony_complete(
     text: &str,
     _config: &JsonParserConfig,
 ) -> anyhow::Result<(Vec<ToolCallResponse>, Option<String>)> {
-    let enc = match get_harmony_encoding().await.as_ref() {
+    let enc = match get_harmony_encoding().as_ref() {
         Ok(e) => e,
         Err(e) => {
             tracing::debug!("Failed to load harmony encoding: {e}. Tool calls will not be parsed.");
@@ -346,8 +338,8 @@ mod tests {
         (call.function.name, args)
     }
 
-    #[tokio::test]
-    async fn test_parse_tool_calls_harmony_basic() {
+    #[test]
+    fn test_parse_tool_calls_harmony_basic() {
         let text = r#"
 <|channel|>analysis<|message|>Need to use function get_current_weather.<|end|>
 <|start|>assistant<|channel|>commentary to=functions.get_current_weather <|constrain|>json
@@ -358,7 +350,7 @@ mod tests {
             tool_call_end_tokens: vec!["<|call|>".to_string()],
             ..Default::default()
         };
-        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).await.unwrap();
+        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).unwrap();
         assert_eq!(
             normal_content,
             Some("Need to use function get_current_weather.".to_string())
@@ -373,9 +365,7 @@ mod tests {
     async fn test_parse_tool_calls_harmony_complete_basic() {
         let text = r#"<|channel|>commentary to=functions.get_current_weather <|constrain|>json<|message|>{"format":"celsius","location":"San Francisco"}"#;
         let (tool_calls, normal_content) =
-            parse_tool_calls_harmony_complete(text, &Default::default())
-                .await
-                .unwrap();
+            parse_tool_calls_harmony_complete(text, &Default::default()).unwrap();
         assert_eq!(normal_content, Some("".to_string()));
         let (name, args) = extract_name_and_args(tool_calls[0].clone());
         assert_eq!(name, "get_current_weather");
@@ -394,13 +384,13 @@ mod tests {
             tool_call_end_tokens: vec!["<|call|>".to_string()],
             ..Default::default()
         };
-        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).await.unwrap();
+        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).unwrap();
         assert_eq!(normal_content, Some(text.trim().to_string()));
         assert_eq!(tool_calls.len(), 0);
     }
 
-    #[tokio::test]
-    async fn test_parse_tool_calls_harmony_with_multi_args() {
+    #[test]
+    fn test_parse_tool_calls_harmony_with_multi_args() {
         let text = r#"
         <|channel|>analysis<|message|>Need to use function get_current_weather.<|end|>
         <|start|>assistant<|channel|>commentary to=functions.get_current_weather <|constrain|>json
@@ -411,7 +401,7 @@ mod tests {
             tool_call_end_tokens: vec!["<|call|>".to_string()],
             ..Default::default()
         };
-        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).await.unwrap();
+        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).unwrap();
         assert_eq!(
             normal_content,
             Some("Need to use function get_current_weather.".to_string())
@@ -423,8 +413,8 @@ mod tests {
         assert_eq!(args["unit"], "fahrenheit");
     }
 
-    #[tokio::test]
-    async fn test_parse_tool_calls_harmony_with_normal_text() {
+    #[test]
+    fn test_parse_tool_calls_harmony_with_normal_text() {
         let text = r#"
         <|channel|>analysis<|message|>Need to use function get_current_weather.<|end|>
         <|start|>assistant<|channel|>commentary to=functions.get_current_weather <|constrain|>json
@@ -435,7 +425,7 @@ mod tests {
             tool_call_end_tokens: vec!["<|call|>".to_string()],
             ..Default::default()
         };
-        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).await.unwrap();
+        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).unwrap();
         assert_eq!(
             normal_content,
             Some("Need to use function get_current_weather.".to_string())
@@ -446,15 +436,15 @@ mod tests {
         assert_eq!(args["location"], "San Francisco");
     }
 
-    #[tokio::test]
-    async fn test_parse_tool_calls_harmony_without_call_token() {
+    #[test]
+    fn test_parse_tool_calls_harmony_without_call_token() {
         let text = r#"<|channel|>analysis<|message|>We need to call get_weather function. The user asks "What's the weather like in San Francisco in Celsius?" So location: "San Francisco, CA" unit: "celsius". Let's call function.<|end|><|start|>assistant<|channel|>commentary to=functions.get_weather <|constrain|>json<|message|>{"location":"San Francisco, CA","unit":"celsius"}"#;
         let config = JsonParserConfig {
             tool_call_start_tokens: vec!["<|start|>assistant<|channel|>commentary".to_string()],
             tool_call_end_tokens: vec!["<|call|>".to_string()],
             ..Default::default()
         };
-        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).await.unwrap();
+        let (tool_calls, normal_content) = parse_tool_calls_harmony(text, &config).unwrap();
         assert_eq!(normal_content, Some("We need to call get_weather function. The user asks \"What's the weather like in San Francisco in Celsius?\" So location: \"San Francisco, CA\" unit: \"celsius\". Let's call function.".to_string()));
         assert_eq!(tool_calls.len(), 1);
         let (name, args) = extract_name_and_args(tool_calls[0].clone());

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};