1use chrono::{DateTime, Utc};
29use serde::{Deserialize, Serialize};
30use sha2::{Digest, Sha256};
31use std::collections::HashMap;
32use uuid::Uuid;
33
34#[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#[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>, #[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#[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>, #[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>, }
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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
189pub struct ExecutionContext {
190 pub runtime_version: Option<String>, #[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#[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#[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 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 pub fn to_canonical_json(&self) -> Result<String, serde_json::Error> {
375 let mut copy = self.clone();
377 copy.metadata.checksum = None;
378
379 let value = serde_json::to_value(©)?;
381 serde_json::to_string(&value)
382 }
383}
384
385#[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 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(©)?;
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}