From 33929fd04b2077983679349a5bb395f7f5fac518 Mon Sep 17 00:00:00 2001 From: Chili Palmer Date: Sat, 18 Jul 2026 23:15:04 +0200 Subject: [PATCH] feat: return normalized token usage from one-shot AI calls (closes #15) Co-Authored-By: Claude Fable 5 --- crates/bds-core/src/engine/ai.rs | 114 ++++++++++++++++-- .../bds-core/src/engine/auto_translation.rs | 4 +- crates/bds-ui/src/app.rs | 31 +++-- crates/bds-ui/src/views/dashboard.rs | 9 +- 4 files changed, 132 insertions(+), 26 deletions(-) diff --git a/crates/bds-core/src/engine/ai.rs b/crates/bds-core/src/engine/ai.rs index 5c0a2fc..ca00dd0 100644 --- a/crates/bds-core/src/engine/ai.rs +++ b/crates/bds-core/src/engine/ai.rs @@ -135,6 +135,14 @@ pub struct MediaTranslationResult { pub caption: String, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] +pub struct TokenUsage { + pub input_tokens: Option, + pub output_tokens: Option, + pub cache_read_tokens: Option, + pub cache_write_tokens: Option, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub enum OneShotResponse { Taxonomy(TaxonomySuggestion), @@ -323,7 +331,7 @@ pub fn run_one_shot( conn: &Connection, offline_mode: bool, request: &OneShotRequest, -) -> EngineResult { +) -> EngineResult<(OneShotResponse, TokenUsage)> { let settings = load_ai_settings(conn, offline_mode)?; let endpoint = active_endpoint(conn, offline_mode)?; let model = select_model(&settings, &endpoint, &request.operation)?; @@ -370,7 +378,24 @@ pub fn run_one_shot( .ok_or_else(|| { EngineError::Parse("chat completions response missing message content".to_string()) })?; - parse_one_shot_response(request, content) + let response = parse_one_shot_response(request, content)?; + Ok((response, parse_token_usage(&body))) +} + +fn parse_token_usage(body: &Value) -> TokenUsage { + let usage = body.get("usage").unwrap_or(&Value::Null); + TokenUsage { + input_tokens: usage.get("prompt_tokens").and_then(Value::as_u64), + output_tokens: usage.get("completion_tokens").and_then(Value::as_u64), + cache_read_tokens: usage + .get("prompt_tokens_details") + .and_then(|details| details.get("cached_tokens")) + .and_then(Value::as_u64), + cache_write_tokens: usage + .get("completion_tokens_details") + .and_then(|details| details.get("cached_tokens")) + .and_then(Value::as_u64), + } } fn build_http_client() -> EngineResult { @@ -823,7 +848,7 @@ mod tests { || request.contains("Authorization: Bearer secret-token") ); http_ok( - r#"{"choices":[{"message":{"content":"{\"title\":\"Better title\",\"excerpt\":\"Short summary\",\"slug\":\"better-title\"}"}}]}"#, + r#"{"choices":[{"message":{"content":"{\"title\":\"Better title\",\"excerpt\":\"Short summary\",\"slug\":\"better-title\"}"}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}"#, ) }, 1, @@ -842,7 +867,7 @@ mod tests { .unwrap(); save_model_preferences(db.conn(), None, Some("gpt-4.1-mini"), None, "").unwrap(); - let response = run_one_shot( + let (response, usage) = run_one_shot( db.conn(), false, &OneShotRequest { @@ -852,6 +877,8 @@ mod tests { ) .unwrap(); + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.output_tokens, Some(5)); assert_eq!( response, OneShotResponse::PostAnalysis(PostAnalysisResult { @@ -900,7 +927,7 @@ mod tests { "categories": ["engineering"] }), }, - r#"{"choices":[{"message":{"content":"{\"tags\":[\"rust\",\"preview\"],\"categories\":[\"engineering\"]}"}}]}"#, + r#"{"choices":[{"message":{"content":"{\"tags\":[\"rust\",\"preview\"],\"categories\":[\"engineering\"]}"}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}"#, ); assert_eq!( @@ -919,7 +946,7 @@ mod tests { operation: OneShotOperation::AnalyzePost, content: json!({"title":"Draft title","excerpt":"","content":"Body"}), }, - r#"{"choices":[{"message":{"content":"{\"title\":\"Better title\",\"excerpt\":\"Short summary\",\"slug\":\"better-title\"}"}}]}"#, + r#"{"choices":[{"message":{"content":"{\"title\":\"Better title\",\"excerpt\":\"Short summary\",\"slug\":\"better-title\"}"}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}"#, ); assert_eq!( @@ -939,7 +966,7 @@ mod tests { operation: OneShotOperation::DetectLanguage, content: json!({"text": "Bonjour tout le monde"}), }, - r#"{"choices":[{"message":{"content":"{\"language_code\":\"fr\"}"}}]}"#, + r#"{"choices":[{"message":{"content":"{\"language_code\":\"fr\"}"}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}"#, ); assert_eq!( @@ -963,7 +990,7 @@ mod tests { "content": "Body" }), }, - r#"{"choices":[{"message":{"content":"{\"title\":\"Hallo\",\"excerpt\":\"Kurzfassung\",\"content\":\"Inhalt\"}"}}]}"#, + r#"{"choices":[{"message":{"content":"{\"title\":\"Hallo\",\"excerpt\":\"Kurzfassung\",\"content\":\"Inhalt\"}"}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}"#, ); assert_eq!( @@ -990,7 +1017,7 @@ mod tests { "image_data_url": "data:image/jpeg;base64,abc123" }), }, - r#"{"choices":[{"message":{"content":"{\"title\":\"Hero image\",\"alt\":\"A scenic mountain\",\"caption\":\"Sunrise over the ridge\"}"}}]}"#, + r#"{"choices":[{"message":{"content":"{\"title\":\"Hero image\",\"alt\":\"A scenic mountain\",\"caption\":\"Sunrise over the ridge\"}"}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}"#, ); assert_eq!( @@ -1016,7 +1043,7 @@ mod tests { "caption": "Morning light" }), }, - r#"{"choices":[{"message":{"content":"{\"title\":\"Montagna\",\"alt\":\"Cresta innevata\",\"caption\":\"Luce del mattino\"}"}}]}"#, + r#"{"choices":[{"message":{"content":"{\"title\":\"Montagna\",\"alt\":\"Cresta innevata\",\"caption\":\"Luce del mattino\"}"}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}"#, ); assert_eq!( @@ -1050,7 +1077,72 @@ mod tests { ) .unwrap(); - run_one_shot(db.conn(), true, &request).unwrap() + let (response, usage) = run_one_shot(db.conn(), true, &request).unwrap(); + assert_eq!( + usage, + TokenUsage { + input_tokens: Some(10), + output_tokens: Some(5), + cache_read_tokens: None, + cache_write_tokens: None, + } + ); + response + } + + #[test] + fn parse_token_usage_normalizes_complete_usage() { + let usage = parse_token_usage(&json!({ + "usage": { + "prompt_tokens": 100, + "completion_tokens": 42, + "prompt_tokens_details": { "cached_tokens": 80 }, + "completion_tokens_details": { "cached_tokens": 7 } + } + })); + assert_eq!( + usage, + TokenUsage { + input_tokens: Some(100), + output_tokens: Some(42), + cache_read_tokens: Some(80), + cache_write_tokens: Some(7), + } + ); + } + + #[test] + fn parse_token_usage_leaves_missing_counters_null() { + let usage = parse_token_usage(&json!({ + "usage": { "prompt_tokens": 100 } + })); + assert_eq!( + usage, + TokenUsage { + input_tokens: Some(100), + output_tokens: None, + cache_read_tokens: None, + cache_write_tokens: None, + } + ); + } + + #[test] + fn parse_token_usage_handles_absent_usage_object() { + assert_eq!(parse_token_usage(&json!({})), TokenUsage::default()); + } + + #[test] + fn parse_token_usage_ignores_non_numeric_counters() { + let usage = parse_token_usage(&json!({ + "usage": { + "prompt_tokens": "many", + "completion_tokens": -3, + "prompt_tokens_details": { "cached_tokens": null }, + "completion_tokens_details": "nope" + } + })); + assert_eq!(usage, TokenUsage::default()); } fn spawn_test_server( diff --git a/crates/bds-core/src/engine/auto_translation.rs b/crates/bds-core/src/engine/auto_translation.rs index a042993..b59984a 100644 --- a/crates/bds-core/src/engine/auto_translation.rs +++ b/crates/bds-core/src/engine/auto_translation.rs @@ -293,7 +293,7 @@ fn translate_post_ai( }), }, )? { - OneShotResponse::Translation(result) => Ok(result), + (OneShotResponse::Translation(result), _usage) => Ok(result), _ => Err(EngineError::Parse( "unexpected post translation response".into(), )), @@ -320,7 +320,7 @@ fn translate_media_ai( }), }, )? { - OneShotResponse::MediaTranslation(result) => Ok(result), + (OneShotResponse::MediaTranslation(result), _usage) => Ok(result), _ => Err(EngineError::Parse( "unexpected media translation response".into(), )), diff --git a/crates/bds-ui/src/app.rs b/crates/bds-ui/src/app.rs index a110565..1344688 100644 --- a/crates/bds-ui/src/app.rs +++ b/crates/bds-ui/src/app.rs @@ -2708,8 +2708,17 @@ impl BdsApp { .then_with(|| left.0.to_lowercase().cmp(&right.0.to_lowercase())) }); tag_items.truncate(40); - let max_count = tag_items.iter().map(|item| item.1).max().unwrap_or(1).max(1); - let min_count = tag_items.iter().map(|item| item.1).min().unwrap_or(max_count); + let max_count = tag_items + .iter() + .map(|item| item.1) + .max() + .unwrap_or(1) + .max(1); + let min_count = tag_items + .iter() + .map(|item| item.1) + .min() + .unwrap_or(max_count); let range = (max_count - min_count).max(1) as f32; let mut tag_cloud = tag_items .into_iter() @@ -5518,7 +5527,7 @@ impl BdsApp { }), }; match ai::run_one_shot(db.conn(), self.offline_mode, &request) { - Ok(ai::OneShotResponse::PostAnalysis(result)) => { + Ok((ai::OneShotResponse::PostAnalysis(result), _usage)) => { self.active_modal = Some(modal::ModalState::AISuggestions { target: modal::AiEntityTarget::Post(post_id.to_string()), fields: vec![ @@ -5577,7 +5586,7 @@ impl BdsApp { }), }; match ai::run_one_shot(db.conn(), self.offline_mode, &request) { - Ok(ai::OneShotResponse::Taxonomy(result)) => { + Ok((ai::OneShotResponse::Taxonomy(result), _usage)) => { self.active_modal = Some(modal::ModalState::AISuggestions { target: modal::AiEntityTarget::Post(post_id.to_string()), fields: vec![ @@ -5624,7 +5633,7 @@ impl BdsApp { }), }; match ai::run_one_shot(db.conn(), self.offline_mode, &request) { - Ok(ai::OneShotResponse::LanguageDetection(result)) => { + Ok((ai::OneShotResponse::LanguageDetection(result), _usage)) => { if let Some(editor) = self.post_editors.get_mut(post_id) { editor.language = result.language_code; editor.mark_dirty(); @@ -5734,7 +5743,7 @@ impl BdsApp { }), }; match ai::run_one_shot(db.conn(), self.offline_mode, &request) { - Ok(ai::OneShotResponse::Translation(result)) => { + Ok((ai::OneShotResponse::Translation(result), _usage)) => { if let Some(editor) = self.post_editors.get_mut(post_id) { editor.switch_language(target_language); editor.title = result.title.clone(); @@ -5798,7 +5807,7 @@ impl BdsApp { }), }; match ai::run_one_shot(db.conn(), self.offline_mode, &request) { - Ok(ai::OneShotResponse::ImageAnalysis(result)) => { + Ok((ai::OneShotResponse::ImageAnalysis(result), _usage)) => { self.active_modal = Some(modal::ModalState::AISuggestions { target: modal::AiEntityTarget::Media(media_id.to_string()), fields: vec![ @@ -5853,7 +5862,7 @@ impl BdsApp { }), }; match ai::run_one_shot(db.conn(), self.offline_mode, &request) { - Ok(ai::OneShotResponse::LanguageDetection(result)) => { + Ok((ai::OneShotResponse::LanguageDetection(result), _usage)) => { if let Some(editor) = self.media_editors.get_mut(media_id) { editor.language = result.language_code; editor.is_dirty = true; @@ -5920,7 +5929,9 @@ impl BdsApp { content: json!({ "text": format!("{}\n{}\n{}", state.title, state.alt, state.caption) }), }, ) { - Ok(ai::OneShotResponse::LanguageDetection(result)) => result.language_code, + Ok((ai::OneShotResponse::LanguageDetection(result), _usage)) => { + result.language_code + } Ok(_) => String::new(), Err(error) => { self.notify(ToastLevel::Error, &error.to_string()); @@ -5946,7 +5957,7 @@ impl BdsApp { }), }; match ai::run_one_shot(db.conn(), self.offline_mode, &request) { - Ok(ai::OneShotResponse::MediaTranslation(result)) => { + Ok((ai::OneShotResponse::MediaTranslation(result), _usage)) => { if let Some(editor) = self.media_editors.get_mut(media_id) { editor.switch_language(target_language); editor.title = result.title.clone(); diff --git a/crates/bds-ui/src/views/dashboard.rs b/crates/bds-ui/src/views/dashboard.rs index 8d0ce32..52682bd 100644 --- a/crates/bds-ui/src/views/dashboard.rs +++ b/crates/bds-ui/src/views/dashboard.rs @@ -265,8 +265,7 @@ fn tag_cloud<'a>(tags: &'a [DashboardTag], locale: UiLocale) -> Element<'a, Mess let words = tags .iter() .map(|tag| { - let bg = - parse_color(tag.color.as_deref()).unwrap_or(Color::from_rgb(0.18, 0.21, 0.28)); + let bg = parse_color(tag.color.as_deref()).unwrap_or(Color::from_rgb(0.18, 0.21, 0.28)); let fg = contrast_color(bg); tooltip( container(text(tag.name.clone()).size(tag.font_size).color(fg)) @@ -324,7 +323,11 @@ fn category_cloud<'a>( } fn post_count_label(locale: UiLocale, count: usize) -> String { - tw(locale, "dashboard.postCount", &[("count", &count.to_string())]) + tw( + locale, + "dashboard.postCount", + &[("count", &count.to_string())], + ) } fn recent_posts<'a>(posts: &'a [DashboardRecentPost], locale: UiLocale) -> Element<'a, Message> {