use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum EstimatorKind {
DiffInDiff,
PropensityScore,
InstrumentalVariables,
OLS,
Bayesian,
BeforeAfter,
Custom(String),
}
impl EstimatorKind {
pub fn as_str(&self) -> &str {
match self {
Self::DiffInDiff => "diff_in_diff",
Self::PropensityScore => "propensity_score",
Self::InstrumentalVariables => "instrumental_variables",
Self::OLS => "ols",
Self::Bayesian => "bayesian",
Self::BeforeAfter => "before_after",
Self::Custom(s) => s,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct EstimatorMeta {
pub kind: EstimatorKind,
pub version: String,
pub parameters: serde_json::Value,
pub random_seed: Option<u64>,
pub environment: Option<EnvironmentFingerprint>,
pub timeout_secs: Option<u64>,
pub failure_mode: Option<String>,
pub request_schema_version: Option<String>,
pub response_schema_version: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct EnvironmentFingerprint {
pub python_version: Option<String>,
pub package_versions: serde_json::Value,
pub platform: Option<String>,
pub env_hash: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct SidecarExecution {
pub estimator: EstimatorMeta,
pub request: serde_json::Value,
pub response: Option<serde_json::Value>,
pub duration_ms: Option<u64>,
pub success: bool,
pub error: Option<String>,
pub started_at: String,
pub completed_at: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn estimator_meta_serde() {
let meta = EstimatorMeta {
kind: EstimatorKind::DiffInDiff,
version: "1.0.0".into(),
parameters: serde_json::json!({"method": "linear"}),
random_seed: Some(42),
environment: Some(EnvironmentFingerprint {
python_version: Some("3.11".into()),
package_versions: serde_json::json!({"dowhy": "0.11"}),
platform: Some("linux-x86_64".into()),
env_hash: None,
}),
timeout_secs: Some(300),
failure_mode: None,
request_schema_version: Some("v1".into()),
response_schema_version: Some("v1".into()),
};
let json = serde_json::to_string(&meta).unwrap();
let back: EstimatorMeta = serde_json::from_str(&json).unwrap();
assert_eq!(back.version, "1.0.0");
assert_eq!(back.random_seed, Some(42));
}
#[test]
fn sidecar_execution_serde() {
let exec = SidecarExecution {
estimator: EstimatorMeta {
kind: EstimatorKind::PropensityScore,
version: "2.0.0".into(),
parameters: serde_json::json!({}),
random_seed: None,
environment: None,
timeout_secs: Some(60),
failure_mode: None,
request_schema_version: None,
response_schema_version: None,
},
request: serde_json::json!({"data": [1, 2, 3]}),
response: Some(serde_json::json!({"estimate": 0.5})),
duration_ms: Some(1500),
success: true,
error: None,
started_at: "2024-01-01T00:00:00Z".into(),
completed_at: Some("2024-01-01T00:00:01Z".into()),
};
let json = serde_json::to_string(&exec).unwrap();
let back: SidecarExecution = serde_json::from_str(&json).unwrap();
assert!(back.success);
assert_eq!(back.duration_ms, Some(1500));
}
}