Skip to main content

lean_ctx/core/context_kernel/
attribution.rs

1//! Receipt-based savings attribution for Context Kernel providers.
2
3use std::collections::HashMap;
4
5use super::types::{ContextPlanV1, ContextReceiptV1, PlanEntry, ReceiptOutcome};
6
7#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
8pub struct AttributionEntry {
9    pub provider: String,
10    pub tokens_contributed: usize,
11    pub tokens_saved: usize,
12    pub efficiency: f64,
13}
14
15#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
16pub struct AttributionReport {
17    pub plan_id: String,
18    pub receipt_id: String,
19    pub total_tokens_delivered: usize,
20    pub total_tokens_saved: usize,
21    pub entries: Vec<AttributionEntry>,
22}
23
24fn entry_savings(entry: &PlanEntry, compression_ratio: f64) -> usize {
25    let delivered_share = (entry.tokens as f64 * compression_ratio) as usize;
26    entry.tokens.saturating_sub(delivered_share)
27}
28
29pub fn compute_attribution(plan: &ContextPlanV1, receipt: &ContextReceiptV1) -> AttributionReport {
30    let compression_ratio = if plan.budget.total_tokens == 0 {
31        0.0
32    } else {
33        receipt.delivered_tokens as f64 / plan.budget.total_tokens as f64
34    };
35    let _outcome: ReceiptOutcome = receipt.outcome;
36    let mut provider_totals: HashMap<String, (usize, usize)> = HashMap::new();
37
38    for entry in &plan.selected {
39        let totals = provider_totals.entry(entry.provider.clone()).or_default();
40        totals.0 = totals.0.saturating_add(entry.tokens);
41        totals.1 = totals
42            .1
43            .saturating_add(entry_savings(entry, compression_ratio));
44    }
45
46    let mut entries: Vec<AttributionEntry> = provider_totals
47        .into_iter()
48        .map(
49            |(provider, (tokens_contributed, tokens_saved))| AttributionEntry {
50                provider,
51                tokens_contributed,
52                tokens_saved,
53                efficiency: tokens_saved as f64 / tokens_contributed.max(1) as f64,
54            },
55        )
56        .collect();
57    entries.sort_by_key(|entry| std::cmp::Reverse(entry.tokens_saved));
58
59    let total_tokens_saved = entries.iter().map(|entry| entry.tokens_saved).sum();
60
61    AttributionReport {
62        plan_id: plan.plan_id.clone(),
63        receipt_id: receipt.receipt_id.clone(),
64        total_tokens_delivered: receipt.delivered_tokens,
65        total_tokens_saved,
66        entries,
67    }
68}
69
70pub fn format_attribution_summary(report: &AttributionReport) -> String {
71    let top_provider = report
72        .entries
73        .first()
74        .map_or_else(|| "none", |entry| entry.provider.as_str());
75
76    format!(
77        "Attribution: {} providers, {} tokens saved, top: {}",
78        report.entries.len(),
79        report.total_tokens_saved,
80        top_provider
81    )
82}
83
84#[cfg(test)]
85mod tests {
86    use std::collections::HashMap;
87
88    use super::{compute_attribution, format_attribution_summary};
89    use crate::core::context_kernel::types::{
90        ContextPlanV1, ContextReceiptV1, PlanBudget, PlanEntry, ReceiptOutcome,
91    };
92
93    fn sample_plan() -> ContextPlanV1 {
94        ContextPlanV1 {
95            plan_id: "plan:test".to_owned(),
96            intent: "test attribution".to_owned(),
97            budget: PlanBudget {
98                total_tokens: 1_000,
99                used_tokens: 800,
100                remaining_tokens: 200,
101            },
102            selected: vec![
103                PlanEntry {
104                    object_id: "file:a".to_owned(),
105                    provider: "files".to_owned(),
106                    view: "full".to_owned(),
107                    tokens: 600,
108                    phi: 1.0,
109                    reason: "relevant".to_owned(),
110                },
111                PlanEntry {
112                    object_id: "fact:b".to_owned(),
113                    provider: "knowledge".to_owned(),
114                    view: "summary".to_owned(),
115                    tokens: 200,
116                    phi: 0.8,
117                    reason: "supporting".to_owned(),
118                },
119            ],
120            excluded: Vec::new(),
121            deferred: Vec::new(),
122            provider_stats: HashMap::new(),
123        }
124    }
125
126    fn sample_receipt() -> ContextReceiptV1 {
127        ContextReceiptV1 {
128            receipt_id: "receipt:test".to_owned(),
129            plan_id: "plan:test".to_owned(),
130            delivered_tokens: 500,
131            cache_hits: 0,
132            cache_misses: 0,
133            outcome: ReceiptOutcome::Accepted,
134            quality_signals: Vec::new(),
135            feedback_attribution: HashMap::new(),
136        }
137    }
138
139    #[test]
140    fn attribution_computes_savings() {
141        let report = compute_attribution(&sample_plan(), &sample_receipt());
142
143        assert_eq!(report.total_tokens_delivered, 500);
144        assert_eq!(report.total_tokens_saved, 400);
145        assert_eq!(report.entries[0].tokens_contributed, 600);
146        assert!((report.entries[0].efficiency - 0.5).abs() < f64::EPSILON);
147    }
148
149    #[test]
150    fn attribution_sorted_by_savings() {
151        let report = compute_attribution(&sample_plan(), &sample_receipt());
152
153        assert_eq!(report.entries[0].provider, "files");
154        assert_eq!(report.entries[0].tokens_saved, 300);
155        assert_eq!(report.entries[1].provider, "knowledge");
156        assert_eq!(report.entries[1].tokens_saved, 100);
157    }
158
159    #[test]
160    fn format_summary_contains_top_provider() {
161        let report = compute_attribution(&sample_plan(), &sample_receipt());
162        let summary = format_attribution_summary(&report);
163
164        assert_eq!(
165            summary,
166            "Attribution: 2 providers, 400 tokens saved, top: files"
167        );
168    }
169}