use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum AdviceSource {
Rule,
Llm,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum AdviceType {
Nl2Sql,
Intent,
Index,
Rewrite,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AiAdviceAuditRecord {
pub source_engine: AdviceSource,
pub llm_model: Option<String>,
pub confidence: f32,
pub advice_type: AdviceType,
pub timestamp: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenefitEstimate {
pub speedup_ratio: f64,
pub confidence: f32,
pub uncertain: bool,
}
impl AiAdviceAuditRecord {
pub fn from_rule(advice_type: AdviceType, confidence: f32) -> Self {
Self {
source_engine: AdviceSource::Rule,
llm_model: None,
confidence,
advice_type,
timestamp: current_timestamp(),
}
}
pub fn from_llm(advice_type: AdviceType, confidence: f32, model: impl Into<String>) -> Self {
Self {
source_engine: AdviceSource::Llm,
llm_model: Some(model.into()),
confidence,
advice_type,
timestamp: current_timestamp(),
}
}
}
impl BenefitEstimate {
pub fn certain(speedup_ratio: f64, confidence: f32) -> Self {
Self {
speedup_ratio,
confidence,
uncertain: false,
}
}
pub fn uncertain(speedup_ratio: f64, confidence: f32) -> Self {
Self {
speedup_ratio,
confidence,
uncertain: true,
}
}
}
fn current_timestamp() -> i64 {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_advice_source_serde() {
let s = AdviceSource::Rule;
let json = serde_json::to_string(&s).unwrap();
let de: AdviceSource = serde_json::from_str(&json).unwrap();
assert_eq!(s, de);
}
#[test]
fn test_advice_type_serde() {
let t = AdviceType::Index;
let json = serde_json::to_string(&t).unwrap();
let de: AdviceType = serde_json::from_str(&json).unwrap();
assert_eq!(t, de);
}
#[test]
fn test_audit_record_from_rule() {
let record = AiAdviceAuditRecord::from_rule(AdviceType::Intent, 0.8);
assert_eq!(record.source_engine, AdviceSource::Rule);
assert!(record.llm_model.is_none());
assert!((record.confidence - 0.8).abs() < 1e-6);
assert_eq!(record.advice_type, AdviceType::Intent);
assert!(record.timestamp > 0);
}
#[test]
fn test_audit_record_from_llm() {
let record = AiAdviceAuditRecord::from_llm(AdviceType::Nl2Sql, 0.9, "gpt-4o-mini");
assert_eq!(record.source_engine, AdviceSource::Llm);
assert_eq!(record.llm_model.as_deref(), Some("gpt-4o-mini"));
assert!((record.confidence - 0.9).abs() < 1e-6);
assert_eq!(record.advice_type, AdviceType::Nl2Sql);
}
#[test]
fn test_benefit_estimate_certain() {
let be = BenefitEstimate::certain(3.5, 0.8);
assert!((be.speedup_ratio - 3.5).abs() < 1e-6);
assert!(!be.uncertain);
}
#[test]
fn test_benefit_estimate_uncertain() {
let be = BenefitEstimate::uncertain(2.0, 0.5);
assert!(be.uncertain);
}
}