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/src/agents/agent.rs
+1 -4
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);
crates/goose/src/agents/capabilities.rs
+1 -32
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();
crates/goose/src/agents/reference.rs
-7
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);
crates/goose/src/agents/summarize.rs
-8
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);
crates/goose/src/agents/truncate.rs
-8
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);