aljaz/goose
1
0
Fork 0
mirror of https://github.com/aljazceru/goose.git synced 2025-12-19 15:14:21 +01:00
Code Issues Packages Projects Releases Wiki Activity

fix(anthropic): include cached tokens in usage count logs (#947)

Browse Source 
Co-authored-by: angiejones <jones.angie@gmail.com>
This commit is contained in:
evalstate
2025-02-11 03:45:26 +00:00
committed by GitHub
parent 9c336607fe
commit 162348dd1f
1 changed files with 23 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
crates/goose/src/providers/formats/anthropic.rs
View File

@@ -187,18 +187,31 @@ pub fn response_to_message(response: Value) -> Result<Message> {
pub fn get_usage(data: &Value) -> Result<Usage> { pub fn get_usage(data: &Value) -> Result<Usage> {
// Extract usage data if available // Extract usage data if available
if let Some(usage) = data.get("usage") { if let Some(usage) = data.get("usage") {
let input_tokens = usage // Sum up all input token types:
// - input_tokens (fresh/uncached)
// - cache_creation_input_tokens (being written to cache)
// - cache_read_input_tokens (read from cache)
let total_input_tokens = usage
.get("input_tokens") .get("input_tokens")
.and_then(|v| v.as_u64()) .and_then(|v| v.as_u64())
.map(|v| v as i32); .unwrap_or(0)
+ usage
.get("cache_creation_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0)
+ usage
.get("cache_read_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0);
let input_tokens = Some(total_input_tokens as i32);
let output_tokens = usage let output_tokens = usage
.get("output_tokens") .get("output_tokens")
.and_then(|v| v.as_u64()) .and_then(|v| v.as_u64())
.map(|v| v as i32); .map(|v| v as i32);
let total_tokens = match (input_tokens, output_tokens) {
(Some(i), Some(o)) => Some(i + o), let total_tokens = output_tokens.map(|o| total_input_tokens as i32 + o);
_ => None,
};
Ok(Usage::new(input_tokens, output_tokens, total_tokens)) Ok(Usage::new(input_tokens, output_tokens, total_tokens))
} else { } else {
@@ -295,9 +308,9 @@ mod tests {
panic!("Expected Text content"); panic!("Expected Text content");
} }
assert_eq!(usage.input_tokens, Some(12)); assert_eq!(usage.input_tokens, Some(24)); // 12 + 12 + 0
assert_eq!(usage.output_tokens, Some(15)); assert_eq!(usage.output_tokens, Some(15));
assert_eq!(usage.total_tokens, Some(27)); assert_eq!(usage.total_tokens, Some(39)); // 24 + 15
Ok(()) Ok(())
} }
@@ -338,9 +351,9 @@ mod tests {
panic!("Expected ToolRequest content"); panic!("Expected ToolRequest content");
} }
assert_eq!(usage.input_tokens, Some(15)); assert_eq!(usage.input_tokens, Some(30)); // 15 + 15 + 0
assert_eq!(usage.output_tokens, Some(20)); assert_eq!(usage.output_tokens, Some(20));
assert_eq!(usage.total_tokens, Some(35)); assert_eq!(usage.total_tokens, Some(50)); // 30 + 20
Ok(()) Ok(())
} }