use std::collections::HashMap;
use std::time::{Duration, SystemTime};
use serde::{Deserialize, Serialize};
use crate::types::{Priority, SourceId};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SourceStats {
pub contributions: usize,
pub kept: usize,
pub tokens: usize,
pub priority: Priority,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DropReason {
OverBudget,
Governance,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DroppedContribution {
pub source: SourceId,
pub reason: DropReason,
pub tokens: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RedactionRecord {
pub source: SourceId,
pub rule: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct Timings {
pub fanout: Duration,
pub tokenize: Duration,
pub prune: Duration,
pub governance: Duration,
pub total: Duration,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BriefReceipt {
pub built_at: SystemTime,
pub total_tokens: usize,
pub by_source: HashMap<SourceId, SourceStats>,
pub dropped: Vec<DroppedContribution>,
pub redactions: Vec<RedactionRecord>,
pub timings: Timings,
}
impl BriefReceipt {
pub fn empty() -> Self {
BriefReceipt {
built_at: SystemTime::now(),
total_tokens: 0,
by_source: HashMap::new(),
dropped: Vec::new(),
redactions: Vec::new(),
timings: Timings::default(),
}
}
}
impl Default for BriefReceipt {
fn default() -> Self {
BriefReceipt::empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty_receipt_has_zero_counts() {
let receipt = BriefReceipt::empty();
assert_eq!(receipt.total_tokens, 0);
assert!(receipt.by_source.is_empty());
assert!(receipt.dropped.is_empty());
assert!(receipt.redactions.is_empty());
assert_eq!(receipt.timings, Timings::default());
}
#[test]
fn dropped_round_trips_through_serde_json() {
let dropped = DroppedContribution {
source: SourceId::new("history"),
reason: DropReason::OverBudget,
tokens: 42,
};
let json = serde_json::to_string(&dropped).expect("serialize");
let back: DroppedContribution = serde_json::from_str(&json).expect("deserialize");
assert_eq!(dropped, back);
}
#[test]
fn source_stats_round_trips_through_serde_json() {
let stats = SourceStats {
contributions: 3,
kept: 2,
tokens: 120,
priority: Priority::Normal,
};
let json = serde_json::to_string(&stats).expect("serialize");
let back: SourceStats = serde_json::from_str(&json).expect("deserialize");
assert_eq!(stats, back);
}
#[test]
fn redaction_record_round_trips_through_serde_json() {
let rec = RedactionRecord {
source: SourceId::new("perception"),
rule: "rule:no_bank_dom".into(),
};
let json = serde_json::to_string(&rec).expect("serialize");
let back: RedactionRecord = serde_json::from_str(&json).expect("deserialize");
assert_eq!(rec, back);
}
}