118 lines · 4.5 KB · Rust
RawDownload
use std::str::FromStr; /// Parses tool call arguments JSON, treating empty strings as empty objects.////// Many LLM providers return empty strings for tool calls with no arguments./// This helper normalizes that behavior by converting empty strings to `{}`.pub fn parse_tool_arguments(arguments: &str) -> Result<serde_json::Value, serde_json::Error> {    if arguments.is_empty() {        Ok(serde_json::Value::Object(Default::default()))    } else {        serde_json::Value::from_str(arguments)    }} /// `partial_json_fixer::fix_json` converts a trailing `\` inside a string into `\\`/// (a literal backslash). When used for incremental parsing (comparing successive/// parses to extract deltas), this produces a spurious backslash character that/// doesn't exist in the final text, corrupting the output.////// This function strips any trailing incomplete escape sequence before fixing,/// so each intermediate parse produces a true prefix of the final string value.pub fn fix_streamed_json(partial_json: &str) -> String {    let json = strip_trailing_incomplete_escape(partial_json);    partial_json_fixer::fix_json(json)} fn strip_trailing_incomplete_escape(json: &str) -> &str {    let trailing_backslashes = json        .as_bytes()        .iter()        .rev()        .take_while(|&&b| b == b'\\')        .count();    if trailing_backslashes % 2 == 1 {        &json[..json.len() - 1]    } else {        json    }} /// Parses a "prompt is too long: N tokens ..." message and extracts the token count.pub fn parse_prompt_too_long(message: &str) -> Option<u64> {    message        .strip_prefix("prompt is too long: ")?        .split_once(" tokens")?        .0        .parse()        .ok()} /// Recognizes OpenAI-style context window overflow errors, which arrive either/// with the `context_length_exceeded` error code or a "Your input exceeds the/// context window of this model" message.pub fn is_context_window_exceeded_message(message: &str) -> bool {    message.contains("context_length_exceeded") || message.contains("exceeds the context window")} #[cfg(test)]mod tests {    use super::*;     #[test]    fn test_fix_streamed_json_strips_incomplete_escape() {        let fixed = fix_streamed_json(r#"{"text": "hello\"#);        let parsed: serde_json::Value = serde_json::from_str(&fixed).expect("valid json");        assert_eq!(parsed["text"], "hello");    }     #[test]    fn test_fix_streamed_json_preserves_complete_escape() {        let fixed = fix_streamed_json(r#"{"text": "hello\\"#);        let parsed: serde_json::Value = serde_json::from_str(&fixed).expect("valid json");        assert_eq!(parsed["text"], "hello\\");    }     #[test]    fn test_fix_streamed_json_strips_escape_after_complete_escape() {        let fixed = fix_streamed_json(r#"{"text": "hello\\\"#);        let parsed: serde_json::Value = serde_json::from_str(&fixed).expect("valid json");        assert_eq!(parsed["text"], "hello\\");    }     #[test]    fn test_fix_streamed_json_no_escape_at_end() {        let fixed = fix_streamed_json(r#"{"text": "hello"#);        let parsed: serde_json::Value = serde_json::from_str(&fixed).expect("valid json");        assert_eq!(parsed["text"], "hello");    }     #[test]    fn test_fix_streamed_json_newline_escape_boundary() {        let fixed = fix_streamed_json(r#"{"text": "line1\"#);        let parsed: serde_json::Value = serde_json::from_str(&fixed).expect("valid json");        assert_eq!(parsed["text"], "line1");         let fixed = fix_streamed_json(r#"{"text": "line1\nline2"#);        let parsed: serde_json::Value = serde_json::from_str(&fixed).expect("valid json");        assert_eq!(parsed["text"], "line1\nline2");    }     #[test]    fn test_fix_streamed_json_incremental_delta_correctness() {        let chunk1 = r#"{"replacement_text": "fn foo() {\"#;        let fixed1 = fix_streamed_json(chunk1);        let parsed1: serde_json::Value = serde_json::from_str(&fixed1).expect("valid json");        let text1 = parsed1["replacement_text"].as_str().expect("string");        assert_eq!(text1, "fn foo() {");         let chunk2 = r#"{"replacement_text": "fn foo() {\n    return bar;\n}"}"#;        let fixed2 = fix_streamed_json(chunk2);        let parsed2: serde_json::Value = serde_json::from_str(&fixed2).expect("valid json");        let text2 = parsed2["replacement_text"].as_str().expect("string");        assert_eq!(text2, "fn foo() {\n    return bar;\n}");         let delta = &text2[text1.len()..];        assert_eq!(delta, "\n    return bar;\n}");    }}