refactor: 抽离 AI pipeline 与调度共享能力逻辑

This commit is contained in:
fawney19
2026-04-10 01:46:14 +08:00
parent b901a6ffc7
commit 5014e2f5fd
255 changed files with 15057 additions and 3115 deletions

View File

@@ -79,12 +79,12 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
.values()
.filter(|item| {
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
if item.created_at_unix_secs < created_from_unix_secs {
if item.created_at_unix_ms < created_from_unix_secs {
return false;
}
}
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
if item.created_at_unix_secs >= created_until_unix_secs {
if item.created_at_unix_ms >= created_until_unix_secs {
return false;
}
}
@@ -108,8 +108,8 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
.cloned()
.collect();
items.sort_by(|left, right| {
left.created_at_unix_secs
.cmp(&right.created_at_unix_secs)
left.created_at_unix_ms
.cmp(&right.created_at_unix_ms)
.then_with(|| left.request_id.cmp(&right.request_id))
});
Ok(items)
@@ -133,8 +133,8 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
.collect();
items.sort_by(|left, right| {
right
.created_at_unix_secs
.cmp(&left.created_at_unix_secs)
.created_at_unix_ms
.cmp(&left.created_at_unix_ms)
.then_with(|| left.id.cmp(&right.id))
});
items.truncate(limit);
@@ -209,10 +209,10 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
usage.validate()?;
let mut by_request_id = self.by_request_id.write().expect("usage repository lock");
let created_at_unix_secs = by_request_id
let created_at_unix_ms = by_request_id
.get(&usage.request_id)
.map(|existing| existing.created_at_unix_secs)
.or(usage.created_at_unix_secs)
.map(|existing| existing.created_at_unix_ms)
.or(usage.created_at_unix_ms)
.unwrap_or(usage.updated_at_unix_secs);
let total_tokens = usage
@@ -320,7 +320,7 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
request_metadata: usage
.request_metadata
.or_else(|| existing.and_then(|existing| existing.request_metadata.clone())),
created_at_unix_secs,
created_at_unix_ms,
updated_at_unix_secs: usage.updated_at_unix_secs,
finalized_at_unix_secs: usage.finalized_at_unix_secs,
};
@@ -339,7 +339,7 @@ mod tests {
};
use serde_json::json;
fn sample_usage(request_id: &str, created_at_unix_secs: i64) -> StoredRequestUsageAudit {
fn sample_usage(request_id: &str, created_at_unix_ms: i64) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
request_id.to_string(),
@@ -374,9 +374,9 @@ mod tests {
Some(120),
"completed".to_string(),
"settled".to_string(),
created_at_unix_secs,
created_at_unix_secs + 1,
Some(created_at_unix_secs + 2),
created_at_unix_ms,
created_at_unix_ms + 1,
Some(created_at_unix_ms + 2),
)
.expect("usage should build")
}
@@ -450,7 +450,7 @@ mod tests {
client_response_body: None,
request_metadata: None,
finalized_at_unix_secs: None,
created_at_unix_secs: Some(100),
created_at_unix_ms: Some(100),
updated_at_unix_secs: 101,
})
.await

View File

@@ -59,7 +59,7 @@ SELECT
client_response_headers,
client_response_body,
request_metadata,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
FROM "usage"
@@ -116,7 +116,7 @@ SELECT
client_response_headers,
client_response_body,
request_metadata,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
FROM "usage"
@@ -203,7 +203,7 @@ SELECT
NULL::jsonb AS client_response_headers,
NULL::jsonb AS client_response_body,
NULL::jsonb AS request_metadata,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
FROM "usage"
@@ -258,7 +258,7 @@ SELECT
NULL::jsonb AS client_response_headers,
NULL::jsonb AS client_response_body,
NULL::jsonb AS request_metadata,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
FROM "usage"
@@ -465,7 +465,7 @@ RETURNING
client_response_headers,
client_response_body,
request_metadata,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
"#;
@@ -713,7 +713,7 @@ impl SqlxUsageReadRepository {
.bind(&usage.client_response_body)
.bind(&usage.request_metadata)
.bind(usage.finalized_at_unix_secs.map(|value| value as f64))
.bind(usage.created_at_unix_secs.map(|value| value as f64))
.bind(usage.created_at_unix_ms.map(|value| value as f64))
.fetch_one(&mut **tx)
.await
.map_postgres_err()?;
@@ -816,7 +816,7 @@ fn map_usage_row(row: &sqlx::postgres::PgRow) -> Result<StoredRequestUsageAudit,
row.try_get("first_byte_time_ms").map_postgres_err()?,
row.try_get("status").map_postgres_err()?,
row.try_get("billing_status").map_postgres_err()?,
row.try_get("created_at_unix_secs").map_postgres_err()?,
row.try_get("created_at_unix_ms").map_postgres_err()?,
row.try_get("updated_at_unix_secs").map_postgres_err()?,
row.try_get("finalized_at_unix_secs").map_postgres_err()?,
)?;
@@ -953,7 +953,7 @@ mod tests {
client_response_body: None,
request_metadata: None,
finalized_at_unix_secs: None,
created_at_unix_secs: Some(100),
created_at_unix_ms: Some(100),
updated_at_unix_secs: 101,
})
.await;