From 6a0b8c25d5170473756da7ed6b557fdecaef0a53 Mon Sep 17 00:00:00 2001 From: David Katz Date: Thu, 18 Dec 2025 11:44:36 -0500 Subject: [PATCH] Integrate pricing with canonical model (#6130) --- crates/goose-cli/src/session/mod.rs | 17 +- crates/goose-cli/src/session/output.rs | 66 +-- crates/goose-server/src/commands/agent.rs | 9 - crates/goose-server/src/openapi.rs | 4 + .../src/routes/config_management.rs | 93 +--- crates/goose/src/providers/canonical/mod.rs | 6 + crates/goose/src/providers/mod.rs | 1 - crates/goose/src/providers/pricing.rs | 408 ------------------ ui/desktop/openapi.json | 98 +++++ ui/desktop/src/api/sdk.gen.ts | 11 +- ui/desktop/src/api/types.gen.ts | 35 ++ .../components/bottom_menu/CostTracker.tsx | 97 +---- .../settings/app/AppSettingsSection.tsx | 134 +----- ui/desktop/src/hooks/useAgent.ts | 10 - ui/desktop/src/hooks/useCostTracking.ts | 77 ++-- ui/desktop/src/utils/costDatabase.ts | 207 --------- ui/desktop/src/utils/pricing.ts | 24 ++ 17 files changed, 263 insertions(+), 1034 deletions(-) delete mode 100644 crates/goose/src/providers/pricing.rs delete mode 100644 ui/desktop/src/utils/costDatabase.ts create mode 100644 ui/desktop/src/utils/pricing.ts diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index de2c57e4..a9a38757 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -33,7 +33,6 @@ use goose::agents::extension::{Envs, ExtensionConfig, PLATFORM_EXTENSIONS}; use goose::agents::types::RetryConfig; use goose::agents::{Agent, SessionConfig, COMPACT_TRIGGERS}; use goose::config::{Config, GooseMode}; -use goose::providers::pricing::initialize_pricing_cache; use goose::session::SessionManager; use input::InputResult; use rmcp::model::PromptMessage; @@ -1416,19 +1415,6 @@ impl CliSession { .get_goose_provider() .unwrap_or_else(|_| "unknown".to_string()); - // Do not get costing information if show cost is disabled - // This will prevent the API call to openrouter.ai - // This is useful if for cases where openrouter.ai may be blocked by corporate firewalls - if show_cost { - // 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: {e}. Pricing data may not be available." - ); - } - } - match self.get_session().await { Ok(metadata) => { let total_tokens = metadata.total_tokens.unwrap_or(0) as usize; @@ -1443,8 +1429,7 @@ impl CliSession { &model_config.model_name, input_tokens, output_tokens, - ) - .await; + ); } } Err(_) => { diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index 04808ee3..9243eff9 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -5,11 +5,9 @@ use goose::config::Config; use goose::conversation::message::{ ActionRequiredData, Message, MessageContent, ToolRequest, ToolResponse, }; -use goose::providers::pricing::get_model_pricing; -use goose::providers::pricing::parse_model_id; +use goose::providers::canonical::maybe_get_canonical_model; use goose::utils::safe_truncate; use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; -use regex::Regex; use rmcp::model::{CallToolRequestParam, JsonObject, PromptArgument}; use serde_json::Value; use std::cell::RefCell; @@ -795,69 +793,25 @@ pub fn display_context_usage(total_tokens: usize, context_limit: usize) { ); } -fn normalize_model_name(model: &str) -> String { - let mut result = model.to_string(); - - // Remove "-latest" suffix - if result.ends_with("-latest") { - result = result.strip_suffix("-latest").unwrap().to_string(); - } - - // Remove date-like suffixes: -YYYYMMDD - let re_date = Regex::new(r"-\d{8}$").unwrap(); - if re_date.is_match(&result) { - result = re_date.replace(&result, "").to_string(); - } - - // Convert version numbers like -3-7- to -3.7- (e.g., claude-3-7-sonnet -> claude-3.7-sonnet) - let re_version = Regex::new(r"-(\d+)-(\d+)-").unwrap(); - if re_version.is_match(&result) { - result = re_version.replace(&result, "-$1.$2-").to_string(); - } - - result -} - -async fn estimate_cost_usd( +fn estimate_cost_usd( provider: &str, model: &str, input_tokens: usize, output_tokens: usize, ) -> Option { - // For OpenRouter, parse the model name to extract real provider/model - let openrouter_data = if provider == "openrouter" { - parse_model_id(model) - } else { - None - }; + let canonical_model = maybe_get_canonical_model(provider, model)?; - let (provider_to_use, model_to_use) = match &openrouter_data { - Some((real_provider, real_model)) => (real_provider.as_str(), real_model.as_str()), - None => (provider, model), - }; + let input_cost_per_token = canonical_model.pricing.prompt?; + let output_cost_per_token = canonical_model.pricing.completion?; - // Use the pricing module's get_model_pricing which handles model name mapping internally - let cleaned_model = normalize_model_name(model_to_use); - let pricing_info = get_model_pricing(provider_to_use, &cleaned_model).await; - - match pricing_info { - Some(pricing) => { - let input_cost = pricing.input_cost * input_tokens as f64; - let output_cost = pricing.output_cost * output_tokens as f64; - Some(input_cost + output_cost) - } - None => None, - } + let input_cost = input_cost_per_token * input_tokens as f64; + let output_cost = output_cost_per_token * output_tokens as f64; + Some(input_cost + output_cost) } /// Display cost information, if price data is available. -pub async fn display_cost_usage( - provider: &str, - model: &str, - input_tokens: usize, - output_tokens: usize, -) { - if let Some(cost) = estimate_cost_usd(provider, model, input_tokens, output_tokens).await { +pub fn display_cost_usage(provider: &str, model: &str, input_tokens: usize, output_tokens: usize) { + if let Some(cost) = estimate_cost_usd(provider, model, input_tokens, output_tokens) { use console::style; eprintln!( "Cost: {} USD ({} tokens: in {}, out {})", diff --git a/crates/goose-server/src/commands/agent.rs b/crates/goose-server/src/commands/agent.rs index 95064098..a68d56ee 100644 --- a/crates/goose-server/src/commands/agent.rs +++ b/crates/goose-server/src/commands/agent.rs @@ -6,8 +6,6 @@ use goose_server::auth::check_token; use tower_http::cors::{Any, CorsLayer}; use tracing::info; -use goose::providers::pricing::initialize_pricing_cache; - // Graceful shutdown signal #[cfg(unix)] async fn shutdown_signal() { @@ -32,13 +30,6 @@ pub async fn run() -> Result<()> { let settings = configuration::Settings::new()?; - 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()); diff --git a/crates/goose-server/src/openapi.rs b/crates/goose-server/src/openapi.rs index e576c87d..751947b0 100644 --- a/crates/goose-server/src/openapi.rs +++ b/crates/goose-server/src/openapi.rs @@ -351,6 +351,7 @@ derive_utoipa!(Icon as IconSchema); super::routes::config_management::remove_custom_provider, super::routes::config_management::check_provider, super::routes::config_management::set_config_provider, + super::routes::config_management::get_pricing, super::routes::agent::start_agent, super::routes::agent::resume_agent, super::routes::agent::get_tools, @@ -417,6 +418,9 @@ derive_utoipa!(Icon as IconSchema); super::routes::config_management::UpdateCustomProviderRequest, super::routes::config_management::CheckProviderRequest, super::routes::config_management::SetProviderRequest, + super::routes::config_management::PricingQuery, + super::routes::config_management::PricingResponse, + super::routes::config_management::PricingData, super::routes::action_required::ConfirmToolActionRequest, super::routes::reply::ChatRequest, super::routes::session::ImportSessionRequest, diff --git a/crates/goose-server/src/routes/config_management.rs b/crates/goose-server/src/routes/config_management.rs index 3dc618ce..9c483f4a 100644 --- a/crates/goose-server/src/routes/config_management.rs +++ b/crates/goose-server/src/routes/config_management.rs @@ -13,10 +13,8 @@ use goose::config::{Config, ConfigError}; use goose::model::ModelConfig; use goose::providers::auto_detect::detect_provider_from_api_key; use goose::providers::base::{ProviderMetadata, ProviderType}; +use goose::providers::canonical::maybe_get_canonical_model; use goose::providers::create_with_default_model; -use goose::providers::pricing::{ - get_all_pricing, get_model_pricing, parse_model_id, refresh_pricing, -}; use goose::providers::providers as get_providers; use goose::{ agents::execute_commands, agents::ExtensionConfig, config::permission::PermissionLevel, @@ -470,7 +468,8 @@ pub struct PricingResponse { #[derive(Deserialize, ToSchema)] pub struct PricingQuery { - pub configured_only: bool, + pub provider: String, + pub model: String, } #[utoipa::path( @@ -484,84 +483,28 @@ pub struct PricingQuery { pub async fn get_pricing( Json(query): Json, ) -> Result, StatusCode> { - let configured_only = query.configured_only; - - // 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 canonical_model = + maybe_get_canonical_model(&query.provider, &query.model).ok_or(StatusCode::NOT_FOUND)?; 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 { - for (metadata, provider_type) in get_providers().await { - // Skip unconfigured providers if filtering - if !check_provider_configured(&metadata, provider_type) { - continue; - } - - for model_info in &metadata.known_models { - // Handle OpenRouter models specially - they store full provider/model names - let (lookup_provider, lookup_model) = if metadata.name == "openrouter" { - // For OpenRouter, parse the model name to extract real provider/model - if let Some((provider, model)) = parse_model_id(&model_info.name) { - (provider, model) - } else { - // Fallback if parsing fails - (metadata.name.clone(), model_info.name.clone()) - } - } else { - // For other providers, use names as-is - (metadata.name.clone(), model_info.name.clone()) - }; - - // Only get pricing from OpenRouter cache - if let Some(pricing) = get_model_pricing(&lookup_provider, &lookup_model).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, - }); - } - // No fallback to hardcoded prices - } - } + if let (Some(input_cost), Some(output_cost)) = ( + canonical_model.pricing.prompt, + canonical_model.pricing.completion, + ) { + pricing_data.push(PricingData { + provider: query.provider.clone(), + model: query.model.clone(), + input_token_cost: input_cost, + output_token_cost: output_cost, + currency: "$".to_string(), + context_length: Some(canonical_model.context_length as u32), + }); } - tracing::debug!( - "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(), + source: "canonical".to_string(), })) } diff --git a/crates/goose/src/providers/canonical/mod.rs b/crates/goose/src/providers/canonical/mod.rs index 1aac285c..b7d0e560 100644 --- a/crates/goose/src/providers/canonical/mod.rs +++ b/crates/goose/src/providers/canonical/mod.rs @@ -20,3 +20,9 @@ impl ModelMapping { } } } + +pub fn maybe_get_canonical_model(provider: &str, model: &str) -> Option { + let registry = CanonicalModelRegistry::bundled().ok()?; + let canonical_id = map_to_canonical_model(provider, model, registry)?; + registry.get(&canonical_id).cloned() +} diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 3c3e6434..04b5a491 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -24,7 +24,6 @@ pub mod oauth; pub mod ollama; pub mod openai; pub mod openrouter; -pub mod pricing; pub mod provider_registry; pub mod provider_test; mod retry; diff --git a/crates/goose/src/providers/pricing.rs b/crates/goose/src/providers/pricing.rs deleted file mode 100644 index 50593b43..00000000 --- a/crates/goose/src/providers/pricing.rs +++ /dev/null @@ -1,408 +0,0 @@ -use anyhow::{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 { - 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>, - /// 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, -} - -/// Cache for OpenRouter pricing data with disk persistence -pub struct PricingCache { - /// In-memory cache - memory_cache: Arc>>, -} - -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> { - 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::(&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 { - Ok(Some(cached)) - } else { - 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?; - Ok(()) - } - - /// Get pricing for a specific model - pub async fn get_model_pricing(&self, provider: &str, model: &str) -> Option { - // 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> = 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(), - }; - - self.save_to_disk(&cached_data).await?; - - { - 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 { - { - let mut cache = self.memory_cache.write().await; - *cache = Some(cached); - } - - return Ok(()); - } - - 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(); -} - -fn create_http_client() -> Result { - Client::builder() - .timeout(Duration::from_secs(30)) - .pool_idle_timeout(Duration::from_secs(90)) - .pool_max_idle_per_host(10) - .build() - .map_err(|e| anyhow!(e)) -} - -/// 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, - pub architecture: Option, -} - -#[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, -} - -/// Response from OpenRouter models endpoint -#[derive(Debug, Deserialize)] -pub struct OpenRouterModelsResponse { - pub data: Vec, -} - -/// Internal function to fetch pricing data -async fn fetch_openrouter_pricing_internal() -> Result> { - let client = create_http_client()?; - let response = 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 { - 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> { - 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-sonnet-4-20250514" -> ("anthropic", "claude-sonnet-4-20250514") -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 { - // OpenRouter prices are already in USD per token - price_str.parse::().ok() -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_parse_model_id() { - assert_eq!( - parse_model_id("anthropic/claude-sonnet-4-20250514"), - Some(( - "anthropic".to_string(), - "claude-sonnet-4-20250514".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 the specific model causing issues - assert_eq!( - parse_model_id("anthropic/claude-sonnet-4-20250514"), - Some(( - "anthropic".to_string(), - "claude-sonnet-4-20250514".to_string() - )) - ); - } - - #[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); - } - - #[tokio::test] - async fn test_claude_sonnet_4_pricing_lookup() { - // Initialize the cache to load from disk - if let Err(e) = initialize_pricing_cache().await { - println!("Failed to initialize pricing cache: {}", e); - return; - } - - // Test lookup for the specific model (use the name that actually exists in cache) - let pricing = get_model_pricing("anthropic", "claude-sonnet-4").await; - - println!( - "Pricing lookup result for anthropic/claude-sonnet-4: {:?}", - pricing - ); - - // Should find pricing data - if let Some(pricing_info) = pricing { - assert!(pricing_info.input_cost > 0.0); - assert!(pricing_info.output_cost > 0.0); - println!( - "Found pricing: input={}, output={}", - pricing_info.input_cost, pricing_info.output_cost - ); - } else { - // Print debug info - let all_pricing = get_all_pricing().await; - if let Some(anthropic_models) = all_pricing.get("anthropic") { - println!("Available anthropic models in cache:"); - for model_name in anthropic_models.keys() { - println!(" {}", model_name); - } - } - panic!("Expected to find pricing for anthropic/claude-sonnet-4"); - } - } -} diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 82be74e7..6bc6dfbc 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -854,6 +854,36 @@ } } }, + "/config/pricing": { + "post": { + "tags": [ + "super::routes::config_management" + ], + "operationId": "get_pricing", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/PricingQuery" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "Model pricing data retrieved successfully", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/PricingResponse" + } + } + } + } + } + } + }, "/config/providers": { "get": { "tags": [ @@ -4223,6 +4253,74 @@ "never_allow" ] }, + "PricingData": { + "type": "object", + "required": [ + "provider", + "model", + "input_token_cost", + "output_token_cost", + "currency" + ], + "properties": { + "context_length": { + "type": "integer", + "format": "int32", + "nullable": true, + "minimum": 0 + }, + "currency": { + "type": "string" + }, + "input_token_cost": { + "type": "number", + "format": "double" + }, + "model": { + "type": "string" + }, + "output_token_cost": { + "type": "number", + "format": "double" + }, + "provider": { + "type": "string" + } + } + }, + "PricingQuery": { + "type": "object", + "required": [ + "provider", + "model" + ], + "properties": { + "model": { + "type": "string" + }, + "provider": { + "type": "string" + } + } + }, + "PricingResponse": { + "type": "object", + "required": [ + "pricing", + "source" + ], + "properties": { + "pricing": { + "type": "array", + "items": { + "$ref": "#/components/schemas/PricingData" + } + }, + "source": { + "type": "string" + } + } + }, "PrincipalType": { "type": "string", "enum": [ diff --git a/ui/desktop/src/api/sdk.gen.ts b/ui/desktop/src/api/sdk.gen.ts index 0ec5ba6b..0aaa6fb4 100644 --- a/ui/desktop/src/api/sdk.gen.ts +++ b/ui/desktop/src/api/sdk.gen.ts @@ -2,7 +2,7 @@ import type { Client, Options as Options2, TDataShape } from './client'; import { client } from './client.gen'; -import type { AddExtensionData, AddExtensionErrors, AddExtensionResponses, AgentAddExtensionData, AgentAddExtensionErrors, AgentAddExtensionResponses, AgentRemoveExtensionData, AgentRemoveExtensionErrors, AgentRemoveExtensionResponses, BackupConfigData, BackupConfigErrors, BackupConfigResponses, CallToolData, CallToolErrors, CallToolResponses, CheckProviderData, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionResponses, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleResponses, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponses, DeleteSessionData, DeleteSessionErrors, DeleteSessionResponses, DetectProviderData, DetectProviderErrors, DetectProviderResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponses, EditMessageData, EditMessageErrors, EditMessageResponses, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeResponses, ExportSessionData, ExportSessionErrors, ExportSessionResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponses, GetSessionData, GetSessionErrors, GetSessionInsightsData, GetSessionInsightsErrors, GetSessionInsightsResponses, GetSessionResponses, GetSlashCommandsData, GetSlashCommandsResponses, GetToolsData, GetToolsErrors, GetToolsResponses, GetTunnelStatusData, GetTunnelStatusResponses, ImportSessionData, ImportSessionErrors, ImportSessionResponses, InitConfigData, InitConfigErrors, InitConfigResponses, InspectRunningJobData, InspectRunningJobErrors, InspectRunningJobResponses, KillRunningJobData, KillRunningJobResponses, ListRecipesData, ListRecipesErrors, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponses, ListSessionsData, ListSessionsErrors, ListSessionsResponses, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, ParseRecipeData, ParseRecipeErrors, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponses, ProvidersData, ProvidersResponses, ReadAllConfigData, ReadAllConfigResponses, ReadConfigData, ReadConfigErrors, ReadConfigResponses, ReadResourceData, ReadResourceErrors, ReadResourceResponses, RecoverConfigData, RecoverConfigErrors, RecoverConfigResponses, RemoveConfigData, RemoveConfigErrors, RemoveConfigResponses, RemoveCustomProviderData, RemoveCustomProviderErrors, RemoveCustomProviderResponses, RemoveExtensionData, RemoveExtensionErrors, RemoveExtensionResponses, ReplyData, ReplyErrors, ReplyResponses, ResumeAgentData, ResumeAgentErrors, ResumeAgentResponses, RunNowHandlerData, RunNowHandlerErrors, RunNowHandlerResponses, SaveRecipeData, SaveRecipeErrors, SaveRecipeResponses, ScanRecipeData, ScanRecipeResponses, ScheduleRecipeData, ScheduleRecipeErrors, ScheduleRecipeResponses, SendTelemetryEventData, SendTelemetryEventResponses, SessionsHandlerData, SessionsHandlerErrors, SessionsHandlerResponses, SetConfigProviderData, SetRecipeSlashCommandData, SetRecipeSlashCommandErrors, SetRecipeSlashCommandResponses, StartAgentData, StartAgentErrors, StartAgentResponses, StartOpenrouterSetupData, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponses, StartTunnelData, StartTunnelErrors, StartTunnelResponses, StatusData, StatusResponses, StopTunnelData, StopTunnelErrors, StopTunnelResponses, UnpauseScheduleData, UnpauseScheduleErrors, UnpauseScheduleResponses, UpdateAgentProviderData, UpdateAgentProviderErrors, UpdateAgentProviderResponses, UpdateCustomProviderData, UpdateCustomProviderErrors, UpdateCustomProviderResponses, UpdateFromSessionData, UpdateFromSessionErrors, UpdateFromSessionResponses, UpdateRouterToolSelectorData, UpdateRouterToolSelectorErrors, UpdateRouterToolSelectorResponses, UpdateScheduleData, UpdateScheduleErrors, UpdateScheduleResponses, UpdateSessionNameData, UpdateSessionNameErrors, UpdateSessionNameResponses, UpdateSessionUserRecipeValuesData, UpdateSessionUserRecipeValuesErrors, UpdateSessionUserRecipeValuesResponses, UpsertConfigData, UpsertConfigErrors, UpsertConfigResponses, UpsertPermissionsData, UpsertPermissionsErrors, UpsertPermissionsResponses, ValidateConfigData, ValidateConfigErrors, ValidateConfigResponses } from './types.gen'; +import type { AddExtensionData, AddExtensionErrors, AddExtensionResponses, AgentAddExtensionData, AgentAddExtensionErrors, AgentAddExtensionResponses, AgentRemoveExtensionData, AgentRemoveExtensionErrors, AgentRemoveExtensionResponses, BackupConfigData, BackupConfigErrors, BackupConfigResponses, CallToolData, CallToolErrors, CallToolResponses, CheckProviderData, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionResponses, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleResponses, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponses, DeleteSessionData, DeleteSessionErrors, DeleteSessionResponses, DetectProviderData, DetectProviderErrors, DetectProviderResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponses, EditMessageData, EditMessageErrors, EditMessageResponses, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeResponses, ExportSessionData, ExportSessionErrors, ExportSessionResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponses, GetPricingData, GetPricingResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponses, GetSessionData, GetSessionErrors, GetSessionInsightsData, GetSessionInsightsErrors, GetSessionInsightsResponses, GetSessionResponses, GetSlashCommandsData, GetSlashCommandsResponses, GetToolsData, GetToolsErrors, GetToolsResponses, GetTunnelStatusData, GetTunnelStatusResponses, ImportSessionData, ImportSessionErrors, ImportSessionResponses, InitConfigData, InitConfigErrors, InitConfigResponses, InspectRunningJobData, InspectRunningJobErrors, InspectRunningJobResponses, KillRunningJobData, KillRunningJobResponses, ListRecipesData, ListRecipesErrors, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponses, ListSessionsData, ListSessionsErrors, ListSessionsResponses, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, ParseRecipeData, ParseRecipeErrors, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponses, ProvidersData, ProvidersResponses, ReadAllConfigData, ReadAllConfigResponses, ReadConfigData, ReadConfigErrors, ReadConfigResponses, ReadResourceData, ReadResourceErrors, ReadResourceResponses, RecoverConfigData, RecoverConfigErrors, RecoverConfigResponses, RemoveConfigData, RemoveConfigErrors, RemoveConfigResponses, RemoveCustomProviderData, RemoveCustomProviderErrors, RemoveCustomProviderResponses, RemoveExtensionData, RemoveExtensionErrors, RemoveExtensionResponses, ReplyData, ReplyErrors, ReplyResponses, ResumeAgentData, ResumeAgentErrors, ResumeAgentResponses, RunNowHandlerData, RunNowHandlerErrors, RunNowHandlerResponses, SaveRecipeData, SaveRecipeErrors, SaveRecipeResponses, ScanRecipeData, ScanRecipeResponses, ScheduleRecipeData, ScheduleRecipeErrors, ScheduleRecipeResponses, SendTelemetryEventData, SendTelemetryEventResponses, SessionsHandlerData, SessionsHandlerErrors, SessionsHandlerResponses, SetConfigProviderData, SetRecipeSlashCommandData, SetRecipeSlashCommandErrors, SetRecipeSlashCommandResponses, StartAgentData, StartAgentErrors, StartAgentResponses, StartOpenrouterSetupData, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponses, StartTunnelData, StartTunnelErrors, StartTunnelResponses, StatusData, StatusResponses, StopTunnelData, StopTunnelErrors, StopTunnelResponses, UnpauseScheduleData, UnpauseScheduleErrors, UnpauseScheduleResponses, UpdateAgentProviderData, UpdateAgentProviderErrors, UpdateAgentProviderResponses, UpdateCustomProviderData, UpdateCustomProviderErrors, UpdateCustomProviderResponses, UpdateFromSessionData, UpdateFromSessionErrors, UpdateFromSessionResponses, UpdateRouterToolSelectorData, UpdateRouterToolSelectorErrors, UpdateRouterToolSelectorResponses, UpdateScheduleData, UpdateScheduleErrors, UpdateScheduleResponses, UpdateSessionNameData, UpdateSessionNameErrors, UpdateSessionNameResponses, UpdateSessionUserRecipeValuesData, UpdateSessionUserRecipeValuesErrors, UpdateSessionUserRecipeValuesResponses, UpsertConfigData, UpsertConfigErrors, UpsertConfigResponses, UpsertPermissionsData, UpsertPermissionsErrors, UpsertPermissionsResponses, ValidateConfigData, ValidateConfigErrors, ValidateConfigResponses } from './types.gen'; export type Options = Options2 & { /** @@ -178,6 +178,15 @@ export const upsertPermissions = (options: } }); +export const getPricing = (options: Options) => (options.client ?? client).post({ + url: '/config/pricing', + ...options, + headers: { + 'Content-Type': 'application/json', + ...options.headers + } +}); + export const providers = (options?: Options) => (options?.client ?? client).get({ url: '/config/providers', ...options }); export const getProviderModels = (options: Options) => (options.client ?? client).get({ url: '/config/providers/{name}/models', ...options }); diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index b19a9b8c..0166872c 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -527,6 +527,25 @@ export type ParseRecipeResponse = { */ export type PermissionLevel = 'always_allow' | 'ask_before' | 'never_allow'; +export type PricingData = { + context_length?: number | null; + currency: string; + input_token_cost: number; + model: string; + output_token_cost: number; + provider: string; +}; + +export type PricingQuery = { + model: string; + provider: string; +}; + +export type PricingResponse = { + pricing: Array; + source: string; +}; + export type PrincipalType = 'Extension' | 'Tool'; export type ProviderDetails = { @@ -1718,6 +1737,22 @@ export type UpsertPermissionsResponses = { export type UpsertPermissionsResponse = UpsertPermissionsResponses[keyof UpsertPermissionsResponses]; +export type GetPricingData = { + body: PricingQuery; + path?: never; + query?: never; + url: '/config/pricing'; +}; + +export type GetPricingResponses = { + /** + * Model pricing data retrieved successfully + */ + 200: PricingResponse; +}; + +export type GetPricingResponse = GetPricingResponses[keyof GetPricingResponses]; + export type ProvidersData = { body?: never; path?: never; diff --git a/ui/desktop/src/components/bottom_menu/CostTracker.tsx b/ui/desktop/src/components/bottom_menu/CostTracker.tsx index 54cf6620..fc930e93 100644 --- a/ui/desktop/src/components/bottom_menu/CostTracker.tsx +++ b/ui/desktop/src/components/bottom_menu/CostTracker.tsx @@ -1,14 +1,9 @@ import { useState, useEffect } from 'react'; import { useModelAndProvider } from '../ModelAndProviderContext'; -import { useConfig } from '../ConfigContext'; import { CoinIcon } from '../icons'; import { Tooltip, TooltipContent, TooltipTrigger } from '../ui/Tooltip'; -import { - getCostForModel, - initializeCostDatabase, - updateAllModelCosts, - fetchAndCachePricing, -} from '../../utils/costDatabase'; +import { fetchModelPricing } from '../../utils/pricing'; +import { PricingData } from '../../api'; interface CostTrackerProps { inputTokens?: number; @@ -24,18 +19,10 @@ interface CostTrackerProps { export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }: CostTrackerProps) { const { currentModel, currentProvider } = useModelAndProvider(); - const { getProviders } = useConfig(); - const [costInfo, setCostInfo] = useState<{ - input_token_cost?: number; - output_token_cost?: number; - currency?: string; - } | null>(null); + const [costInfo, setCostInfo] = useState(null); const [isLoading, setIsLoading] = useState(true); const [showPricing, setShowPricing] = useState(true); const [pricingFailed, setPricingFailed] = useState(false); - const [modelNotFound, setModelNotFound] = useState(false); - const [hasAttemptedFetch, setHasAttemptedFetch] = useState(false); - const [initialLoadComplete, setInitialLoadComplete] = useState(false); // Check if pricing is enabled useEffect(() => { @@ -44,33 +31,11 @@ export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }: setShowPricing(stored !== 'false'); }; - // Check on mount checkPricingSetting(); - - // Listen for storage changes window.addEventListener('storage', checkPricingSetting); return () => window.removeEventListener('storage', checkPricingSetting); }, []); - // Set initial load complete after a short delay - useEffect(() => { - const timer = setTimeout(() => { - setInitialLoadComplete(true); - }, 3000); // Give 3 seconds for initial load - - return () => window.clearTimeout(timer); - }, []); - - // Debug log props removed - - // Initialize cost database on mount - useEffect(() => { - initializeCostDatabase(); - - // Update costs for all models in background - updateAllModelCosts().catch(() => {}); - }, [getProviders]); - useEffect(() => { const loadCostInfo = async () => { if (!currentModel || !currentProvider) { @@ -78,49 +43,20 @@ export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }: return; } + setIsLoading(true); try { - // First check sync cache - let costData = getCostForModel(currentProvider, currentModel); - + const costData = await fetchModelPricing(currentProvider, currentModel); if (costData) { - // We have cached data setCostInfo(costData); setPricingFailed(false); - setModelNotFound(false); - setIsLoading(false); - setHasAttemptedFetch(true); } else { - // Need to fetch from backend - setIsLoading(true); - const result = await fetchAndCachePricing(currentProvider, currentModel); - setHasAttemptedFetch(true); - - if (result && result.costInfo) { - setCostInfo(result.costInfo); - setPricingFailed(false); - setModelNotFound(false); - } else if (result && result.error === 'model_not_found') { - // Model not found in pricing database, but API call succeeded - setModelNotFound(true); - setPricingFailed(false); - } else { - // API call failed or other error - const freeProviders = ['ollama', 'local', 'localhost']; - if (!freeProviders.includes(currentProvider.toLowerCase())) { - setPricingFailed(true); - setModelNotFound(false); - } - } - setIsLoading(false); + setPricingFailed(true); + setCostInfo(null); } } catch { - setHasAttemptedFetch(true); - // Only set pricing failed if we're not dealing with a known free provider - const freeProviders = ['ollama', 'local', 'localhost']; - if (!freeProviders.includes(currentProvider.toLowerCase())) { - setPricingFailed(true); - setModelNotFound(false); - } + setPricingFailed(true); + setCostInfo(null); + } finally { setIsLoading(false); } }; @@ -221,10 +157,9 @@ export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }: // Otherwise show as unavailable const getUnavailableTooltip = () => { - if (pricingFailed && hasAttemptedFetch && initialLoadComplete) { - return `Pricing data unavailable - OpenRouter connection failed. Click refresh in settings to retry.`; + if (pricingFailed) { + return `Pricing data unavailable for ${currentModel}`; } - // If we reach here, it must be modelNotFound (since we only get here after attempting fetch) return `Cost data not available for ${currentModel} (${inputTokens.toLocaleString()} input, ${outputTokens.toLocaleString()} output tokens)`; }; @@ -249,12 +184,8 @@ export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }: // Build tooltip content const getTooltipContent = (): string => { // Handle error states first - if (pricingFailed && hasAttemptedFetch && initialLoadComplete) { - return `Pricing data unavailable - OpenRouter connection failed. Click refresh in settings to retry.`; - } - - if (modelNotFound && hasAttemptedFetch && initialLoadComplete) { - return `Pricing not available for ${currentProvider}/${currentModel}. This model may not be supported by the pricing service.`; + if (pricingFailed) { + return `Pricing data unavailable for ${currentProvider}/${currentModel}`; } // Handle session costs diff --git a/ui/desktop/src/components/settings/app/AppSettingsSection.tsx b/ui/desktop/src/components/settings/app/AppSettingsSection.tsx index b5579907..cf0d04bf 100644 --- a/ui/desktop/src/components/settings/app/AppSettingsSection.tsx +++ b/ui/desktop/src/components/settings/app/AppSettingsSection.tsx @@ -1,13 +1,12 @@ import { useState, useEffect, useRef } from 'react'; import { Switch } from '../../ui/switch'; import { Button } from '../../ui/button'; -import { Settings, RefreshCw, ExternalLink } from 'lucide-react'; +import { Settings } from 'lucide-react'; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from '../../ui/dialog'; import UpdateSection from './UpdateSection'; import TunnelSection from '../tunnel/TunnelSection'; import { COST_TRACKING_ENABLED, UPDATES_ENABLED } from '../../../updates'; -import { getApiUrl } from '../../../config'; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '../../ui/card'; import ThemeSelector from '../../GooseSidebar/ThemeSelector'; import BlockLogoBlack from './icons/block-lockup_black.png'; @@ -26,9 +25,6 @@ export default function AppSettingsSection({ scrollToSection }: AppSettingsSecti const [isMacOS, setIsMacOS] = useState(false); const [isDockSwitchDisabled, setIsDockSwitchDisabled] = useState(false); const [showNotificationModal, setShowNotificationModal] = useState(false); - const [pricingStatus, setPricingStatus] = useState<'loading' | 'success' | 'error'>('loading'); - const [lastFetchTime, setLastFetchTime] = useState(null); - const [isRefreshing, setIsRefreshing] = useState(false); const [showPricing, setShowPricing] = useState(true); const [isDarkMode, setIsDarkMode] = useState(false); const updateSectionRef = useRef(null); @@ -66,71 +62,6 @@ export default function AppSettingsSection({ scrollToSection }: AppSettingsSecti setShowPricing(stored !== 'false'); }, []); - // Check pricing status on mount - useEffect(() => { - checkPricingStatus(); - }, []); - - const checkPricingStatus = async () => { - try { - const apiUrl = getApiUrl('/config/pricing'); - const secretKey = await window.electron.getSecretKey(); - - const headers: HeadersInit = { 'Content-Type': 'application/json' }; - if (secretKey) { - headers['X-Secret-Key'] = secretKey; - } - - const response = await fetch(apiUrl, { - method: 'POST', - headers, - body: JSON.stringify({ configured_only: true }), - }); - - if (response.ok) { - await response.json(); - setPricingStatus('success'); - setLastFetchTime(new Date()); - } else { - setPricingStatus('error'); - } - } catch { - setPricingStatus('error'); - } - }; - - const handleRefreshPricing = async () => { - setIsRefreshing(true); - try { - const apiUrl = getApiUrl('/config/pricing'); - const secretKey = await window.electron.getSecretKey(); - - const headers: HeadersInit = { 'Content-Type': 'application/json' }; - if (secretKey) { - headers['X-Secret-Key'] = secretKey; - } - - const response = await fetch(apiUrl, { - method: 'POST', - headers, - body: JSON.stringify({ configured_only: false }), - }); - - if (response.ok) { - setPricingStatus('success'); - setLastFetchTime(new Date()); - // Trigger a reload of the cost database - window.dispatchEvent(new CustomEvent('pricing-updated')); - } else { - setPricingStatus('error'); - } - } catch { - setPricingStatus('error'); - } finally { - setIsRefreshing(false); - } - }; - // Handle scrolling to update section useEffect(() => { if (scrollToSection === 'update' && updateSectionRef.current) { @@ -326,69 +257,6 @@ export default function AppSettingsSection({ scrollToSection }: AppSettingsSecti )} - {/* Pricing Status - only show if cost tracking is enabled */} - {COST_TRACKING_ENABLED && showPricing && ( - <> -
- Pricing Source: - - OpenRouter Docs - - -
- -
- Status: -
- - {pricingStatus === 'success' - ? '✓ Connected' - : pricingStatus === 'error' - ? '✗ Failed' - : '... Checking'} - - -
-
- - {lastFetchTime && ( -
- Last updated: - {lastFetchTime.toLocaleTimeString()} -
- )} - - {pricingStatus === 'error' && ( -

- Unable to fetch pricing data. Costs will not be displayed. -

- )} - - )} diff --git a/ui/desktop/src/hooks/useAgent.ts b/ui/desktop/src/hooks/useAgent.ts index 37fea45d..ec9608e8 100644 --- a/ui/desktop/src/hooks/useAgent.ts +++ b/ui/desktop/src/hooks/useAgent.ts @@ -2,7 +2,6 @@ import { useCallback, useRef, useState } from 'react'; import { useConfig } from '../components/ConfigContext'; import { ChatType } from '../types/chat'; import { initializeSystem } from '../utils/providerUtils'; -import { initializeCostDatabase } from '../utils/costDatabase'; import { backupConfig, initConfig, @@ -13,7 +12,6 @@ import { startAgent, validateConfig, } from '../api'; -import { COST_TRACKING_ENABLED } from '../updates'; export enum AgentState { UNINITIALIZED = 'uninitialized', @@ -235,14 +233,6 @@ export function useAgent(): UseAgentReturn { recipe: recipeForInit, }); - if (COST_TRACKING_ENABLED) { - try { - await initializeCostDatabase(); - } catch (error) { - console.error('Failed to initialize cost database:', error); - } - } - const recipe = initContext.recipe || agentSession.recipe; const conversation = agentSession.conversation || []; // If we're loading a recipe from initContext (new recipe load), start with empty messages diff --git a/ui/desktop/src/hooks/useCostTracking.ts b/ui/desktop/src/hooks/useCostTracking.ts index 4e8b390b..974c8ee8 100644 --- a/ui/desktop/src/hooks/useCostTracking.ts +++ b/ui/desktop/src/hooks/useCostTracking.ts @@ -1,6 +1,6 @@ import { useEffect, useRef, useState } from 'react'; import { useModelAndProvider } from '../components/ModelAndProviderContext'; -import { getCostForModel } from '../utils/costDatabase'; +import { fetchModelPricing } from '../utils/pricing'; import { Session } from '../api'; interface UseCostTrackingProps { @@ -32,46 +32,53 @@ export const useCostTracking = ({ // Handle model changes and accumulate costs useEffect(() => { - if ( - prevModelRef.current !== undefined && - prevProviderRef.current !== undefined && - (prevModelRef.current !== currentModel || prevProviderRef.current !== currentProvider) - ) { - // Model/provider has changed, save the costs for the previous model - const prevKey = `${prevProviderRef.current}/${prevModelRef.current}`; + const handleModelChange = async () => { + if ( + prevModelRef.current !== undefined && + prevProviderRef.current !== undefined && + (prevModelRef.current !== currentModel || prevProviderRef.current !== currentProvider) + ) { + // Model/provider has changed, save the costs for the previous model + const prevKey = `${prevProviderRef.current}/${prevModelRef.current}`; - // Get pricing info for the previous model - const prevCostInfo = getCostForModel(prevProviderRef.current, prevModelRef.current); + // Get pricing info for the previous model + const prevCostInfo = await fetchModelPricing( + prevProviderRef.current, + prevModelRef.current + ); - if (prevCostInfo) { - const prevInputCost = - (sessionInputTokens || localInputTokens) * (prevCostInfo.input_token_cost || 0); - const prevOutputCost = - (sessionOutputTokens || localOutputTokens) * (prevCostInfo.output_token_cost || 0); - const prevTotalCost = prevInputCost + prevOutputCost; + if (prevCostInfo) { + const prevInputCost = + (sessionInputTokens || localInputTokens) * (prevCostInfo.input_token_cost || 0); + const prevOutputCost = + (sessionOutputTokens || localOutputTokens) * (prevCostInfo.output_token_cost || 0); + const prevTotalCost = prevInputCost + prevOutputCost; - // Save the accumulated costs for this model - setSessionCosts((prev) => ({ - ...prev, - [prevKey]: { - inputTokens: sessionInputTokens || localInputTokens, - outputTokens: sessionOutputTokens || localOutputTokens, - totalCost: prevTotalCost, - }, - })); + // Save the accumulated costs for this model + setSessionCosts((prev) => ({ + ...prev, + [prevKey]: { + inputTokens: sessionInputTokens || localInputTokens, + outputTokens: sessionOutputTokens || localOutputTokens, + totalCost: prevTotalCost, + }, + })); + } + + console.log( + 'Model changed from', + `${prevProviderRef.current}/${prevModelRef.current}`, + 'to', + `${currentProvider}/${currentModel}`, + '- saved costs and restored session token counters' + ); } - console.log( - 'Model changed from', - `${prevProviderRef.current}/${prevModelRef.current}`, - 'to', - `${currentProvider}/${currentModel}`, - '- saved costs and restored session token counters' - ); - } + prevModelRef.current = currentModel || undefined; + prevProviderRef.current = currentProvider || undefined; + }; - prevModelRef.current = currentModel || undefined; - prevProviderRef.current = currentProvider || undefined; + handleModelChange(); }, [ currentModel, currentProvider, diff --git a/ui/desktop/src/utils/costDatabase.ts b/ui/desktop/src/utils/costDatabase.ts deleted file mode 100644 index 5b3bde1d..00000000 --- a/ui/desktop/src/utils/costDatabase.ts +++ /dev/null @@ -1,207 +0,0 @@ -import { getApiUrl } from '../config'; -import { safeJsonParse } from './conversionUtils'; - -export interface ModelCostInfo { - input_token_cost: number; // Cost per token for input (in USD) - output_token_cost: number; // Cost per token for output (in USD) - currency: string; // Currency symbol -} - -// In-memory cache for current session only -const sessionPricingCache = new Map(); - -/** - * Fetch pricing data from backend for specific provider/model - */ -async function fetchPricingForModel( - provider: string, - model: string -): Promise { - // For OpenRouter models, we need to use the parsed provider and model for the API lookup - let lookupProvider = provider; - let lookupModel = model; - - if (provider.toLowerCase() === 'openrouter') { - const parsed = parseOpenRouterModel(model); - if (parsed) { - lookupProvider = parsed[0]; - lookupModel = parsed[1]; - } - } - - const apiUrl = getApiUrl('/config/pricing'); - const secretKey = await window.electron.getSecretKey(); - - const headers: HeadersInit = { 'Content-Type': 'application/json' }; - if (secretKey) { - headers['X-Secret-Key'] = secretKey; - } - - const response = await fetch(apiUrl, { - method: 'POST', - headers, - body: JSON.stringify({ configured_only: false }), - }); - - if (!response.ok) { - throw new Error(`API request failed with status ${response.status}`); - } - - const data = await safeJsonParse<{ - pricing: Array<{ - provider: string; - model: string; - input_token_cost: number; - output_token_cost: number; - currency: string; - }>; - }>(response, 'Failed to parse pricing data'); - - // Find the specific model pricing using the lookup provider/model - const pricing = data.pricing?.find( - (p: { - provider: string; - model: string; - input_token_cost: number; - output_token_cost: number; - currency: string; - }) => { - const providerMatch = p.provider.toLowerCase() === lookupProvider.toLowerCase(); - - // More flexible model matching - handle versioned models - let modelMatch = p.model === lookupModel; - - // If exact match fails, try matching without version suffix - if (!modelMatch && lookupModel.includes('-20')) { - // Remove date suffix like -20241022 - const modelWithoutDate = lookupModel.replace(/-20\d{6}$/, ''); - modelMatch = p.model === modelWithoutDate; - - // Also try with dots instead of dashes (claude-3-5-sonnet vs claude-3.5-sonnet) - if (!modelMatch) { - const modelWithDots = modelWithoutDate.replace(/-(\d)-/g, '.$1.'); - modelMatch = p.model === modelWithDots; - } - } - - return providerMatch && modelMatch; - } - ); - - if (pricing) { - return { - input_token_cost: pricing.input_token_cost, - output_token_cost: pricing.output_token_cost, - currency: pricing.currency || '$', - }; - } - - // API call succeeded but model not found in pricing data - return null; -} - -/** - * Initialize the cost database - no-op since we fetch on demand now - */ -export async function initializeCostDatabase(): Promise { - // Clear session cache on init - sessionPricingCache.clear(); -} - -/** - * Update model costs from providers - no-op since we fetch on demand - */ -export async function updateAllModelCosts(): Promise { - // No-op - we fetch on demand now -} - -/** - * Parse OpenRouter model ID to extract provider and model - * e.g., "anthropic/claude-sonnet-4" -> ["anthropic", "claude-sonnet-4"] - */ -function parseOpenRouterModel(modelId: string): [string, string] | null { - const parts = modelId.split('/'); - if (parts.length === 2) { - return [parts[0], parts[1]]; - } - return null; -} - -/** - * Get cost information for a specific model with session caching - */ -export function getCostForModel(provider: string, model: string): ModelCostInfo | null { - const cacheKey = `${provider}/${model}`; - - // Check session cache first - if (sessionPricingCache.has(cacheKey)) { - return sessionPricingCache.get(cacheKey) || null; - } - - // For OpenRouter models, also check if we have cached data under the parsed provider/model - if (provider.toLowerCase() === 'openrouter') { - const parsed = parseOpenRouterModel(model); - if (parsed) { - const [parsedProvider, parsedModel] = parsed; - const parsedCacheKey = `${parsedProvider}/${parsedModel}`; - if (sessionPricingCache.has(parsedCacheKey)) { - const cachedData = sessionPricingCache.get(parsedCacheKey) || null; - // Also cache it under the original OpenRouter key for future lookups - sessionPricingCache.set(cacheKey, cachedData); - return cachedData; - } - } - } - - // For local/free providers, return zero cost immediately - const freeProviders = ['ollama', 'local', 'localhost']; - if (freeProviders.includes(provider.toLowerCase())) { - const zeroCost = { - input_token_cost: 0, - output_token_cost: 0, - currency: '$', - }; - sessionPricingCache.set(cacheKey, zeroCost); - return zeroCost; - } - - // Need to fetch - return null and let component handle async fetch - return null; -} - -/** - * Fetch and cache pricing for a model - */ -export async function fetchAndCachePricing( - provider: string, - model: string -): Promise<{ costInfo: ModelCostInfo | null; error?: string } | null> { - try { - const cacheKey = `${provider}/${model}`; - const costInfo = await fetchPricingForModel(provider, model); - - // Cache the result in session cache under the original key - sessionPricingCache.set(cacheKey, costInfo); - - // For OpenRouter models, also cache under the parsed provider/model key - // This helps with cross-referencing between frontend requests and backend responses - if (provider.toLowerCase() === 'openrouter') { - const parsed = parseOpenRouterModel(model); - if (parsed) { - const [parsedProvider, parsedModel] = parsed; - const parsedCacheKey = `${parsedProvider}/${parsedModel}`; - sessionPricingCache.set(parsedCacheKey, costInfo); - } - } - - if (costInfo) { - return { costInfo }; - } else { - // Model not found in pricing data - return { costInfo: null, error: 'model_not_found' }; - } - } catch { - // This is a real API/network error - return null; - } -} diff --git a/ui/desktop/src/utils/pricing.ts b/ui/desktop/src/utils/pricing.ts new file mode 100644 index 00000000..509b9240 --- /dev/null +++ b/ui/desktop/src/utils/pricing.ts @@ -0,0 +1,24 @@ +import { getPricing, PricingData } from '../api'; + +/** + * Fetch pricing for a specific provider/model from the backend + */ +export async function fetchModelPricing( + provider: string, + model: string +): Promise { + try { + const response = await getPricing({ + body: { provider, model }, + throwOnError: false, + }); + + if (!response.data) { + return null; + } + + return response.data.pricing?.[0] ?? null; + } catch { + return null; + } +}