fix(cost): resolve databricks_v2 pricing and surface cost in standard usage update (#9925)
Co-authored-by: Douwe M Osinga <douwe@sidewalklabs.com>
This commit is contained in:
@@ -35,7 +35,10 @@ pub fn canonical_name(provider: &str, model: &str) -> String {
|
||||
}
|
||||
|
||||
fn is_meta_provider(provider: &str) -> bool {
|
||||
matches!(provider, "databricks" | "tetrate" | "bedrock" | "azure")
|
||||
matches!(
|
||||
provider,
|
||||
"databricks" | "databricks_v2" | "tetrate" | "bedrock" | "azure"
|
||||
)
|
||||
}
|
||||
|
||||
pub fn map_provider_name(provider: &str) -> &str {
|
||||
@@ -46,6 +49,7 @@ pub fn map_provider_name(provider: &str) -> &str {
|
||||
"aws_bedrock" => "amazon-bedrock",
|
||||
"gcp_vertex_ai" => "google-vertex",
|
||||
"gemini_oauth" => "google",
|
||||
"databricks_v2" => "databricks",
|
||||
"zhipu" => "zhipuai",
|
||||
"novita" => "novita-ai",
|
||||
"opencode_go" => "opencode-go",
|
||||
@@ -121,6 +125,21 @@ pub fn map_to_canonical_model(
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback for meta-providers: some native aliases are keyed under the
|
||||
// meta-provider itself (e.g. "databricks/databricks-gpt-oss-120b") and do
|
||||
// not infer back to a first-party provider. Only try this after inference,
|
||||
// so models that DO infer (e.g. databricks-claude-* -> anthropic/*) keep
|
||||
// resolving to the richer first-party catalog entry.
|
||||
if is_meta_provider(provider) {
|
||||
if let Some(canonical) = registry.get(registry_provider, model) {
|
||||
return Some(canonical.id.clone());
|
||||
}
|
||||
let normalized_model = strip_version_suffix(model);
|
||||
if let Some(canonical) = registry.get(registry_provider, &normalized_model) {
|
||||
return Some(canonical.id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
@@ -537,4 +556,36 @@ mod tests {
|
||||
Some("google-vertex/claude-haiku-4.5".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
// Databricks-native open-weight ids are keyed under the meta-provider itself
|
||||
// (e.g. "databricks/databricks-gpt-oss-120b") and do not infer back to another
|
||||
// provider, so they must resolve via the direct meta-provider lookup. These
|
||||
// particular ids are unversioned, so the assertions are not catalog-version brittle.
|
||||
#[test]
|
||||
fn test_databricks_native_open_weight_ids_resolve() {
|
||||
let r = super::super::CanonicalModelRegistry::bundled().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
map_to_canonical_model("databricks_v2", "databricks-gpt-oss-120b", r),
|
||||
Some("databricks/databricks-gpt-oss-120b".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
map_to_canonical_model("databricks_v2", "databricks-gpt-oss-20b", r),
|
||||
Some("databricks/databricks-gpt-oss-20b".to_string())
|
||||
);
|
||||
// Legacy provider name resolves identically.
|
||||
assert_eq!(
|
||||
map_to_canonical_model("databricks", "databricks-gpt-oss-120b", r),
|
||||
Some("databricks/databricks-gpt-oss-120b".to_string())
|
||||
);
|
||||
|
||||
// Regression guard: the meta-provider lookup must remain a *fallback*
|
||||
// after inference. databricks-claude-* aliases infer back to the richer
|
||||
// first-party "anthropic/*" entry (which carries thinking_mode used for
|
||||
// adaptive thinking), not the metadata-poor "databricks/databricks-*" one.
|
||||
assert_eq!(
|
||||
map_to_canonical_model("databricks", "databricks-claude-opus-4-7", r),
|
||||
Some("anthropic/claude-opus-4.7".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -42,7 +42,7 @@ use crate::utils::sanitize_unicode_tags;
|
||||
use agent_client_protocol::schema::{
|
||||
AgentCapabilities, Annotations, AuthMethod, AuthMethodAgent, AuthenticateRequest,
|
||||
AuthenticateResponse, BlobResourceContents, CancelNotification, CloseSessionRequest,
|
||||
CloseSessionResponse, ConfigOptionUpdate, Content, ContentBlock, ContentChunk,
|
||||
CloseSessionResponse, ConfigOptionUpdate, Content, ContentBlock, ContentChunk, Cost,
|
||||
CurrentModeUpdate, EmbeddedResource, EmbeddedResourceResource, FileSystemCapabilities,
|
||||
ForkSessionRequest, ForkSessionResponse, ImageContent, Implementation, InitializeRequest,
|
||||
InitializeResponse, ListSessionsRequest, ListSessionsResponse, LoadSessionRequest,
|
||||
@@ -834,7 +834,13 @@ pub(super) fn build_usage_updates(session: &Session) -> Option<UsageUpdates> {
|
||||
accumulated_cost: session.accumulated_cost,
|
||||
}),
|
||||
},
|
||||
standard: UsageUpdate::new(used, ctx_limit),
|
||||
standard: {
|
||||
let mut standard = UsageUpdate::new(used, ctx_limit);
|
||||
if let Some(amount) = session.accumulated_cost {
|
||||
standard = standard.cost(Cost::new(amount, "USD"));
|
||||
}
|
||||
standard
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user