1use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
11#[serde(rename_all = "snake_case")]
12pub enum EstimatorKind {
13 DiffInDiff,
15 PropensityScore,
17 InstrumentalVariables,
19 OLS,
21 Bayesian,
23 BeforeAfter,
25 Custom(String),
27}
28
29impl EstimatorKind {
30 pub fn as_str(&self) -> &str {
32 match self {
33 Self::DiffInDiff => "diff_in_diff",
34 Self::PropensityScore => "propensity_score",
35 Self::InstrumentalVariables => "instrumental_variables",
36 Self::OLS => "ols",
37 Self::Bayesian => "bayesian",
38 Self::BeforeAfter => "before_after",
39 Self::Custom(s) => s,
40 }
41 }
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
48pub struct EstimatorMeta {
49 pub kind: EstimatorKind,
51 pub version: String,
53 pub parameters: serde_json::Value,
55 pub random_seed: Option<u64>,
57 pub environment: Option<EnvironmentFingerprint>,
59 pub timeout_secs: Option<u64>,
61 pub failure_mode: Option<String>,
63 pub request_schema_version: Option<String>,
65 pub response_schema_version: Option<String>,
67}
68
69#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
73pub struct EnvironmentFingerprint {
74 pub python_version: Option<String>,
76 pub package_versions: serde_json::Value,
78 pub platform: Option<String>,
80 pub env_hash: Option<String>,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
88pub struct SidecarExecution {
89 pub estimator: EstimatorMeta,
91 pub request: serde_json::Value,
93 pub response: Option<serde_json::Value>,
95 pub duration_ms: Option<u64>,
97 pub success: bool,
99 pub error: Option<String>,
101 pub started_at: String,
103 pub completed_at: Option<String>,
105}
106
107#[cfg(test)]
108mod tests {
109 use super::*;
110
111 #[test]
112 fn estimator_meta_serde() {
113 let meta = EstimatorMeta {
114 kind: EstimatorKind::DiffInDiff,
115 version: "1.0.0".into(),
116 parameters: serde_json::json!({"method": "linear"}),
117 random_seed: Some(42),
118 environment: Some(EnvironmentFingerprint {
119 python_version: Some("3.11".into()),
120 package_versions: serde_json::json!({"dowhy": "0.11"}),
121 platform: Some("linux-x86_64".into()),
122 env_hash: None,
123 }),
124 timeout_secs: Some(300),
125 failure_mode: None,
126 request_schema_version: Some("v1".into()),
127 response_schema_version: Some("v1".into()),
128 };
129
130 let json = serde_json::to_string(&meta).unwrap();
131 let back: EstimatorMeta = serde_json::from_str(&json).unwrap();
132 assert_eq!(back.version, "1.0.0");
133 assert_eq!(back.random_seed, Some(42));
134 }
135
136 #[test]
137 fn sidecar_execution_serde() {
138 let exec = SidecarExecution {
139 estimator: EstimatorMeta {
140 kind: EstimatorKind::PropensityScore,
141 version: "2.0.0".into(),
142 parameters: serde_json::json!({}),
143 random_seed: None,
144 environment: None,
145 timeout_secs: Some(60),
146 failure_mode: None,
147 request_schema_version: None,
148 response_schema_version: None,
149 },
150 request: serde_json::json!({"data": [1, 2, 3]}),
151 response: Some(serde_json::json!({"estimate": 0.5})),
152 duration_ms: Some(1500),
153 success: true,
154 error: None,
155 started_at: "2024-01-01T00:00:00Z".into(),
156 completed_at: Some("2024-01-01T00:00:01Z".into()),
157 };
158
159 let json = serde_json::to_string(&exec).unwrap();
160 let back: SidecarExecution = serde_json::from_str(&json).unwrap();
161 assert!(back.success);
162 assert_eq!(back.duration_ms, Some(1500));
163 }
164}