Skip to main content

briefcase_core/
models.rs

1//! Core data models for AI decision tracking and observability
2//!
3//! This module contains the fundamental data structures used throughout the Briefcase AI system:
4//! - [`Input`] and [`Output`] for capturing AI function parameters and results
5//! - [`DecisionSnapshot`] for recording complete AI decision context
6//! - [`ModelParameters`] for tracking AI model configuration
7//! - [`ExecutionContext`] for environmental reproducibility
8//! - [`Snapshot`] for grouping related decisions
9//!
10//! ## Example Usage
11//!
12//! ```rust
13//! use briefcase_core::models::*;
14//! use serde_json::json;
15//!
16//! // Create input and output data
17//! let input = Input::new("prompt", json!("What is AI?"), "string");
18//! let output = Output::new("response", json!("AI is..."), "string")
19//!     .with_confidence(0.92);
20//!
21//! // Create a decision snapshot
22//! let decision = DecisionSnapshot::new("gpt_query")
23//!     .add_input(input)
24//!     .add_output(output)
25//!     .with_execution_time(234.5);
26//! ```
27
28use chrono::{DateTime, Utc};
29use serde::{Deserialize, Serialize};
30use sha2::{Digest, Sha256};
31use std::collections::HashMap;
32use uuid::Uuid;
33
34/// Input parameter to an AI decision point
35///
36/// Represents a single input parameter passed to an AI function or model.
37/// Inputs are captured with their name, value, and type information for
38/// complete reproducibility.
39///
40/// # Examples
41///
42/// ```rust
43/// use briefcase_core::models::Input;
44/// use serde_json::json;
45///
46/// let input = Input::new("user_query", json!("Hello world"), "string");
47/// assert_eq!(input.name, "user_query");
48/// assert_eq!(input.data_type, "string");
49/// ```
50#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
51pub struct Input {
52    pub name: String,
53    pub value: serde_json::Value,
54    pub data_type: String,
55    #[serde(default = "default_schema_version")]
56    pub schema_version: String,
57}
58
59impl Input {
60    pub fn new(
61        name: impl Into<String>,
62        value: serde_json::Value,
63        data_type: impl Into<String>,
64    ) -> Self {
65        Self {
66            name: name.into(),
67            value,
68            data_type: data_type.into(),
69            schema_version: default_schema_version(),
70        }
71    }
72}
73
74/// Output result from an AI decision point
75///
76/// Represents the result produced by an AI function or model.
77/// Outputs can include confidence scores and are captured with
78/// type information for analysis and replay.
79///
80/// # Examples
81///
82/// ```rust
83/// use briefcase_core::models::Output;
84/// use serde_json::json;
85///
86/// let output = Output::new("response", json!("Hello back!"), "string")
87///     .with_confidence(0.95);
88/// assert_eq!(output.confidence, Some(0.95));
89/// ```
90#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
91pub struct Output {
92    pub name: String,
93    pub value: serde_json::Value,
94    pub data_type: String,
95    pub confidence: Option<f64>, // 0.0 - 1.0
96    #[serde(default = "default_schema_version")]
97    pub schema_version: String,
98}
99
100impl Output {
101    pub fn new(
102        name: impl Into<String>,
103        value: serde_json::Value,
104        data_type: impl Into<String>,
105    ) -> Self {
106        Self {
107            name: name.into(),
108            value,
109            data_type: data_type.into(),
110            confidence: None,
111            schema_version: default_schema_version(),
112        }
113    }
114
115    pub fn with_confidence(mut self, confidence: f64) -> Self {
116        self.confidence = Some(confidence.clamp(0.0, 1.0));
117        self
118    }
119}
120
121/// AI model parameters for reproducibility and tracking
122///
123/// Captures the configuration of an AI model including its name, version,
124/// provider, and all parameters needed for reproducible execution.
125///
126/// # Examples
127///
128/// ```rust
129/// use briefcase_core::models::ModelParameters;
130/// use serde_json::json;
131///
132/// let params = ModelParameters::new("gpt-4")
133///     .with_provider("openai")
134///     .with_parameter("temperature", json!(0.7))
135///     .with_hyperparameter("max_tokens", json!(1000));
136/// ```
137#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
138pub struct ModelParameters {
139    pub model_name: String,
140    pub model_version: Option<String>,
141    pub provider: Option<String>, // openai, anthropic, etc.
142    #[serde(default)]
143    pub parameters: HashMap<String, serde_json::Value>,
144    #[serde(default)]
145    pub hyperparameters: HashMap<String, serde_json::Value>,
146    pub weights_hash: Option<String>, // For reproducibility verification
147}
148
149impl ModelParameters {
150    pub fn new(model_name: impl Into<String>) -> Self {
151        Self {
152            model_name: model_name.into(),
153            model_version: None,
154            provider: None,
155            parameters: HashMap::new(),
156            hyperparameters: HashMap::new(),
157            weights_hash: None,
158        }
159    }
160
161    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
162        self.provider = Some(provider.into());
163        self
164    }
165
166    pub fn with_version(mut self, version: impl Into<String>) -> Self {
167        self.model_version = Some(version.into());
168        self
169    }
170
171    pub fn with_parameter(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
172        self.parameters.insert(key.into(), value);
173        self
174    }
175
176    pub fn with_hyperparameter(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
177        self.hyperparameters.insert(key.into(), value);
178        self
179    }
180
181    pub fn with_weights_hash(mut self, hash: impl Into<String>) -> Self {
182        self.weights_hash = Some(hash.into());
183        self
184    }
185}
186
187/// Execution context for deterministic replay
188#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
189pub struct ExecutionContext {
190    pub runtime_version: Option<String>, // Python 3.11, Node 20, etc.
191    #[serde(default)]
192    pub dependencies: HashMap<String, String>,
193    pub random_seed: Option<i64>,
194    #[serde(default)]
195    pub environment_variables: HashMap<String, String>,
196    #[serde(default)]
197    pub hardware_info: HashMap<String, serde_json::Value>,
198}
199
200impl ExecutionContext {
201    pub fn new() -> Self {
202        Self::default()
203    }
204
205    pub fn with_runtime_version(mut self, version: impl Into<String>) -> Self {
206        self.runtime_version = Some(version.into());
207        self
208    }
209
210    pub fn with_dependency(mut self, name: impl Into<String>, version: impl Into<String>) -> Self {
211        self.dependencies.insert(name.into(), version.into());
212        self
213    }
214
215    pub fn with_random_seed(mut self, seed: i64) -> Self {
216        self.random_seed = Some(seed);
217        self
218    }
219
220    pub fn with_env_var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
221        self.environment_variables.insert(key.into(), value.into());
222        self
223    }
224}
225
226/// Metadata for a snapshot
227#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
228pub struct SnapshotMetadata {
229    pub snapshot_id: Uuid,
230    pub timestamp: DateTime<Utc>,
231    #[serde(default = "default_schema_version")]
232    pub schema_version: String,
233    pub sdk_version: String,
234    pub created_by: Option<String>,
235    pub checksum: Option<String>,
236}
237
238impl SnapshotMetadata {
239    pub fn new() -> Self {
240        Self {
241            snapshot_id: Uuid::new_v4(),
242            timestamp: Utc::now(),
243            schema_version: default_schema_version(),
244            sdk_version: env!("CARGO_PKG_VERSION").to_string(),
245            created_by: None,
246            checksum: None,
247        }
248    }
249
250    pub fn with_created_by(mut self, created_by: impl Into<String>) -> Self {
251        self.created_by = Some(created_by.into());
252        self
253    }
254
255    pub fn compute_checksum(&mut self, data: &[u8]) {
256        let mut hasher = Sha256::new();
257        hasher.update(data);
258        let result = hasher.finalize();
259        self.checksum = Some(format!("{:x}", result));
260    }
261}
262
263impl Default for SnapshotMetadata {
264    fn default() -> Self {
265        Self::new()
266    }
267}
268
269/// A single AI decision capture
270#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
271pub struct DecisionSnapshot {
272    pub metadata: SnapshotMetadata,
273    pub context: ExecutionContext,
274    pub function_name: String,
275    pub module_name: Option<String>,
276    pub inputs: Vec<Input>,
277    pub outputs: Vec<Output>,
278    pub model_parameters: Option<ModelParameters>,
279    pub execution_time_ms: Option<f64>,
280    pub error: Option<String>,
281    pub error_type: Option<String>,
282    #[serde(default)]
283    pub tags: HashMap<String, String>,
284    #[serde(default)]
285    pub custom_data: HashMap<String, serde_json::Value>,
286}
287
288impl DecisionSnapshot {
289    pub fn new(function_name: impl Into<String>) -> Self {
290        let metadata = SnapshotMetadata::new();
291        let snapshot = Self {
292            metadata,
293            context: ExecutionContext::new(),
294            function_name: function_name.into(),
295            module_name: None,
296            inputs: Vec::new(),
297            outputs: Vec::new(),
298            model_parameters: None,
299            execution_time_ms: None,
300            error: None,
301            error_type: None,
302            tags: HashMap::new(),
303            custom_data: HashMap::new(),
304        };
305
306        // Update checksum after creation
307        let mut result = snapshot;
308        result.update_checksum();
309        result
310    }
311
312    pub fn with_module(mut self, module_name: impl Into<String>) -> Self {
313        self.module_name = Some(module_name.into());
314        self.update_checksum();
315        self
316    }
317
318    pub fn with_context(mut self, context: ExecutionContext) -> Self {
319        self.context = context;
320        self.update_checksum();
321        self
322    }
323
324    pub fn add_input(mut self, input: Input) -> Self {
325        self.inputs.push(input);
326        self.update_checksum();
327        self
328    }
329
330    pub fn add_output(mut self, output: Output) -> Self {
331        self.outputs.push(output);
332        self.update_checksum();
333        self
334    }
335
336    pub fn with_model_parameters(mut self, params: ModelParameters) -> Self {
337        self.model_parameters = Some(params);
338        self.update_checksum();
339        self
340    }
341
342    pub fn with_execution_time(mut self, time_ms: f64) -> Self {
343        self.execution_time_ms = Some(time_ms);
344        self.update_checksum();
345        self
346    }
347
348    pub fn with_error(mut self, error: impl Into<String>, error_type: Option<String>) -> Self {
349        self.error = Some(error.into());
350        self.error_type = error_type;
351        self.update_checksum();
352        self
353    }
354
355    pub fn add_tag(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
356        self.tags.insert(key.into(), value.into());
357        self.update_checksum();
358        self
359    }
360
361    pub fn add_custom_data(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
362        self.custom_data.insert(key.into(), value);
363        self.update_checksum();
364        self
365    }
366
367    fn update_checksum(&mut self) {
368        if let Ok(json_bytes) = serde_json::to_vec(self) {
369            self.metadata.compute_checksum(&json_bytes);
370        }
371    }
372
373    /// Serialize to canonical JSON for consistent checksums
374    pub fn to_canonical_json(&self) -> Result<String, serde_json::Error> {
375        // Create a copy without the checksum for canonical representation
376        let mut copy = self.clone();
377        copy.metadata.checksum = None;
378
379        // Use sorted keys for consistent JSON
380        let value = serde_json::to_value(&copy)?;
381        serde_json::to_string(&value)
382    }
383}
384
385/// Root snapshot containing multiple decisions (e.g., a session)
386#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
387pub struct Snapshot {
388    pub metadata: SnapshotMetadata,
389    pub decisions: Vec<DecisionSnapshot>,
390    pub snapshot_type: SnapshotType,
391}
392
393impl Snapshot {
394    pub fn new(snapshot_type: SnapshotType) -> Self {
395        let metadata = SnapshotMetadata::new();
396        let snapshot = Self {
397            metadata,
398            decisions: Vec::new(),
399            snapshot_type,
400        };
401
402        let mut result = snapshot;
403        result.update_checksum();
404        result
405    }
406
407    pub fn add_decision(&mut self, decision: DecisionSnapshot) {
408        self.decisions.push(decision);
409        self.update_checksum();
410    }
411
412    pub fn with_created_by(mut self, created_by: impl Into<String>) -> Self {
413        self.metadata.created_by = Some(created_by.into());
414        self.update_checksum();
415        self
416    }
417
418    fn update_checksum(&mut self) {
419        if let Ok(json_bytes) = serde_json::to_vec(self) {
420            self.metadata.compute_checksum(&json_bytes);
421        }
422    }
423
424    /// Serialize to canonical JSON for consistent checksums
425    pub fn to_canonical_json(&self) -> Result<String, serde_json::Error> {
426        let mut copy = self.clone();
427        copy.metadata.checksum = None;
428
429        let value = serde_json::to_value(&copy)?;
430        serde_json::to_string(&value)
431    }
432}
433
434#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
435pub enum SnapshotType {
436    Decision,
437    Batch,
438    Session,
439}
440
441fn default_schema_version() -> String {
442    "1.0".to_string()
443}
444
445#[cfg(test)]
446mod tests {
447    use super::*;
448    use serde_json::json;
449
450    #[test]
451    fn test_input_creation() {
452        let input = Input::new("test", json!("value"), "string");
453        assert_eq!(input.name, "test");
454        assert_eq!(input.value, json!("value"));
455        assert_eq!(input.data_type, "string");
456        assert_eq!(input.schema_version, "1.0");
457    }
458
459    #[test]
460    fn test_output_with_confidence() {
461        let output = Output::new("result", json!(42), "number").with_confidence(0.95);
462        assert_eq!(output.confidence, Some(0.95));
463    }
464
465    #[test]
466    fn test_output_confidence_clamping() {
467        let output = Output::new("result", json!(42), "number").with_confidence(1.5);
468        assert_eq!(output.confidence, Some(1.0));
469
470        let output = Output::new("result", json!(42), "number").with_confidence(-0.5);
471        assert_eq!(output.confidence, Some(0.0));
472    }
473
474    #[test]
475    fn test_model_parameters_builder() {
476        let params = ModelParameters::new("gpt-4")
477            .with_provider("openai")
478            .with_version("1.0")
479            .with_parameter("temperature", json!(0.7))
480            .with_hyperparameter("max_tokens", json!(1000));
481
482        assert_eq!(params.model_name, "gpt-4");
483        assert_eq!(params.provider, Some("openai".to_string()));
484        assert_eq!(params.model_version, Some("1.0".to_string()));
485        assert_eq!(params.parameters.get("temperature"), Some(&json!(0.7)));
486        assert_eq!(params.hyperparameters.get("max_tokens"), Some(&json!(1000)));
487    }
488
489    #[test]
490    fn test_execution_context_builder() {
491        let context = ExecutionContext::new()
492            .with_runtime_version("Python 3.11")
493            .with_dependency("numpy", "1.24.0")
494            .with_random_seed(42)
495            .with_env_var("DEBUG", "true");
496
497        assert_eq!(context.runtime_version, Some("Python 3.11".to_string()));
498        assert_eq!(
499            context.dependencies.get("numpy"),
500            Some(&"1.24.0".to_string())
501        );
502        assert_eq!(context.random_seed, Some(42));
503        assert_eq!(
504            context.environment_variables.get("DEBUG"),
505            Some(&"true".to_string())
506        );
507    }
508
509    #[test]
510    fn test_decision_snapshot_creation() {
511        let snapshot = DecisionSnapshot::new("test_function");
512        assert_eq!(snapshot.function_name, "test_function");
513        assert!(snapshot.metadata.checksum.is_some());
514    }
515
516    #[test]
517    fn test_decision_snapshot_builder() {
518        let input = Input::new("x", json!(1), "number");
519        let output = Output::new("result", json!(2), "number");
520        let params = ModelParameters::new("gpt-4");
521
522        let snapshot = DecisionSnapshot::new("add")
523            .with_module("math")
524            .add_input(input)
525            .add_output(output)
526            .with_model_parameters(params)
527            .with_execution_time(100.5)
528            .add_tag("version", "1.0");
529
530        assert_eq!(snapshot.function_name, "add");
531        assert_eq!(snapshot.module_name, Some("math".to_string()));
532        assert_eq!(snapshot.inputs.len(), 1);
533        assert_eq!(snapshot.outputs.len(), 1);
534        assert!(snapshot.model_parameters.is_some());
535        assert_eq!(snapshot.execution_time_ms, Some(100.5));
536        assert_eq!(snapshot.tags.get("version"), Some(&"1.0".to_string()));
537    }
538
539    #[test]
540    fn test_snapshot_creation() {
541        let mut snapshot = Snapshot::new(SnapshotType::Session);
542        let decision = DecisionSnapshot::new("test");
543
544        snapshot.add_decision(decision);
545
546        assert_eq!(snapshot.snapshot_type, SnapshotType::Session);
547        assert_eq!(snapshot.decisions.len(), 1);
548        assert!(snapshot.metadata.checksum.is_some());
549    }
550
551    #[test]
552    fn test_json_serialization_roundtrip() {
553        let input = Input::new("test", json!({"key": "value"}), "object");
554        let output = Output::new("result", json!([1, 2, 3]), "array");
555        let params = ModelParameters::new("gpt-4")
556            .with_provider("openai")
557            .with_parameter("temperature", json!(0.8));
558
559        let snapshot = DecisionSnapshot::new("process")
560            .add_input(input)
561            .add_output(output)
562            .with_model_parameters(params);
563
564        let json = serde_json::to_string(&snapshot).unwrap();
565        let deserialized: DecisionSnapshot = serde_json::from_str(&json).unwrap();
566
567        assert_eq!(snapshot, deserialized);
568    }
569
570    #[test]
571    fn test_canonical_json_consistency() {
572        let snapshot = DecisionSnapshot::new("test")
573            .add_tag("key1", "value1")
574            .add_tag("key2", "value2");
575
576        let json1 = snapshot.to_canonical_json().unwrap();
577        let json2 = snapshot.to_canonical_json().unwrap();
578
579        assert_eq!(json1, json2);
580    }
581
582    #[test]
583    fn test_checksum_updates() {
584        let mut snapshot = DecisionSnapshot::new("test");
585        let initial_checksum = snapshot.metadata.checksum.clone();
586
587        snapshot = snapshot.add_tag("new", "tag");
588        let updated_checksum = snapshot.metadata.checksum.clone();
589
590        assert_ne!(initial_checksum, updated_checksum);
591    }
592}