somatize_core/
schema.rs

1//! Schema — dtype and shape for compile-time type checking between filters.
2//!
3//! The compiler validates that connected filters have compatible schemas
4//! before execution begins, catching shape/type mismatches early.
5
6use serde::{Deserialize, Serialize};
7use std::fmt;
8
9/// Primitive data types that Soma values can contain.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
11#[non_exhaustive]
12pub enum DataType {
13    /// 64-bit floating point.
14    Float64,
15    /// 32-bit floating point.
16    Float32,
17    /// 64-bit signed integer.
18    Int64,
19    /// Boolean.
20    Bool,
21    /// UTF-8 string.
22    Utf8,
23    /// Raw bytes.
24    Bytes,
25    /// Structured JSON (any shape).
26    Json,
27    /// A conversation: a list of [`crate::message::Message`].
28    ///
29    /// Distinct from `Json` so the compiler can reject an edge that hands a
30    /// tensor, or an arbitrary document, to a node expecting a conversation.
31    Messages,
32}
33
34impl DataType {
35    /// Could a value of this type ever be read as `target`?
36    ///
37    /// This is coarser than equality on purpose. `is_compatible_with`
38    /// demands an exact dtype match, which makes every mismatch equally
39    /// suspicious — `f32` meeting `f64` and a conversation meeting a tensor
40    /// both come out as "not compatible", so neither can be more than a
41    /// warning without breaking the first case.
42    ///
43    /// This answers the stronger question: is there *any* reading under
44    /// which this connection makes sense? Numeric widths differ but describe
45    /// the same thing; a tensor and a conversation do not. The second kind
46    /// is what fails a multi-agent handoff, and it is worth refusing to
47    /// compile rather than warning about.
48    ///
49    /// The permitted coercions are exactly the ones the runtime performs:
50    /// - anything → `Json` (every `Value` has a `to_plain_json`)
51    /// - `Json` → anything (it is the dynamic type; the reader checks)
52    /// - `Utf8` → `Messages` (a bare prompt becomes a user turn)
53    /// - `Messages` → `Utf8` (a conversation's prose)
54    /// - numeric ↔ numeric (widths differ; meaning does not)
55    pub fn can_coerce_to(&self, target: &DataType) -> bool {
56        use DataType::*;
57
58        if self == target {
59            return true;
60        }
61        // Json is the dynamic type: it absorbs and yields anything.
62        if matches!(self, Json) || matches!(target, Json) {
63            return true;
64        }
65        if self.is_numeric() && target.is_numeric() {
66            return true;
67        }
68        matches!((self, target), (Utf8, Messages) | (Messages, Utf8))
69    }
70
71    /// Numeric in the sense [`Self::can_coerce_to`] uses: `Float64`,
72    /// `Float32`, `Int64` — and `Bool`, which tensors carry as 0/1.
73    pub fn is_numeric(&self) -> bool {
74        matches!(
75            self,
76            Self::Float64 | Self::Float32 | Self::Int64 | Self::Bool
77        )
78    }
79}
80
81impl fmt::Display for DataType {
82    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
83        match self {
84            Self::Float64 => write!(f, "f64"),
85            Self::Float32 => write!(f, "f32"),
86            Self::Int64 => write!(f, "i64"),
87            Self::Bool => write!(f, "bool"),
88            Self::Utf8 => write!(f, "str"),
89            Self::Bytes => write!(f, "bytes"),
90            Self::Json => write!(f, "json"),
91            Self::Messages => write!(f, "messages"),
92        }
93    }
94}
95
96/// Describes the shape and type of a Value, without holding the actual data.
97///
98/// Used by:
99/// - Filters: declare what they accept (input) and produce (output)
100/// - Compiler: validate type compatibility between connected filters
101/// - VirtualValue: know schema without materializing
102/// - Cache metadata: describe stored entries
103#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
104pub struct Schema {
105    /// The primitive data type.
106    pub dtype: DataType,
107
108    /// Shape dimensions. Empty for scalars, `[n]` for vectors, `[r,c]` for matrices, etc.
109    /// `None` means shape is dynamic/unknown.
110    pub shape: Option<Vec<Dimension>>,
111}
112
113/// A single dimension in a tensor shape.
114#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
115pub enum Dimension {
116    /// Fixed size (e.g., 128 features).
117    Fixed(usize),
118    /// Dynamic size (e.g., batch dimension). Named for documentation.
119    Dynamic(String),
120}
121
122impl fmt::Display for Dimension {
123    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
124        match self {
125            Self::Fixed(n) => write!(f, "{n}"),
126            Self::Dynamic(name) => write!(f, "{name}"),
127        }
128    }
129}
130
131impl Schema {
132    /// Create a schema for a 1D tensor (vector) of known length.
133    pub fn vector(dtype: DataType, len: usize) -> Self {
134        Self {
135            dtype,
136            shape: Some(vec![Dimension::Fixed(len)]),
137        }
138    }
139
140    /// Create a schema for a 2D tensor (matrix) with known dimensions.
141    pub fn matrix(dtype: DataType, rows: usize, cols: usize) -> Self {
142        Self {
143            dtype,
144            shape: Some(vec![Dimension::Fixed(rows), Dimension::Fixed(cols)]),
145        }
146    }
147
148    /// Create a schema for a tensor with a dynamic batch dimension.
149    pub fn batched(dtype: DataType, feature_dims: &[usize]) -> Self {
150        let mut dims = vec![Dimension::Dynamic("batch".into())];
151        dims.extend(feature_dims.iter().map(|&d| Dimension::Fixed(d)));
152        Self {
153            dtype,
154            shape: Some(dims),
155        }
156    }
157
158    /// Create a schema for a scalar value.
159    pub fn scalar(dtype: DataType) -> Self {
160        Self {
161            dtype,
162            shape: Some(vec![]),
163        }
164    }
165
166    /// Create a schema for JSON data (shape is irrelevant).
167    pub fn json() -> Self {
168        Self {
169            dtype: DataType::Json,
170            shape: None,
171        }
172    }
173
174    /// Create a schema for UTF-8 text (shape is irrelevant).
175    ///
176    /// This is what a prompt or a completion carries. An edge typed `text`
177    /// will not accept a tensor, which is how a mis-wired handoff between two
178    /// agent nodes becomes a compile error rather than a runtime surprise.
179    pub fn text() -> Self {
180        Self {
181            dtype: DataType::Utf8,
182            shape: None,
183        }
184    }
185
186    /// Create a schema for a conversation.
187    pub fn messages() -> Self {
188        Self {
189            dtype: DataType::Messages,
190            shape: None,
191        }
192    }
193
194    /// Create a schema for raw bytes.
195    pub fn bytes() -> Self {
196        Self {
197            dtype: DataType::Bytes,
198            shape: None,
199        }
200    }
201
202    /// Create a schema with fully dynamic (unknown) shape.
203    pub fn dynamic(dtype: DataType) -> Self {
204        Self { dtype, shape: None }
205    }
206
207    /// Is connecting these two definitely a mistake?
208    ///
209    /// True when no reading of `self` could satisfy `other` — a tensor
210    /// arriving where a conversation is expected, say. The compiler refuses
211    /// to build such a graph, rather than warning and letting it fail
212    /// mid-run once tokens have been spent.
213    pub fn is_incompatible_with(&self, other: &Schema) -> bool {
214        !self.dtype.can_coerce_to(&other.dtype)
215    }
216
217    /// Check if this schema is compatible with another (can be connected in a pipeline).
218    ///
219    /// Compatibility rules:
220    /// - Same dtype required (no implicit coercion)
221    /// - If both shapes are known, fixed dimensions must match
222    /// - Dynamic dimensions are compatible with any size
223    /// - Unknown shape (None) is compatible with anything of the same dtype
224    pub fn is_compatible_with(&self, other: &Schema) -> bool {
225        if self.dtype != other.dtype {
226            return false;
227        }
228
229        match (&self.shape, &other.shape) {
230            (None, _) | (_, None) => true, // unknown shape is flexible
231            (Some(a), Some(b)) => {
232                if a.len() != b.len() {
233                    return false;
234                }
235                a.iter().zip(b.iter()).all(|(da, db)| match (da, db) {
236                    (Dimension::Fixed(x), Dimension::Fixed(y)) => x == y,
237                    _ => true, // dynamic is compatible with anything
238                })
239            }
240        }
241    }
242
243    /// Number of known dimensions (rank).
244    pub fn rank(&self) -> Option<usize> {
245        self.shape.as_ref().map(|s| s.len())
246    }
247}
248
249impl fmt::Display for Schema {
250    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
251        write!(f, "{}", self.dtype)?;
252        if let Some(shape) = &self.shape {
253            if shape.is_empty() {
254                write!(f, " (scalar)")?;
255            } else {
256                let dims: Vec<String> = shape.iter().map(|d| d.to_string()).collect();
257                write!(f, "[{}]", dims.join(", "))?;
258            }
259        }
260        Ok(())
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    #[test]
269    fn schema_display() {
270        assert_eq!(
271            Schema::scalar(DataType::Float64).to_string(),
272            "f64 (scalar)"
273        );
274        assert_eq!(
275            Schema::vector(DataType::Float64, 128).to_string(),
276            "f64[128]"
277        );
278        assert_eq!(
279            Schema::matrix(DataType::Float64, 100, 50).to_string(),
280            "f64[100, 50]"
281        );
282        assert_eq!(
283            Schema::batched(DataType::Float32, &[128]).to_string(),
284            "f32[batch, 128]"
285        );
286        assert_eq!(Schema::json().to_string(), "json");
287    }
288
289    #[test]
290    fn compatible_same_schema() {
291        let s = Schema::vector(DataType::Float64, 128);
292        assert!(s.is_compatible_with(&s));
293    }
294
295    #[test]
296    fn compatible_dynamic_with_fixed() {
297        let dynamic = Schema::batched(DataType::Float64, &[128]);
298        let fixed = Schema::matrix(DataType::Float64, 32, 128);
299        assert!(dynamic.is_compatible_with(&fixed));
300        assert!(fixed.is_compatible_with(&dynamic));
301    }
302
303    #[test]
304    fn compatible_unknown_shape() {
305        let unknown = Schema::dynamic(DataType::Float64);
306        let known = Schema::vector(DataType::Float64, 128);
307        assert!(unknown.is_compatible_with(&known));
308        assert!(known.is_compatible_with(&unknown));
309    }
310
311    #[test]
312    fn incompatible_different_dtype() {
313        let f64_schema = Schema::vector(DataType::Float64, 128);
314        let i64_schema = Schema::vector(DataType::Int64, 128);
315        assert!(!f64_schema.is_compatible_with(&i64_schema));
316    }
317
318    #[test]
319    fn incompatible_different_fixed_dims() {
320        let a = Schema::vector(DataType::Float64, 128);
321        let b = Schema::vector(DataType::Float64, 256);
322        assert!(!a.is_compatible_with(&b));
323    }
324
325    #[test]
326    fn incompatible_different_rank() {
327        let vec = Schema::vector(DataType::Float64, 128);
328        let mat = Schema::matrix(DataType::Float64, 128, 64);
329        assert!(!vec.is_compatible_with(&mat));
330    }
331
332    #[test]
333    fn json_compatible_with_json() {
334        assert!(Schema::json().is_compatible_with(&Schema::json()));
335    }
336
337    #[test]
338    fn json_incompatible_with_tensor() {
339        assert!(!Schema::json().is_compatible_with(&Schema::vector(DataType::Float64, 10)));
340    }
341
342    #[test]
343    fn serde_roundtrip() {
344        let schemas = vec![
345            Schema::scalar(DataType::Float64),
346            Schema::vector(DataType::Float32, 100),
347            Schema::batched(DataType::Float64, &[128, 64]),
348            Schema::json(),
349            Schema::dynamic(DataType::Int64),
350        ];
351        for s in schemas {
352            let json = serde_json::to_string(&s).unwrap();
353            let deserialized: Schema = serde_json::from_str(&json).unwrap();
354            assert_eq!(s, deserialized);
355        }
356    }
357
358    #[test]
359    fn rank() {
360        assert_eq!(Schema::scalar(DataType::Float64).rank(), Some(0));
361        assert_eq!(Schema::vector(DataType::Float64, 10).rank(), Some(1));
362        assert_eq!(Schema::matrix(DataType::Float64, 10, 5).rank(), Some(2));
363        assert_eq!(Schema::json().rank(), None);
364    }
365}
somatize_core/schema.rs

somatize_core/
schema.rs
