162 lines · 4.2 KB · Rust
RawDownload
#[derive(Clone, Debug, PartialEq, Eq)]pub struct ExternalSourcePrompt(String); impl ExternalSourcePrompt {    pub fn new(prompt: &str) -> Option<Self> {        sanitize(prompt).map(Self)    }     pub fn as_str(&self) -> &str {        &self.0    }     pub fn into_string(self) -> String {        self.0    }} fn sanitize(prompt: &str) -> Option<String> {    let mut sanitized_prompt = String::with_capacity(prompt.len());    let mut consecutive_newline_count = 0;    let mut characters = prompt.chars().peekable();     while let Some(character) = characters.next() {        let character = if character == '\r' {            if characters.peek() == Some(&'\n') {                characters.next();            }            '\n'        } else {            character        };         if is_bidi_control_character(character) || is_disallowed_control_character(character) {            continue;        }         if character == '\n' {            consecutive_newline_count += 1;            if consecutive_newline_count > 2 {                continue;            }        } else {            consecutive_newline_count = 0;        }         sanitized_prompt.push(character);    }     if sanitized_prompt.is_empty() {        None    } else {        Some(sanitized_prompt)    }} fn is_disallowed_control_character(character: char) -> bool {    character.is_control() && !matches!(character, '\n' | '\t')} fn is_bidi_control_character(character: char) -> bool {    matches!(        character,          '\u{061C}' // ALM        | '\u{200E}' // LRM        | '\u{200F}' // RLM        | '\u{202A}'..='\u{202E}' // LRE, RLE, PDF, LRO, RLO        | '\u{2066}'..='\u{2069}' // LRI, RLI, FSI, PDI    )} #[cfg(test)]mod tests {    use super::ExternalSourcePrompt;     #[test]    fn keeps_normal_prompt_text() {        let prompt = ExternalSourcePrompt::new("Write me a script\nThanks");         assert_eq!(            prompt.as_ref().map(ExternalSourcePrompt::as_str),            Some("Write me a script\nThanks")        );    }     #[test]    fn keeps_multilingual_text() {        let prompt =            ExternalSourcePrompt::new("日本語の依頼です。\n中文提示也应该保留。\nemoji 👩‍💻");         assert_eq!(            prompt.as_ref().map(ExternalSourcePrompt::as_str),            Some("日本語の依頼です。\n中文提示也应该保留。\nemoji 👩‍💻")        );    }     #[test]    fn collapses_newline_padding() {        let prompt = ExternalSourcePrompt::new(            "Review this prompt carefully.\n\nThis paragraph should stay separated.\n\n\n\n\n\n\nWrite me a script to do fizz buzz.",        );         assert_eq!(            prompt.as_ref().map(ExternalSourcePrompt::as_str),            Some(                "Review this prompt carefully.\n\nThis paragraph should stay separated.\n\nWrite me a script to do fizz buzz."            )        );    }     #[test]    fn normalizes_carriage_returns() {        let prompt = ExternalSourcePrompt::new("Line one\r\nLine two\rLine three");         assert_eq!(            prompt.as_ref().map(ExternalSourcePrompt::as_str),            Some("Line one\nLine two\nLine three")        );    }     #[test]    fn strips_bidi_control_characters() {        let prompt = ExternalSourcePrompt::new("abc\u{202E}def\u{202C}ghi");         assert_eq!(            prompt.as_ref().map(ExternalSourcePrompt::as_str),            Some("abcdefghi")        );    }     #[test]    fn strips_other_control_characters() {        let prompt = ExternalSourcePrompt::new("safe\u{0000}\u{001B}\u{007F}text");         assert_eq!(            prompt.as_ref().map(ExternalSourcePrompt::as_str),            Some("safetext")        );    }     #[test]    fn keeps_tabs() {        let prompt = ExternalSourcePrompt::new("keep\tindentation");         assert_eq!(            prompt.as_ref().map(ExternalSourcePrompt::as_str),            Some("keep\tindentation")        );    }     #[test]    fn drops_empty_prompt() {        assert_eq!(ExternalSourcePrompt::new(""), None);    }     #[test]    fn drops_prompt_with_only_removed_characters() {        assert_eq!(            ExternalSourcePrompt::new("\u{202E}\u{202C}\u{0000}\u{001B}"),            None        );    }}