tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

refactor: clean up log usage (#1704)

Browse Source
This commit is contained in:
Salman Mohammed
2025-03-17 15:18:21 -04:00
committed by GitHub
parent 32ab726d5b
commit ea0960f645
9 changed files with 20 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/bench.rs
+1 -11
View File
@@ -48,17 +48,7 @@ impl BenchAgent for BenchSession {
}
async fn get_token_usage(&self) -> Option<i32> {
// Get token usage from the provider
if let Ok(usage) = self.session.get_usage().await {
// Sum up total tokens across all usage entries
let total_tokens = usage
.iter()
.map(|u| u.usage.total_tokens.unwrap_or(0))
.sum();
Some(total_tokens)
} else {
None
}
self.session.get_total_token_usage().ok().flatten()
}
}
crates/goose-cli/src/lib.rs
-1
View File
@@ -1,7 +1,6 @@
use etcetera::AppStrategyArgs;
use once_cell::sync::Lazy;
pub mod commands;
pub mod log_usage;
pub mod logging;
pub mod session;
crates/goose-cli/src/log_usage.rs
-107
View File
@@ -1,107 +0,0 @@
use etcetera::AppStrategy;
use goose::providers::base::ProviderUsage;
#[derive(Debug, serde::Serialize, serde::Deserialize)]
struct SessionLog {
session_file: String,
usage: Vec<ProviderUsage>,
}
pub fn log_usage(
#[cfg(target_os = "windows")] home_dir: etcetera::app_strategy::Windows,
#[cfg(any(target_os = "macos", target_os = "linux"))] home_dir: etcetera::app_strategy::Xdg,
session_file: String,
usage: Vec<ProviderUsage>,
) {
let log = SessionLog {
session_file,
usage,
};
// choose_app_strategy().state_dir()
// - macOS/Linux: ~/.local/state/goose/logs/
// - Windows: ~\AppData\Roaming\Block\goose\data\logs
// - Windows has no convention for state_dir, use data_dir instead
let log_dir = home_dir
.in_state_dir("logs")
.unwrap_or_else(|| home_dir.in_data_dir("logs"));
if let Err(e) = std::fs::create_dir_all(&log_dir) {
eprintln!("Failed to create log directory: {}", e);
return;
}
let log_file = log_dir.join("goose.log");
let serialized = match serde_json::to_string(&log) {
Ok(s) => s,
Err(e) => {
eprintln!("Failed to serialize usage log: {}", e);
return;
}
};
// Append to log file
if let Err(e) = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(log_file)
.and_then(|mut file| {
std::io::Write::write_all(&mut file, serialized.as_bytes())?;
std::io::Write::write_all(&mut file, b"\n")?;
Ok(())
})
{
eprintln!("Failed to write to usage log file: {}", e);
}
}
#[cfg(test)]
mod tests {
use etcetera::{choose_app_strategy, AppStrategy};
use goose::providers::base::{ProviderUsage, Usage};
use crate::log_usage::{log_usage, SessionLog};
#[test]
fn test_session_logging() {
use tempfile::tempdir;
// Create a temporary directory
let temp_dir = tempdir().unwrap();
let temp_home = temp_dir.path().to_path_buf();
// Temporarily set `HOME` to the temp directory
std::env::set_var("HOME", temp_home.as_os_str());
let home_dir = choose_app_strategy(crate::APP_STRATEGY.clone()).unwrap();
let log_file = home_dir
.in_state_dir("logs")
.unwrap_or_else(|| home_dir.in_data_dir("logs"))
.join("goose.log");
log_usage(
home_dir,
"path.txt".to_string(),
vec![ProviderUsage::new(
"model".to_string(),
Usage::new(Some(10), Some(20), Some(30)),
)],
);
// Check if log file exists and contains the expected content
assert!(log_file.exists(), "Log file should exist");
let log_content = std::fs::read_to_string(&log_file).unwrap();
let log: SessionLog = serde_json::from_str(&log_content).unwrap();
assert!(log.session_file.contains("path.txt"));
assert_eq!(log.usage[0].usage.input_tokens, Some(10));
assert_eq!(log.usage[0].usage.output_tokens, Some(20));
assert_eq!(log.usage[0].usage.total_tokens, Some(30));
assert_eq!(log.usage[0].model, "model");
// Remove the log file after test
std::fs::remove_file(&log_file).ok();
std::env::remove_var("HOME");
}
}
crates/goose-cli/src/session/mod.rs
+17 -19
View File
@@ -16,7 +16,6 @@ use goose::agents::extension::{Envs, ExtensionConfig};
use goose::agents::{Agent, SessionConfig};
use goose::config::Config;
use goose::message::{Message, MessageContent};
use goose::providers::base::ProviderUsage;
use goose::session;
use mcp_core::handler::ToolError;
use mcp_core::prompt::PromptMessage;
@@ -29,8 +28,6 @@ use std::sync::Arc;
use std::time::Instant;
use tokio;
use crate::log_usage::log_usage;
pub struct Session {
agent: Box<dyn Agent>,
messages: Vec<Message>,
@@ -413,19 +410,10 @@ impl Session {
}
}
// Log usage and cleanup
if let Ok(home_dir) = choose_app_strategy(crate::APP_STRATEGY.clone()) {
let usage = self.agent.usage().await;
log_usage(
home_dir,
self.session_file.to_string_lossy().to_string(),
usage,
);
println!(
"\nClosing session. Recorded to {}",
self.session_file.display()
);
}
println!(
"\nClosing session. Recorded to {}",
self.session_file.display()
);
Ok(())
}
@@ -645,8 +633,18 @@ impl Session {
self.messages.clone()
}
/// Get the token usage from the agent
pub async fn get_usage(&self) -> Result<Vec<ProviderUsage>> {
Ok(self.agent.usage().await)
/// Get the session metadata
pub fn get_metadata(&self) -> Result<session::SessionMetadata> {
if !self.session_file.exists() {
return Err(anyhow::anyhow!("Session file does not exist"));
}
session::read_metadata(&self.session_file)
}
// Get the session's total token usage
pub fn get_total_token_usage(&self) -> Result<Option<i32>> {
let metadata = self.get_metadata()?;
Ok(metadata.total_tokens)
}
}