mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 18:37:46 +08:00
feat: 扩展 Rust gateway 全功能模块，新增 billing/crypto/wallet crate 及完整数据层
- 新增 aether-billing、aether-crypto、aether-wallet 独立 crate - aether-data 扩展 repository 层：announcements、auth_modules、billing、 candidate_selection、gemini_file_mappings、global_models、management_tokens、 oauth_providers、proxy_nodes、quota、users、wallet 等模块 - aether-gateway 新增 api/auth/billing/control/middleware/scheduler/usage/ video_tasks/hooks/maintenance/model_fetch/provider_transport 等功能模块 - 重构 executor decision 和 gateway state 为模块目录结构 - 新增 gateway router、frontdoor 路由层及对应测试 - Python 侧 API 路由重构，新增 compat/support 模块 - 前端 Logo 组件更新及 Provider 管理页面调整
This commit is contained in:
fawney19
@@ -0,0 +1,775 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_billing::map_usage_from_response;
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
|
||||
use aether_data::repository::usage::UpsertUsageRecord;
|
||||
use aether_data::DataLayerError;
|
||||
use base64::Engine as _;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::reporting::{GatewayStreamReportRequest, GatewaySyncReportRequest};
|
||||
use super::{UsageEvent, UsageEventData, UsageEventType};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum UsageLifecycleState {
|
||||
Pending,
|
||||
Streaming,
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
pub(crate) fn build_pending_usage_record(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<UpsertUsageRecord, DataLayerError> {
|
||||
build_upsert_usage_record(
|
||||
plan,
|
||||
report_context,
|
||||
UsageEventData::default(),
|
||||
UsageLifecycleState::Pending,
|
||||
None,
|
||||
updated_at_unix_secs,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_streaming_usage_record(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
status_code: u16,
|
||||
response_headers: &BTreeMap<String, String>,
|
||||
telemetry: Option<&ExecutionTelemetry>,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<UpsertUsageRecord, DataLayerError> {
|
||||
build_upsert_usage_record(
|
||||
plan,
|
||||
report_context,
|
||||
UsageEventData {
|
||||
status_code: Some(status_code),
|
||||
response_time_ms: telemetry.and_then(|value| value.elapsed_ms),
|
||||
first_byte_time_ms: telemetry.and_then(|value| value.ttfb_ms),
|
||||
response_headers: Some(headers_to_json(response_headers)),
|
||||
client_response_headers: Some(headers_to_json(response_headers)),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
UsageLifecycleState::Streaming,
|
||||
None,
|
||||
updated_at_unix_secs,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_sync_terminal_usage_event(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
event_type: UsageEventType,
|
||||
) -> Result<UsageEvent, DataLayerError> {
|
||||
let mut data = UsageEventData {
|
||||
status_code: Some(payload.status_code),
|
||||
response_time_ms: payload
|
||||
.telemetry
|
||||
.as_ref()
|
||||
.and_then(|value| value.elapsed_ms),
|
||||
first_byte_time_ms: payload.telemetry.as_ref().and_then(|value| value.ttfb_ms),
|
||||
response_headers: Some(headers_to_json(&payload.headers)),
|
||||
client_response_headers: Some(headers_to_json(&payload.headers)),
|
||||
response_body: payload
|
||||
.body_json
|
||||
.clone()
|
||||
.or_else(|| decode_body_for_storage(payload.body_base64.as_deref())),
|
||||
client_response_body: payload.client_body_json.clone(),
|
||||
error_message: resolve_error_message(
|
||||
payload.status_code,
|
||||
payload.body_json.as_ref(),
|
||||
payload.body_base64.as_deref(),
|
||||
),
|
||||
error_category: resolve_error_category(payload.status_code, event_type),
|
||||
..UsageEventData::default()
|
||||
};
|
||||
|
||||
let standardized_response_body = payload
|
||||
.body_json
|
||||
.as_ref()
|
||||
.cloned()
|
||||
.or_else(|| data.response_body.clone());
|
||||
if let Some(tokens) = standardized_response_body.as_ref() {
|
||||
apply_standardized_usage(
|
||||
resolve_billing_api_format(plan, report_context),
|
||||
tokens,
|
||||
&mut data,
|
||||
);
|
||||
}
|
||||
if data.total_tokens.is_none() {
|
||||
if let Some(tokens) = payload
|
||||
.body_json
|
||||
.as_ref()
|
||||
.and_then(extract_token_counts_from_json)
|
||||
.or_else(|| {
|
||||
payload
|
||||
.body_base64
|
||||
.as_deref()
|
||||
.and_then(extract_token_counts_from_base64_payload)
|
||||
})
|
||||
{
|
||||
data.input_tokens = Some(tokens.0);
|
||||
data.output_tokens = Some(tokens.1);
|
||||
data.total_tokens = Some(tokens.2);
|
||||
}
|
||||
}
|
||||
|
||||
build_usage_event(plan, report_context, data, event_type)
|
||||
}
|
||||
|
||||
pub(crate) fn build_stream_terminal_usage_event(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
payload: &GatewayStreamReportRequest,
|
||||
event_type: UsageEventType,
|
||||
) -> Result<UsageEvent, DataLayerError> {
|
||||
let mut data = UsageEventData {
|
||||
status_code: Some(payload.status_code),
|
||||
response_time_ms: payload
|
||||
.telemetry
|
||||
.as_ref()
|
||||
.and_then(|value| value.elapsed_ms),
|
||||
first_byte_time_ms: payload.telemetry.as_ref().and_then(|value| value.ttfb_ms),
|
||||
response_headers: Some(headers_to_json(&payload.headers)),
|
||||
client_response_headers: Some(headers_to_json(&payload.headers)),
|
||||
response_body: decode_body_for_storage(payload.body_base64.as_deref()),
|
||||
error_message: resolve_error_message(
|
||||
payload.status_code,
|
||||
None,
|
||||
payload.body_base64.as_deref(),
|
||||
),
|
||||
error_category: resolve_error_category(payload.status_code, event_type),
|
||||
..UsageEventData::default()
|
||||
};
|
||||
|
||||
let standardized_response_body = data.response_body.clone();
|
||||
if let Some(response_body) = standardized_response_body.as_ref() {
|
||||
apply_standardized_usage(
|
||||
resolve_billing_api_format(plan, report_context),
|
||||
response_body,
|
||||
&mut data,
|
||||
);
|
||||
}
|
||||
if data.total_tokens.is_none() {
|
||||
if let Some(tokens) = payload
|
||||
.body_base64
|
||||
.as_deref()
|
||||
.and_then(extract_token_counts_from_base64_payload)
|
||||
{
|
||||
data.input_tokens = Some(tokens.0);
|
||||
data.output_tokens = Some(tokens.1);
|
||||
data.total_tokens = Some(tokens.2);
|
||||
}
|
||||
}
|
||||
|
||||
build_usage_event(plan, report_context, data, event_type)
|
||||
}
|
||||
|
||||
pub(crate) fn build_upsert_usage_record_from_event(
|
||||
event: &UsageEvent,
|
||||
) -> Result<UpsertUsageRecord, DataLayerError> {
|
||||
let lifecycle_state = match event.event_type {
|
||||
UsageEventType::Pending => UsageLifecycleState::Pending,
|
||||
UsageEventType::Streaming => UsageLifecycleState::Streaming,
|
||||
UsageEventType::Completed => UsageLifecycleState::Completed,
|
||||
UsageEventType::Failed => UsageLifecycleState::Failed,
|
||||
UsageEventType::Cancelled => UsageLifecycleState::Cancelled,
|
||||
};
|
||||
|
||||
let now_unix_secs = event.timestamp_ms / 1_000;
|
||||
build_upsert_usage_record(
|
||||
&ExecutionPlan {
|
||||
request_id: event.request_id.clone(),
|
||||
candidate_id: None,
|
||||
provider_name: Some(event.data.provider_name.clone()),
|
||||
provider_id: event.data.provider_id.clone().unwrap_or_default(),
|
||||
endpoint_id: event.data.provider_endpoint_id.clone().unwrap_or_default(),
|
||||
key_id: event.data.provider_api_key_id.clone().unwrap_or_default(),
|
||||
method: "POST".to_string(),
|
||||
url: String::new(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: aether_contracts::RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: event.data.is_stream.unwrap_or(false),
|
||||
client_api_format: event.data.api_format.clone().unwrap_or_default(),
|
||||
provider_api_format: event
|
||||
.data
|
||||
.endpoint_api_format
|
||||
.clone()
|
||||
.or_else(|| event.data.api_format.clone())
|
||||
.unwrap_or_default(),
|
||||
model_name: Some(event.data.model.clone()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
},
|
||||
None,
|
||||
event.data.clone(),
|
||||
lifecycle_state,
|
||||
Some(now_unix_secs),
|
||||
now_unix_secs,
|
||||
)
|
||||
}
|
||||
|
||||
fn build_usage_event(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
data: UsageEventData,
|
||||
event_type: UsageEventType,
|
||||
) -> Result<UsageEvent, DataLayerError> {
|
||||
let base = build_base_usage_data(plan, report_context);
|
||||
Ok(UsageEvent::new(
|
||||
event_type,
|
||||
plan.request_id.clone(),
|
||||
merge_usage_data(base, data),
|
||||
))
|
||||
}
|
||||
|
||||
fn build_upsert_usage_record(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
data: UsageEventData,
|
||||
lifecycle_state: UsageLifecycleState,
|
||||
finalized_at_unix_secs: Option<u64>,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<UpsertUsageRecord, DataLayerError> {
|
||||
let data = merge_usage_data(build_base_usage_data(plan, report_context), data);
|
||||
let (status, billing_status) = lifecycle_status_and_billing(lifecycle_state);
|
||||
|
||||
Ok(UpsertUsageRecord {
|
||||
request_id: plan.request_id.clone(),
|
||||
user_id: data.user_id,
|
||||
api_key_id: data.api_key_id,
|
||||
username: data.username,
|
||||
api_key_name: data.api_key_name,
|
||||
provider_name: data.provider_name,
|
||||
model: data.model,
|
||||
target_model: data.target_model,
|
||||
provider_id: empty_to_none(data.provider_id),
|
||||
provider_endpoint_id: empty_to_none(data.provider_endpoint_id),
|
||||
provider_api_key_id: empty_to_none(data.provider_api_key_id),
|
||||
request_type: data.request_type,
|
||||
api_format: data.api_format,
|
||||
api_family: data.api_family,
|
||||
endpoint_kind: data.endpoint_kind,
|
||||
endpoint_api_format: data.endpoint_api_format,
|
||||
provider_api_family: data.provider_api_family,
|
||||
provider_endpoint_kind: data.provider_endpoint_kind,
|
||||
has_format_conversion: data.has_format_conversion,
|
||||
is_stream: data.is_stream,
|
||||
input_tokens: data.input_tokens,
|
||||
output_tokens: data.output_tokens,
|
||||
total_tokens: data.total_tokens,
|
||||
cache_creation_input_tokens: data.cache_creation_input_tokens,
|
||||
cache_read_input_tokens: data.cache_read_input_tokens,
|
||||
cache_creation_cost_usd: data.cache_creation_cost_usd,
|
||||
cache_read_cost_usd: data.cache_read_cost_usd,
|
||||
output_price_per_1m: data.output_price_per_1m,
|
||||
total_cost_usd: data.total_cost_usd,
|
||||
actual_total_cost_usd: data.actual_total_cost_usd,
|
||||
status_code: data.status_code,
|
||||
error_message: data.error_message,
|
||||
error_category: data.error_category,
|
||||
response_time_ms: data.response_time_ms,
|
||||
first_byte_time_ms: data.first_byte_time_ms,
|
||||
status: status.to_string(),
|
||||
billing_status: billing_status.to_string(),
|
||||
request_headers: data.request_headers,
|
||||
request_body: data.request_body,
|
||||
provider_request_headers: data.provider_request_headers,
|
||||
provider_request_body: data.provider_request_body,
|
||||
response_headers: data.response_headers,
|
||||
response_body: data.response_body,
|
||||
client_response_headers: data.client_response_headers,
|
||||
client_response_body: data.client_response_body,
|
||||
request_metadata: data.request_metadata,
|
||||
finalized_at_unix_secs,
|
||||
created_at_unix_secs: Some(updated_at_unix_secs),
|
||||
updated_at_unix_secs,
|
||||
})
|
||||
}
|
||||
|
||||
fn build_base_usage_data(plan: &ExecutionPlan, report_context: Option<&Value>) -> UsageEventData {
|
||||
let context = report_context.and_then(Value::as_object);
|
||||
let api_format = context_string(context, "client_api_format")
|
||||
.or_else(|| non_empty_string(Some(plan.client_api_format.clone())));
|
||||
let endpoint_api_format = context_string(context, "provider_api_format")
|
||||
.or_else(|| non_empty_string(Some(plan.provider_api_format.clone())));
|
||||
let model = context_string(context, "model")
|
||||
.or_else(|| plan.model_name.clone())
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let provider_name = context_string(context, "provider_name")
|
||||
.or_else(|| plan.provider_name.clone())
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
|
||||
UsageEventData {
|
||||
user_id: context_string(context, "user_id"),
|
||||
api_key_id: context_string(context, "api_key_id"),
|
||||
username: context_string(context, "username"),
|
||||
api_key_name: context_string(context, "api_key_name"),
|
||||
provider_name,
|
||||
model,
|
||||
target_model: context_string(context, "mapped_model"),
|
||||
provider_id: context_string(context, "provider_id")
|
||||
.or_else(|| non_empty_string(Some(plan.provider_id.clone()))),
|
||||
provider_endpoint_id: context_string(context, "endpoint_id")
|
||||
.or_else(|| non_empty_string(Some(plan.endpoint_id.clone()))),
|
||||
provider_api_key_id: context_string(context, "key_id")
|
||||
.or_else(|| non_empty_string(Some(plan.key_id.clone()))),
|
||||
request_type: Some(infer_request_type(api_format.as_deref())),
|
||||
api_format: api_format.clone(),
|
||||
api_family: api_format
|
||||
.as_deref()
|
||||
.and_then(infer_api_family)
|
||||
.map(ToOwned::to_owned),
|
||||
endpoint_kind: api_format
|
||||
.as_deref()
|
||||
.and_then(infer_endpoint_kind)
|
||||
.map(ToOwned::to_owned),
|
||||
endpoint_api_format: endpoint_api_format.clone(),
|
||||
provider_api_family: endpoint_api_format
|
||||
.as_deref()
|
||||
.and_then(infer_api_family)
|
||||
.map(ToOwned::to_owned),
|
||||
provider_endpoint_kind: endpoint_api_format
|
||||
.as_deref()
|
||||
.and_then(infer_endpoint_kind)
|
||||
.map(ToOwned::to_owned),
|
||||
has_format_conversion: context_bool(context, "needs_conversion"),
|
||||
is_stream: Some(plan.stream),
|
||||
request_headers: context_value(context, "original_headers"),
|
||||
request_body: context_value(context, "original_request_body")
|
||||
.or_else(|| plan.body.json_body.clone()),
|
||||
provider_request_headers: context_value(context, "provider_request_headers")
|
||||
.or_else(|| Some(headers_to_json(&plan.headers))),
|
||||
provider_request_body: context_value(context, "provider_request_body")
|
||||
.or_else(|| plan.body.json_body.clone()),
|
||||
request_metadata: Some(Value::Object(Map::from_iter([
|
||||
(
|
||||
"request_id".to_string(),
|
||||
Value::String(plan.request_id.clone()),
|
||||
),
|
||||
(
|
||||
"candidate_id".to_string(),
|
||||
plan.candidate_id
|
||||
.clone()
|
||||
.map(Value::String)
|
||||
.unwrap_or(Value::Null),
|
||||
),
|
||||
]))),
|
||||
..UsageEventData::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_usage_data(base: UsageEventData, override_data: UsageEventData) -> UsageEventData {
|
||||
UsageEventData {
|
||||
user_id: override_data.user_id.or(base.user_id),
|
||||
api_key_id: override_data.api_key_id.or(base.api_key_id),
|
||||
username: override_data.username.or(base.username),
|
||||
api_key_name: override_data.api_key_name.or(base.api_key_name),
|
||||
provider_name: if override_data.provider_name.trim().is_empty() {
|
||||
base.provider_name
|
||||
} else {
|
||||
override_data.provider_name
|
||||
},
|
||||
model: if override_data.model.trim().is_empty() {
|
||||
base.model
|
||||
} else {
|
||||
override_data.model
|
||||
},
|
||||
target_model: override_data.target_model.or(base.target_model),
|
||||
provider_id: override_data.provider_id.or(base.provider_id),
|
||||
provider_endpoint_id: override_data
|
||||
.provider_endpoint_id
|
||||
.or(base.provider_endpoint_id),
|
||||
provider_api_key_id: override_data
|
||||
.provider_api_key_id
|
||||
.or(base.provider_api_key_id),
|
||||
request_type: override_data.request_type.or(base.request_type),
|
||||
api_format: override_data.api_format.or(base.api_format),
|
||||
api_family: override_data.api_family.or(base.api_family),
|
||||
endpoint_kind: override_data.endpoint_kind.or(base.endpoint_kind),
|
||||
endpoint_api_format: override_data
|
||||
.endpoint_api_format
|
||||
.or(base.endpoint_api_format),
|
||||
provider_api_family: override_data
|
||||
.provider_api_family
|
||||
.or(base.provider_api_family),
|
||||
provider_endpoint_kind: override_data
|
||||
.provider_endpoint_kind
|
||||
.or(base.provider_endpoint_kind),
|
||||
has_format_conversion: override_data
|
||||
.has_format_conversion
|
||||
.or(base.has_format_conversion),
|
||||
is_stream: override_data.is_stream.or(base.is_stream),
|
||||
input_tokens: override_data.input_tokens.or(base.input_tokens),
|
||||
output_tokens: override_data.output_tokens.or(base.output_tokens),
|
||||
total_tokens: override_data.total_tokens.or(base.total_tokens),
|
||||
cache_creation_input_tokens: override_data
|
||||
.cache_creation_input_tokens
|
||||
.or(base.cache_creation_input_tokens),
|
||||
cache_read_input_tokens: override_data
|
||||
.cache_read_input_tokens
|
||||
.or(base.cache_read_input_tokens),
|
||||
cache_creation_cost_usd: override_data
|
||||
.cache_creation_cost_usd
|
||||
.or(base.cache_creation_cost_usd),
|
||||
cache_read_cost_usd: override_data
|
||||
.cache_read_cost_usd
|
||||
.or(base.cache_read_cost_usd),
|
||||
output_price_per_1m: override_data
|
||||
.output_price_per_1m
|
||||
.or(base.output_price_per_1m),
|
||||
total_cost_usd: override_data.total_cost_usd.or(base.total_cost_usd),
|
||||
actual_total_cost_usd: override_data
|
||||
.actual_total_cost_usd
|
||||
.or(base.actual_total_cost_usd),
|
||||
status_code: override_data.status_code.or(base.status_code),
|
||||
error_message: override_data.error_message.or(base.error_message),
|
||||
error_category: override_data.error_category.or(base.error_category),
|
||||
response_time_ms: override_data.response_time_ms.or(base.response_time_ms),
|
||||
first_byte_time_ms: override_data.first_byte_time_ms.or(base.first_byte_time_ms),
|
||||
request_headers: override_data.request_headers.or(base.request_headers),
|
||||
request_body: override_data.request_body.or(base.request_body),
|
||||
provider_request_headers: override_data
|
||||
.provider_request_headers
|
||||
.or(base.provider_request_headers),
|
||||
provider_request_body: override_data
|
||||
.provider_request_body
|
||||
.or(base.provider_request_body),
|
||||
response_headers: override_data.response_headers.or(base.response_headers),
|
||||
response_body: override_data.response_body.or(base.response_body),
|
||||
client_response_headers: override_data
|
||||
.client_response_headers
|
||||
.or(base.client_response_headers),
|
||||
client_response_body: override_data
|
||||
.client_response_body
|
||||
.or(base.client_response_body),
|
||||
request_metadata: merge_json_value(base.request_metadata, override_data.request_metadata),
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_json_value(base: Option<Value>, override_value: Option<Value>) -> Option<Value> {
|
||||
match (base, override_value) {
|
||||
(Some(Value::Object(mut base)), Some(Value::Object(override_object))) => {
|
||||
for (key, value) in override_object {
|
||||
base.insert(key, value);
|
||||
}
|
||||
Some(Value::Object(base))
|
||||
}
|
||||
(Some(base), None) => Some(base),
|
||||
(_, Some(override_value)) => Some(override_value),
|
||||
(None, None) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn lifecycle_status_and_billing(state: UsageLifecycleState) -> (&'static str, &'static str) {
|
||||
match state {
|
||||
UsageLifecycleState::Pending => ("pending", "pending"),
|
||||
UsageLifecycleState::Streaming => ("streaming", "pending"),
|
||||
UsageLifecycleState::Completed => ("completed", "pending"),
|
||||
UsageLifecycleState::Failed => ("failed", "void"),
|
||||
UsageLifecycleState::Cancelled => ("cancelled", "void"),
|
||||
}
|
||||
}
|
||||
|
||||
fn context_string(context: Option<&Map<String, Value>>, key: &str) -> Option<String> {
|
||||
context
|
||||
.and_then(|value| value.get(key))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn context_bool(context: Option<&Map<String, Value>>, key: &str) -> Option<bool> {
|
||||
context
|
||||
.and_then(|value| value.get(key))
|
||||
.and_then(Value::as_bool)
|
||||
}
|
||||
|
||||
fn context_value(context: Option<&Map<String, Value>>, key: &str) -> Option<Value> {
|
||||
context.and_then(|value| value.get(key)).cloned()
|
||||
}
|
||||
|
||||
fn non_empty_string(value: Option<String>) -> Option<String> {
|
||||
value
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn infer_request_type(api_format: Option<&str>) -> String {
|
||||
match infer_endpoint_kind(api_format.unwrap_or_default()) {
|
||||
Some("video") => "video".to_string(),
|
||||
Some("image") => "image".to_string(),
|
||||
_ => "chat".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn infer_api_family(api_format: &str) -> Option<&str> {
|
||||
api_format.split_once(':').map(|(family, _)| family)
|
||||
}
|
||||
|
||||
fn infer_endpoint_kind(api_format: &str) -> Option<&str> {
|
||||
api_format.split_once(':').map(|(_, kind)| kind)
|
||||
}
|
||||
|
||||
fn resolve_billing_api_format(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<String> {
|
||||
let context = report_context.and_then(Value::as_object);
|
||||
context_string(context, "provider_api_format")
|
||||
.or_else(|| non_empty_string(Some(plan.provider_api_format.clone())))
|
||||
.or_else(|| context_string(context, "client_api_format"))
|
||||
.or_else(|| non_empty_string(Some(plan.client_api_format.clone())))
|
||||
}
|
||||
|
||||
fn apply_standardized_usage(
|
||||
api_format: Option<String>,
|
||||
response_body: &Value,
|
||||
data: &mut UsageEventData,
|
||||
) {
|
||||
let Some(api_format) = api_format.as_deref() else {
|
||||
return;
|
||||
};
|
||||
if !response_body.is_object() {
|
||||
return;
|
||||
}
|
||||
let usage = map_usage_from_response(response_body, api_format);
|
||||
if usage.input_tokens > 0 {
|
||||
data.input_tokens = Some(usage.input_tokens as u64);
|
||||
}
|
||||
if usage.output_tokens > 0 {
|
||||
data.output_tokens = Some(usage.output_tokens as u64);
|
||||
}
|
||||
if usage.cache_creation_tokens > 0 {
|
||||
data.cache_creation_input_tokens = Some(usage.cache_creation_tokens as u64);
|
||||
}
|
||||
if usage.cache_read_tokens > 0 {
|
||||
data.cache_read_input_tokens = Some(usage.cache_read_tokens as u64);
|
||||
}
|
||||
let total_tokens = usage
|
||||
.input_tokens
|
||||
.saturating_add(usage.output_tokens)
|
||||
.max(0) as u64;
|
||||
if total_tokens > 0 {
|
||||
data.total_tokens = Some(total_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
fn headers_to_json(headers: &BTreeMap<String, String>) -> Value {
|
||||
Value::Object(Map::from_iter(
|
||||
headers
|
||||
.iter()
|
||||
.map(|(key, value)| (key.clone(), Value::String(value.clone()))),
|
||||
))
|
||||
}
|
||||
|
||||
fn resolve_error_category(status_code: u16, event_type: UsageEventType) -> Option<String> {
|
||||
match event_type {
|
||||
UsageEventType::Cancelled => Some("cancelled".to_string()),
|
||||
UsageEventType::Failed if status_code >= 500 => Some("server_error".to_string()),
|
||||
UsageEventType::Failed if status_code >= 400 => Some("client_error".to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_error_message(
|
||||
status_code: u16,
|
||||
body_json: Option<&Value>,
|
||||
body_base64: Option<&str>,
|
||||
) -> Option<String> {
|
||||
let explicit_error_message = body_json
|
||||
.and_then(extract_explicit_error_message_from_json)
|
||||
.or_else(|| {
|
||||
body_base64
|
||||
.and_then(|value| decode_body_for_storage(Some(value)))
|
||||
.as_ref()
|
||||
.and_then(extract_explicit_error_message_from_json)
|
||||
});
|
||||
if explicit_error_message.is_some() {
|
||||
return explicit_error_message;
|
||||
}
|
||||
if status_code < 400 {
|
||||
return None;
|
||||
}
|
||||
|
||||
body_json
|
||||
.and_then(extract_generic_error_message_from_json)
|
||||
.or_else(|| {
|
||||
body_base64
|
||||
.and_then(|value| decode_body_for_storage(Some(value)))
|
||||
.as_ref()
|
||||
.and_then(extract_generic_error_message_from_json)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_explicit_error_message_from_json(value: &Value) -> Option<String> {
|
||||
value
|
||||
.get("error")
|
||||
.and_then(|error| error.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn extract_generic_error_message_from_json(value: &Value) -> Option<String> {
|
||||
extract_explicit_error_message_from_json(value).or_else(|| {
|
||||
value
|
||||
.get("message")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn decode_body_for_storage(body_base64: Option<&str>) -> Option<Value> {
|
||||
let body_base64 = body_base64?;
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.ok()?;
|
||||
if let Ok(json_body) = serde_json::from_slice::<Value>(&bytes) {
|
||||
return Some(json_body);
|
||||
}
|
||||
if let Ok(text) = String::from_utf8(bytes) {
|
||||
return Some(Value::String(text));
|
||||
}
|
||||
Some(Value::String(body_base64.to_string()))
|
||||
}
|
||||
|
||||
fn extract_token_counts_from_base64_payload(body_base64: &str) -> Option<(u64, u64, u64)> {
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.ok()?;
|
||||
if let Ok(json_body) = serde_json::from_slice::<Value>(&bytes) {
|
||||
return extract_token_counts_from_json(&json_body);
|
||||
}
|
||||
let text = String::from_utf8(bytes).ok()?;
|
||||
extract_token_counts_from_sse_text(&text)
|
||||
}
|
||||
|
||||
fn extract_token_counts_from_sse_text(text: &str) -> Option<(u64, u64, u64)> {
|
||||
let mut last_seen = None;
|
||||
for line in text.lines() {
|
||||
let line = line.trim();
|
||||
if line.is_empty() || !line.starts_with("data:") {
|
||||
continue;
|
||||
}
|
||||
let payload = line.trim_start_matches("data:").trim();
|
||||
if payload.is_empty() || payload == "[DONE]" {
|
||||
continue;
|
||||
}
|
||||
if let Ok(json_body) = serde_json::from_str::<Value>(payload) {
|
||||
if let Some(tokens) = extract_token_counts_from_json(&json_body) {
|
||||
last_seen = Some(tokens);
|
||||
}
|
||||
}
|
||||
}
|
||||
last_seen
|
||||
}
|
||||
|
||||
fn extract_token_counts_from_json(value: &Value) -> Option<(u64, u64, u64)> {
|
||||
if let Some(usage) = value.get("usage").and_then(Value::as_object) {
|
||||
let input = usage
|
||||
.get("input_tokens")
|
||||
.or_else(|| usage.get("prompt_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or_default();
|
||||
let output = usage
|
||||
.get("output_tokens")
|
||||
.or_else(|| usage.get("completion_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or_default();
|
||||
let total = usage
|
||||
.get("total_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(input + output);
|
||||
return Some((input, output, total));
|
||||
}
|
||||
|
||||
if let Some(usage) = value.get("usageMetadata").and_then(Value::as_object) {
|
||||
let input = usage
|
||||
.get("promptTokenCount")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or_default();
|
||||
let output = usage
|
||||
.get("candidatesTokenCount")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or_default();
|
||||
let total = usage
|
||||
.get("totalTokenCount")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(input + output);
|
||||
return Some((input, output, total));
|
||||
}
|
||||
|
||||
if let Some(response) = value.get("response") {
|
||||
return extract_token_counts_from_json(response);
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn empty_to_none(value: Option<String>) -> Option<String> {
|
||||
value
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{build_upsert_usage_record_from_event, extract_token_counts_from_json};
|
||||
use crate::gateway::usage::{UsageEvent, UsageEventData, UsageEventType};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn extracts_openai_usage_tokens() {
|
||||
let tokens = extract_token_counts_from_json(&json!({
|
||||
"usage": {
|
||||
"input_tokens": 3,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 8
|
||||
}
|
||||
}))
|
||||
.expect("tokens should exist");
|
||||
|
||||
assert_eq!(tokens, (3, 5, 8));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_upsert_record_from_terminal_event() {
|
||||
let record = build_upsert_usage_record_from_event(&UsageEvent {
|
||||
event_type: UsageEventType::Completed,
|
||||
request_id: "req-1".to_string(),
|
||||
timestamp_ms: 1_700_000_000_000,
|
||||
data: UsageEventData {
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
status_code: Some(200),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
})
|
||||
.expect("record should build");
|
||||
|
||||
assert_eq!(record.request_id, "req-1");
|
||||
assert_eq!(record.status, "completed");
|
||||
assert_eq!(record.billing_status, "pending");
|
||||
assert_eq!(record.total_tokens, Some(30));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user