1use af_context::{ProviderAttemptId, ToolCallId};
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(rename_all = "snake_case")]
8pub enum MeteringSource {
9 Reported,
11 Estimated,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(rename_all = "snake_case")]
18pub enum MeteringOutcome {
19 Completed,
21 Failed,
23 Cancelled,
25 Steered,
27 TimedOut,
29 NotDispatched,
31 Unknown,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
37#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
38pub enum MeteringDetails {
39 Model {
41 model: String,
43 provider: Option<String>,
45 provider_attempt_id: ProviderAttemptId,
47 source: MeteringSource,
49 outcome: MeteringOutcome,
51 },
52 Tool {
54 call_id: ToolCallId,
56 name: String,
58 source: MeteringSource,
60 outcome: MeteringOutcome,
62 },
63}
64
65impl MeteringDetails {
66 pub fn completes(&self, prepared: &Self) -> bool {
68 match (self, prepared) {
69 (
70 Self::Model {
71 model,
72 provider,
73 provider_attempt_id,
74 ..
75 },
76 Self::Model {
77 model: expected_model,
78 provider: expected_provider,
79 provider_attempt_id: expected_attempt,
80 ..
81 },
82 ) => {
83 model == expected_model
84 && provider_attempt_id == expected_attempt
85 && expected_provider
86 .as_ref()
87 .is_none_or(|expected| provider.as_ref() == Some(expected))
88 }
89 (
90 Self::Tool { call_id, name, .. },
91 Self::Tool {
92 call_id: expected_call,
93 name: expected_name,
94 ..
95 },
96 ) => call_id == expected_call && name == expected_name,
97 _ => false,
98 }
99 }
100
101 pub fn validate(&self) -> Result<(), crate::EventError> {
103 let valid =
104 |s: &str| !s.trim().is_empty() && s.len() <= 512 && !s.chars().any(char::is_control);
105 let ok = match self {
106 Self::Model {
107 model,
108 provider,
109 provider_attempt_id,
110 ..
111 } => {
112 valid(model)
113 && provider.as_deref().is_none_or(valid)
114 && valid(provider_attempt_id.as_str())
115 }
116 Self::Tool { call_id, name, .. } => valid(call_id.as_str()) && valid(name),
117 };
118 if ok {
119 Ok(())
120 } else {
121 Err(crate::EventError::InvalidMetering)
122 }
123 }
124}
125
126impl MeteringSource {
127 pub fn as_str(self) -> &'static str {
129 match self {
130 Self::Reported => "reported",
131 Self::Estimated => "estimated",
132 }
133 }
134}
135impl MeteringOutcome {
136 pub fn as_str(self) -> &'static str {
138 match self {
139 Self::Completed => "completed",
140 Self::Failed => "failed",
141 Self::Cancelled => "cancelled",
142 Self::Steered => "steered",
143 Self::TimedOut => "timed_out",
144 Self::NotDispatched => "not_dispatched",
145 Self::Unknown => "unknown",
146 }
147 }
148}
149
150#[cfg(test)]
151mod tests {
152 use super::*;
153 #[test]
154 fn attribution_is_bounded_and_preserves_unknown_provider_and_outcomes() {
155 for outcome in [
156 MeteringOutcome::Completed,
157 MeteringOutcome::Failed,
158 MeteringOutcome::Cancelled,
159 MeteringOutcome::Steered,
160 MeteringOutcome::TimedOut,
161 MeteringOutcome::NotDispatched,
162 MeteringOutcome::Unknown,
163 ] {
164 let mut details = MeteringDetails::Model {
165 model: "m".into(),
166 provider: None,
167 provider_attempt_id: "attempt".parse().unwrap(),
168 source: MeteringSource::Estimated,
169 outcome,
170 };
171 details.validate().unwrap();
172 assert_eq!(serde_json::to_value(outcome).unwrap(), outcome.as_str());
173 assert_eq!(
174 serde_json::from_value::<MeteringDetails>(serde_json::to_value(&details).unwrap())
175 .unwrap(),
176 details
177 );
178 if let MeteringDetails::Model { provider, .. } = &mut details {
179 *provider = Some("\n".into());
180 }
181 assert!(details.validate().is_err());
182 }
183 for name in ["".to_owned(), "a".repeat(513), "a\nb".into()] {
184 assert!(MeteringDetails::Tool {
185 name,
186 call_id: "call".parse().unwrap(),
187 source: MeteringSource::Estimated,
188 outcome: MeteringOutcome::Unknown
189 }
190 .validate()
191 .is_err());
192 }
193 }
194}