mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +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:
@@ -44,12 +44,21 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
|
||||
provider_api_key_id: Option<&str>,
|
||||
global_model_name: &str,
|
||||
) -> 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 = (
|
||||
provider_id.to_string(),
|
||||
global_model_name.to_string(),
|
||||
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) {
|
||||
return Ok(Some(value.clone()));
|
||||
}
|
||||
@@ -68,6 +77,88 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
|
||||
})
|
||||
.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)]
|
||||
@@ -110,4 +201,77 @@ mod tests {
|
||||
assert_eq!(stored.provider_id, "provider-1");
|
||||
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
|
||||
FROM providers p
|
||||
INNER JOIN global_models gm
|
||||
ON gm.name = $2
|
||||
AND gm.is_active = TRUE
|
||||
ON gm.is_active = TRUE
|
||||
LEFT JOIN models m
|
||||
ON m.global_model_id = gm.id
|
||||
AND m.provider_id = p.id
|
||||
@@ -33,7 +32,71 @@ LEFT JOIN provider_api_keys pak
|
||||
ON pak.id = $3
|
||||
AND pak.provider_id = p.id
|
||||
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
|
||||
"#;
|
||||
|
||||
@@ -62,6 +125,22 @@ impl SqlxBillingReadRepository {
|
||||
.map_postgres_err()?;
|
||||
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]
|
||||
@@ -74,6 +153,15 @@ impl BillingReadRepository for SqlxBillingReadRepository {
|
||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||
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> {
|
||||
|
||||
Reference in New Issue
Block a user