sqz_engine/
json_projection.rs

1/// Schema-Aware JSON Projection — strips JSON to only the fields
2/// relevant to the current context.
3///
4/// Unlike `strip_nulls` (removes null values) or `keep_fields` (requires
5/// an explicit field list), projection automatically removes low-value
6/// fields by default:
7///
8/// - Bookkeeping fields by exact name: `__v`, `etag`, request and
9///   correlation ids
10/// - Empty collections and strings: `[]`, `{}`, `""`
11///
12/// Prefix rules, same-day timestamp dedup, and depth summaries are opt-in.
13/// Names like `_id`, `_source`, `_embedded`, `_links`, `cursor`, or
14/// `traceback` carry the payload, the record id, or the next page in real
15/// APIs, so the defaults leave them alone.
16
17use crate::error::Result;
18
19/// Configuration for JSON projection.
20#[derive(Debug, Clone)]
21pub struct ProjectionConfig {
22    /// Remove fields matching these prefixes (case-insensitive).
23    pub strip_prefixes: Vec<String>,
24    /// Remove fields matching these exact names (case-insensitive).
25    pub strip_names: Vec<String>,
26    /// Maximum nesting depth to preserve. Objects deeper than this
27    /// are replaced with `{...N keys}`. Default: 64.
28    pub max_depth: usize,
29    /// Remove empty arrays and objects. Default: true.
30    pub strip_empty: bool,
31    /// Keep one timestamp per calendar day per object. Default: false,
32    /// since `closed_at` on the day an issue opened still matters.
33    pub dedup_timestamps: bool,
34}
35
36impl Default for ProjectionConfig {
37    fn default() -> Self {
38        Self {
39            strip_prefixes: Vec::new(),
40            strip_names: vec![
41                "__v".to_string(),
42                "etag".to_string(),
43                "request_id".to_string(),
44                "x_request_id".to_string(),
45                "correlation_id".to_string(),
46            ],
47            // Guards pathological nesting only; real payloads go deeper than 5.
48            max_depth: 64,
49            strip_empty: true,
50            dedup_timestamps: false,
51        }
52    }
53}
54
55/// Result of JSON projection.
56#[derive(Debug, Clone)]
57pub struct ProjectionResult {
58    /// The projected JSON string.
59    pub data: String,
60    /// Number of fields removed.
61    pub fields_removed: usize,
62    /// Estimated tokens saved.
63    pub tokens_saved: u32,
64}
65
66/// Apply schema-aware projection to a JSON string.
67///
68/// Returns the projected JSON and stats, or the original string unchanged
69/// if it's not valid JSON or projection doesn't help.
70pub fn project_json(input: &str, config: &ProjectionConfig) -> Result<ProjectionResult> {
71    let trimmed = input.trim();
72    let mut value: serde_json::Value = match serde_json::from_str(trimmed) {
73        Ok(v) => v,
74        Err(_) => {
75            return Ok(ProjectionResult {
76                data: input.to_string(),
77                fields_removed: 0,
78                tokens_saved: 0,
79            });
80        }
81    };
82
83    let original_tokens = estimate_tokens(input);
84    let mut fields_removed = 0;
85
86    project_value(&mut value, config, 0, &mut fields_removed);
87
88    let projected = serde_json::to_string(&value)
89        .unwrap_or_else(|_| input.to_string());
90
91    let projected_tokens = estimate_tokens(&projected);
92    let tokens_saved = original_tokens.saturating_sub(projected_tokens);
93
94    // Only return projected version if it's actually smaller or fields were removed
95    if projected.len() < input.len() || fields_removed > 0 {
96        Ok(ProjectionResult {
97            data: projected,
98            fields_removed,
99            tokens_saved,
100        })
101    } else {
102        Ok(ProjectionResult {
103            data: input.to_string(),
104            fields_removed: 0,
105            tokens_saved: 0,
106        })
107    }
108}
109
110/// Recursively project a JSON value, removing low-value fields.
111fn project_value(
112    value: &mut serde_json::Value,
113    config: &ProjectionConfig,
114    depth: usize,
115    removed: &mut usize,
116) {
117    match value {
118        serde_json::Value::Object(map) => {
119            // At max depth, replace deep objects with a summary
120            if depth >= config.max_depth {
121                let key_count = map.len();
122                if key_count > 0 {
123                    map.clear();
124                    map.insert(
125                        "_sqz_summary".to_string(),
126                        serde_json::Value::String(format!("{{...{key_count} keys}}")),
127                    );
128                    *removed += key_count;
129                }
130                return;
131            }
132
133            let keys_to_remove: Vec<String> = map
134                .keys()
135                .filter(|k| should_strip_field(k, config))
136                .cloned()
137                .collect();
138
139            for key in &keys_to_remove {
140                map.remove(key);
141                *removed += 1;
142            }
143
144            // Strip empty collections
145            if config.strip_empty {
146                let empty_keys: Vec<String> = map
147                    .iter()
148                    .filter(|(_, v)| is_empty_collection(v))
149                    .map(|(k, _)| k.clone())
150                    .collect();
151                for key in &empty_keys {
152                    map.remove(key);
153                    *removed += 1;
154                }
155            }
156
157            // Dedup timestamps: if multiple timestamp fields exist on the same
158            // day, keep only the most descriptive one
159            if config.dedup_timestamps {
160                dedup_timestamps(map, removed);
161            }
162
163            // Recurse into remaining values
164            for v in map.values_mut() {
165                project_value(v, config, depth + 1, removed);
166            }
167        }
168        serde_json::Value::Array(arr) => {
169            for item in arr.iter_mut() {
170                project_value(item, config, depth + 1, removed);
171            }
172        }
173        _ => {}
174    }
175}
176
177/// Check if a field name should be stripped based on the config.
178fn should_strip_field(name: &str, config: &ProjectionConfig) -> bool {
179    let lower = name.to_lowercase();
180
181    // Check exact name matches
182    for strip_name in &config.strip_names {
183        if lower == strip_name.to_lowercase() {
184            return true;
185        }
186    }
187
188    // Check prefix matches
189    for prefix in &config.strip_prefixes {
190        let prefix_lower = prefix.to_lowercase();
191        if lower.starts_with(&prefix_lower) && lower != prefix_lower {
192            // Don't strip if the field IS the prefix (e.g., don't strip "id" for prefix "_")
193            // But DO strip "_id", "debug_info", etc.
194            return true;
195        }
196    }
197
198    false
199}
200
201/// Check if a value is an empty collection.
202fn is_empty_collection(value: &serde_json::Value) -> bool {
203    match value {
204        serde_json::Value::Array(arr) => arr.is_empty(),
205        serde_json::Value::Object(map) => map.is_empty(),
206        serde_json::Value::String(s) => s.is_empty(),
207        _ => false,
208    }
209}
210
211/// Remove redundant timestamp fields. If multiple fields end in `_at` or `_date`
212/// and have the same date prefix (YYYY-MM-DD), keep only the first one.
213fn dedup_timestamps(
214    map: &mut serde_json::Map<String, serde_json::Value>,
215    removed: &mut usize,
216) {
217    let timestamp_fields: Vec<(String, String)> = map
218        .iter()
219        .filter_map(|(k, v)| {
220            if (k.ends_with("_at") || k.ends_with("_date") || k.ends_with("_time"))
221                && v.is_string()
222            {
223                let date_prefix = v.as_str().unwrap_or("").chars().take(10).collect::<String>();
224                if date_prefix.len() == 10 && date_prefix.contains('-') {
225                    return Some((k.clone(), date_prefix));
226                }
227            }
228            None
229        })
230        .collect();
231
232    if timestamp_fields.len() <= 1 {
233        return;
234    }
235
236    // Group by date prefix
237    let mut seen_dates: std::collections::HashSet<String> = std::collections::HashSet::new();
238    let mut first_field_per_date: std::collections::HashMap<String, String> = std::collections::HashMap::new();
239    let mut to_remove = Vec::new();
240
241    for (field, date) in &timestamp_fields {
242        if seen_dates.contains(date) {
243            // This date was already seen — remove this field unless it's the "primary" one
244            let primary = first_field_per_date.get(date).unwrap();
245            // Keep the more descriptive field (created_at > updated_at > modified_at)
246            let dominated = if field.contains("created") {
247                // created_at is more important — remove the previous one
248                to_remove.push(primary.clone());
249                false
250            } else {
251                true
252            };
253            if dominated {
254                to_remove.push(field.clone());
255            }
256        } else {
257            seen_dates.insert(date.clone());
258            first_field_per_date.insert(date.clone(), field.clone());
259        }
260    }
261
262    for field in &to_remove {
263        map.remove(field);
264        *removed += 1;
265    }
266}
267
268fn estimate_tokens(text: &str) -> u32 {
269    ((text.len() as f64) / 4.0).ceil() as u32
270}
271
272// ── Tests ─────────────────────────────────────────────────────────────────
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277    use serde_json::json;
278
279    fn project_default(input: serde_json::Value) -> (serde_json::Value, usize) {
280        let result =
281            project_json(&serde_json::to_string(&input).unwrap(), &ProjectionConfig::default()).unwrap();
282        (serde_json::from_str(&result.data).unwrap(), result.fields_removed)
283    }
284
285    #[test]
286    fn test_strips_bookkeeping_fields_only() {
287        let (parsed, removed) = project_default(json!({
288            "id": 1,
289            "name": "Alice",
290            "_id": "abc123",
291            "__v": 3,
292            "etag": "W/\"1\"",
293            "debug_info": "verbose stuff",
294            "internal_state": "hidden"
295        }));
296        assert!(parsed.get("id").is_some(), "id should be kept");
297        assert!(parsed.get("_id").is_some(), "a Mongo _id is the record id");
298        assert!(parsed.get("debug_info").is_some(), "no prefix guessing by default");
299        assert!(parsed.get("internal_state").is_some());
300        assert!(parsed.get("__v").is_none(), "__v should be stripped");
301        assert!(parsed.get("etag").is_none(), "etag should be stripped");
302        assert_eq!(removed, 2);
303    }
304
305    #[test]
306    fn test_keeps_payload_fields_real_apis_put_behind_underscores() {
307        let (parsed, _) = project_default(json!({
308            "hits": {"hits": [{"_id": "doc-1", "_score": 9.5, "_source": {"title": "Runbook"}}]},
309            "_embedded": {"orders": [{"id": 7}]},
310            "_links": {"next": {"href": "/orders?page=2"}},
311            "cursor": "c2",
312            "traceback": "File \"app.py\", line 3",
313            "__typename": "Issue"
314        }));
315        assert_eq!(parsed["hits"]["hits"][0]["_source"]["title"], "Runbook");
316        assert_eq!(parsed["_embedded"]["orders"][0]["id"], 7);
317        assert_eq!(parsed["_links"]["next"]["href"], "/orders?page=2");
318        assert_eq!(parsed["cursor"], "c2");
319        assert!(parsed.get("traceback").is_some());
320        assert_eq!(parsed["__typename"], "Issue");
321    }
322
323    #[test]
324    fn test_keeps_deeply_nested_text() {
325        // Notion search: the title is six levels down.
326        let (parsed, _) = project_default(json!({
327            "results": [{"properties": {"title": {"title": [{"plain_text": "Design note"}]}}}]
328        }));
329        assert_eq!(
330            parsed["results"][0]["properties"]["title"]["title"][0]["plain_text"],
331            "Design note"
332        );
333    }
334
335    #[test]
336    fn test_keeps_same_day_timestamps_by_default() {
337        let (parsed, removed) = project_default(json!({
338            "created_at": "2024-01-15T10:00:00Z",
339            "closed_at": "2024-01-15T14:30:00Z"
340        }));
341        assert!(parsed.get("closed_at").is_some());
342        assert_eq!(removed, 0);
343    }
344
345    #[test]
346    fn test_strips_empty_collections() {
347        let input = json!({
348            "name": "Bob",
349            "tags": [],
350            "metadata": {},
351            "bio": ""
352        });
353        let config = ProjectionConfig::default();
354        let result = project_json(&serde_json::to_string(&input).unwrap(), &config).unwrap();
355        let parsed: serde_json::Value = serde_json::from_str(&result.data).unwrap();
356        assert!(parsed.get("name").is_some());
357        assert!(parsed.get("tags").is_none(), "empty array should be stripped");
358        assert!(parsed.get("metadata").is_none(), "empty object should be stripped");
359        assert!(parsed.get("bio").is_none(), "empty string should be stripped");
360    }
361
362    #[test]
363    fn test_max_depth_truncation() {
364        let input = json!({
365            "a": {"b": {"c": {"d": {"e": {"f": "deep"}}}}}
366        });
367        let config = ProjectionConfig {
368            max_depth: 3,
369            ..Default::default()
370        };
371        let result = project_json(&serde_json::to_string(&input).unwrap(), &config).unwrap();
372        // At depth 3, the nested object should be replaced with a summary
373        let parsed: serde_json::Value = serde_json::from_str(&result.data).unwrap();
374        // Navigate to depth 3: a -> b -> c (this is where truncation happens)
375        let at_depth = &parsed["a"]["b"]["c"];
376        assert!(
377            at_depth.get("_sqz_summary").is_some() || result.fields_removed > 0,
378            "deep nesting should be truncated at max_depth: {:?}", parsed
379        );
380    }
381
382    #[test]
383    fn test_dedup_timestamps() {
384        let input = json!({
385            "name": "Alice",
386            "created_at": "2024-01-15T10:00:00Z",
387            "updated_at": "2024-01-15T14:30:00Z",
388            "modified_at": "2024-01-15T14:30:00Z"
389        });
390        let config = ProjectionConfig {
391            dedup_timestamps: true,
392            ..Default::default()
393        };
394        let result = project_json(&serde_json::to_string(&input).unwrap(), &config).unwrap();
395        let parsed: serde_json::Value = serde_json::from_str(&result.data).unwrap();
396        // Should keep created_at, drop one of the redundant same-day timestamps
397        assert!(parsed.get("created_at").is_some(), "created_at should be kept");
398        assert!(result.fields_removed > 0, "redundant timestamps should be removed");
399    }
400
401    #[test]
402    fn test_non_json_passthrough() {
403        let input = "not json at all";
404        let config = ProjectionConfig::default();
405        let result = project_json(input, &config).unwrap();
406        assert_eq!(result.data, input);
407        assert_eq!(result.fields_removed, 0);
408    }
409
410    #[test]
411    fn test_preserves_important_fields() {
412        let input = json!({
413            "id": 42,
414            "name": "Alice",
415            "email": "alice@example.com",
416            "role": "admin",
417            "status": "active"
418        });
419        let config = ProjectionConfig::default();
420        let result = project_json(&serde_json::to_string(&input).unwrap(), &config).unwrap();
421        let parsed: serde_json::Value = serde_json::from_str(&result.data).unwrap();
422        assert_eq!(parsed["id"], 42);
423        assert_eq!(parsed["name"], "Alice");
424        assert_eq!(parsed["role"], "admin");
425    }
426
427    #[test]
428    fn test_custom_strip_prefixes() {
429        let input = json!({
430            "name": "Alice",
431            "tmp_cache": "data",
432            "tmp_buffer": "more data"
433        });
434        let config = ProjectionConfig {
435            strip_prefixes: vec!["tmp_".to_string()],
436            ..Default::default()
437        };
438        let result = project_json(&serde_json::to_string(&input).unwrap(), &config).unwrap();
439        let parsed: serde_json::Value = serde_json::from_str(&result.data).unwrap();
440        assert!(parsed.get("name").is_some());
441        assert!(parsed.get("tmp_cache").is_none());
442        assert!(parsed.get("tmp_buffer").is_none());
443    }
444
445    #[test]
446    fn test_nested_projection() {
447        let input = json!({
448            "user": {
449                "id": 1,
450                "name": "Alice",
451                "_internal_id": "xyz",
452                "debug_flags": [1, 2, 3]
453            }
454        });
455        let config = ProjectionConfig {
456            strip_prefixes: vec!["_".to_string(), "debug".to_string()],
457            ..Default::default()
458        };
459        let result = project_json(&serde_json::to_string(&input).unwrap(), &config).unwrap();
460        let parsed: serde_json::Value = serde_json::from_str(&result.data).unwrap();
461        assert!(parsed["user"].get("id").is_some());
462        assert!(parsed["user"].get("_internal_id").is_none());
463        assert!(parsed["user"].get("debug_flags").is_none());
464    }
465
466    use proptest::prelude::*;
467
468    proptest! {
469        /// Projection never produces invalid JSON from valid JSON input.
470        #[test]
471        fn prop_projection_produces_valid_json(
472            key1 in "[a-z]{3,10}",
473            key2 in "[a-z]{3,10}",
474            val in "[a-z0-9 ]{1,50}",
475        ) {
476            let input = format!(r#"{{"{key1}":"{val}","{key2}":42}}"#);
477            let config = ProjectionConfig::default();
478            let result = project_json(&input, &config).unwrap();
479            let parsed: std::result::Result<serde_json::Value, _> = serde_json::from_str(&result.data);
480            prop_assert!(parsed.is_ok(), "projection output must be valid JSON");
481        }
482
483        /// Fields removed count is non-negative.
484        #[test]
485        fn prop_fields_removed_non_negative(
486            val in "[a-z]{1,20}",
487        ) {
488            let input = format!(r#"{{"name":"{val}","_debug":"x","__v":1}}"#);
489            let config = ProjectionConfig::default();
490            let result = project_json(&input, &config).unwrap();
491            // fields_removed is usize, always >= 0
492            let _ = result.fields_removed;
493        }
494    }
495}
sqz_engine/json_projection.rs

sqz_engine/
json_projection.rs
