Skip to main content

briefcase_core/models/
base.rs

1//! Core data models for AI decision tracking and observability
2
3use super::agent::AgentMetadata;
4use super::hardware::HardwareMetadata;
5use super::scorecard::{ExperimentMetadata, Scorecard};
6use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8use sha2::{Digest, Sha256};
9use std::collections::HashMap;
10use uuid::Uuid;
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
13pub struct Input {
14    pub name: String,
15    pub value: serde_json::Value,
16    pub data_type: String,
17    #[serde(default = "default_schema_version")]
18    pub schema_version: String,
19}
20
21impl Input {
22    pub fn new(
23        name: impl Into<String>,
24        value: serde_json::Value,
25        data_type: impl Into<String>,
26    ) -> Self {
27        Self {
28            name: name.into(),
29            value,
30            data_type: data_type.into(),
31            schema_version: default_schema_version(),
32        }
33    }
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
37pub struct Output {
38    pub name: String,
39    pub value: serde_json::Value,
40    pub data_type: String,
41    pub confidence: Option<f64>,
42    #[serde(default = "default_schema_version")]
43    pub schema_version: String,
44}
45
46impl Output {
47    pub fn new(
48        name: impl Into<String>,
49        value: serde_json::Value,
50        data_type: impl Into<String>,
51    ) -> Self {
52        Self {
53            name: name.into(),
54            value,
55            data_type: data_type.into(),
56            confidence: None,
57            schema_version: default_schema_version(),
58        }
59    }
60    pub fn with_confidence(mut self, confidence: f64) -> Self {
61        self.confidence = Some(confidence.clamp(0.0, 1.0));
62        self
63    }
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
67pub struct ModelParameters {
68    pub model_name: String,
69    pub model_version: Option<String>,
70    pub provider: Option<String>,
71    #[serde(default)]
72    pub parameters: HashMap<String, serde_json::Value>,
73    #[serde(default)]
74    pub hyperparameters: HashMap<String, serde_json::Value>,
75    pub weights_hash: Option<String>,
76}
77
78impl ModelParameters {
79    pub fn new(model_name: impl Into<String>) -> Self {
80        Self {
81            model_name: model_name.into(),
82            model_version: None,
83            provider: None,
84            parameters: HashMap::new(),
85            hyperparameters: HashMap::new(),
86            weights_hash: None,
87        }
88    }
89    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
90        self.provider = Some(provider.into());
91        self
92    }
93    pub fn with_version(mut self, version: impl Into<String>) -> Self {
94        self.model_version = Some(version.into());
95        self
96    }
97    pub fn with_parameter(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
98        self.parameters.insert(key.into(), value);
99        self
100    }
101    pub fn with_hyperparameter(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
102        self.hyperparameters.insert(key.into(), value);
103        self
104    }
105}
106
107#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
108pub struct ExecutionContext {
109    pub runtime_version: Option<String>,
110    #[serde(default)]
111    pub dependencies: HashMap<String, String>,
112    pub random_seed: Option<i64>,
113    #[serde(default)]
114    pub environment_variables: HashMap<String, String>,
115    #[serde(default)]
116    pub hardware_info: HashMap<String, serde_json::Value>,
117}
118
119impl ExecutionContext {
120    pub fn new() -> Self {
121        Self::default()
122    }
123    pub fn with_runtime_version(mut self, v: impl Into<String>) -> Self {
124        self.runtime_version = Some(v.into());
125        self
126    }
127    pub fn with_random_seed(mut self, s: i64) -> Self {
128        self.random_seed = Some(s);
129        self
130    }
131    pub fn with_dependency(mut self, n: impl Into<String>, v: impl Into<String>) -> Self {
132        self.dependencies.insert(n.into(), v.into());
133        self
134    }
135}
136
137#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
138pub struct SnapshotMetadata {
139    pub snapshot_id: Uuid,
140    pub timestamp: DateTime<Utc>,
141    pub schema_version: String,
142    pub sdk_version: String,
143    pub created_by: Option<String>,
144    pub checksum: Option<String>,
145}
146
147impl Default for SnapshotMetadata {
148    fn default() -> Self {
149        Self::new()
150    }
151}
152
153impl SnapshotMetadata {
154    pub fn new() -> Self {
155        Self {
156            snapshot_id: Uuid::new_v4(),
157            timestamp: Utc::now(),
158            schema_version: "1.0".to_string(),
159            sdk_version: "2.5.9".to_string(),
160            created_by: None,
161            checksum: None,
162        }
163    }
164    pub fn compute_checksum(&mut self, data: &[u8]) {
165        let mut hasher = Sha256::new();
166        hasher.update(data);
167        self.checksum = Some(format!("{:x}", hasher.finalize()));
168    }
169}
170
171#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
172pub struct DecisionSnapshot {
173    pub metadata: SnapshotMetadata,
174    pub context: ExecutionContext,
175    pub function_name: String,
176    pub module_name: Option<String>,
177    pub inputs: Vec<Input>,
178    pub outputs: Vec<Output>,
179    pub model_parameters: Option<ModelParameters>,
180    pub execution_time_ms: Option<f64>,
181    pub error: Option<String>,
182    pub error_type: Option<String>,
183    pub scorecard: Option<Scorecard>,
184    pub experiment: Option<ExperimentMetadata>,
185    pub agent: Option<AgentMetadata>,
186    pub hardware: Option<HardwareMetadata>,
187    #[serde(default)]
188    pub tags: HashMap<String, String>,
189}
190
191impl DecisionSnapshot {
192    pub fn new(function_name: impl Into<String>) -> Self {
193        let mut s = Self {
194            metadata: SnapshotMetadata::new(),
195            context: ExecutionContext::new(),
196            function_name: function_name.into(),
197            module_name: None,
198            inputs: Vec::new(),
199            outputs: Vec::new(),
200            model_parameters: None,
201            execution_time_ms: None,
202            error: None,
203            error_type: None,
204            scorecard: None,
205            experiment: None,
206            agent: None,
207            hardware: None,
208            tags: HashMap::new(),
209        };
210        s.update_checksum();
211        s
212    }
213    pub fn with_module(mut self, m: impl Into<String>) -> Self {
214        self.module_name = Some(m.into());
215        self.update_checksum();
216        self
217    }
218    pub fn with_context(mut self, c: ExecutionContext) -> Self {
219        self.context = c;
220        self.update_checksum();
221        self
222    }
223    pub fn add_input(mut self, i: Input) -> Self {
224        self.inputs.push(i);
225        self.update_checksum();
226        self
227    }
228    pub fn add_output(mut self, o: Output) -> Self {
229        self.outputs.push(o);
230        self.update_checksum();
231        self
232    }
233    pub fn with_model_parameters(mut self, p: ModelParameters) -> Self {
234        self.model_parameters = Some(p);
235        self.update_checksum();
236        self
237    }
238    pub fn with_execution_time(mut self, t: f64) -> Self {
239        self.execution_time_ms = Some(t);
240        self.update_checksum();
241        self
242    }
243    pub fn with_error(mut self, e: impl Into<String>, et: Option<String>) -> Self {
244        self.error = Some(e.into());
245        self.error_type = et;
246        self.update_checksum();
247        self
248    }
249    pub fn with_scorecard(mut self, s: Scorecard) -> Self {
250        self.scorecard = Some(s);
251        self.update_checksum();
252        self
253    }
254    pub fn with_agent(mut self, a: AgentMetadata) -> Self {
255        self.agent = Some(a);
256        self.update_checksum();
257        self
258    }
259    pub fn with_hardware(mut self, h: HardwareMetadata) -> Self {
260        self.hardware = Some(h);
261        self.update_checksum();
262        self
263    }
264    pub fn add_tag(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
265        self.tags.insert(k.into(), v.into());
266        self.update_checksum();
267        self
268    }
269
270    fn update_checksum(&mut self) {
271        if let Ok(b) = serde_json::to_vec(self) {
272            self.metadata.compute_checksum(&b);
273        }
274    }
275    pub fn fingerprint(&self) -> String {
276        let mut hasher = Sha256::new();
277        hasher.update(self.function_name.as_bytes());
278        for i in &self.inputs {
279            hasher.update(i.name.as_bytes());
280            hasher.update(
281                serde_json::to_string(&i.value)
282                    .unwrap_or_default()
283                    .as_bytes(),
284            );
285        }
286        if let Some(p) = &self.model_parameters {
287            hasher.update(p.model_name.as_bytes());
288        }
289        format!("{:x}", hasher.finalize())
290    }
291}
292
293#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
294pub struct Snapshot {
295    pub metadata: SnapshotMetadata,
296    pub decisions: Vec<DecisionSnapshot>,
297    pub snapshot_type: SnapshotType,
298}
299
300impl Snapshot {
301    pub fn new(t: SnapshotType) -> Self {
302        let mut s = Self {
303            metadata: SnapshotMetadata::new(),
304            decisions: Vec::new(),
305            snapshot_type: t,
306        };
307        s.update_checksum();
308        s
309    }
310    pub fn add_decision(&mut self, d: DecisionSnapshot) {
311        self.decisions.push(d);
312        self.update_checksum();
313    }
314    pub fn with_created_by(mut self, c: impl Into<String>) -> Self {
315        self.metadata.created_by = Some(c.into());
316        self.update_checksum();
317        self
318    }
319    fn update_checksum(&mut self) {
320        if let Ok(b) = serde_json::to_vec(self) {
321            self.metadata.compute_checksum(&b);
322        }
323    }
324}
325
326#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
327pub enum SnapshotType {
328    Decision,
329    Batch,
330    Session,
331}
332
333fn default_schema_version() -> String {
334    "1.0".to_string()
335}