agentic_forge_core/query/
intent.rs1use 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}