llm_toolkit/intent/
frame.rs

1use crate::extract::FlexibleExtractor;
2use crate::extract::core::ContentExtractor;
3use crate::intent::{IntentError, IntentExtractor};
4use std::str::FromStr;
5
6/// A frame-based intent extractor that uses separate tags for input wrapping and extraction.
7pub struct IntentFrame {
8    input_tag: String,
9    extractor_tag: String,
10}
11
12impl IntentFrame {
13    /// Creates a new `IntentFrame` with specified input and extractor tags.
14    pub fn new(input_tag: &str, extractor_tag: &str) -> Self {
15        Self {
16            input_tag: input_tag.to_string(),
17            extractor_tag: extractor_tag.to_string(),
18        }
19    }
20
21    /// Wraps the given text with the input tag.
22    pub fn wrap(&self, text: &str) -> String {
23        format!("<{0}>{1}</{0}>", self.input_tag, text)
24    }
25}
26
27impl<T> IntentExtractor<T> for IntentFrame
28where
29    T: FromStr,
30{
31    fn extract_intent(&self, text: &str) -> Result<T, IntentError> {
32        // Use FlexibleExtractor to get the string inside the extractor tag
33        let extractor = FlexibleExtractor::new();
34        let extracted_str = extractor
35            .extract_tagged(text, &self.extractor_tag)
36            .ok_or_else(|| IntentError::TagNotFound {
37                tag: self.extractor_tag.clone(),
38            })?;
39
40        // Parse the string into the user's type
41        T::from_str(&extracted_str).map_err(|_| IntentError::ParseFailed {
42            value: extracted_str.to_string(),
43        })
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    #[derive(Debug, PartialEq)]
52    enum TestIntent {
53        Login,
54        Logout,
55    }
56
57    impl FromStr for TestIntent {
58        type Err = String;
59
60        fn from_str(s: &str) -> Result<Self, Self::Err> {
61            match s {
62                "Login" => Ok(TestIntent::Login),
63                "Logout" => Ok(TestIntent::Logout),
64                _ => Err(format!("Unknown intent: {}", s)),
65            }
66        }
67    }
68
69    #[test]
70    fn test_wrap_method() {
71        let frame = IntentFrame::new("input", "output");
72        let wrapped = frame.wrap("test content");
73        assert_eq!(wrapped, "<input>test content</input>");
74    }
75
76    #[test]
77    fn test_extract_intent_success() {
78        let frame = IntentFrame::new("input", "intent");
79        let text = "<intent>Login</intent>";
80        let result: Result<TestIntent, _> = IntentExtractor::extract_intent(&frame, text);
81        assert_eq!(result.unwrap(), TestIntent::Login);
82    }
83
84    #[test]
85    fn test_extract_intent_tag_not_found() {
86        let frame = IntentFrame::new("input", "intent");
87        let text = "No intent tag here";
88        let result: Result<TestIntent, _> = IntentExtractor::extract_intent(&frame, text);
89
90        match result {
91            Err(IntentError::TagNotFound { tag }) => {
92                assert_eq!(tag, "intent");
93            }
94            _ => panic!("Expected TagNotFound error"),
95        }
96    }
97
98    #[test]
99    fn test_extract_intent_parse_failed() {
100        let frame = IntentFrame::new("input", "intent");
101        let text = "<intent>Invalid</intent>";
102        let result: Result<TestIntent, _> = IntentExtractor::extract_intent(&frame, text);
103
104        match result {
105            Err(IntentError::ParseFailed { value }) => {
106                assert_eq!(value, "Invalid");
107            }
108            _ => panic!("Expected ParseFailed error"),
109        }
110    }
111}
llm_toolkit/intent/frame.rs

llm_toolkit/intent/
frame.rs
