From ea0960f645f7a659d3c72dc285ece5ff89788983 Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Mon, 17 Mar 2025 15:18:21 -0400 Subject: [PATCH] refactor: clean up log usage (#1704) --- crates/goose-cli/src/commands/bench.rs | 12 +-- crates/goose-cli/src/lib.rs | 1 - crates/goose-cli/src/log_usage.rs | 107 ------------------------ crates/goose-cli/src/session/mod.rs | 36 ++++---- crates/goose/src/agents/agent.rs | 5 +- crates/goose/src/agents/capabilities.rs | 33 +------- crates/goose/src/agents/reference.rs | 7 -- crates/goose/src/agents/summarize.rs | 8 -- crates/goose/src/agents/truncate.rs | 8 -- 9 files changed, 20 insertions(+), 197 deletions(-) delete mode 100644 crates/goose-cli/src/log_usage.rs diff --git a/crates/goose-cli/src/commands/bench.rs b/crates/goose-cli/src/commands/bench.rs index 90466e44..7b5ee16e 100644 --- a/crates/goose-cli/src/commands/bench.rs +++ b/crates/goose-cli/src/commands/bench.rs @@ -48,17 +48,7 @@ impl BenchAgent for BenchSession { } async fn get_token_usage(&self) -> Option { - // 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() } } diff --git a/crates/goose-cli/src/lib.rs b/crates/goose-cli/src/lib.rs index 207af817..d972f9aa 100644 --- a/crates/goose-cli/src/lib.rs +++ b/crates/goose-cli/src/lib.rs @@ -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; diff --git a/crates/goose-cli/src/log_usage.rs b/crates/goose-cli/src/log_usage.rs deleted file mode 100644 index c87671af..00000000 --- a/crates/goose-cli/src/log_usage.rs +++ /dev/null @@ -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, -} - -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, -) { - 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"); - } -} diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index ad41061f..ce7c44bc 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -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, messages: Vec, @@ -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> { - Ok(self.agent.usage().await) + /// Get the session metadata + pub fn get_metadata(&self) -> Result { + 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> { + let metadata = self.get_metadata()?; + Ok(metadata.total_tokens) } } diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index d724fb4a..fbde767c 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use super::extension::{ExtensionConfig, ExtensionResult}; use crate::message::Message; -use crate::providers::base::{Provider, ProviderUsage}; +use crate::providers::base::Provider; use crate::session; use mcp_core::prompt::Prompt; use mcp_core::protocol::GetPromptResult; @@ -47,9 +47,6 @@ pub trait Agent: Send + Sync { /// Pass through a JSON-RPC request to a specific extension async fn passthrough(&self, extension: &str, request: Value) -> ExtensionResult; - /// Get the total usage of the agent - async fn usage(&self) -> Vec; - /// Add custom text to be included in the system prompt async fn extend_system_prompt(&mut self, extension: String); diff --git a/crates/goose/src/agents/capabilities.rs b/crates/goose/src/agents/capabilities.rs index 18105c24..44627b3b 100644 --- a/crates/goose/src/agents/capabilities.rs +++ b/crates/goose/src/agents/capabilities.rs @@ -13,7 +13,7 @@ use tracing::{debug, instrument}; use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult}; use crate::config::Config; use crate::prompt_template; -use crate::providers::base::{Provider, ProviderUsage}; +use crate::providers::base::Provider; use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait}; use mcp_client::transport::{SseTransport, StdioTransport, Transport}; use mcp_core::{prompt::Prompt, Content, Tool, ToolCall, ToolError, ToolResult}; @@ -32,7 +32,6 @@ pub struct Capabilities { instructions: HashMap, resource_capable_extensions: HashSet, provider: Arc>, - provider_usage: Mutex>, system_prompt_override: Option, system_prompt_extensions: Vec, } @@ -92,7 +91,6 @@ impl Capabilities { instructions: HashMap::new(), resource_capable_extensions: HashSet::new(), provider: Arc::new(provider), - provider_usage: Mutex::new(Vec::new()), system_prompt_override: None, system_prompt_extensions: Vec::new(), } @@ -207,12 +205,6 @@ impl Capabilities { Arc::clone(&self.provider) } - /// Record provider usage - // TODO consider moving this off to the provider or as a form of logging - pub async fn record_usage(&self, usage: ProviderUsage) { - self.provider_usage.lock().await.push(usage); - } - /// Get aggregated usage statistics pub async fn remove_extension(&mut self, name: &str) -> ExtensionResult<()> { let sanitized_name = normalize(name.to_string()); @@ -227,29 +219,6 @@ impl Capabilities { Ok(self.clients.keys().cloned().collect()) } - pub async fn get_usage(&self) -> Vec { - let provider_usage = self.provider_usage.lock().await.clone(); - let mut usage_map: HashMap = HashMap::new(); - - provider_usage.iter().for_each(|usage| { - usage_map - .entry(usage.model.clone()) - .and_modify(|e| { - e.usage.input_tokens = Some( - e.usage.input_tokens.unwrap_or(0) + usage.usage.input_tokens.unwrap_or(0), - ); - e.usage.output_tokens = Some( - e.usage.output_tokens.unwrap_or(0) + usage.usage.output_tokens.unwrap_or(0), - ); - e.usage.total_tokens = Some( - e.usage.total_tokens.unwrap_or(0) + usage.usage.total_tokens.unwrap_or(0), - ); - }) - .or_insert_with(|| usage.clone()); - }); - usage_map.into_values().collect() - } - /// Get all tools from all clients with proper prefixing pub async fn get_prefixed_tools(&mut self) -> ExtensionResult> { let mut tools = Vec::new(); diff --git a/crates/goose/src/agents/reference.rs b/crates/goose/src/agents/reference.rs index 089202de..324df39b 100644 --- a/crates/goose/src/agents/reference.rs +++ b/crates/goose/src/agents/reference.rs @@ -13,7 +13,6 @@ use crate::agents::capabilities::Capabilities; use crate::agents::extension::{ExtensionConfig, ExtensionResult}; use crate::message::{Message, ToolRequest}; use crate::providers::base::Provider; -use crate::providers::base::ProviderUsage; use crate::token_counter::TokenCounter; use crate::{register_agent, session}; use anyhow::{anyhow, Result}; @@ -146,7 +145,6 @@ impl Agent for ReferenceAgent { &messages, &tools, ).await?; - capabilities.record_usage(usage.clone()).await; // record usage for the session in the session file if let Some(session) = session.clone() { @@ -204,11 +202,6 @@ impl Agent for ReferenceAgent { })) } - async fn usage(&self) -> Vec { - let capabilities = self.capabilities.lock().await; - capabilities.get_usage().await - } - async fn extend_system_prompt(&mut self, extension: String) { let mut capabilities = self.capabilities.lock().await; capabilities.add_system_prompt_extension(extension); diff --git a/crates/goose/src/agents/summarize.rs b/crates/goose/src/agents/summarize.rs index 0a5bb368..07857266 100644 --- a/crates/goose/src/agents/summarize.rs +++ b/crates/goose/src/agents/summarize.rs @@ -18,7 +18,6 @@ use crate::config::Config; use crate::memory_condense::condense_messages; use crate::message::{Message, ToolRequest}; use crate::providers::base::Provider; -use crate::providers::base::ProviderUsage; use crate::providers::errors::ProviderError; use crate::register_agent; use crate::session; @@ -243,8 +242,6 @@ impl Agent for SummarizeAgent { &tools, ).await { Ok((response, usage)) => { - capabilities.record_usage(usage.clone()).await; - // record usage for the session in the session file if let Some(session) = session.clone() { // TODO: track session_id in langfuse tracing @@ -419,11 +416,6 @@ impl Agent for SummarizeAgent { })) } - async fn usage(&self) -> Vec { - let capabilities = self.capabilities.lock().await; - capabilities.get_usage().await - } - async fn extend_system_prompt(&mut self, extension: String) { let mut capabilities = self.capabilities.lock().await; capabilities.add_system_prompt_extension(extension); diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index a20212e4..298fca8b 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -17,7 +17,6 @@ use crate::agents::ToolPermissionStore; use crate::config::Config; use crate::message::{Message, ToolRequest}; use crate::providers::base::Provider; -use crate::providers::base::ProviderUsage; use crate::providers::errors::ProviderError; use crate::providers::toolshim::{ augment_message_with_tool_calls, modify_system_prompt_for_tool_json, OllamaInterpreter, @@ -258,8 +257,6 @@ impl Agent for TruncateAgent { response = augment_message_with_tool_calls(&interpreter, response, &toolshim_tools).await?; } - capabilities.record_usage(usage.clone()).await; - // record usage for the session in the session file if let Some(session) = session.clone() { // TODO: track session_id in langfuse tracing @@ -473,11 +470,6 @@ impl Agent for TruncateAgent { })) } - async fn usage(&self) -> Vec { - let capabilities = self.capabilities.lock().await; - capabilities.get_usage().await - } - async fn extend_system_prompt(&mut self, extension: String) { let mut capabilities = self.capabilities.lock().await; capabilities.add_system_prompt_extension(extension);