tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Integrate pricing with canonical model (#6130)

Browse Source
This commit is contained in:
David Katz
2025-12-18 11:44:36 -05:00
committed by GitHub
parent 473f269daa
commit 6a0b8c25d5
17 changed files with 263 additions and 1034 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/mod.rs
+1 -16
View File
@@ -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(_) => {
crates/goose-cli/src/session/output.rs
+10 -56
View File
@@ -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<f64> {
// 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 {})",
crates/goose-server/src/commands/agent.rs
-9
View File
@@ -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());
crates/goose-server/src/openapi.rs
+4
View File
@@ -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,
crates/goose-server/src/routes/config_management.rs
+18 -75
View File
@@ -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<PricingQuery>,
) -> Result<Json<PricingResponse>, 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(),
}))
}
crates/goose/src/providers/canonical/mod.rs
+6
View File
@@ -20,3 +20,9 @@ impl ModelMapping {
}
}
}
pub fn maybe_get_canonical_model(provider: &str, model: &str) -> Option<CanonicalModel> {
let registry = CanonicalModelRegistry::bundled().ok()?;
let canonical_id = map_to_canonical_model(provider, model, registry)?;
registry.get(&canonical_id).cloned()
}
crates/goose/src/providers/mod.rs
-1
View File
@@ -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;
crates/goose/src/providers/pricing.rs
-408
View File
@@ -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<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 {
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<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(),
};
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> {
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<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 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<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-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<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-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");
}
}
}