Skip to main content

agentic_forge_core/query/
intent.rs

1//! Intent declaration and scoped extraction.
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::collections::HashMap;
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
8pub enum ExtractionIntent {
9    Exists,
10    #[default]
11    IdsOnly,
12    Summary,
13    Fields(Vec<String>),
14    Full,
15}
16
17impl ExtractionIntent {
18    pub fn estimated_tokens(&self) -> u64 {
19        match self {
20            Self::Exists => 1,
21            Self::IdsOnly => 10,
22            Self::Summary => 50,
23            Self::Fields(f) => 20 * f.len() as u64,
24            Self::Full => 500,
25        }
26    }
27
28    pub fn from_label(s: &str) -> Self {
29        match s {
30            "exists" => Self::Exists,
31            "ids" | "ids_only" => Self::IdsOnly,
32            "summary" => Self::Summary,
33            "full" => Self::Full,
34            _ => Self::Full,
35        }
36    }
37
38    pub fn is_minimal(&self) -> bool {
39        matches!(self, Self::Exists | Self::IdsOnly)
40    }
41
42    pub fn includes_content(&self) -> bool {
43        matches!(self, Self::Summary | Self::Fields(_) | Self::Full)
44    }
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub enum ScopedResult {
49    Bool(bool),
50    Id(String),
51    Ids(Vec<String>),
52    Summary(String),
53    Fields(HashMap<String, Value>),
54    Full(Value),
55    Count(usize),
56}
57
58impl ScopedResult {
59    pub fn estimated_tokens(&self) -> u64 {
60        match self {
61            Self::Bool(_) => 1,
62            Self::Id(_) => 5,
63            Self::Ids(ids) => ids.len() as u64 * 5,
64            Self::Summary(_) => 50,
65            Self::Fields(f) => f.len() as u64 * 20,
66            Self::Full(v) => serde_json::to_string(v)
67                .map(|s| s.len() as u64 / 4)
68                .unwrap_or(500),
69            Self::Count(_) => 2,
70        }
71    }
72}
73
74pub trait Scopeable {
75    fn id_str(&self) -> String;
76    fn summarize(&self) -> String;
77    fn extract_fields(&self, fields: &[String]) -> HashMap<String, Value>;
78    fn to_json(&self) -> Value;
79}
80
81pub fn apply_intent<T: Scopeable>(intent: &ExtractionIntent, item: &T) -> ScopedResult {
82    match intent {
83        ExtractionIntent::Exists => ScopedResult::Bool(true),
84        ExtractionIntent::IdsOnly => ScopedResult::Id(item.id_str()),
85        ExtractionIntent::Summary => ScopedResult::Summary(item.summarize()),
86        ExtractionIntent::Fields(f) => ScopedResult::Fields(item.extract_fields(f)),
87        ExtractionIntent::Full => ScopedResult::Full(item.to_json()),
88    }
89}
90
91pub fn apply_intent_many<T: Scopeable>(intent: &ExtractionIntent, items: &[T]) -> ScopedResult {
92    match intent {
93        ExtractionIntent::Exists => ScopedResult::Bool(!items.is_empty()),
94        ExtractionIntent::IdsOnly => ScopedResult::Ids(items.iter().map(|i| i.id_str()).collect()),
95        ExtractionIntent::Summary => ScopedResult::Count(items.len()),
96        _ => ScopedResult::Full(serde_json::json!(items
97            .iter()
98            .map(|i| i.to_json())
99            .collect::<Vec<_>>())),
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    struct MockItem {
108        id: String,
109        name: String,
110    }
111    impl Scopeable for MockItem {
112        fn id_str(&self) -> String {
113            self.id.clone()
114        }
115        fn summarize(&self) -> String {
116            format!("{}: {}", self.id, self.name)
117        }
118        fn extract_fields(&self, fields: &[String]) -> HashMap<String, Value> {
119            let mut map = HashMap::new();
120            for f in fields {
121                match f.as_str() {
122                    "name" => {
123                        map.insert("name".into(), Value::String(self.name.clone()));
124                    }
125                    _ => {}
126                }
127            }
128            map
129        }
130        fn to_json(&self) -> Value {
131            serde_json::json!({"id": self.id, "name": self.name})
132        }
133    }
134
135    #[test]
136    fn test_intent_exists() {
137        let item = MockItem {
138            id: "1".into(),
139            name: "Test".into(),
140        };
141        let result = apply_intent(&ExtractionIntent::Exists, &item);
142        assert!(matches!(result, ScopedResult::Bool(true)));
143        assert_eq!(result.estimated_tokens(), 1);
144    }
145
146    #[test]
147    fn test_intent_ids_only() {
148        let item = MockItem {
149            id: "42".into(),
150            name: "Test".into(),
151        };
152        let result = apply_intent(&ExtractionIntent::IdsOnly, &item);
153        assert!(matches!(result, ScopedResult::Id(ref s) if s == "42"));
154    }
155
156    #[test]
157    fn test_intent_summary() {
158        let item = MockItem {
159            id: "1".into(),
160            name: "Test".into(),
161        };
162        let result = apply_intent(&ExtractionIntent::Summary, &item);
163        assert!(matches!(result, ScopedResult::Summary(_)));
164    }
165
166    #[test]
167    fn test_intent_fields() {
168        let item = MockItem {
169            id: "1".into(),
170            name: "Hello".into(),
171        };
172        let result = apply_intent(&ExtractionIntent::Fields(vec!["name".into()]), &item);
173        if let ScopedResult::Fields(map) = result {
174            assert_eq!(map.get("name"), Some(&Value::String("Hello".into())));
175        } else {
176            panic!("Expected Fields");
177        }
178    }
179
180    #[test]
181    fn test_intent_full() {
182        let item = MockItem {
183            id: "1".into(),
184            name: "Full".into(),
185        };
186        let result = apply_intent(&ExtractionIntent::Full, &item);
187        assert!(matches!(result, ScopedResult::Full(_)));
188    }
189
190    #[test]
191    fn test_intent_default_is_minimal() {
192        let intent = ExtractionIntent::default();
193        assert!(intent.is_minimal());
194        assert!(!intent.includes_content());
195    }
196
197    #[test]
198    fn test_scoped_query_cheaper_than_full() {
199        let ids_cost = ExtractionIntent::IdsOnly.estimated_tokens();
200        let full_cost = ExtractionIntent::Full.estimated_tokens();
201        assert!(
202            ids_cost < full_cost / 10,
203            "IDs should be >10x cheaper than Full"
204        );
205    }
206
207    #[test]
208    fn test_intent_from_str() {
209        assert_eq!(
210            ExtractionIntent::from_label("exists"),
211            ExtractionIntent::Exists
212        );
213        assert_eq!(
214            ExtractionIntent::from_label("ids"),
215            ExtractionIntent::IdsOnly
216        );
217        assert_eq!(
218            ExtractionIntent::from_label("summary"),
219            ExtractionIntent::Summary
220        );
221        assert_eq!(ExtractionIntent::from_label("full"), ExtractionIntent::Full);
222        assert_eq!(
223            ExtractionIntent::from_label("unknown"),
224            ExtractionIntent::Full
225        );
226    }
227
228    #[test]
229    fn test_apply_intent_many_ids() {
230        let items = vec![
231            MockItem {
232                id: "1".into(),
233                name: "A".into(),
234            },
235            MockItem {
236                id: "2".into(),
237                name: "B".into(),
238            },
239        ];
240        let result = apply_intent_many(&ExtractionIntent::IdsOnly, &items);
241        if let ScopedResult::Ids(ids) = result {
242            assert_eq!(ids, vec!["1", "2"]);
243        } else {
244            panic!("Expected Ids");
245        }
246    }
247
248    #[test]
249    fn test_apply_intent_many_exists_empty() {
250        let items: Vec<MockItem> = vec![];
251        let result = apply_intent_many(&ExtractionIntent::Exists, &items);
252        assert!(matches!(result, ScopedResult::Bool(false)));
253    }
254}