sqlite_graphrag/agent_surface/
shape.rs

1//! Stateless reshaping primitives applied to the result array of an envelope.
2//!
3//! Every function here takes and returns plain `serde_json` data so a single
4//! call site in [`crate::output`] can serve every subcommand. None of them
5//! knows what a memory, an entity or a hit is.
6
7use super::filter::{lookup, scalar_text, FilterExpr};
8use serde_json::{Map, Value};
9
10/// Keeps only the elements accepted by every predicate.
11pub fn filter(items: Vec<Value>, filters: &[FilterExpr]) -> Vec<Value> {
12    if filters.is_empty() {
13        return items;
14    }
15    items
16        .into_iter()
17        .filter(|item| super::filter::matches_all(filters, item))
18        .collect()
19}
20
21/// Sorts elements ascending by the scalar found at the dotted `key`.
22///
23/// Numeric values compare numerically, everything else compares as text.
24/// Elements without the key keep their relative order at the end of the list,
25/// so a partially populated payload never loses rows to sorting.
26pub fn sort(mut items: Vec<Value>, key: &str) -> Vec<Value> {
27    let path: Vec<String> = key.split('.').map(str::to_string).collect();
28    items.sort_by(|a, b| {
29        let left = lookup(a, &path);
30        let right = lookup(b, &path);
31        match (left, right) {
32            (None, None) => std::cmp::Ordering::Equal,
33            (None, Some(_)) => std::cmp::Ordering::Greater,
34            (Some(_), None) => std::cmp::Ordering::Less,
35            (Some(l), Some(r)) => compare(l, r),
36        }
37    });
38    items
39}
40
41/// Total order over two JSON scalars used by [`sort`].
42fn compare(left: &Value, right: &Value) -> std::cmp::Ordering {
43    if let (Some(l), Some(r)) = (left.as_f64(), right.as_f64()) {
44        return l.partial_cmp(&r).unwrap_or(std::cmp::Ordering::Equal);
45    }
46    match (scalar_text(left), scalar_text(right)) {
47        (Some(l), Some(r)) => l.cmp(&r),
48        (Some(_), None) => std::cmp::Ordering::Less,
49        (None, Some(_)) => std::cmp::Ordering::Greater,
50        (None, None) => std::cmp::Ordering::Equal,
51    }
52}
53
54/// Drops later elements whose scalar at `key` was already seen.
55///
56/// Elements lacking the key are always kept: dropping them would silently
57/// collapse rows that were never proven duplicate.
58pub fn dedupe(items: Vec<Value>, key: &str) -> Vec<Value> {
59    let path: Vec<String> = key.split('.').map(str::to_string).collect();
60    let mut seen = std::collections::HashSet::new();
61    let mut out = Vec::with_capacity(items.len());
62    for item in items {
63        match lookup(&item, &path).and_then(scalar_text) {
64            Some(text) => {
65                if seen.insert(text) {
66                    out.push(item);
67                }
68            }
69            None => out.push(item),
70        }
71    }
72    out
73}
74
75/// Truncates the list to at most `max` elements. `0` means "no cap".
76pub fn limit(mut items: Vec<Value>, max: usize) -> Vec<Value> {
77    if max > 0 && items.len() > max {
78        items.truncate(max);
79    }
80    items
81}
82
83/// Rewrites each object element to carry only `keys`, in the requested order.
84///
85/// Keys absent from an element are skipped rather than emitted as `null`, so a
86/// projection never invents fields. Non-object elements pass through unchanged.
87pub fn project(items: Vec<Value>, keys: &[String]) -> Vec<Value> {
88    if keys.is_empty() {
89        return items;
90    }
91    items
92        .into_iter()
93        .map(|item| project_one(item, keys))
94        .collect()
95}
96
97/// Projects a single value; see [`project`].
98pub fn project_one(item: Value, keys: &[String]) -> Value {
99    if keys.is_empty() {
100        return item;
101    }
102    let Value::Object(_) = &item else {
103        return item;
104    };
105    let mut out = Map::new();
106    for key in keys {
107        let path: Vec<String> = key.split('.').map(str::to_string).collect();
108        if let Some(found) = lookup(&item, &path) {
109            out.insert(key.clone(), found.clone());
110        }
111    }
112    Value::Object(out)
113}
114
115/// Shortens every string longer than `max` characters, recursively.
116///
117/// Returns `true` when at least one string was cut, so the caller can flag the
118/// envelope. Truncation counts characters, not bytes, and therefore never
119/// splits a UTF-8 sequence.
120pub fn truncate_strings(value: &mut Value, max: usize) -> bool {
121    if max == 0 {
122        return false;
123    }
124    match value {
125        Value::String(s) => {
126            if s.chars().count() > max {
127                let cut = s
128                    .char_indices()
129                    .nth(max)
130                    .map_or(s.len(), |(byte_idx, _)| byte_idx);
131                s.truncate(cut);
132                true
133            } else {
134                false
135            }
136        }
137        Value::Array(items) => {
138            let mut hit = false;
139            for item in items {
140                hit |= truncate_strings(item, max);
141            }
142            hit
143        }
144        Value::Object(map) => {
145            let mut hit = false;
146            for (_, item) in map.iter_mut() {
147                hit |= truncate_strings(item, max);
148            }
149            hit
150        }
151        Value::Null | Value::Bool(_) | Value::Number(_) => false,
152    }
153}
sqlite_graphrag/agent_surface/shape.rs

sqlite_graphrag/agent_surface/
shape.rs
