Skip to content

Commit

Permalink
fix(anthropic): include cached tokens in usage count logs (#947)
Browse files Browse the repository at this point in the history
Co-authored-by: angiejones <jones.angie@gmail.com>
  • Loading branch information
evalstate and angiejones authored Feb 11, 2025
1 parent 9c33660 commit 162348d
Showing 1 changed file with 23 additions and 10 deletions.
33 changes: 23 additions & 10 deletions crates/goose/src/providers/formats/anthropic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -187,18 +187,31 @@ pub fn response_to_message(response: Value) -> Result<Message> {
pub fn get_usage(data: &Value) -> Result<Usage> {
// Extract usage data if available
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")
.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
.get("output_tokens")
.and_then(|v| v.as_u64())
.map(|v| v as i32);
let total_tokens = match (input_tokens, output_tokens) {
(Some(i), Some(o)) => Some(i + o),
_ => None,
};

let total_tokens = output_tokens.map(|o| total_input_tokens as i32 + o);

Ok(Usage::new(input_tokens, output_tokens, total_tokens))
} else {
Expand Down Expand Up @@ -295,9 +308,9 @@ mod tests {
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.total_tokens, Some(27));
assert_eq!(usage.total_tokens, Some(39)); // 24 + 15

Ok(())
}
Expand Down Expand Up @@ -338,9 +351,9 @@ mod tests {
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.total_tokens, Some(35));
assert_eq!(usage.total_tokens, Some(50)); // 30 + 20

Ok(())
}
Expand Down

0 comments on commit 162348d

Please sign in to comment.