feat: return normalized token usage from one-shot AI calls (closes #15)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -135,6 +135,14 @@ pub struct MediaTranslationResult {
|
|||||||
pub caption: String,
|
pub caption: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
|
||||||
|
pub struct TokenUsage {
|
||||||
|
pub input_tokens: Option<u64>,
|
||||||
|
pub output_tokens: Option<u64>,
|
||||||
|
pub cache_read_tokens: Option<u64>,
|
||||||
|
pub cache_write_tokens: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
pub enum OneShotResponse {
|
pub enum OneShotResponse {
|
||||||
Taxonomy(TaxonomySuggestion),
|
Taxonomy(TaxonomySuggestion),
|
||||||
@@ -323,7 +331,7 @@ pub fn run_one_shot(
|
|||||||
conn: &Connection,
|
conn: &Connection,
|
||||||
offline_mode: bool,
|
offline_mode: bool,
|
||||||
request: &OneShotRequest,
|
request: &OneShotRequest,
|
||||||
) -> EngineResult<OneShotResponse> {
|
) -> EngineResult<(OneShotResponse, TokenUsage)> {
|
||||||
let settings = load_ai_settings(conn, offline_mode)?;
|
let settings = load_ai_settings(conn, offline_mode)?;
|
||||||
let endpoint = active_endpoint(conn, offline_mode)?;
|
let endpoint = active_endpoint(conn, offline_mode)?;
|
||||||
let model = select_model(&settings, &endpoint, &request.operation)?;
|
let model = select_model(&settings, &endpoint, &request.operation)?;
|
||||||
@@ -370,7 +378,24 @@ pub fn run_one_shot(
|
|||||||
.ok_or_else(|| {
|
.ok_or_else(|| {
|
||||||
EngineError::Parse("chat completions response missing message content".to_string())
|
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<Client> {
|
fn build_http_client() -> EngineResult<Client> {
|
||||||
@@ -823,7 +848,7 @@ mod tests {
|
|||||||
|| request.contains("Authorization: Bearer secret-token")
|
|| request.contains("Authorization: Bearer secret-token")
|
||||||
);
|
);
|
||||||
http_ok(
|
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,
|
1,
|
||||||
@@ -842,7 +867,7 @@ mod tests {
|
|||||||
.unwrap();
|
.unwrap();
|
||||||
save_model_preferences(db.conn(), None, Some("gpt-4.1-mini"), None, "").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(),
|
db.conn(),
|
||||||
false,
|
false,
|
||||||
&OneShotRequest {
|
&OneShotRequest {
|
||||||
@@ -852,6 +877,8 @@ mod tests {
|
|||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(usage.input_tokens, Some(10));
|
||||||
|
assert_eq!(usage.output_tokens, Some(5));
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
response,
|
response,
|
||||||
OneShotResponse::PostAnalysis(PostAnalysisResult {
|
OneShotResponse::PostAnalysis(PostAnalysisResult {
|
||||||
@@ -900,7 +927,7 @@ mod tests {
|
|||||||
"categories": ["engineering"]
|
"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!(
|
assert_eq!(
|
||||||
@@ -919,7 +946,7 @@ mod tests {
|
|||||||
operation: OneShotOperation::AnalyzePost,
|
operation: OneShotOperation::AnalyzePost,
|
||||||
content: json!({"title":"Draft title","excerpt":"","content":"Body"}),
|
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!(
|
assert_eq!(
|
||||||
@@ -939,7 +966,7 @@ mod tests {
|
|||||||
operation: OneShotOperation::DetectLanguage,
|
operation: OneShotOperation::DetectLanguage,
|
||||||
content: json!({"text": "Bonjour tout le monde"}),
|
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!(
|
assert_eq!(
|
||||||
@@ -963,7 +990,7 @@ mod tests {
|
|||||||
"content": "Body"
|
"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!(
|
assert_eq!(
|
||||||
@@ -990,7 +1017,7 @@ mod tests {
|
|||||||
"image_data_url": "data:image/jpeg;base64,abc123"
|
"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!(
|
assert_eq!(
|
||||||
@@ -1016,7 +1043,7 @@ mod tests {
|
|||||||
"caption": "Morning light"
|
"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!(
|
assert_eq!(
|
||||||
@@ -1050,7 +1077,72 @@ mod tests {
|
|||||||
)
|
)
|
||||||
.unwrap();
|
.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(
|
fn spawn_test_server(
|
||||||
|
|||||||
@@ -293,7 +293,7 @@ fn translate_post_ai(
|
|||||||
}),
|
}),
|
||||||
},
|
},
|
||||||
)? {
|
)? {
|
||||||
OneShotResponse::Translation(result) => Ok(result),
|
(OneShotResponse::Translation(result), _usage) => Ok(result),
|
||||||
_ => Err(EngineError::Parse(
|
_ => Err(EngineError::Parse(
|
||||||
"unexpected post translation response".into(),
|
"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(
|
_ => Err(EngineError::Parse(
|
||||||
"unexpected media translation response".into(),
|
"unexpected media translation response".into(),
|
||||||
)),
|
)),
|
||||||
|
|||||||
@@ -2708,8 +2708,17 @@ impl BdsApp {
|
|||||||
.then_with(|| left.0.to_lowercase().cmp(&right.0.to_lowercase()))
|
.then_with(|| left.0.to_lowercase().cmp(&right.0.to_lowercase()))
|
||||||
});
|
});
|
||||||
tag_items.truncate(40);
|
tag_items.truncate(40);
|
||||||
let max_count = tag_items.iter().map(|item| item.1).max().unwrap_or(1).max(1);
|
let max_count = tag_items
|
||||||
let min_count = tag_items.iter().map(|item| item.1).min().unwrap_or(max_count);
|
.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 range = (max_count - min_count).max(1) as f32;
|
||||||
let mut tag_cloud = tag_items
|
let mut tag_cloud = tag_items
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -5518,7 +5527,7 @@ impl BdsApp {
|
|||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
match ai::run_one_shot(db.conn(), self.offline_mode, &request) {
|
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 {
|
self.active_modal = Some(modal::ModalState::AISuggestions {
|
||||||
target: modal::AiEntityTarget::Post(post_id.to_string()),
|
target: modal::AiEntityTarget::Post(post_id.to_string()),
|
||||||
fields: vec![
|
fields: vec![
|
||||||
@@ -5577,7 +5586,7 @@ impl BdsApp {
|
|||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
match ai::run_one_shot(db.conn(), self.offline_mode, &request) {
|
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 {
|
self.active_modal = Some(modal::ModalState::AISuggestions {
|
||||||
target: modal::AiEntityTarget::Post(post_id.to_string()),
|
target: modal::AiEntityTarget::Post(post_id.to_string()),
|
||||||
fields: vec![
|
fields: vec![
|
||||||
@@ -5624,7 +5633,7 @@ impl BdsApp {
|
|||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
match ai::run_one_shot(db.conn(), self.offline_mode, &request) {
|
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) {
|
if let Some(editor) = self.post_editors.get_mut(post_id) {
|
||||||
editor.language = result.language_code;
|
editor.language = result.language_code;
|
||||||
editor.mark_dirty();
|
editor.mark_dirty();
|
||||||
@@ -5734,7 +5743,7 @@ impl BdsApp {
|
|||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
match ai::run_one_shot(db.conn(), self.offline_mode, &request) {
|
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) {
|
if let Some(editor) = self.post_editors.get_mut(post_id) {
|
||||||
editor.switch_language(target_language);
|
editor.switch_language(target_language);
|
||||||
editor.title = result.title.clone();
|
editor.title = result.title.clone();
|
||||||
@@ -5798,7 +5807,7 @@ impl BdsApp {
|
|||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
match ai::run_one_shot(db.conn(), self.offline_mode, &request) {
|
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 {
|
self.active_modal = Some(modal::ModalState::AISuggestions {
|
||||||
target: modal::AiEntityTarget::Media(media_id.to_string()),
|
target: modal::AiEntityTarget::Media(media_id.to_string()),
|
||||||
fields: vec![
|
fields: vec![
|
||||||
@@ -5853,7 +5862,7 @@ impl BdsApp {
|
|||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
match ai::run_one_shot(db.conn(), self.offline_mode, &request) {
|
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) {
|
if let Some(editor) = self.media_editors.get_mut(media_id) {
|
||||||
editor.language = result.language_code;
|
editor.language = result.language_code;
|
||||||
editor.is_dirty = true;
|
editor.is_dirty = true;
|
||||||
@@ -5920,7 +5929,9 @@ impl BdsApp {
|
|||||||
content: json!({ "text": format!("{}\n{}\n{}", state.title, state.alt, state.caption) }),
|
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(),
|
Ok(_) => String::new(),
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
self.notify(ToastLevel::Error, &error.to_string());
|
self.notify(ToastLevel::Error, &error.to_string());
|
||||||
@@ -5946,7 +5957,7 @@ impl BdsApp {
|
|||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
match ai::run_one_shot(db.conn(), self.offline_mode, &request) {
|
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) {
|
if let Some(editor) = self.media_editors.get_mut(media_id) {
|
||||||
editor.switch_language(target_language);
|
editor.switch_language(target_language);
|
||||||
editor.title = result.title.clone();
|
editor.title = result.title.clone();
|
||||||
|
|||||||
@@ -265,8 +265,7 @@ fn tag_cloud<'a>(tags: &'a [DashboardTag], locale: UiLocale) -> Element<'a, Mess
|
|||||||
let words = tags
|
let words = tags
|
||||||
.iter()
|
.iter()
|
||||||
.map(|tag| {
|
.map(|tag| {
|
||||||
let bg =
|
let bg = parse_color(tag.color.as_deref()).unwrap_or(Color::from_rgb(0.18, 0.21, 0.28));
|
||||||
parse_color(tag.color.as_deref()).unwrap_or(Color::from_rgb(0.18, 0.21, 0.28));
|
|
||||||
let fg = contrast_color(bg);
|
let fg = contrast_color(bg);
|
||||||
tooltip(
|
tooltip(
|
||||||
container(text(tag.name.clone()).size(tag.font_size).color(fg))
|
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 {
|
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> {
|
fn recent_posts<'a>(posts: &'a [DashboardRecentPost], locale: UiLocale) -> Element<'a, Message> {
|
||||||
|
|||||||
Reference in New Issue
Block a user