mockforge_intelligence/handlers/behavioral_cloning/
probability.rs

1//! Probability models and sampling from recorded exchanges.
2
3use crate::behavioral_cloning::storage::{BehavioralCloningStateProvider, BehavioralCloningStore};
4use crate::behavioral_cloning::{EndpointProbabilityModel, ProbabilisticModel};
5use axum::extract::{Path, State};
6use axum::response::Json;
7use serde::{Deserialize, Serialize};
8use serde_json::{json, Value};
9
10/// Request to build a probability model for an endpoint
11#[derive(Debug, Deserialize)]
12pub struct BuildProbabilityModelRequest {
13    /// Endpoint path
14    pub endpoint: String,
15    /// HTTP method
16    pub method: String,
17    /// Optional limit on number of samples to analyze
18    #[serde(default)]
19    pub sample_limit: Option<u32>,
20}
21
22/// Response for probability model
23#[derive(Debug, Serialize)]
24pub struct ProbabilityModelResponse {
25    /// Success flag
26    pub success: bool,
27    /// Probability model
28    pub model: EndpointProbabilityModel,
29}
30
31/// Build a probability model for an endpoint
32///
33/// POST /api/v1/behavioral-cloning/probability-models
34pub async fn build_probability_model<S: BehavioralCloningStateProvider>(
35    State(state): State<S>,
36    Json(request): Json<BuildProbabilityModelRequest>,
37) -> Result<Json<Value>, String> {
38    // Open database connection
39    let db = state.open_store().await?;
40
41    // Get exchanges for the endpoint
42    let limit = request.sample_limit.map(|l| l as i32);
43    let exchanges = db
44        .get_exchanges_for_endpoint(&request.endpoint, &request.method, limit)
45        .await
46        .map_err(|e| format!("Failed to query exchanges: {}", e))?;
47
48    if exchanges.is_empty() {
49        return Err(format!(
50            "No recorded traffic found for {} {}",
51            request.method, request.endpoint
52        ));
53    }
54
55    let data = super::probability_data::ModelData::from_exchanges(&exchanges);
56
57    // Build probability model
58    let model = ProbabilisticModel::build_probability_model_from_data(
59        &request.endpoint,
60        &request.method,
61        &data.status_codes,
62        &data.latencies_ms,
63        &data.error_responses,
64        &data.request_payloads,
65        &data.response_payloads,
66    );
67
68    // Store model in database
69    db.insert_endpoint_probability_model(&model)
70        .await
71        .map_err(|e| format!("Failed to store probability model: {}", e))?;
72
73    Ok(Json(json!({
74        "success": true,
75        "model": model
76    })))
77}
78
79/// Get a probability model for an endpoint
80///
81/// GET /api/v1/behavioral-cloning/probability-models/{endpoint}/{method}
82pub async fn get_probability_model<S: BehavioralCloningStateProvider>(
83    Path((endpoint, method)): Path<(String, String)>,
84    State(state): State<S>,
85) -> Result<Json<Value>, String> {
86    let db = state.open_store().await?;
87
88    let model = db
89        .get_endpoint_probability_model(&endpoint, &method)
90        .await
91        .map_err(|e| format!("Failed to query probability model: {}", e))?
92        .ok_or_else(|| format!("No probability model found for {} {}", method, endpoint))?;
93
94    Ok(Json(json!({
95        "success": true,
96        "model": model
97    })))
98}
99
100/// List all probability models
101///
102/// GET /api/v1/behavioral-cloning/probability-models
103pub async fn list_probability_models<S: BehavioralCloningStateProvider>(
104    State(state): State<S>,
105) -> Result<Json<Value>, String> {
106    let db = state.open_store().await?;
107
108    let models = db
109        .get_all_endpoint_probability_models()
110        .await
111        .map_err(|e| format!("Failed to query probability models: {}", e))?;
112
113    Ok(Json(json!({
114        "success": true,
115        "models": models,
116        "count": models.len()
117    })))
118}
119
120/// Sample a status code from a probability model
121///
122/// POST /api/v1/behavioral-cloning/probability-models/{endpoint}/{method}/sample/status-code
123pub async fn sample_status_code<S: BehavioralCloningStateProvider>(
124    Path((endpoint, method)): Path<(String, String)>,
125    State(state): State<S>,
126) -> Result<Json<Value>, String> {
127    let db = state.open_store().await?;
128
129    let model = db
130        .get_endpoint_probability_model(&endpoint, &method)
131        .await
132        .map_err(|e| format!("Failed to query model: {}", e))?
133        .ok_or_else(|| format!("No probability model found for {} {}", method, endpoint))?;
134
135    let sampled_code = ProbabilisticModel::sample_status_code(&model);
136
137    Ok(Json(json!({
138        "success": true,
139        "endpoint": endpoint,
140        "method": method,
141        "status_code": sampled_code
142    })))
143}
144
145/// Sample latency from a probability model
146///
147/// POST /api/v1/behavioral-cloning/probability-models/{endpoint}/{method}/sample/latency
148pub async fn sample_latency<S: BehavioralCloningStateProvider>(
149    Path((endpoint, method)): Path<(String, String)>,
150    State(state): State<S>,
151) -> Result<Json<Value>, String> {
152    let db = state.open_store().await?;
153
154    let model = db
155        .get_endpoint_probability_model(&endpoint, &method)
156        .await
157        .map_err(|e| format!("Failed to query model: {}", e))?
158        .ok_or_else(|| format!("No probability model found for {} {}", method, endpoint))?;
159
160    let sampled_latency = ProbabilisticModel::sample_latency(&model);
161
162    Ok(Json(json!({
163        "success": true,
164        "endpoint": endpoint,
165        "method": method,
166        "latency_ms": sampled_latency
167    })))
168}
mockforge_intelligence/handlers/behavioral_cloning/probability.rs

mockforge_intelligence/handlers/behavioral_cloning/
probability.rs
