aljaz/goose
1
0
Fork 0
mirror of https://github.com/aljazceru/goose.git synced 2026-02-02 05:04:23 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add input and output token tracking to SessionMetadata (#1979)

Browse Source
This commit is contained in:
Young Kim
2025-04-01 11:08:28 -07:00
committed by GitHub
parent 45e0f7f717
commit 2fbf601446
5 changed files with 20 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
Cargo.lock generated
View File

@@ -4015,9 +4015,9 @@ checksum = "b410bbe7e14ab526a0e86877eb47c6996a2bd7746f027ba551028c925390e4e9"
[[package]]
name = "openssl"
version = "0.10.70"
version = "0.10.71"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61cfb4e166a8bb8c9b55c500bc2308550148ece889be90f609377e58140f42c6"
checksum = "5e14130c6a98cd258fdcb0fb6d744152343ff729cbfcb28c656a9d12b999fbcd"
dependencies = [
"bitflags 2.9.0",
"cfg-if",
@@ -4056,9 +4056,9 @@ dependencies = [
[[package]]
name = "openssl-sys"
version = "0.9.105"
version = "0.9.106"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8b22d5b84be05a8d6947c7cb71f7c849aa0f112acd4bf51c2a7c1c988ac0a9dc"
checksum = "8bb61ea9811cc39e3c2069f40b8b8e2e70d8569b361f879786cc7ed48b777cdd"
dependencies = [
"cc",
"libc",

2
crates/goose/src/agents/reference.rs
View File

@@ -169,6 +169,8 @@ impl Agent for ReferenceAgent {
let mut metadata = session::read_metadata(&session_file)?;
metadata.working_dir = session.working_dir;
metadata.total_tokens = usage.usage.total_tokens;
metadata.input_tokens = usage.usage.input_tokens;
metadata.output_tokens = usage.usage.output_tokens;
// The message count is the number of messages in the session + 1 for the response
// The message count does not include the tool response till next iteration
metadata.message_count = messages.len() + 1;

2
crates/goose/src/agents/summarize.rs
View File

@@ -266,6 +266,8 @@ impl Agent for SummarizeAgent {
let mut metadata = session::read_metadata(&session_file)?;
metadata.working_dir = session.working_dir;
metadata.total_tokens = usage.usage.total_tokens;
metadata.input_tokens = usage.usage.input_tokens;
metadata.output_tokens = usage.usage.output_tokens;
// The message count is the number of messages in the session + 1 for the response
// The message count does not include the tool response till next iteration
metadata.message_count = messages.len() + 1;

2
crates/goose/src/agents/truncate.rs
View File

@@ -280,6 +280,8 @@ impl Agent for TruncateAgent {
let mut metadata = session::read_metadata(&session_file)?;
metadata.working_dir = session.working_dir;
metadata.total_tokens = usage.usage.total_tokens;
metadata.input_tokens = usage.usage.input_tokens;
metadata.output_tokens = usage.usage.output_tokens;
// The message count is the number of messages in the session + 1 for the response
// The message count does not include the tool response till next iteration
metadata.message_count = messages.len() + 1;

10
crates/goose/src/session/storage.rs
View File

@@ -27,6 +27,10 @@ pub struct SessionMetadata {
pub message_count: usize,
/// The total number of tokens used in the session. Retrieved from the provider's last usage.
pub total_tokens: Option<i32>,
/// The number of input tokens used in the session. Retrieved from the provider's last usage.
pub input_tokens: Option<i32>,
/// The number of output tokens used in the session. Retrieved from the provider's last usage.
pub output_tokens: Option<i32>,
}
// Custom deserializer to handle old sessions without working_dir
@@ -40,6 +44,8 @@ impl<'de> Deserialize<'de> for SessionMetadata {
description: String,
message_count: usize,
total_tokens: Option<i32>,
input_tokens: Option<i32>,
output_tokens: Option<i32>,
working_dir: Option<PathBuf>,
}
@@ -49,6 +55,8 @@ impl<'de> Deserialize<'de> for SessionMetadata {
description: helper.description,
message_count: helper.message_count,
total_tokens: helper.total_tokens,
input_tokens: helper.input_tokens,
output_tokens: helper.output_tokens,
working_dir: helper.working_dir.unwrap_or_else(get_home_dir),
})
}
@@ -61,6 +69,8 @@ impl SessionMetadata {
description: String::new(),
message_count: 0,
total_tokens: None,
input_tokens: None,
output_tokens: None,
}
}
}