mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat(billing): 引入 model_id 精确计费查找路径,传播模型 ID 至用量事件与 report context
- UsageEventData 新增 model_id / global_model_id 字段,write.rs seed 结构体同步补充 - report_context 新增 model_id / global_model_id / global_model_name,各 payload 构建时从 candidate 传入 - event_enrichment 优先按 model_id 精确查找计费上下文,回退为按名称多轮查找(保留 NoRule 结果降级逻辑) - BillingReadRepository 新增 find_model_context_by_model_id,memory/sql 分别实现;SQL 查询重构支持按 provider_model_name 和 mappings 匹配并按优先级排序 - pricing.rs 修复:model_tiered_pricing 为空 tiers 时回退到 default_tiered_pricing - request_metadata 允许字段列表补充 model_id / global_model_id / global_model_name - admin usage 路由信息输出及脱敏字段列表同步新增三个 model 相关字段
This commit is contained in:
@@ -91,6 +91,9 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
|||||||
endpoint_id: &candidate.endpoint_id,
|
endpoint_id: &candidate.endpoint_id,
|
||||||
key_id: &candidate.key_id,
|
key_id: &candidate.key_id,
|
||||||
key_name: Some(&candidate.key_name),
|
key_name: Some(&candidate.key_name),
|
||||||
|
model_id: Some(&candidate.model_id),
|
||||||
|
global_model_id: Some(&candidate.global_model_id),
|
||||||
|
global_model_name: Some(&candidate.global_model_name),
|
||||||
provider_api_format: spec_metadata.api_format,
|
provider_api_format: spec_metadata.api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
|||||||
@@ -16,6 +16,9 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
|
|||||||
pub(crate) endpoint_id: &'a str,
|
pub(crate) endpoint_id: &'a str,
|
||||||
pub(crate) key_id: &'a str,
|
pub(crate) key_id: &'a str,
|
||||||
pub(crate) key_name: Option<&'a str>,
|
pub(crate) key_name: Option<&'a str>,
|
||||||
|
pub(crate) model_id: Option<&'a str>,
|
||||||
|
pub(crate) global_model_id: Option<&'a str>,
|
||||||
|
pub(crate) global_model_name: Option<&'a str>,
|
||||||
pub(crate) provider_api_format: &'a str,
|
pub(crate) provider_api_format: &'a str,
|
||||||
pub(crate) client_api_format: &'a str,
|
pub(crate) client_api_format: &'a str,
|
||||||
pub(crate) mapped_model: Option<&'a str>,
|
pub(crate) mapped_model: Option<&'a str>,
|
||||||
@@ -136,6 +139,21 @@ pub(crate) fn build_local_execution_report_context(
|
|||||||
if let Some(key_name) = parts.key_name {
|
if let Some(key_name) = parts.key_name {
|
||||||
object.insert("key_name".to_string(), Value::String(key_name.to_string()));
|
object.insert("key_name".to_string(), Value::String(key_name.to_string()));
|
||||||
}
|
}
|
||||||
|
if let Some(model_id) = parts.model_id {
|
||||||
|
object.insert("model_id".to_string(), Value::String(model_id.to_string()));
|
||||||
|
}
|
||||||
|
if let Some(global_model_id) = parts.global_model_id {
|
||||||
|
object.insert(
|
||||||
|
"global_model_id".to_string(),
|
||||||
|
Value::String(global_model_id.to_string()),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if let Some(global_model_name) = parts.global_model_name {
|
||||||
|
object.insert(
|
||||||
|
"global_model_name".to_string(),
|
||||||
|
Value::String(global_model_name.to_string()),
|
||||||
|
);
|
||||||
|
}
|
||||||
if let Some(mapped_model) = parts.mapped_model {
|
if let Some(mapped_model) = parts.mapped_model {
|
||||||
object.insert(
|
object.insert(
|
||||||
"mapped_model".to_string(),
|
"mapped_model".to_string(),
|
||||||
|
|||||||
@@ -76,6 +76,9 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
|||||||
endpoint_id: &candidate.endpoint_id,
|
endpoint_id: &candidate.endpoint_id,
|
||||||
key_id: &candidate.key_id,
|
key_id: &candidate.key_id,
|
||||||
key_name: None,
|
key_name: None,
|
||||||
|
model_id: Some(&candidate.model_id),
|
||||||
|
global_model_id: Some(&candidate.global_model_id),
|
||||||
|
global_model_name: Some(&candidate.global_model_name),
|
||||||
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
||||||
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
||||||
mapped_model: None,
|
mapped_model: None,
|
||||||
|
|||||||
@@ -68,6 +68,9 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
|||||||
endpoint_id: &candidate.endpoint_id,
|
endpoint_id: &candidate.endpoint_id,
|
||||||
key_id: &candidate.key_id,
|
key_id: &candidate.key_id,
|
||||||
key_name: None,
|
key_name: None,
|
||||||
|
model_id: Some(&candidate.model_id),
|
||||||
|
global_model_id: Some(&candidate.global_model_id),
|
||||||
|
global_model_name: Some(&candidate.global_model_name),
|
||||||
provider_api_format: spec_metadata.api_format,
|
provider_api_format: spec_metadata.api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
|||||||
@@ -59,6 +59,9 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
|||||||
endpoint_id: &candidate.endpoint_id,
|
endpoint_id: &candidate.endpoint_id,
|
||||||
key_id: &candidate.key_id,
|
key_id: &candidate.key_id,
|
||||||
key_name: None,
|
key_name: None,
|
||||||
|
model_id: Some(&candidate.model_id),
|
||||||
|
global_model_id: Some(&candidate.global_model_id),
|
||||||
|
global_model_name: Some(&candidate.global_model_name),
|
||||||
provider_api_format: spec_metadata.api_format,
|
provider_api_format: spec_metadata.api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
|||||||
@@ -65,6 +65,9 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
|||||||
endpoint_id: &candidate.endpoint_id,
|
endpoint_id: &candidate.endpoint_id,
|
||||||
key_id: &candidate.key_id,
|
key_id: &candidate.key_id,
|
||||||
key_name: Some(&candidate.key_name),
|
key_name: Some(&candidate.key_name),
|
||||||
|
model_id: Some(&candidate.model_id),
|
||||||
|
global_model_id: Some(&candidate.global_model_id),
|
||||||
|
global_model_name: Some(&candidate.global_model_name),
|
||||||
provider_api_format: &resolved.provider_api_format,
|
provider_api_format: &resolved.provider_api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
|||||||
@@ -82,6 +82,9 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
|||||||
endpoint_id: &candidate.endpoint_id,
|
endpoint_id: &candidate.endpoint_id,
|
||||||
key_id: &candidate.key_id,
|
key_id: &candidate.key_id,
|
||||||
key_name: Some(&candidate.key_name),
|
key_name: Some(&candidate.key_name),
|
||||||
|
model_id: Some(&candidate.model_id),
|
||||||
|
global_model_id: Some(&candidate.global_model_id),
|
||||||
|
global_model_name: Some(&candidate.global_model_name),
|
||||||
provider_api_format: &resolved.provider_api_format,
|
provider_api_format: &resolved.provider_api_format,
|
||||||
client_api_format: "openai:chat",
|
client_api_format: "openai:chat",
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
|||||||
@@ -86,6 +86,9 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
|||||||
endpoint_id: &candidate.endpoint_id,
|
endpoint_id: &candidate.endpoint_id,
|
||||||
key_id: &candidate.key_id,
|
key_id: &candidate.key_id,
|
||||||
key_name: Some(&candidate.key_name),
|
key_name: Some(&candidate.key_name),
|
||||||
|
model_id: Some(&candidate.model_id),
|
||||||
|
global_model_id: Some(&candidate.global_model_id),
|
||||||
|
global_model_name: Some(&candidate.global_model_name),
|
||||||
provider_api_format: &resolved.provider_api_format,
|
provider_api_format: &resolved.provider_api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
|||||||
@@ -157,6 +157,21 @@ impl MinimalCandidateSelectionRowSource for GatewayDataState {
|
|||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl BillingModelContextLookup for GatewayDataState {
|
impl BillingModelContextLookup for GatewayDataState {
|
||||||
|
async fn find_billing_model_context_by_model_id(
|
||||||
|
&self,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
|
GatewayDataState::find_billing_model_context_by_model_id(
|
||||||
|
self,
|
||||||
|
provider_id,
|
||||||
|
provider_api_key_id,
|
||||||
|
model_id,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
async fn find_billing_model_context(
|
async fn find_billing_model_context(
|
||||||
&self,
|
&self,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
|
|||||||
@@ -1241,6 +1241,22 @@ impl GatewayDataState {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn find_billing_model_context_by_model_id(
|
||||||
|
&self,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
|
match &self.billing_reader {
|
||||||
|
Some(repository) => {
|
||||||
|
repository
|
||||||
|
.find_model_context_by_model_id(provider_id, provider_api_key_id, model_id)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
None => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn read_request_candidate_trace(
|
pub(crate) async fn read_request_candidate_trace(
|
||||||
&self,
|
&self,
|
||||||
request_id: &str,
|
request_id: &str,
|
||||||
|
|||||||
@@ -285,6 +285,9 @@ fn admin_usage_strip_routing_metadata(metadata: &mut serde_json::Map<String, Val
|
|||||||
metadata.remove("candidate_id");
|
metadata.remove("candidate_id");
|
||||||
metadata.remove("candidate_index");
|
metadata.remove("candidate_index");
|
||||||
metadata.remove("key_name");
|
metadata.remove("key_name");
|
||||||
|
metadata.remove("model_id");
|
||||||
|
metadata.remove("global_model_id");
|
||||||
|
metadata.remove("global_model_name");
|
||||||
metadata.remove("planner_kind");
|
metadata.remove("planner_kind");
|
||||||
metadata.remove("route_family");
|
metadata.remove("route_family");
|
||||||
metadata.remove("route_kind");
|
metadata.remove("route_kind");
|
||||||
@@ -795,6 +798,17 @@ fn admin_usage_routing_json(
|
|||||||
"key_name",
|
"key_name",
|
||||||
provider_key_name.or_else(|| item.routing_key_name()),
|
provider_key_name.or_else(|| item.routing_key_name()),
|
||||||
);
|
);
|
||||||
|
maybe_insert_string_field(&mut routing, "model_id", item.routing_model_id());
|
||||||
|
maybe_insert_string_field(
|
||||||
|
&mut routing,
|
||||||
|
"global_model_id",
|
||||||
|
item.routing_global_model_id(),
|
||||||
|
);
|
||||||
|
maybe_insert_string_field(
|
||||||
|
&mut routing,
|
||||||
|
"global_model_name",
|
||||||
|
item.routing_global_model_name(),
|
||||||
|
);
|
||||||
maybe_insert_string_field(&mut routing, "planner_kind", item.routing_planner_kind());
|
maybe_insert_string_field(&mut routing, "planner_kind", item.routing_planner_kind());
|
||||||
maybe_insert_string_field(&mut routing, "route_family", item.routing_route_family());
|
maybe_insert_string_field(&mut routing, "route_family", item.routing_route_family());
|
||||||
maybe_insert_string_field(&mut routing, "route_kind", item.routing_route_kind());
|
maybe_insert_string_field(&mut routing, "route_kind", item.routing_route_kind());
|
||||||
@@ -2395,7 +2409,10 @@ mod tests {
|
|||||||
fn detail_payload_exposes_typed_routing_section() {
|
fn detail_payload_exposes_typed_routing_section() {
|
||||||
let item = StoredRequestUsageAudit {
|
let item = StoredRequestUsageAudit {
|
||||||
request_metadata: Some(json!({
|
request_metadata: Some(json!({
|
||||||
"trace_id": "trace-routing-detail"
|
"trace_id": "trace-routing-detail",
|
||||||
|
"model_id": "model-1",
|
||||||
|
"global_model_id": "global-model-1",
|
||||||
|
"global_model_name": "gpt-5"
|
||||||
})),
|
})),
|
||||||
candidate_id: Some("cand-1".to_string()),
|
candidate_id: Some("cand-1".to_string()),
|
||||||
candidate_index: Some(2),
|
candidate_index: Some(2),
|
||||||
@@ -2423,6 +2440,9 @@ mod tests {
|
|||||||
assert_eq!(payload["routing"]["candidate_id"], "cand-1");
|
assert_eq!(payload["routing"]["candidate_id"], "cand-1");
|
||||||
assert_eq!(payload["routing"]["candidate_index"], 2);
|
assert_eq!(payload["routing"]["candidate_index"], 2);
|
||||||
assert_eq!(payload["routing"]["key_name"], "resolved-primary");
|
assert_eq!(payload["routing"]["key_name"], "resolved-primary");
|
||||||
|
assert_eq!(payload["routing"]["model_id"], "model-1");
|
||||||
|
assert_eq!(payload["routing"]["global_model_id"], "global-model-1");
|
||||||
|
assert_eq!(payload["routing"]["global_model_name"], "gpt-5");
|
||||||
assert_eq!(payload["routing"]["planner_kind"], "claude_cli_sync");
|
assert_eq!(payload["routing"]["planner_kind"], "claude_cli_sync");
|
||||||
assert_eq!(payload["routing"]["route_family"], "claude");
|
assert_eq!(payload["routing"]["route_family"], "claude");
|
||||||
assert_eq!(payload["routing"]["route_kind"], "cli");
|
assert_eq!(payload["routing"]["route_kind"], "cli");
|
||||||
@@ -2437,6 +2457,9 @@ mod tests {
|
|||||||
assert!(payload["metadata"]["candidate_id"].is_null());
|
assert!(payload["metadata"]["candidate_id"].is_null());
|
||||||
assert!(payload["metadata"]["candidate_index"].is_null());
|
assert!(payload["metadata"]["candidate_index"].is_null());
|
||||||
assert!(payload["metadata"]["key_name"].is_null());
|
assert!(payload["metadata"]["key_name"].is_null());
|
||||||
|
assert!(payload["metadata"]["model_id"].is_null());
|
||||||
|
assert!(payload["metadata"]["global_model_id"].is_null());
|
||||||
|
assert!(payload["metadata"]["global_model_name"].is_null());
|
||||||
assert!(payload["metadata"]["planner_kind"].is_null());
|
assert!(payload["metadata"]["planner_kind"].is_null());
|
||||||
assert!(payload["metadata"]["trace_id"].is_null());
|
assert!(payload["metadata"]["trace_id"].is_null());
|
||||||
assert!(payload["metadata"]["route_family"].is_null());
|
assert!(payload["metadata"]["route_family"].is_null());
|
||||||
|
|||||||
@@ -4,10 +4,23 @@ use aether_usage_runtime::{UsageEvent, UsageEventType};
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use serde_json::{Map, Value};
|
use serde_json::{Map, Value};
|
||||||
|
|
||||||
use crate::{BillingModelPricingSnapshot, BillingService, BillingUsageInput};
|
use crate::{
|
||||||
|
BillingComputation, BillingModelPricingSnapshot, BillingService, BillingSnapshotStatus,
|
||||||
|
BillingUsageInput,
|
||||||
|
};
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
pub trait BillingModelContextLookup: Send + Sync {
|
pub trait BillingModelContextLookup: Send + Sync {
|
||||||
|
async fn find_billing_model_context_by_model_id(
|
||||||
|
&self,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
|
let _ = (provider_id, provider_api_key_id, model_id);
|
||||||
|
Ok(None)
|
||||||
|
}
|
||||||
|
|
||||||
async fn find_billing_model_context(
|
async fn find_billing_model_context(
|
||||||
&self,
|
&self,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
@@ -35,23 +48,78 @@ pub async fn enrich_usage_event_with_billing(
|
|||||||
else {
|
else {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
};
|
};
|
||||||
let model_name = event.data.model.trim();
|
if let Some(model_id) = event
|
||||||
if model_name.is_empty() {
|
.data
|
||||||
|
.model_id
|
||||||
|
.as_deref()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
{
|
||||||
|
if let Some(context) = data
|
||||||
|
.find_billing_model_context_by_model_id(
|
||||||
|
provider_id,
|
||||||
|
event.data.provider_api_key_id.as_deref(),
|
||||||
|
model_id,
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
let pricing = map_pricing_context(context);
|
||||||
|
let computation = calculate_billing_computation(&pricing, event)?;
|
||||||
|
apply_billing_computation(event, computation)?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut first_no_rule = None;
|
||||||
|
for lookup_name in billing_model_lookup_names(&event.data) {
|
||||||
|
let Some(context) = data
|
||||||
|
.find_billing_model_context(
|
||||||
|
provider_id,
|
||||||
|
event.data.provider_api_key_id.as_deref(),
|
||||||
|
lookup_name,
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
|
let pricing = map_pricing_context(context);
|
||||||
|
let computation = calculate_billing_computation(&pricing, event)?;
|
||||||
|
if matches!(
|
||||||
|
computation.cost_result.status,
|
||||||
|
BillingSnapshotStatus::NoRule
|
||||||
|
) {
|
||||||
|
first_no_rule.get_or_insert(computation);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
apply_billing_computation(event, computation)?;
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
let Some(context) = data
|
if let Some(computation) = first_no_rule {
|
||||||
.find_billing_model_context(
|
apply_billing_computation(event, computation)?;
|
||||||
provider_id,
|
}
|
||||||
event.data.provider_api_key_id.as_deref(),
|
Ok(())
|
||||||
model_name,
|
}
|
||||||
)
|
|
||||||
.await?
|
|
||||||
else {
|
|
||||||
return Ok(());
|
|
||||||
};
|
|
||||||
|
|
||||||
let pricing = map_pricing_context(context);
|
fn billing_model_lookup_names(data: &aether_usage_runtime::UsageEventData) -> Vec<&str> {
|
||||||
|
let mut names = Vec::new();
|
||||||
|
for value in [data.target_model.as_deref(), Some(data.model.as_str())]
|
||||||
|
.into_iter()
|
||||||
|
.flatten()
|
||||||
|
{
|
||||||
|
let value = value.trim();
|
||||||
|
if !value.is_empty() && !names.contains(&value) {
|
||||||
|
names.push(value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
names
|
||||||
|
}
|
||||||
|
|
||||||
|
fn calculate_billing_computation(
|
||||||
|
pricing: &BillingModelPricingSnapshot,
|
||||||
|
event: &UsageEvent,
|
||||||
|
) -> Result<BillingComputation, DataLayerError> {
|
||||||
let input = BillingUsageInput {
|
let input = BillingUsageInput {
|
||||||
task_type: event
|
task_type: event
|
||||||
.data
|
.data
|
||||||
@@ -85,11 +153,17 @@ pub async fn enrich_usage_event_with_billing(
|
|||||||
cache_ttl_minutes: pricing.provider_api_key_cache_ttl_minutes,
|
cache_ttl_minutes: pricing.provider_api_key_cache_ttl_minutes,
|
||||||
};
|
};
|
||||||
|
|
||||||
let computation = BillingService::new()
|
BillingService::new()
|
||||||
.calculate(&pricing, &input)
|
.calculate(pricing, &input)
|
||||||
.map_err(|err| {
|
.map_err(|err| {
|
||||||
DataLayerError::UnexpectedValue(format!("billing calculation failed: {err}"))
|
DataLayerError::UnexpectedValue(format!("billing calculation failed: {err}"))
|
||||||
})?;
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_billing_computation(
|
||||||
|
event: &mut UsageEvent,
|
||||||
|
computation: BillingComputation,
|
||||||
|
) -> Result<(), DataLayerError> {
|
||||||
event.data.total_cost_usd = Some(computation.cost_result.cost);
|
event.data.total_cost_usd = Some(computation.cost_result.cost);
|
||||||
event.data.actual_total_cost_usd = Some(computation.actual_total_cost);
|
event.data.actual_total_cost_usd = Some(computation.actual_total_cost);
|
||||||
merge_billing_snapshot_metadata(
|
merge_billing_snapshot_metadata(
|
||||||
@@ -97,8 +171,7 @@ pub async fn enrich_usage_event_with_billing(
|
|||||||
&computation.cost_result.snapshot,
|
&computation.cost_result.snapshot,
|
||||||
computation.rate_multiplier,
|
computation.rate_multiplier,
|
||||||
computation.is_free_tier,
|
computation.is_free_tier,
|
||||||
)?;
|
)
|
||||||
Ok(())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn map_pricing_context(context: StoredBillingModelContext) -> BillingModelPricingSnapshot {
|
fn map_pricing_context(context: StoredBillingModelContext) -> BillingModelPricingSnapshot {
|
||||||
@@ -153,11 +226,22 @@ mod tests {
|
|||||||
use super::{enrich_usage_event_with_billing, BillingModelContextLookup};
|
use super::{enrich_usage_event_with_billing, BillingModelContextLookup};
|
||||||
|
|
||||||
struct TestLookup {
|
struct TestLookup {
|
||||||
context: Option<StoredBillingModelContext>,
|
name_context: Option<StoredBillingModelContext>,
|
||||||
|
model_id_context: Option<StoredBillingModelContext>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl BillingModelContextLookup for TestLookup {
|
impl BillingModelContextLookup for TestLookup {
|
||||||
|
async fn find_billing_model_context_by_model_id(
|
||||||
|
&self,
|
||||||
|
_provider_id: &str,
|
||||||
|
_provider_api_key_id: Option<&str>,
|
||||||
|
_model_id: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
|
||||||
|
{
|
||||||
|
Ok(self.model_id_context.clone())
|
||||||
|
}
|
||||||
|
|
||||||
async fn find_billing_model_context(
|
async fn find_billing_model_context(
|
||||||
&self,
|
&self,
|
||||||
_provider_id: &str,
|
_provider_id: &str,
|
||||||
@@ -165,14 +249,14 @@ mod tests {
|
|||||||
_global_model_name: &str,
|
_global_model_name: &str,
|
||||||
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
|
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
|
||||||
{
|
{
|
||||||
Ok(self.context.clone())
|
Ok(self.name_context.clone())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn enriches_completed_usage_event_with_billing_snapshot() {
|
async fn enriches_completed_usage_event_with_billing_snapshot() {
|
||||||
let lookup = TestLookup {
|
let lookup = TestLookup {
|
||||||
context: Some(
|
name_context: Some(
|
||||||
StoredBillingModelContext::new(
|
StoredBillingModelContext::new(
|
||||||
"provider-1".to_string(),
|
"provider-1".to_string(),
|
||||||
Some("pay_as_you_go".to_string()),
|
Some("pay_as_you_go".to_string()),
|
||||||
@@ -192,6 +276,7 @@ mod tests {
|
|||||||
)
|
)
|
||||||
.expect("billing context should build"),
|
.expect("billing context should build"),
|
||||||
),
|
),
|
||||||
|
model_id_context: None,
|
||||||
};
|
};
|
||||||
let mut event = UsageEvent::new(
|
let mut event = UsageEvent::new(
|
||||||
UsageEventType::Completed,
|
UsageEventType::Completed,
|
||||||
@@ -229,4 +314,82 @@ mod tests {
|
|||||||
Some("complete")
|
Some("complete")
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn enriches_by_provider_model_id_before_name_fallback() {
|
||||||
|
let blank_name_context = StoredBillingModelContext::new(
|
||||||
|
"provider-1".to_string(),
|
||||||
|
Some("pay_as_you_go".to_string()),
|
||||||
|
Some("key-1".to_string()),
|
||||||
|
None,
|
||||||
|
Some(60),
|
||||||
|
"global-model-blank".to_string(),
|
||||||
|
"claude-sonnet-4-6".to_string(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some("model-blank".to_string()),
|
||||||
|
Some("claude-sonnet-4-6".to_string()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.expect("blank billing context should build");
|
||||||
|
let priced_model_context = StoredBillingModelContext::new(
|
||||||
|
"provider-1".to_string(),
|
||||||
|
Some("pay_as_you_go".to_string()),
|
||||||
|
Some("key-1".to_string()),
|
||||||
|
None,
|
||||||
|
Some(60),
|
||||||
|
"global-model-priced".to_string(),
|
||||||
|
"claude-sonnet-4-6".to_string(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some("model-priced".to_string()),
|
||||||
|
Some("claude-sonnet-4-6".to_string()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some(
|
||||||
|
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]}),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.expect("priced billing context should build");
|
||||||
|
let lookup = TestLookup {
|
||||||
|
name_context: Some(blank_name_context),
|
||||||
|
model_id_context: Some(priced_model_context),
|
||||||
|
};
|
||||||
|
let mut event = UsageEvent::new(
|
||||||
|
UsageEventType::Completed,
|
||||||
|
"req-billing-model-id-1",
|
||||||
|
UsageEventData {
|
||||||
|
provider_name: "NekoCode".to_string(),
|
||||||
|
model: "claude-sonnet-4-6".to_string(),
|
||||||
|
model_id: Some("model-priced".to_string()),
|
||||||
|
provider_id: Some("provider-1".to_string()),
|
||||||
|
provider_api_key_id: Some("key-1".to_string()),
|
||||||
|
request_type: Some("chat".to_string()),
|
||||||
|
input_tokens: Some(1_000),
|
||||||
|
output_tokens: Some(500),
|
||||||
|
status_code: Some(200),
|
||||||
|
..UsageEventData::default()
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
enrich_usage_event_with_billing(&lookup, &mut event)
|
||||||
|
.await
|
||||||
|
.expect("billing should succeed");
|
||||||
|
|
||||||
|
assert!(event.data.total_cost_usd.unwrap_or_default() > 0.0);
|
||||||
|
assert_eq!(
|
||||||
|
event
|
||||||
|
.data
|
||||||
|
.request_metadata
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|value| value.get("billing_snapshot"))
|
||||||
|
.and_then(|value| value.get("status"))
|
||||||
|
.and_then(Value::as_str),
|
||||||
|
Some("complete")
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -24,6 +24,7 @@ impl BillingModelPricingSnapshot {
|
|||||||
pub fn effective_tiered_pricing(&self) -> Option<&Value> {
|
pub fn effective_tiered_pricing(&self) -> Option<&Value> {
|
||||||
self.model_tiered_pricing
|
self.model_tiered_pricing
|
||||||
.as_ref()
|
.as_ref()
|
||||||
|
.filter(|value| has_tiered_pricing_tiers(value))
|
||||||
.or(self.default_tiered_pricing.as_ref())
|
.or(self.default_tiered_pricing.as_ref())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -58,6 +59,67 @@ impl BillingModelPricingSnapshot {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn has_tiered_pricing_tiers(value: &Value) -> bool {
|
||||||
|
value
|
||||||
|
.get("tiers")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.is_some_and(|tiers| !tiers.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
use super::BillingModelPricingSnapshot;
|
||||||
|
|
||||||
|
fn snapshot(
|
||||||
|
model_tiered_pricing: Option<serde_json::Value>,
|
||||||
|
default_tiered_pricing: Option<serde_json::Value>,
|
||||||
|
) -> BillingModelPricingSnapshot {
|
||||||
|
BillingModelPricingSnapshot {
|
||||||
|
provider_id: "provider-1".to_string(),
|
||||||
|
provider_billing_type: None,
|
||||||
|
provider_api_key_id: None,
|
||||||
|
provider_api_key_rate_multipliers: None,
|
||||||
|
provider_api_key_cache_ttl_minutes: None,
|
||||||
|
global_model_id: "global-model-1".to_string(),
|
||||||
|
global_model_name: "gpt-5".to_string(),
|
||||||
|
global_model_config: None,
|
||||||
|
default_price_per_request: None,
|
||||||
|
default_tiered_pricing,
|
||||||
|
model_id: Some("model-1".to_string()),
|
||||||
|
model_provider_model_name: Some("gpt-5-upstream".to_string()),
|
||||||
|
model_config: None,
|
||||||
|
model_price_per_request: None,
|
||||||
|
model_tiered_pricing,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn empty_provider_tiered_pricing_inherits_global_default() {
|
||||||
|
let default_pricing =
|
||||||
|
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]});
|
||||||
|
let pricing = snapshot(Some(json!({})), Some(default_pricing.clone()));
|
||||||
|
|
||||||
|
assert_eq!(pricing.effective_tiered_pricing(), Some(&default_pricing));
|
||||||
|
|
||||||
|
let pricing = snapshot(Some(json!({"tiers": []})), Some(default_pricing.clone()));
|
||||||
|
|
||||||
|
assert_eq!(pricing.effective_tiered_pricing(), Some(&default_pricing));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn populated_provider_tiered_pricing_overrides_global_default() {
|
||||||
|
let provider_pricing =
|
||||||
|
json!({"tiers":[{"up_to":null,"input_price_per_1m":1.0,"output_price_per_1m":2.0}]});
|
||||||
|
let default_pricing =
|
||||||
|
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]});
|
||||||
|
let pricing = snapshot(Some(provider_pricing.clone()), Some(default_pricing));
|
||||||
|
|
||||||
|
assert_eq!(pricing.effective_tiered_pricing(), Some(&provider_pricing));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||||
pub struct BillingUsageInput {
|
pub struct BillingUsageInput {
|
||||||
pub task_type: String,
|
pub task_type: String,
|
||||||
|
|||||||
@@ -150,4 +150,14 @@ pub trait BillingReadRepository: Send + Sync {
|
|||||||
provider_api_key_id: Option<&str>,
|
provider_api_key_id: Option<&str>,
|
||||||
global_model_name: &str,
|
global_model_name: &str,
|
||||||
) -> Result<Option<StoredBillingModelContext>, crate::DataLayerError>;
|
) -> Result<Option<StoredBillingModelContext>, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn find_model_context_by_model_id(
|
||||||
|
&self,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, crate::DataLayerError> {
|
||||||
|
let _ = (provider_id, provider_api_key_id, model_id);
|
||||||
|
Ok(None)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -500,6 +500,18 @@ impl StoredRequestUsageAudit {
|
|||||||
.or_else(|| self.request_metadata_string("key_name"))
|
.or_else(|| self.request_metadata_string("key_name"))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn routing_model_id(&self) -> Option<&str> {
|
||||||
|
self.request_metadata_string("model_id")
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn routing_global_model_id(&self) -> Option<&str> {
|
||||||
|
self.request_metadata_string("global_model_id")
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn routing_global_model_name(&self) -> Option<&str> {
|
||||||
|
self.request_metadata_string("global_model_name")
|
||||||
|
}
|
||||||
|
|
||||||
pub fn routing_planner_kind(&self) -> Option<&str> {
|
pub fn routing_planner_kind(&self) -> Option<&str> {
|
||||||
self.planner_kind
|
self.planner_kind
|
||||||
.as_deref()
|
.as_deref()
|
||||||
|
|||||||
@@ -44,12 +44,21 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
|
|||||||
provider_api_key_id: Option<&str>,
|
provider_api_key_id: Option<&str>,
|
||||||
global_model_name: &str,
|
global_model_name: &str,
|
||||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
|
let by_key = self.by_key.read().expect("billing repository lock");
|
||||||
|
if let Some(value) = find_context_by_provider_model_name(
|
||||||
|
&by_key,
|
||||||
|
provider_id,
|
||||||
|
provider_api_key_id,
|
||||||
|
global_model_name,
|
||||||
|
) {
|
||||||
|
return Ok(Some(value));
|
||||||
|
}
|
||||||
|
|
||||||
let key = (
|
let key = (
|
||||||
provider_id.to_string(),
|
provider_id.to_string(),
|
||||||
global_model_name.to_string(),
|
global_model_name.to_string(),
|
||||||
provider_api_key_id.map(ToOwned::to_owned),
|
provider_api_key_id.map(ToOwned::to_owned),
|
||||||
);
|
);
|
||||||
let by_key = self.by_key.read().expect("billing repository lock");
|
|
||||||
if let Some(value) = by_key.get(&key) {
|
if let Some(value) = by_key.get(&key) {
|
||||||
return Ok(Some(value.clone()));
|
return Ok(Some(value.clone()));
|
||||||
}
|
}
|
||||||
@@ -68,6 +77,88 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
|
|||||||
})
|
})
|
||||||
.map(|(_, value)| value.clone()))
|
.map(|(_, value)| value.clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn find_model_context_by_model_id(
|
||||||
|
&self,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
|
let by_key = self.by_key.read().expect("billing repository lock");
|
||||||
|
if let Some(value) =
|
||||||
|
find_context_by_model_id_and_key(&by_key, provider_id, provider_api_key_id, model_id)
|
||||||
|
{
|
||||||
|
return Ok(Some(value));
|
||||||
|
}
|
||||||
|
if let Some(value) = find_context_by_model_id_and_key(&by_key, provider_id, None, model_id)
|
||||||
|
{
|
||||||
|
return Ok(Some(value));
|
||||||
|
}
|
||||||
|
Ok(by_key
|
||||||
|
.iter()
|
||||||
|
.find(|((stored_provider_id, _, _), value)| {
|
||||||
|
stored_provider_id == provider_id && value.model_id.as_deref() == Some(model_id)
|
||||||
|
})
|
||||||
|
.map(|(_, value)| value.clone()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn find_context_by_provider_model_name(
|
||||||
|
by_key: &BillingContextMap,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
provider_model_name: &str,
|
||||||
|
) -> Option<StoredBillingModelContext> {
|
||||||
|
find_context_by_provider_model_name_and_key(
|
||||||
|
by_key,
|
||||||
|
provider_id,
|
||||||
|
provider_api_key_id,
|
||||||
|
provider_model_name,
|
||||||
|
)
|
||||||
|
.or_else(|| {
|
||||||
|
find_context_by_provider_model_name_and_key(by_key, provider_id, None, provider_model_name)
|
||||||
|
})
|
||||||
|
.or_else(|| {
|
||||||
|
by_key
|
||||||
|
.iter()
|
||||||
|
.find(|((stored_provider_id, _, _), value)| {
|
||||||
|
stored_provider_id == provider_id
|
||||||
|
&& value.model_provider_model_name.as_deref() == Some(provider_model_name)
|
||||||
|
})
|
||||||
|
.map(|(_, value)| value.clone())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn find_context_by_provider_model_name_and_key(
|
||||||
|
by_key: &BillingContextMap,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
provider_model_name: &str,
|
||||||
|
) -> Option<StoredBillingModelContext> {
|
||||||
|
by_key
|
||||||
|
.iter()
|
||||||
|
.find(|((stored_provider_id, _, stored_key_id), value)| {
|
||||||
|
stored_provider_id == provider_id
|
||||||
|
&& stored_key_id.as_deref() == provider_api_key_id
|
||||||
|
&& value.model_provider_model_name.as_deref() == Some(provider_model_name)
|
||||||
|
})
|
||||||
|
.map(|(_, value)| value.clone())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn find_context_by_model_id_and_key(
|
||||||
|
by_key: &BillingContextMap,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Option<StoredBillingModelContext> {
|
||||||
|
by_key
|
||||||
|
.iter()
|
||||||
|
.find(|((stored_provider_id, _, stored_key_id), value)| {
|
||||||
|
stored_provider_id == provider_id
|
||||||
|
&& stored_key_id.as_deref() == provider_api_key_id
|
||||||
|
&& value.model_id.as_deref() == Some(model_id)
|
||||||
|
})
|
||||||
|
.map(|(_, value)| value.clone())
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -110,4 +201,77 @@ mod tests {
|
|||||||
assert_eq!(stored.provider_id, "provider-1");
|
assert_eq!(stored.provider_id, "provider-1");
|
||||||
assert_eq!(stored.global_model_name, "gpt-5");
|
assert_eq!(stored.global_model_name, "gpt-5");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn resolves_by_provider_model_name_before_global_name_collision() {
|
||||||
|
let global_named_context = StoredBillingModelContext::new(
|
||||||
|
"provider-1".to_string(),
|
||||||
|
Some("pay_as_you_go".to_string()),
|
||||||
|
Some("key-1".to_string()),
|
||||||
|
None,
|
||||||
|
Some(60),
|
||||||
|
"global-model-blank".to_string(),
|
||||||
|
"claude-sonnet-4-6".to_string(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some("model-blank".to_string()),
|
||||||
|
Some("blank-upstream".to_string()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.expect("blank billing context should build");
|
||||||
|
let provider_priced_context = StoredBillingModelContext::new(
|
||||||
|
"provider-1".to_string(),
|
||||||
|
Some("pay_as_you_go".to_string()),
|
||||||
|
Some("key-1".to_string()),
|
||||||
|
None,
|
||||||
|
Some(60),
|
||||||
|
"global-model-priced".to_string(),
|
||||||
|
"claude-opus-4-6".to_string(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]})),
|
||||||
|
Some("model-priced".to_string()),
|
||||||
|
Some("claude-sonnet-4-6".to_string()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.expect("priced billing context should build");
|
||||||
|
let repository = InMemoryBillingReadRepository::seed(vec![
|
||||||
|
global_named_context,
|
||||||
|
provider_priced_context,
|
||||||
|
]);
|
||||||
|
|
||||||
|
let stored = repository
|
||||||
|
.find_model_context("provider-1", Some("key-1"), "claude-sonnet-4-6")
|
||||||
|
.await
|
||||||
|
.expect("lookup should succeed")
|
||||||
|
.expect("context should exist");
|
||||||
|
|
||||||
|
assert_eq!(stored.global_model_name, "claude-opus-4-6");
|
||||||
|
assert_eq!(
|
||||||
|
stored.model_provider_model_name.as_deref(),
|
||||||
|
Some("claude-sonnet-4-6")
|
||||||
|
);
|
||||||
|
assert!(stored.default_tiered_pricing.is_some());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn resolves_by_model_id() {
|
||||||
|
let repository = InMemoryBillingReadRepository::seed(vec![sample_context()]);
|
||||||
|
let stored = repository
|
||||||
|
.find_model_context_by_model_id("provider-1", Some("key-1"), "model-1")
|
||||||
|
.await
|
||||||
|
.expect("lookup should succeed")
|
||||||
|
.expect("context should exist");
|
||||||
|
|
||||||
|
assert_eq!(stored.global_model_name, "gpt-5");
|
||||||
|
assert_eq!(
|
||||||
|
stored.model_provider_model_name.as_deref(),
|
||||||
|
Some("gpt-5-upstream")
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -23,8 +23,7 @@ SELECT
|
|||||||
m.tiered_pricing AS model_tiered_pricing
|
m.tiered_pricing AS model_tiered_pricing
|
||||||
FROM providers p
|
FROM providers p
|
||||||
INNER JOIN global_models gm
|
INNER JOIN global_models gm
|
||||||
ON gm.name = $2
|
ON gm.is_active = TRUE
|
||||||
AND gm.is_active = TRUE
|
|
||||||
LEFT JOIN models m
|
LEFT JOIN models m
|
||||||
ON m.global_model_id = gm.id
|
ON m.global_model_id = gm.id
|
||||||
AND m.provider_id = p.id
|
AND m.provider_id = p.id
|
||||||
@@ -33,7 +32,71 @@ LEFT JOIN provider_api_keys pak
|
|||||||
ON pak.id = $3
|
ON pak.id = $3
|
||||||
AND pak.provider_id = p.id
|
AND pak.provider_id = p.id
|
||||||
WHERE p.id = $1
|
WHERE p.id = $1
|
||||||
ORDER BY COALESCE(m.is_available, FALSE) DESC, m.created_at ASC
|
AND (
|
||||||
|
gm.name = $2
|
||||||
|
OR m.provider_model_name = $2
|
||||||
|
OR (
|
||||||
|
m.provider_model_mappings IS NOT NULL
|
||||||
|
AND (
|
||||||
|
m.provider_model_mappings @> jsonb_build_array(jsonb_build_object('name', $2::TEXT))
|
||||||
|
OR m.provider_model_mappings @> jsonb_build_array(to_jsonb($2::TEXT))
|
||||||
|
OR m.provider_model_mappings @> jsonb_build_object('name', $2::TEXT)
|
||||||
|
OR m.provider_model_mappings = to_jsonb($2::TEXT)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
ORDER BY
|
||||||
|
CASE
|
||||||
|
WHEN m.provider_model_name = $2 THEN 0
|
||||||
|
WHEN m.provider_model_mappings IS NOT NULL
|
||||||
|
AND (
|
||||||
|
m.provider_model_mappings @> jsonb_build_array(jsonb_build_object('name', $2::TEXT))
|
||||||
|
OR m.provider_model_mappings @> jsonb_build_array(to_jsonb($2::TEXT))
|
||||||
|
OR m.provider_model_mappings @> jsonb_build_object('name', $2::TEXT)
|
||||||
|
OR m.provider_model_mappings = to_jsonb($2::TEXT)
|
||||||
|
) THEN 1
|
||||||
|
WHEN gm.name = $2 THEN 2
|
||||||
|
ELSE 3
|
||||||
|
END ASC,
|
||||||
|
COALESCE(m.is_available, FALSE) DESC,
|
||||||
|
CASE
|
||||||
|
WHEN m.tiered_pricing IS NOT NULL OR m.price_per_request IS NOT NULL THEN 0
|
||||||
|
WHEN gm.default_tiered_pricing IS NOT NULL OR gm.default_price_per_request IS NOT NULL THEN 1
|
||||||
|
ELSE 2
|
||||||
|
END ASC,
|
||||||
|
m.created_at ASC
|
||||||
|
LIMIT 1
|
||||||
|
"#;
|
||||||
|
|
||||||
|
const FIND_MODEL_CONTEXT_BY_MODEL_ID_SQL: &str = r#"
|
||||||
|
SELECT
|
||||||
|
p.id AS provider_id,
|
||||||
|
CAST(p.billing_type AS TEXT) AS provider_billing_type,
|
||||||
|
pak.id AS provider_api_key_id,
|
||||||
|
pak.rate_multipliers AS provider_api_key_rate_multipliers,
|
||||||
|
pak.cache_ttl_minutes AS provider_api_key_cache_ttl_minutes,
|
||||||
|
gm.id AS global_model_id,
|
||||||
|
gm.name AS global_model_name,
|
||||||
|
gm.config AS global_model_config,
|
||||||
|
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
|
||||||
|
gm.default_tiered_pricing AS default_tiered_pricing,
|
||||||
|
m.id AS model_id,
|
||||||
|
m.provider_model_name AS model_provider_model_name,
|
||||||
|
m.config AS model_config,
|
||||||
|
CAST(m.price_per_request AS DOUBLE PRECISION) AS model_price_per_request,
|
||||||
|
m.tiered_pricing AS model_tiered_pricing
|
||||||
|
FROM providers p
|
||||||
|
INNER JOIN models m
|
||||||
|
ON m.id = $2
|
||||||
|
AND m.provider_id = p.id
|
||||||
|
AND m.is_active = TRUE
|
||||||
|
INNER JOIN global_models gm
|
||||||
|
ON gm.id = m.global_model_id
|
||||||
|
AND gm.is_active = TRUE
|
||||||
|
LEFT JOIN provider_api_keys pak
|
||||||
|
ON pak.id = $3
|
||||||
|
AND pak.provider_id = p.id
|
||||||
|
WHERE p.id = $1
|
||||||
LIMIT 1
|
LIMIT 1
|
||||||
"#;
|
"#;
|
||||||
|
|
||||||
@@ -62,6 +125,22 @@ impl SqlxBillingReadRepository {
|
|||||||
.map_postgres_err()?;
|
.map_postgres_err()?;
|
||||||
row.as_ref().map(map_row).transpose()
|
row.as_ref().map(map_row).transpose()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn find_model_context_by_model_id(
|
||||||
|
&self,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
|
let row = sqlx::query(FIND_MODEL_CONTEXT_BY_MODEL_ID_SQL)
|
||||||
|
.bind(provider_id)
|
||||||
|
.bind(model_id)
|
||||||
|
.bind(provider_api_key_id)
|
||||||
|
.fetch_optional(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_postgres_err()?;
|
||||||
|
row.as_ref().map(map_row).transpose()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
@@ -74,6 +153,15 @@ impl BillingReadRepository for SqlxBillingReadRepository {
|
|||||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
Self::find_model_context(self, provider_id, provider_api_key_id, global_model_name).await
|
Self::find_model_context(self, provider_id, provider_api_key_id, global_model_name).await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn find_model_context_by_model_id(
|
||||||
|
&self,
|
||||||
|
provider_id: &str,
|
||||||
|
provider_api_key_id: Option<&str>,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||||
|
Self::find_model_context_by_model_id(self, provider_id, provider_api_key_id, model_id).await
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn map_row(row: &sqlx::postgres::PgRow) -> Result<StoredBillingModelContext, DataLayerError> {
|
fn map_row(row: &sqlx::postgres::PgRow) -> Result<StoredBillingModelContext, DataLayerError> {
|
||||||
|
|||||||
@@ -33,6 +33,10 @@ pub struct UsageEventData {
|
|||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub target_model: Option<String>,
|
pub target_model: Option<String>,
|
||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub model_id: Option<String>,
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub global_model_id: Option<String>,
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub provider_id: Option<String>,
|
pub provider_id: Option<String>,
|
||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub provider_endpoint_id: Option<String>,
|
pub provider_endpoint_id: Option<String>,
|
||||||
|
|||||||
@@ -75,6 +75,9 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
|
|||||||
copy_non_null_value(source, target, "billing_snapshot");
|
copy_non_null_value(source, target, "billing_snapshot");
|
||||||
copy_non_empty_string(source, target, "billing_snapshot_schema_version");
|
copy_non_empty_string(source, target, "billing_snapshot_schema_version");
|
||||||
copy_non_empty_string(source, target, "billing_snapshot_status");
|
copy_non_empty_string(source, target, "billing_snapshot_status");
|
||||||
|
copy_non_empty_string(source, target, "model_id");
|
||||||
|
copy_non_empty_string(source, target, "global_model_id");
|
||||||
|
copy_non_empty_string(source, target, "global_model_name");
|
||||||
copy_non_null_value(source, target, "dimensions");
|
copy_non_null_value(source, target, "dimensions");
|
||||||
copy_non_null_value(source, target, "billing_rule_snapshot");
|
copy_non_null_value(source, target, "billing_rule_snapshot");
|
||||||
copy_non_null_value(source, target, "scheduling_audit");
|
copy_non_null_value(source, target, "scheduling_audit");
|
||||||
@@ -97,6 +100,9 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
|
|||||||
remove_non_null_value(&mut source, target, "billing_snapshot");
|
remove_non_null_value(&mut source, target, "billing_snapshot");
|
||||||
remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version");
|
remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version");
|
||||||
remove_non_empty_string(&mut source, target, "billing_snapshot_status");
|
remove_non_empty_string(&mut source, target, "billing_snapshot_status");
|
||||||
|
remove_non_empty_string(&mut source, target, "model_id");
|
||||||
|
remove_non_empty_string(&mut source, target, "global_model_id");
|
||||||
|
remove_non_empty_string(&mut source, target, "global_model_name");
|
||||||
remove_non_null_value(&mut source, target, "dimensions");
|
remove_non_null_value(&mut source, target, "dimensions");
|
||||||
remove_non_null_value(&mut source, target, "billing_rule_snapshot");
|
remove_non_null_value(&mut source, target, "billing_rule_snapshot");
|
||||||
remove_non_null_value(&mut source, target, "scheduling_audit");
|
remove_non_null_value(&mut source, target, "scheduling_audit");
|
||||||
@@ -378,6 +384,9 @@ mod tests {
|
|||||||
"billing_snapshot": {"status": "complete"},
|
"billing_snapshot": {"status": "complete"},
|
||||||
"billing_snapshot_schema_version": "2.0",
|
"billing_snapshot_schema_version": "2.0",
|
||||||
"billing_snapshot_status": "complete",
|
"billing_snapshot_status": "complete",
|
||||||
|
"model_id": "model-1",
|
||||||
|
"global_model_id": "global-model-1",
|
||||||
|
"global_model_name": "gpt-5",
|
||||||
"dimensions": {"total_input_context": 10},
|
"dimensions": {"total_input_context": 10},
|
||||||
"rate_multiplier": 1.25,
|
"rate_multiplier": 1.25,
|
||||||
"is_free_tier": false,
|
"is_free_tier": false,
|
||||||
@@ -405,6 +414,9 @@ mod tests {
|
|||||||
"billing_snapshot": {"status": "complete"},
|
"billing_snapshot": {"status": "complete"},
|
||||||
"billing_snapshot_schema_version": "2.0",
|
"billing_snapshot_schema_version": "2.0",
|
||||||
"billing_snapshot_status": "complete",
|
"billing_snapshot_status": "complete",
|
||||||
|
"model_id": "model-1",
|
||||||
|
"global_model_id": "global-model-1",
|
||||||
|
"global_model_name": "gpt-5",
|
||||||
"dimensions": {"total_input_context": 10},
|
"dimensions": {"total_input_context": 10},
|
||||||
"rate_multiplier": 1.25,
|
"rate_multiplier": 1.25,
|
||||||
"is_free_tier": false,
|
"is_free_tier": false,
|
||||||
@@ -455,6 +467,9 @@ mod tests {
|
|||||||
"client_requested_stream": false,
|
"client_requested_stream": false,
|
||||||
"upstream_is_stream": true,
|
"upstream_is_stream": true,
|
||||||
"provider_id": "provider-1",
|
"provider_id": "provider-1",
|
||||||
|
"model_id": "model-1",
|
||||||
|
"global_model_id": "global-model-1",
|
||||||
|
"global_model_name": "gpt-5",
|
||||||
"billing_snapshot": {"status": "complete"}
|
"billing_snapshot": {"status": "complete"}
|
||||||
})
|
})
|
||||||
.as_object()
|
.as_object()
|
||||||
@@ -468,6 +483,9 @@ mod tests {
|
|||||||
json!({
|
json!({
|
||||||
"client_requested_stream": false,
|
"client_requested_stream": false,
|
||||||
"upstream_is_stream": true,
|
"upstream_is_stream": true,
|
||||||
|
"model_id": "model-1",
|
||||||
|
"global_model_id": "global-model-1",
|
||||||
|
"global_model_name": "gpt-5",
|
||||||
"billing_snapshot": {"status": "complete"}
|
"billing_snapshot": {"status": "complete"}
|
||||||
})
|
})
|
||||||
);
|
);
|
||||||
|
|||||||
@@ -74,6 +74,8 @@ pub struct LifecycleUsageSeed {
|
|||||||
pub provider_name: String,
|
pub provider_name: String,
|
||||||
pub model: String,
|
pub model: String,
|
||||||
pub target_model: Option<String>,
|
pub target_model: Option<String>,
|
||||||
|
pub model_id: Option<String>,
|
||||||
|
pub global_model_id: Option<String>,
|
||||||
pub provider_id: Option<String>,
|
pub provider_id: Option<String>,
|
||||||
pub provider_endpoint_id: Option<String>,
|
pub provider_endpoint_id: Option<String>,
|
||||||
pub provider_api_key_id: Option<String>,
|
pub provider_api_key_id: Option<String>,
|
||||||
@@ -110,6 +112,8 @@ pub struct TerminalUsageContextSeed {
|
|||||||
pub provider_name: String,
|
pub provider_name: String,
|
||||||
pub model: String,
|
pub model: String,
|
||||||
pub target_model: Option<String>,
|
pub target_model: Option<String>,
|
||||||
|
pub model_id: Option<String>,
|
||||||
|
pub global_model_id: Option<String>,
|
||||||
pub provider_id: Option<String>,
|
pub provider_id: Option<String>,
|
||||||
pub provider_endpoint_id: Option<String>,
|
pub provider_endpoint_id: Option<String>,
|
||||||
pub provider_api_key_id: Option<String>,
|
pub provider_api_key_id: Option<String>,
|
||||||
@@ -168,6 +172,8 @@ pub struct TerminalUsageSeed {
|
|||||||
pub provider_name: String,
|
pub provider_name: String,
|
||||||
pub model: String,
|
pub model: String,
|
||||||
pub target_model: Option<String>,
|
pub target_model: Option<String>,
|
||||||
|
pub model_id: Option<String>,
|
||||||
|
pub global_model_id: Option<String>,
|
||||||
pub provider_id: Option<String>,
|
pub provider_id: Option<String>,
|
||||||
pub provider_endpoint_id: Option<String>,
|
pub provider_endpoint_id: Option<String>,
|
||||||
pub provider_api_key_id: Option<String>,
|
pub provider_api_key_id: Option<String>,
|
||||||
@@ -262,6 +268,8 @@ pub fn build_lifecycle_usage_seed(
|
|||||||
provider_name,
|
provider_name,
|
||||||
model,
|
model,
|
||||||
target_model: context_string(context, "mapped_model"),
|
target_model: context_string(context, "mapped_model"),
|
||||||
|
model_id: context_string(context, "model_id"),
|
||||||
|
global_model_id: context_string(context, "global_model_id"),
|
||||||
provider_id: empty_to_none(
|
provider_id: empty_to_none(
|
||||||
context_string(context, "provider_id")
|
context_string(context, "provider_id")
|
||||||
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
|
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
|
||||||
@@ -486,6 +494,8 @@ fn build_terminal_usage_event_from_seed_impl(
|
|||||||
provider_name,
|
provider_name,
|
||||||
model,
|
model,
|
||||||
target_model,
|
target_model,
|
||||||
|
model_id,
|
||||||
|
global_model_id,
|
||||||
provider_id,
|
provider_id,
|
||||||
provider_endpoint_id,
|
provider_endpoint_id,
|
||||||
provider_api_key_id,
|
provider_api_key_id,
|
||||||
@@ -548,6 +558,8 @@ fn build_terminal_usage_event_from_seed_impl(
|
|||||||
provider_name,
|
provider_name,
|
||||||
model,
|
model,
|
||||||
target_model,
|
target_model,
|
||||||
|
model_id,
|
||||||
|
global_model_id,
|
||||||
provider_id,
|
provider_id,
|
||||||
provider_endpoint_id,
|
provider_endpoint_id,
|
||||||
provider_api_key_id,
|
provider_api_key_id,
|
||||||
@@ -657,6 +669,8 @@ pub fn build_terminal_usage_context_seed(
|
|||||||
.or_else(|| non_empty_str(plan.model_name.as_deref()))
|
.or_else(|| non_empty_str(plan.model_name.as_deref()))
|
||||||
.unwrap_or_else(|| "unknown".to_string()),
|
.unwrap_or_else(|| "unknown".to_string()),
|
||||||
target_model: context_string(context, "mapped_model"),
|
target_model: context_string(context, "mapped_model"),
|
||||||
|
model_id: context_string(context, "model_id"),
|
||||||
|
global_model_id: context_string(context, "global_model_id"),
|
||||||
provider_id: context_string(context, "provider_id")
|
provider_id: context_string(context, "provider_id")
|
||||||
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
|
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
|
||||||
provider_endpoint_id: context_string(context, "endpoint_id")
|
provider_endpoint_id: context_string(context, "endpoint_id")
|
||||||
@@ -813,6 +827,8 @@ pub fn build_sync_terminal_usage_seed(
|
|||||||
provider_name: context_seed.provider_name,
|
provider_name: context_seed.provider_name,
|
||||||
model: context_seed.model,
|
model: context_seed.model,
|
||||||
target_model: context_seed.target_model,
|
target_model: context_seed.target_model,
|
||||||
|
model_id: context_seed.model_id,
|
||||||
|
global_model_id: context_seed.global_model_id,
|
||||||
provider_id: context_seed.provider_id,
|
provider_id: context_seed.provider_id,
|
||||||
provider_endpoint_id: context_seed.provider_endpoint_id,
|
provider_endpoint_id: context_seed.provider_endpoint_id,
|
||||||
provider_api_key_id: context_seed.provider_api_key_id,
|
provider_api_key_id: context_seed.provider_api_key_id,
|
||||||
@@ -881,6 +897,8 @@ pub fn build_stream_terminal_usage_seed(
|
|||||||
provider_name: context_seed.provider_name,
|
provider_name: context_seed.provider_name,
|
||||||
model: context_seed.model,
|
model: context_seed.model,
|
||||||
target_model: context_seed.target_model,
|
target_model: context_seed.target_model,
|
||||||
|
model_id: context_seed.model_id,
|
||||||
|
global_model_id: context_seed.global_model_id,
|
||||||
provider_id: context_seed.provider_id,
|
provider_id: context_seed.provider_id,
|
||||||
provider_endpoint_id: context_seed.provider_endpoint_id,
|
provider_endpoint_id: context_seed.provider_endpoint_id,
|
||||||
provider_api_key_id: context_seed.provider_api_key_id,
|
provider_api_key_id: context_seed.provider_api_key_id,
|
||||||
@@ -1242,6 +1260,8 @@ fn build_usage_event_data_seed_with_detail(
|
|||||||
provider_name,
|
provider_name,
|
||||||
model,
|
model,
|
||||||
target_model: context_string(context, "mapped_model"),
|
target_model: context_string(context, "mapped_model"),
|
||||||
|
model_id: context_string(context, "model_id"),
|
||||||
|
global_model_id: context_string(context, "global_model_id"),
|
||||||
provider_id: context_string(context, "provider_id")
|
provider_id: context_string(context, "provider_id")
|
||||||
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
|
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
|
||||||
provider_endpoint_id: context_string(context, "endpoint_id")
|
provider_endpoint_id: context_string(context, "endpoint_id")
|
||||||
@@ -3523,6 +3543,8 @@ mod tests {
|
|||||||
provider_name: "OpenAI".to_string(),
|
provider_name: "OpenAI".to_string(),
|
||||||
model: "gpt-5".to_string(),
|
model: "gpt-5".to_string(),
|
||||||
target_model: None,
|
target_model: None,
|
||||||
|
model_id: None,
|
||||||
|
global_model_id: None,
|
||||||
provider_id: Some("provider-1".to_string()),
|
provider_id: Some("provider-1".to_string()),
|
||||||
provider_endpoint_id: Some("endpoint-1".to_string()),
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||||||
provider_api_key_id: Some("upstream-key-1".to_string()),
|
provider_api_key_id: Some("upstream-key-1".to_string()),
|
||||||
@@ -3646,6 +3668,8 @@ mod tests {
|
|||||||
provider_name: "OpenAI".to_string(),
|
provider_name: "OpenAI".to_string(),
|
||||||
model: "gpt-5".to_string(),
|
model: "gpt-5".to_string(),
|
||||||
target_model: None,
|
target_model: None,
|
||||||
|
model_id: None,
|
||||||
|
global_model_id: None,
|
||||||
provider_id: Some("provider-1".to_string()),
|
provider_id: Some("provider-1".to_string()),
|
||||||
provider_endpoint_id: Some("endpoint-1".to_string()),
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||||||
provider_api_key_id: Some("upstream-key-1".to_string()),
|
provider_api_key_id: Some("upstream-key-1".to_string()),
|
||||||
|
|||||||
Reference in New Issue
Block a user