aljaz/goose
1
0
Fork 0
mirror of https://github.com/aljazceru/goose.git synced 2025-12-18 14:44:21 +01:00
Code Issues 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

12
crates/goose-cli/src/commands/bench.rs
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()
}
}

1
crates/goose-cli/src/lib.rs
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;

107
crates/goose-cli/src/log_usage.rs
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");
}
}

36
crates/goose-cli/src/session/mod.rs
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)
}
}

5
crates/goose/src/agents/agent.rs
View File

@@ -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<Value>;
/// Get the total usage of the agent
async fn usage(&self) -> Vec<ProviderUsage>;
/// Add custom text to be included in the system prompt
async fn extend_system_prompt(&mut self, extension: String);

33
crates/goose/src/agents/capabilities.rs
View File

@@ -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<String, String>,
resource_capable_extensions: HashSet<String>,
provider: Arc<Box<dyn Provider>>,
provider_usage: Mutex<Vec<ProviderUsage>>,
system_prompt_override: Option<String>,
system_prompt_extensions: Vec<String>,
}
@@ -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<ProviderUsage> {
let provider_usage = self.provider_usage.lock().await.clone();
let mut usage_map: HashMap<String, ProviderUsage> = 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<Vec<Tool>> {
let mut tools = Vec::new();

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

@@ -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<ProviderUsage> {
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);

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

@@ -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<ProviderUsage> {
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);

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

@@ -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<ProviderUsage> {
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);