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:
@@ -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))
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user