somatize_core/
util.rs

1//! Shared utility functions.
2
3use std::time::{SystemTime, UNIX_EPOCH};
4
5/// Generate a hex-encoded nanosecond timestamp ID with a prefix.
6///
7/// Used for unique run IDs, plan IDs, etc.
8pub fn timestamp_id(prefix: &str) -> String {
9    let nanos = SystemTime::now()
10        .duration_since(UNIX_EPOCH)
11        .unwrap_or_default()
12        .as_nanos();
13    format!("{prefix}_{nanos:x}")
14}
15
16/// Pull a JSON object out of a reply that may be fenced or prefaced.
17///
18/// Models wrap JSON in ```json fences, or preface it with a sentence,
19/// often enough that requiring a bare object would fail on working
20/// output.
21///
22/// Two implementations of this used to live in two crates. The naive one
23/// took everything between the first `{` and the *last* `}`, so a reply
24/// holding an object followed by any later brace — a second example, a
25/// closing fence with a brace in it, prose about `{}` — parsed as
26/// nothing. This is the scanner: the first *balanced* object, respecting
27/// strings and escapes.
28pub fn extract_json(text: &str) -> Option<serde_json::Value> {
29    // The whole reply, when the model simply answered with JSON.
30    if let Ok(value) = serde_json::from_str(text.trim()) {
31        return Some(value);
32    }
33
34    let start = text.find('{')?;
35    let mut depth = 0usize;
36    let mut in_string = false;
37    let mut escaped = false;
38
39    for (i, &b) in text.as_bytes().iter().enumerate().skip(start) {
40        if escaped {
41            escaped = false;
42            continue;
43        }
44        match b {
45            b'\\' if in_string => escaped = true,
46            b'"' => in_string = !in_string,
47            b'{' if !in_string => depth += 1,
48            b'}' if !in_string => {
49                depth -= 1;
50                if depth == 0 {
51                    return serde_json::from_str(&text[start..=i]).ok();
52                }
53            }
54            _ => {}
55        }
56    }
57    None
58}
59
60/// Enough of a long string to recognise, without pasting the whole thing
61/// into an error message.
62pub fn truncate(text: &str, max: usize) -> String {
63    if text.chars().count() <= max {
64        return text.to_string();
65    }
66    text.chars().take(max).collect::<String>() + "…"
67}
68
69#[cfg(test)]
70mod util_tests {
71    use super::*;
72
73    /// The naive version took the first `{` to the *last* `}`, which is
74    /// not the same object as soon as anything follows it.
75    #[test]
76    fn extract_json_takes_the_first_balanced_object() {
77        let reply = "Here you go:\n```json\n{\"a\": 1}\n```\nand another: {\"b\": 2}";
78        let got = extract_json(reply).expect("should find the first object");
79        assert_eq!(got, serde_json::json!({"a": 1}));
80    }
81
82    #[test]
83    fn extract_json_is_not_fooled_by_braces_in_strings() {
84        let reply = r#"prose {"note": "a } inside", "n": 3} trailing"#;
85        let got = extract_json(reply).expect("should find the object");
86        assert_eq!(got, serde_json::json!({"note": "a } inside", "n": 3}));
87    }
88
89    #[test]
90    fn extract_json_accepts_a_bare_object() {
91        assert_eq!(
92            extract_json("  {\"a\": 1}  "),
93            Some(serde_json::json!({"a": 1}))
94        );
95    }
96
97    #[test]
98    fn truncate_keeps_short_text_whole() {
99        assert_eq!(truncate("hola", 10), "hola");
100        assert_eq!(truncate("hola", 2), "ho…");
101    }
102}
somatize_core/util.rs

somatize_core/
util.rs
