lean_ctx/core/context_kernel/
attribution.rs1use 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}