Skip to content

Commit a408d54

Browse files
committed
fix(stats): support provider-prefixed pricing env keys
1 parent 7954d02 commit a408d54

1 file changed

Lines changed: 32 additions & 1 deletion

File tree

src/cortex-cli/src/stats_cmd.rs

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,9 +118,11 @@ fn load_custom_pricing() -> std::collections::HashMap<String, ModelPricing> {
118118
// Try to load from environment variables in format:
119119
// CORTEX_PRICING_<MODEL>=<input_price>,<output_price>
120120
// Example: CORTEX_PRICING_GPT4O=2.5,10.0
121+
// Use double underscore to encode provider-prefixed model IDs:
122+
// CORTEX_PRICING_ANTHROPIC__CLAUDE_3_5_SONNET=3.0,15.0
121123
for (key, value) in std::env::vars() {
122124
if let Some(model_suffix) = key.strip_prefix("CORTEX_PRICING_") {
123-
let model_name = model_suffix.to_lowercase().replace('_', "-");
125+
let model_name = pricing_env_model_name(model_suffix);
124126
let parts: Vec<&str> = value.split(',').collect();
125127
if parts.len() == 2
126128
&& let (Ok(input), Ok(output)) = (
@@ -142,6 +144,14 @@ fn load_custom_pricing() -> std::collections::HashMap<String, ModelPricing> {
142144
custom
143145
}
144146

147+
fn pricing_env_model_name(model_suffix: &str) -> String {
148+
model_suffix
149+
.to_lowercase()
150+
.replace("__", "\0")
151+
.replace('_', "-")
152+
.replace('\0', "/")
153+
}
154+
145155
impl StatsCli {
146156
/// Run the stats command.
147157
pub async fn run(self) -> Result<()> {
@@ -735,6 +745,27 @@ mod tests {
735745
assert!((cost - 12.5).abs() < 0.001);
736746
}
737747

748+
#[test]
749+
fn test_pricing_env_model_name_supports_provider_prefixed_models() {
750+
assert_eq!(
751+
pricing_env_model_name("ANTHROPIC__CLAUDE_3_5_SONNET"),
752+
"anthropic/claude-3-5-sonnet"
753+
);
754+
assert_eq!(
755+
pricing_env_model_name("OPENAI__GPT_4O_MINI"),
756+
"openai/gpt-4o-mini"
757+
);
758+
}
759+
760+
#[test]
761+
fn test_pricing_env_model_name_preserves_legacy_underscore_mapping() {
762+
assert_eq!(pricing_env_model_name("GPT_4O_MINI"), "gpt-4o-mini");
763+
assert_eq!(
764+
pricing_env_model_name("CLAUDE_3_5_HAIKU"),
765+
"claude-3-5-haiku"
766+
);
767+
}
768+
738769
#[test]
739770
fn test_validate_days_range() {
740771
// Valid values

0 commit comments

Comments
 (0)