Add usage client family read model field

This commit is contained in:
RWDai
2026-05-18 19:28:36 +08:00
parent 0a26accca4
commit 6d994917a0
6 changed files with 54 additions and 12 deletions

View File

@@ -1,8 +1,8 @@
mod types;
pub use types::{
parse_usage_body_ref, usage_body_ref, PendingUsageCleanupSummary,
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
parse_usage_body_ref, usage_body_ref, usage_request_metadata_client_family,
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,

View File

@@ -34,6 +34,8 @@ pub struct StoredRequestUsageAudit {
pub provider_endpoint_kind: Option<String>,
pub has_format_conversion: bool,
pub is_stream: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_family: Option<String>,
pub input_tokens: u64,
pub output_tokens: u64,
pub total_tokens: u64,
@@ -211,6 +213,7 @@ impl StoredRequestUsageAudit {
provider_endpoint_kind,
has_format_conversion,
is_stream,
client_family: None,
input_tokens: parse_u64(input_tokens, "usage.input_tokens")?,
output_tokens: parse_u64(output_tokens, "usage.output_tokens")?,
total_tokens: parse_u64(total_tokens, "usage.total_tokens")?,
@@ -311,6 +314,10 @@ impl StoredRequestUsageAudit {
.filter(|value| !value.is_empty())
}
pub fn request_metadata_client_family(&self) -> Option<&str> {
usage_request_metadata_client_family(self.request_metadata.as_ref())
}
fn billing_snapshot_resolved_number(&self, key: &str) -> Option<f64> {
self.request_metadata_object()
.and_then(|metadata| metadata.get("billing_snapshot"))
@@ -1794,6 +1801,18 @@ pub struct UsageCleanupPreviewCounts {
pub log: u64,
}
pub fn usage_request_metadata_client_family(value: Option<&Value>) -> Option<&str> {
let metadata = value.and_then(Value::as_object)?;
metadata
.get("client_session_affinity")
.and_then(Value::as_object)
.and_then(|affinity| affinity.get("client_family"))
.and_then(Value::as_str)
.or_else(|| metadata.get("client_family").and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn parse_u64(value: i32, field_name: &str) -> Result<u64, crate::DataLayerError> {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))