feat: Add comprehensive cost tracking display for LLM usage (#2992)

Co-authored-by: jack <jack@deck.local>
Co-authored-by: Bradley Axen <baxen@squareup.com>
This commit is contained in:
jack
2025-06-26 10:46:14 +01:00
committed by GitHub
parent 9d48ed980f
commit d2ff4f3746
24 changed files with 1985 additions and 47 deletions
+11
View File
@@ -10,12 +10,23 @@ use goose::scheduler_factory::SchedulerFactory;
use tower_http::cors::{Any, CorsLayer};
use tracing::info;
use goose::providers::pricing::initialize_pricing_cache;
pub async fn run() -> Result<()> {
// Initialize logging
crate::logging::setup_logging(Some("goosed"))?;
let settings = configuration::Settings::new()?;
// Initialize pricing cache on startup
tracing::info!("Initializing pricing cache...");
if let Err(e) = initialize_pricing_cache().await {
tracing::warn!(
"Failed to initialize pricing cache: {}. Pricing data may not be available.",
e
);
}
let secret_key =
std::env::var("GOOSE_SERVER__SECRET_KEY").unwrap_or_else(|_| "test".to_string());
@@ -13,6 +13,7 @@ use goose::config::{extensions::name_to_key, PermissionManager};
use goose::config::{ExtensionConfigManager, ExtensionEntry};
use goose::model::ModelConfig;
use goose::providers::base::ProviderMetadata;
use goose::providers::pricing::{get_all_pricing, get_model_pricing, refresh_pricing};
use goose::providers::providers as get_providers;
use goose::{agents::ExtensionConfig, config::permission::PermissionLevel};
use http::{HeaderMap, StatusCode};
@@ -314,6 +315,128 @@ pub async fn providers(
Ok(Json(providers_response))
}
#[derive(Serialize, ToSchema)]
pub struct PricingData {
pub provider: String,
pub model: String,
pub input_token_cost: f64,
pub output_token_cost: f64,
pub currency: String,
pub context_length: Option<u32>,
}
#[derive(Serialize, ToSchema)]
pub struct PricingResponse {
pub pricing: Vec<PricingData>,
pub source: String,
}
#[derive(Deserialize, ToSchema)]
pub struct PricingQuery {
/// If true, only return pricing for configured providers. If false, return all.
pub configured_only: Option<bool>,
}
#[utoipa::path(
post,
path = "/config/pricing",
request_body = PricingQuery,
responses(
(status = 200, description = "Model pricing data retrieved successfully", body = PricingResponse)
)
)]
pub async fn get_pricing(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(query): Json<PricingQuery>,
) -> Result<Json<PricingResponse>, StatusCode> {
verify_secret_key(&headers, &state)?;
let configured_only = query.configured_only.unwrap_or(true);
// If refresh requested (configured_only = false), refresh the cache
if !configured_only {
if let Err(e) = refresh_pricing().await {
tracing::error!("Failed to refresh pricing data: {}", e);
}
}
let mut pricing_data = Vec::new();
if !configured_only {
// Get ALL pricing data from the cache
let all_pricing = get_all_pricing().await;
for (provider, models) in all_pricing {
for (model, pricing) in models {
pricing_data.push(PricingData {
provider: provider.clone(),
model: model.clone(),
input_token_cost: pricing.input_cost,
output_token_cost: pricing.output_cost,
currency: "$".to_string(),
context_length: pricing.context_length,
});
}
}
} else {
// Get only configured providers' pricing
let providers_metadata = get_providers();
for metadata in providers_metadata {
// Skip unconfigured providers if filtering
if !check_provider_configured(&metadata) {
continue;
}
for model_info in &metadata.known_models {
// Try to get pricing from cache
if let Some(pricing) = get_model_pricing(&metadata.name, &model_info.name).await {
pricing_data.push(PricingData {
provider: metadata.name.clone(),
model: model_info.name.clone(),
input_token_cost: pricing.input_cost,
output_token_cost: pricing.output_cost,
currency: "$".to_string(),
context_length: pricing.context_length,
});
}
// Check if the model has embedded pricing data
else if let (Some(input_cost), Some(output_cost)) =
(model_info.input_token_cost, model_info.output_token_cost)
{
pricing_data.push(PricingData {
provider: metadata.name.clone(),
model: model_info.name.clone(),
input_token_cost: input_cost,
output_token_cost: output_cost,
currency: model_info
.currency
.clone()
.unwrap_or_else(|| "$".to_string()),
context_length: Some(model_info.context_limit as u32),
});
}
}
}
}
tracing::info!(
"Returning pricing for {} models{}",
pricing_data.len(),
if configured_only {
" (configured providers only)"
} else {
" (all cached models)"
}
);
Ok(Json(PricingResponse {
pricing: pricing_data,
source: "openrouter".to_string(),
}))
}
#[utoipa::path(
post,
path = "/config/init",
@@ -471,6 +594,7 @@ pub fn routes(state: Arc<AppState>) -> Router {
.route("/config/extensions", post(add_extension))
.route("/config/extensions/{name}", delete(remove_extension))
.route("/config/providers", get(providers))
.route("/config/pricing", post(get_pricing))
.route("/config/init", post(init_config))
.route("/config/backup", post(backup_config))
.route("/config/permissions", post(upsert_permissions))
+1
View File
@@ -17,6 +17,7 @@ mcp-core = { path = "../mcp-core" }
anyhow = "1.0"
thiserror = "1.0"
futures = "0.3"
dirs = "5.0"
reqwest = { version = "0.12.9", features = [
"rustls-tls-native-roots",
"json",
+9 -3
View File
@@ -5,7 +5,7 @@ use reqwest::{Client, StatusCode};
use serde_json::Value;
use std::time::Duration;
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage};
use super::base::{ConfigKey, ModelInfo, Provider, ProviderMetadata, ProviderUsage};
use super::errors::ProviderError;
use super::formats::anthropic::{create_request, get_usage, response_to_message};
use super::utils::{emit_debug_trace, get_model};
@@ -122,12 +122,18 @@ impl AnthropicProvider {
#[async_trait]
impl Provider for AnthropicProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::new(
ProviderMetadata::with_models(
"anthropic",
"Anthropic",
"Claude and other models from Anthropic",
ANTHROPIC_DEFAULT_MODEL,
ANTHROPIC_KNOWN_MODELS.to_vec(),
vec![
ModelInfo::with_cost("claude-3-5-sonnet-20241022", 200000, 0.000003, 0.000015),
ModelInfo::with_cost("claude-3-5-haiku-20241022", 200000, 0.000001, 0.000005),
ModelInfo::with_cost("claude-3-opus-20240229", 200000, 0.000015, 0.000075),
ModelInfo::with_cost("claude-3-sonnet-20240229", 200000, 0.000003, 0.000015),
ModelInfo::with_cost("claude-3-haiku-20240307", 200000, 0.00000025, 0.00000125),
],
ANTHROPIC_DOC_URL,
vec![
ConfigKey::new("ANTHROPIC_API_KEY", true, true, None),
+78
View File
@@ -32,6 +32,41 @@ pub struct ModelInfo {
pub name: String,
/// The maximum context length this model supports
pub context_limit: usize,
/// Cost per token for input (optional)
pub input_token_cost: Option<f64>,
/// Cost per token for output (optional)
pub output_token_cost: Option<f64>,
/// Currency for the costs (default: "$")
pub currency: Option<String>,
}
impl ModelInfo {
/// Create a new ModelInfo with just name and context limit
pub fn new(name: impl Into<String>, context_limit: usize) -> Self {
Self {
name: name.into(),
context_limit,
input_token_cost: None,
output_token_cost: None,
currency: None,
}
}
/// Create a new ModelInfo with cost information (per token)
pub fn with_cost(
name: impl Into<String>,
context_limit: usize,
input_cost: f64,
output_cost: f64,
) -> Self {
Self {
name: name.into(),
context_limit,
input_token_cost: Some(input_cost),
output_token_cost: Some(output_cost),
currency: Some("$".to_string()),
}
}
}
/// Metadata about a provider's configuration requirements and capabilities
@@ -74,6 +109,9 @@ impl ProviderMetadata {
.map(|&name| ModelInfo {
name: name.to_string(),
context_limit: ModelConfig::new(name.to_string()).context_limit(),
input_token_cost: None,
output_token_cost: None,
currency: None,
})
.collect(),
model_doc_link: model_doc_link.to_string(),
@@ -81,6 +119,27 @@ impl ProviderMetadata {
}
}
/// Create a new ProviderMetadata with ModelInfo objects that include cost data
pub fn with_models(
name: &str,
display_name: &str,
description: &str,
default_model: &str,
models: Vec<ModelInfo>,
model_doc_link: &str,
config_keys: Vec<ConfigKey>,
) -> Self {
Self {
name: name.to_string(),
display_name: display_name.to_string(),
description: description.to_string(),
default_model: default_model.to_string(),
known_models: models,
model_doc_link: model_doc_link.to_string(),
config_keys,
}
}
pub fn empty() -> Self {
Self {
name: "".to_string(),
@@ -313,6 +372,9 @@ mod tests {
let info = ModelInfo {
name: "test-model".to_string(),
context_limit: 1000,
input_token_cost: None,
output_token_cost: None,
currency: None,
};
assert_eq!(info.context_limit, 1000);
@@ -320,6 +382,9 @@ mod tests {
let info2 = ModelInfo {
name: "test-model".to_string(),
context_limit: 1000,
input_token_cost: None,
output_token_cost: None,
currency: None,
};
assert_eq!(info, info2);
@@ -327,7 +392,20 @@ mod tests {
let info3 = ModelInfo {
name: "test-model".to_string(),
context_limit: 2000,
input_token_cost: None,
output_token_cost: None,
currency: None,
};
assert_ne!(info, info3);
}
#[test]
fn test_model_info_with_cost() {
let info = ModelInfo::with_cost("gpt-4o", 128000, 0.0000025, 0.00001);
assert_eq!(info.name, "gpt-4o");
assert_eq!(info.context_limit, 128000);
assert_eq!(info.input_token_cost, Some(0.0000025));
assert_eq!(info.output_token_cost, Some(0.00001));
assert_eq!(info.currency, Some("$".to_string()));
}
}
+1
View File
@@ -18,6 +18,7 @@ pub mod oauth;
pub mod ollama;
pub mod openai;
pub mod openrouter;
pub mod pricing;
pub mod sagemaker_tgi;
pub mod snowflake;
pub mod toolshim;
+11 -3
View File
@@ -5,7 +5,7 @@ use serde_json::Value;
use std::collections::HashMap;
use std::time::Duration;
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::base::{ConfigKey, ModelInfo, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::embedding::{EmbeddingCapable, EmbeddingRequest, EmbeddingResponse};
use super::errors::ProviderError;
use super::formats::openai::{create_request, get_usage, response_to_message};
@@ -126,12 +126,20 @@ impl OpenAiProvider {
#[async_trait]
impl Provider for OpenAiProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::new(
ProviderMetadata::with_models(
"openai",
"OpenAI",
"GPT-4 and other OpenAI models, including OpenAI compatible ones",
OPEN_AI_DEFAULT_MODEL,
OPEN_AI_KNOWN_MODELS.to_vec(),
vec![
ModelInfo::with_cost("gpt-4o", 128000, 0.0000025, 0.00001),
ModelInfo::with_cost("gpt-4o-mini", 128000, 0.00000015, 0.0000006),
ModelInfo::with_cost("gpt-4-turbo", 128000, 0.00001, 0.00003),
ModelInfo::with_cost("gpt-3.5-turbo", 16385, 0.0000005, 0.0000015),
ModelInfo::with_cost("o1", 200000, 0.000015, 0.00006),
ModelInfo::with_cost("o3", 200000, 0.000015, 0.00006), // Using o1 pricing as placeholder
ModelInfo::with_cost("o4-mini", 128000, 0.000003, 0.000012), // Using o1-mini pricing as placeholder
],
OPEN_AI_DOC_URL,
vec![
ConfigKey::new("OPENAI_API_KEY", true, true, None),
+387
View File
@@ -0,0 +1,387 @@
use anyhow::Result;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
/// Disk cache configuration
const CACHE_FILE_NAME: &str = "pricing_cache.json";
const CACHE_TTL_DAYS: u64 = 7; // Cache for 7 days
/// Get the cache directory path
fn get_cache_dir() -> Result<PathBuf> {
let cache_dir = if let Ok(goose_dir) = std::env::var("GOOSE_CACHE_DIR") {
PathBuf::from(goose_dir)
} else {
dirs::cache_dir()
.ok_or_else(|| anyhow::anyhow!("Could not determine cache directory"))?
.join("goose")
};
Ok(cache_dir)
}
/// Cached pricing data structure for disk storage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CachedPricingData {
/// Nested HashMap: provider -> model -> pricing info
pub pricing: HashMap<String, HashMap<String, PricingInfo>>,
/// Unix timestamp when data was fetched
pub fetched_at: u64,
}
/// Simplified pricing info for efficient storage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PricingInfo {
pub input_cost: f64, // Cost per token
pub output_cost: f64, // Cost per token
pub context_length: Option<u32>,
}
/// Cache for OpenRouter pricing data with disk persistence
pub struct PricingCache {
/// In-memory cache
memory_cache: Arc<RwLock<Option<CachedPricingData>>>,
}
impl PricingCache {
pub fn new() -> Self {
Self {
memory_cache: Arc::new(RwLock::new(None)),
}
}
/// Load pricing from disk cache
async fn load_from_disk(&self) -> Result<Option<CachedPricingData>> {
let cache_path = get_cache_dir()?.join(CACHE_FILE_NAME);
if !cache_path.exists() {
return Ok(None);
}
match tokio::fs::read(&cache_path).await {
Ok(data) => {
match serde_json::from_slice::<CachedPricingData>(&data) {
Ok(cached) => {
// Check if cache is still valid
let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
let age_days = (now - cached.fetched_at) / (24 * 60 * 60);
if age_days < CACHE_TTL_DAYS {
tracing::info!(
"Loaded pricing data from disk cache (age: {} days)",
age_days
);
Ok(Some(cached))
} else {
tracing::info!("Disk cache expired (age: {} days)", age_days);
Ok(None)
}
}
Err(e) => {
tracing::warn!("Failed to parse pricing cache: {}", e);
Ok(None)
}
}
}
Err(e) => {
tracing::warn!("Failed to read pricing cache: {}", e);
Ok(None)
}
}
}
/// Save pricing data to disk
async fn save_to_disk(&self, data: &CachedPricingData) -> Result<()> {
let cache_dir = get_cache_dir()?;
tokio::fs::create_dir_all(&cache_dir).await?;
let cache_path = cache_dir.join(CACHE_FILE_NAME);
let json_data = serde_json::to_vec_pretty(data)?;
tokio::fs::write(&cache_path, json_data).await?;
tracing::info!("Saved pricing data to disk cache");
Ok(())
}
/// Get pricing for a specific model
pub async fn get_model_pricing(&self, provider: &str, model: &str) -> Option<PricingInfo> {
// Try memory cache first
{
let cache = self.memory_cache.read().await;
if let Some(cached) = &*cache {
return cached
.pricing
.get(&provider.to_lowercase())
.and_then(|models| models.get(model))
.cloned();
}
}
// Try loading from disk
if let Ok(Some(disk_cache)) = self.load_from_disk().await {
// Update memory cache
{
let mut cache = self.memory_cache.write().await;
*cache = Some(disk_cache.clone());
}
return disk_cache
.pricing
.get(&provider.to_lowercase())
.and_then(|models| models.get(model))
.cloned();
}
None
}
/// Force refresh pricing data from OpenRouter
pub async fn refresh(&self) -> Result<()> {
let pricing = fetch_openrouter_pricing_internal().await?;
// Convert to our efficient structure
let mut structured_pricing: HashMap<String, HashMap<String, PricingInfo>> = HashMap::new();
for (model_id, model) in pricing {
if let Some((provider, model_name)) = parse_model_id(&model_id) {
if let (Some(input_cost), Some(output_cost)) = (
convert_pricing(&model.pricing.prompt),
convert_pricing(&model.pricing.completion),
) {
let provider_lower = provider.to_lowercase();
let provider_models = structured_pricing.entry(provider_lower).or_default();
provider_models.insert(
model_name,
PricingInfo {
input_cost,
output_cost,
context_length: model.context_length,
},
);
}
}
}
let cached_data = CachedPricingData {
pricing: structured_pricing,
fetched_at: SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(),
};
// Log how many models we fetched
let total_models: usize = cached_data
.pricing
.values()
.map(|models| models.len())
.sum();
tracing::info!(
"Fetched pricing for {} providers with {} total models from OpenRouter",
cached_data.pricing.len(),
total_models
);
// Save to disk
self.save_to_disk(&cached_data).await?;
// Update memory cache
{
let mut cache = self.memory_cache.write().await;
*cache = Some(cached_data);
}
Ok(())
}
/// Initialize cache (load from disk or fetch if needed)
pub async fn initialize(&self) -> Result<()> {
// Try loading from disk first
if let Ok(Some(cached)) = self.load_from_disk().await {
// Log how many models we have cached
let total_models: usize = cached.pricing.values().map(|models| models.len()).sum();
tracing::info!(
"Loaded {} providers with {} total models from disk cache",
cached.pricing.len(),
total_models
);
// Update memory cache
{
let mut cache = self.memory_cache.write().await;
*cache = Some(cached);
}
return Ok(());
}
// If no disk cache, fetch from OpenRouter
tracing::info!("No valid disk cache found, fetching from OpenRouter");
self.refresh().await
}
}
impl Default for PricingCache {
fn default() -> Self {
Self::new()
}
}
// Global cache instance
lazy_static::lazy_static! {
static ref PRICING_CACHE: PricingCache = PricingCache::new();
static ref HTTP_CLIENT: Client = Client::builder()
.timeout(Duration::from_secs(30))
.pool_idle_timeout(Duration::from_secs(90))
.pool_max_idle_per_host(10)
.build()
.unwrap();
}
/// OpenRouter model pricing information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OpenRouterModel {
pub id: String,
pub name: String,
pub pricing: OpenRouterPricing,
pub context_length: Option<u32>,
pub architecture: Option<Architecture>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OpenRouterPricing {
pub prompt: String, // Cost per token for input (in USD)
pub completion: String, // Cost per token for output (in USD)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Architecture {
pub modality: String,
pub tokenizer: String,
pub instruct_type: Option<String>,
}
/// Response from OpenRouter models endpoint
#[derive(Debug, Deserialize)]
pub struct OpenRouterModelsResponse {
pub data: Vec<OpenRouterModel>,
}
/// Internal function to fetch pricing data
async fn fetch_openrouter_pricing_internal() -> Result<HashMap<String, OpenRouterModel>> {
let response = HTTP_CLIENT
.get("https://openrouter.ai/api/v1/models")
.send()
.await?;
if !response.status().is_success() {
anyhow::bail!(
"Failed to fetch OpenRouter models: HTTP {}",
response.status()
);
}
let models_response: OpenRouterModelsResponse = response.json().await?;
// Create a map for easy lookup
let mut pricing_map = HashMap::new();
for model in models_response.data {
pricing_map.insert(model.id.clone(), model);
}
Ok(pricing_map)
}
/// Initialize pricing cache on startup
pub async fn initialize_pricing_cache() -> Result<()> {
PRICING_CACHE.initialize().await
}
/// Get pricing for a specific model
pub async fn get_model_pricing(provider: &str, model: &str) -> Option<PricingInfo> {
PRICING_CACHE.get_model_pricing(provider, model).await
}
/// Force refresh pricing data
pub async fn refresh_pricing() -> Result<()> {
PRICING_CACHE.refresh().await
}
/// Get all cached pricing data
pub async fn get_all_pricing() -> HashMap<String, HashMap<String, PricingInfo>> {
let cache = PRICING_CACHE.memory_cache.read().await;
if let Some(cached) = &*cache {
cached.pricing.clone()
} else {
// Try loading from disk
if let Ok(Some(disk_cache)) = PRICING_CACHE.load_from_disk().await {
// Update memory cache
drop(cache);
let mut write_cache = PRICING_CACHE.memory_cache.write().await;
*write_cache = Some(disk_cache.clone());
disk_cache.pricing
} else {
HashMap::new()
}
}
}
/// Convert OpenRouter model ID to provider/model format
/// e.g., "anthropic/claude-3.5-sonnet" -> ("anthropic", "claude-3.5-sonnet")
pub fn parse_model_id(model_id: &str) -> Option<(String, String)> {
let parts: Vec<&str> = model_id.splitn(2, '/').collect();
if parts.len() == 2 {
// Normalize provider names to match our internal naming
let provider = match parts[0] {
"openai" => "openai",
"anthropic" => "anthropic",
"google" => "google",
"meta-llama" => "ollama", // Meta models often run via Ollama
"mistralai" => "mistral",
"cohere" => "cohere",
"perplexity" => "perplexity",
"deepseek" => "deepseek",
"groq" => "groq",
"nvidia" => "nvidia",
"microsoft" => "azure",
"replicate" => "replicate",
"huggingface" => "huggingface",
_ => parts[0],
};
Some((provider.to_string(), parts[1].to_string()))
} else {
None
}
}
/// Convert OpenRouter pricing to cost per token (already in that format)
pub fn convert_pricing(price_str: &str) -> Option<f64> {
// OpenRouter prices are already in USD per token
price_str.parse::<f64>().ok()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_model_id() {
assert_eq!(
parse_model_id("anthropic/claude-3.5-sonnet"),
Some(("anthropic".to_string(), "claude-3.5-sonnet".to_string()))
);
assert_eq!(
parse_model_id("openai/gpt-4"),
Some(("openai".to_string(), "gpt-4".to_string()))
);
assert_eq!(parse_model_id("invalid-format"), None);
}
#[test]
fn test_convert_pricing() {
assert_eq!(convert_pricing("0.000003"), Some(0.000003));
assert_eq!(convert_pricing("0.015"), Some(0.015));
assert_eq!(convert_pricing("invalid"), None);
}
}