refactor: 大规模模块拆分与代码精简,新增 ai-pipeline/data-contracts 独立 crate

- 新增 aether-ai-pipeline 和 aether-data-contracts crate,将 pipeline 逻辑与数据契约从 gateway 中解耦
- 重构 admin handlers:拆分单体模块为 auth/billing/endpoint/features/model/observability/provider/system 等独立子模块
- 合并 chat/cli 重复代码路径:精简 conversion、finalize、planner 中的 sync/chat/cli 分支
- 重构 scheduler/executor/data 层,引入 facade 模式降低模块间耦合
- 移除冗余的 intent 模块,将 plan_fallback/policy/stream_path/sync_path 迁移至 executor
- 前端适配:调整 admin API 调用和 provider 模型测试对话框
This commit is contained in:
fawney19
2026-04-07 02:50:19 +08:00
parent 763ff03a7b
commit 5d96d6673b
732 changed files with 28593 additions and 20666 deletions

View File

@@ -0,0 +1,13 @@
[package]
name = "aether-data-contracts"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Shared data contracts and repository traits for Aether Rust services"
[dependencies]
async-trait.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true

View File

@@ -0,0 +1,30 @@
#[derive(Debug, thiserror::Error)]
pub enum DataLayerError {
#[error("invalid configuration: {0}")]
InvalidConfiguration(String),
#[error("invalid input: {0}")]
InvalidInput(String),
#[error("postgres error: {0}")]
Postgres(String),
#[error("redis error: {0}")]
Redis(String),
#[error("operation timed out: {0}")]
TimedOut(String),
#[error("unexpected database value: {0}")]
UnexpectedValue(String),
}
impl DataLayerError {
pub fn postgres(error: impl std::fmt::Display) -> Self {
Self::Postgres(error.to_string())
}
pub fn redis(error: impl std::fmt::Display) -> Self {
Self::Redis(error.to_string())
}
}

View File

@@ -0,0 +1,4 @@
mod error;
pub mod repository;
pub use error::DataLayerError;

View File

@@ -0,0 +1,7 @@
mod types;
pub use types::{
AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingPresetApplyResult,
AdminBillingRuleRecord, AdminBillingRuleWriteInput, BillingReadRepository,
StoredBillingModelContext,
};

View File

@@ -0,0 +1,153 @@
use async_trait::async_trait;
use serde_json::Value;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredBillingModelContext {
pub provider_id: String,
pub provider_billing_type: Option<String>,
pub provider_api_key_id: Option<String>,
pub provider_api_key_rate_multipliers: Option<Value>,
pub provider_api_key_cache_ttl_minutes: Option<i64>,
pub global_model_id: String,
pub global_model_name: String,
pub global_model_config: Option<Value>,
pub default_price_per_request: Option<f64>,
pub default_tiered_pricing: Option<Value>,
pub model_id: Option<String>,
pub model_provider_model_name: Option<String>,
pub model_config: Option<Value>,
pub model_price_per_request: Option<f64>,
pub model_tiered_pricing: Option<Value>,
}
impl StoredBillingModelContext {
#[allow(clippy::too_many_arguments)]
pub fn new(
provider_id: String,
provider_billing_type: Option<String>,
provider_api_key_id: Option<String>,
provider_api_key_rate_multipliers: Option<Value>,
provider_api_key_cache_ttl_minutes: Option<i64>,
global_model_id: String,
global_model_name: String,
global_model_config: Option<Value>,
default_price_per_request: Option<f64>,
default_tiered_pricing: Option<Value>,
model_id: Option<String>,
model_provider_model_name: Option<String>,
model_config: Option<Value>,
model_price_per_request: Option<f64>,
model_tiered_pricing: Option<Value>,
) -> Result<Self, crate::DataLayerError> {
if provider_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"billing.provider_id is empty".to_string(),
));
}
if global_model_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"billing.global_model_id is empty".to_string(),
));
}
if global_model_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"billing.global_model_name is empty".to_string(),
));
}
Ok(Self {
provider_id,
provider_billing_type,
provider_api_key_id,
provider_api_key_rate_multipliers,
provider_api_key_cache_ttl_minutes,
global_model_id,
global_model_name,
global_model_config,
default_price_per_request,
default_tiered_pricing,
model_id,
model_provider_model_name,
model_config,
model_price_per_request,
model_tiered_pricing,
})
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminBillingRuleRecord {
pub id: String,
pub name: String,
pub task_type: String,
pub global_model_id: Option<String>,
pub model_id: Option<String>,
pub expression: String,
pub variables: Value,
pub dimension_mappings: Value,
pub is_enabled: bool,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AdminBillingRuleWriteInput {
pub name: String,
pub task_type: String,
pub global_model_id: Option<String>,
pub model_id: Option<String>,
pub expression: String,
pub variables: Value,
pub dimension_mappings: Value,
pub is_enabled: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminBillingCollectorRecord {
pub id: String,
pub api_format: String,
pub task_type: String,
pub dimension_name: String,
pub source_type: String,
pub source_path: Option<String>,
pub value_type: String,
pub transform_expression: Option<String>,
pub default_value: Option<String>,
pub priority: i32,
pub is_enabled: bool,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AdminBillingCollectorWriteInput {
pub api_format: String,
pub task_type: String,
pub dimension_name: String,
pub source_type: String,
pub source_path: Option<String>,
pub value_type: String,
pub transform_expression: Option<String>,
pub default_value: Option<String>,
pub priority: i32,
pub is_enabled: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminBillingPresetApplyResult {
pub preset: String,
pub mode: String,
pub created: u64,
pub updated: u64,
pub skipped: u64,
pub errors: Vec<String>,
}
#[async_trait]
pub trait BillingReadRepository: Send + Sync {
async fn find_model_context(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, crate::DataLayerError>;
}

View File

@@ -0,0 +1,6 @@
mod types;
pub use types::{
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};

View File

@@ -0,0 +1,83 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderModelMapping {
pub name: String,
pub priority: i32,
pub api_formats: Option<Vec<String>>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredMinimalCandidateSelectionRow {
pub provider_id: String,
pub provider_name: String,
pub provider_type: String,
pub provider_priority: i32,
pub provider_is_active: bool,
pub endpoint_id: String,
pub endpoint_api_format: String,
pub endpoint_api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub endpoint_is_active: bool,
pub key_id: String,
pub key_name: String,
pub key_auth_type: String,
pub key_is_active: bool,
pub key_api_formats: Option<Vec<String>>,
pub key_allowed_models: Option<Vec<String>>,
pub key_capabilities: Option<serde_json::Value>,
pub key_internal_priority: i32,
pub key_global_priority_by_format: Option<serde_json::Value>,
pub model_id: String,
pub global_model_id: String,
pub global_model_name: String,
pub global_model_mappings: Option<Vec<String>>,
pub global_model_supports_streaming: Option<bool>,
pub model_provider_model_name: String,
pub model_provider_model_mappings: Option<Vec<StoredProviderModelMapping>>,
pub model_supports_streaming: Option<bool>,
pub model_is_active: bool,
pub model_is_available: bool,
}
impl StoredMinimalCandidateSelectionRow {
pub fn supports_streaming(&self) -> bool {
self.model_supports_streaming
.or(self.global_model_supports_streaming)
.unwrap_or(true)
}
pub fn key_supports_api_format(&self, api_format: &str) -> bool {
let target = api_format.trim();
match self.key_api_formats.as_deref() {
None => true,
Some(formats) => formats
.iter()
.any(|value| value.eq_ignore_ascii_case(target)),
}
}
}
#[async_trait]
pub trait MinimalCandidateSelectionReadRepository: Send + Sync {
async fn list_for_exact_api_format(
&self,
api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
async fn list_for_exact_api_format_and_global_model(
&self,
api_format: &str,
global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
}
pub trait MinimalCandidateSelectionRepository:
MinimalCandidateSelectionReadRepository + Send + Sync
{
}
impl<T> MinimalCandidateSelectionRepository for T where
T: MinimalCandidateSelectionReadRepository + Send + Sync
{
}

View File

@@ -0,0 +1,9 @@
mod types;
pub use types::{
build_decision_trace, derive_request_candidate_final_status, DecisionTrace,
DecisionTraceCandidate, PublicHealthStatusCount, PublicHealthTimelineBucket,
RequestCandidateFinalStatus, RequestCandidateReadRepository, RequestCandidateRepository,
RequestCandidateStatus, RequestCandidateTrace, RequestCandidateWriteRepository,
StoredRequestCandidate, UpsertRequestCandidateRecord,
};

View File

@@ -0,0 +1,522 @@
use std::collections::BTreeMap;
use async_trait::async_trait;
use crate::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RequestCandidateStatus {
Available,
Unused,
Pending,
Streaming,
Success,
Failed,
Cancelled,
Skipped,
}
impl RequestCandidateStatus {
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"available" => Ok(Self::Available),
"unused" => Ok(Self::Unused),
"pending" => Ok(Self::Pending),
"streaming" => Ok(Self::Streaming),
"success" => Ok(Self::Success),
"failed" => Ok(Self::Failed),
"cancelled" => Ok(Self::Cancelled),
"skipped" => Ok(Self::Skipped),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unsupported request_candidates.status: {other}"
))),
}
}
pub fn is_attempted(self, started_at_unix_secs: Option<u64>) -> bool {
match self {
Self::Available | Self::Unused | Self::Skipped => false,
Self::Pending => started_at_unix_secs.is_some(),
Self::Streaming | Self::Success | Self::Failed | Self::Cancelled => true,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredRequestCandidate {
pub id: String,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub candidate_index: u32,
pub retry_index: u32,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub status: RequestCandidateStatus,
pub skip_reason: Option<String>,
pub is_cached: bool,
pub status_code: Option<u16>,
pub error_type: Option<String>,
pub error_message: Option<String>,
pub latency_ms: Option<u64>,
pub concurrent_requests: Option<u32>,
pub extra_data: Option<serde_json::Value>,
pub required_capabilities: Option<serde_json::Value>,
pub created_at_unix_secs: u64,
pub started_at_unix_secs: Option<u64>,
pub finished_at_unix_secs: Option<u64>,
}
impl StoredRequestCandidate {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
request_id: String,
user_id: Option<String>,
api_key_id: Option<String>,
username: Option<String>,
api_key_name: Option<String>,
candidate_index: i32,
retry_index: i32,
provider_id: Option<String>,
endpoint_id: Option<String>,
key_id: Option<String>,
status: RequestCandidateStatus,
skip_reason: Option<String>,
is_cached: bool,
status_code: Option<i32>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<i32>,
concurrent_requests: Option<i32>,
extra_data: Option<serde_json::Value>,
required_capabilities: Option<serde_json::Value>,
created_at_unix_secs: i64,
started_at_unix_secs: Option<i64>,
finished_at_unix_secs: Option<i64>,
) -> Result<Self, crate::DataLayerError> {
let candidate_index = u32::try_from(candidate_index).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.candidate_index: {candidate_index}"
))
})?;
let retry_index = u32::try_from(retry_index).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.retry_index: {retry_index}"
))
})?;
let status_code = status_code
.map(|value| {
u16::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.status_code: {value}"
))
})
})
.transpose()?;
let latency_ms = latency_ms
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.latency_ms: {value}"
))
})
})
.transpose()?;
let concurrent_requests = concurrent_requests
.map(|value| {
u32::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.concurrent_requests: {value}"
))
})
})
.transpose()?;
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.created_at_unix_secs: {created_at_unix_secs}"
))
})?;
let started_at_unix_secs = started_at_unix_secs
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.started_at_unix_secs: {value}"
))
})
})
.transpose()?;
let finished_at_unix_secs = finished_at_unix_secs
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.finished_at_unix_secs: {value}"
))
})
})
.transpose()?;
Ok(Self {
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
status,
skip_reason,
is_cached,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests,
extra_data,
required_capabilities,
created_at_unix_secs,
started_at_unix_secs,
finished_at_unix_secs,
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RequestCandidateFinalStatus {
Success,
Failed,
Cancelled,
Streaming,
Pending,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct RequestCandidateTrace {
pub request_id: String,
pub total_candidates: usize,
pub final_status: RequestCandidateFinalStatus,
pub total_latency_ms: u64,
pub candidates: Vec<StoredRequestCandidate>,
}
impl RequestCandidateTrace {
pub fn from_candidates(
request_id: impl Into<String>,
all_candidates: Vec<StoredRequestCandidate>,
attempted_only: bool,
) -> Option<Self> {
if all_candidates.is_empty() {
return None;
}
let candidates = if attempted_only {
all_candidates
.iter()
.filter(|candidate| {
candidate
.status
.is_attempted(candidate.started_at_unix_secs)
})
.cloned()
.collect::<Vec<_>>()
} else {
all_candidates.clone()
};
let total_latency_ms = candidates
.iter()
.filter(|candidate| {
matches!(
candidate.status,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
) && candidate.latency_ms.is_some()
})
.map(|candidate| candidate.latency_ms.unwrap_or(0))
.sum();
let final_status_source = if attempted_only && candidates.is_empty() {
&all_candidates
} else {
&candidates
};
Some(Self {
request_id: request_id.into(),
total_candidates: candidates.len(),
final_status: derive_request_candidate_final_status(final_status_source),
total_latency_ms,
candidates,
})
}
}
pub fn derive_request_candidate_final_status(
candidates: &[StoredRequestCandidate],
) -> RequestCandidateFinalStatus {
let has_success = candidates.iter().any(|candidate| {
candidate.status == RequestCandidateStatus::Success
|| matches!(candidate.status_code, Some(status_code) if (200..300).contains(&status_code))
});
if has_success {
return RequestCandidateFinalStatus::Success;
}
if candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Streaming)
{
return RequestCandidateFinalStatus::Streaming;
}
if candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Pending)
{
return RequestCandidateFinalStatus::Pending;
}
let has_cancelled = candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Cancelled);
let has_failed = candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Failed);
if has_cancelled && !has_failed {
return RequestCandidateFinalStatus::Cancelled;
}
RequestCandidateFinalStatus::Failed
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DecisionTraceCandidate {
#[serde(flatten)]
pub candidate: StoredRequestCandidate,
pub provider_name: Option<String>,
pub provider_website: Option<String>,
pub provider_type: Option<String>,
pub endpoint_api_format: Option<String>,
pub endpoint_api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub provider_key_name: Option<String>,
pub provider_key_auth_type: Option<String>,
pub provider_key_capabilities: Option<serde_json::Value>,
pub provider_key_is_active: Option<bool>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DecisionTrace {
pub request_id: String,
pub total_candidates: usize,
pub final_status: RequestCandidateFinalStatus,
pub total_latency_ms: u64,
pub candidates: Vec<DecisionTraceCandidate>,
}
pub fn build_decision_trace(
trace: RequestCandidateTrace,
providers: Vec<StoredProviderCatalogProvider>,
endpoints: Vec<StoredProviderCatalogEndpoint>,
keys: Vec<StoredProviderCatalogKey>,
) -> DecisionTrace {
let provider_map = providers
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
let endpoint_map = endpoints
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
let key_map = keys
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
DecisionTrace {
request_id: trace.request_id,
total_candidates: trace.total_candidates,
final_status: trace.final_status,
total_latency_ms: trace.total_latency_ms,
candidates: trace
.candidates
.into_iter()
.map(|candidate| {
enrich_decision_trace_candidate(candidate, &provider_map, &endpoint_map, &key_map)
})
.collect(),
}
}
fn enrich_decision_trace_candidate(
candidate: StoredRequestCandidate,
provider_map: &BTreeMap<String, StoredProviderCatalogProvider>,
endpoint_map: &BTreeMap<String, StoredProviderCatalogEndpoint>,
key_map: &BTreeMap<String, StoredProviderCatalogKey>,
) -> DecisionTraceCandidate {
let provider = candidate
.provider_id
.as_ref()
.and_then(|provider_id| provider_map.get(provider_id));
let endpoint = candidate
.endpoint_id
.as_ref()
.and_then(|endpoint_id| endpoint_map.get(endpoint_id));
let provider_key = candidate
.key_id
.as_ref()
.and_then(|key_id| key_map.get(key_id));
DecisionTraceCandidate {
provider_name: provider.map(|item| item.name.clone()),
provider_website: provider.and_then(|item| item.website.clone()),
provider_type: provider.map(|item| item.provider_type.clone()),
endpoint_api_format: endpoint.map(|item| item.api_format.clone()),
endpoint_api_family: endpoint.and_then(|item| item.api_family.clone()),
endpoint_kind: endpoint.and_then(|item| item.endpoint_kind.clone()),
provider_key_name: provider_key
.map(|item| item.name.clone())
.or_else(|| candidate.api_key_name.clone()),
provider_key_auth_type: provider_key.map(|item| item.auth_type.clone()),
provider_key_capabilities: provider_key.and_then(|item| item.capabilities.clone()),
provider_key_is_active: provider_key.map(|item| item.is_active),
candidate,
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PublicHealthStatusCount {
pub endpoint_id: String,
pub status: RequestCandidateStatus,
pub count: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PublicHealthTimelineBucket {
pub endpoint_id: String,
pub segment_idx: u32,
pub total_count: u64,
pub success_count: u64,
pub failed_count: u64,
pub min_created_at_unix_secs: Option<u64>,
pub max_created_at_unix_secs: Option<u64>,
}
#[async_trait]
pub trait RequestCandidateReadRepository: Send + Sync {
async fn list_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
async fn list_by_provider_id(
&self,
provider_id: &str,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
async fn list_finalized_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
async fn count_finalized_statuses_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
) -> Result<Vec<PublicHealthStatusCount>, crate::DataLayerError>;
async fn aggregate_finalized_timeline_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
until_unix_secs: u64,
segments: u32,
) -> Result<Vec<PublicHealthTimelineBucket>, crate::DataLayerError>;
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UpsertRequestCandidateRecord {
pub id: String,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub candidate_index: u32,
pub retry_index: u32,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub status: RequestCandidateStatus,
pub skip_reason: Option<String>,
pub is_cached: Option<bool>,
pub status_code: Option<u16>,
pub error_type: Option<String>,
pub error_message: Option<String>,
pub latency_ms: Option<u64>,
pub concurrent_requests: Option<u32>,
pub extra_data: Option<serde_json::Value>,
pub required_capabilities: Option<serde_json::Value>,
pub created_at_unix_secs: Option<u64>,
pub started_at_unix_secs: Option<u64>,
pub finished_at_unix_secs: Option<u64>,
}
impl UpsertRequestCandidateRecord {
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
if self.id.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"request candidate upsert id cannot be empty".to_string(),
));
}
if self.request_id.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"request candidate upsert request_id cannot be empty".to_string(),
));
}
Ok(())
}
}
#[async_trait]
pub trait RequestCandidateWriteRepository: Send + Sync {
async fn upsert(
&self,
candidate: UpsertRequestCandidateRecord,
) -> Result<StoredRequestCandidate, crate::DataLayerError>;
async fn delete_created_before(
&self,
created_before_unix_secs: u64,
limit: usize,
) -> Result<usize, crate::DataLayerError>;
}
pub trait RequestCandidateRepository:
RequestCandidateReadRepository + RequestCandidateWriteRepository + Send + Sync
{
}
impl<T> RequestCandidateRepository for T where
T: RequestCandidateReadRepository + RequestCandidateWriteRepository + Send + Sync
{
}

View File

@@ -0,0 +1,10 @@
mod types;
pub use types::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
GlobalModelReadRepository, GlobalModelWriteRepository, PublicCatalogModelListQuery,
PublicCatalogModelSearchQuery, PublicGlobalModelQuery, StoredAdminGlobalModel,
StoredAdminGlobalModelPage, StoredAdminProviderModel, StoredProviderActiveGlobalModel,
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};

View File

@@ -0,0 +1,688 @@
use async_trait::async_trait;
use serde_json::Value;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredPublicGlobalModel {
pub id: String,
pub name: String,
pub display_name: Option<String>,
pub is_active: bool,
pub default_price_per_request: Option<f64>,
pub default_tiered_pricing: Option<Value>,
pub supported_capabilities: Option<Value>,
pub config: Option<Value>,
pub usage_count: u64,
}
impl StoredPublicGlobalModel {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
name: String,
display_name: Option<String>,
is_active: bool,
default_price_per_request: Option<f64>,
default_tiered_pricing: Option<Value>,
supported_capabilities: Option<Value>,
config: Option<Value>,
usage_count: u64,
) -> Result<Self, crate::DataLayerError> {
if id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.id is empty".to_string(),
));
}
if name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.name is empty".to_string(),
));
}
Ok(Self {
id,
name,
display_name,
is_active,
default_price_per_request,
default_tiered_pricing,
supported_capabilities,
config,
usage_count,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct PublicGlobalModelQuery {
pub offset: usize,
pub limit: usize,
pub is_active: Option<bool>,
pub search: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredPublicCatalogModel {
pub id: String,
pub provider_id: String,
pub provider_name: String,
pub provider_model_name: String,
pub name: String,
pub display_name: String,
pub description: Option<String>,
pub icon_url: Option<String>,
pub input_price_per_1m: Option<f64>,
pub output_price_per_1m: Option<f64>,
pub cache_creation_price_per_1m: Option<f64>,
pub cache_read_price_per_1m: Option<f64>,
pub supports_vision: Option<bool>,
pub supports_function_calling: Option<bool>,
pub supports_streaming: Option<bool>,
pub is_active: bool,
}
impl StoredPublicCatalogModel {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
provider_id: String,
provider_name: String,
provider_model_name: String,
name: String,
display_name: String,
description: Option<String>,
icon_url: Option<String>,
input_price_per_1m: Option<f64>,
output_price_per_1m: Option<f64>,
cache_creation_price_per_1m: Option<f64>,
cache_read_price_per_1m: Option<f64>,
supports_vision: Option<bool>,
supports_function_calling: Option<bool>,
supports_streaming: Option<bool>,
is_active: bool,
) -> Result<Self, crate::DataLayerError> {
if id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.id is empty".to_string(),
));
}
if provider_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.provider_id is empty".to_string(),
));
}
if provider_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"providers.name is empty".to_string(),
));
}
if provider_model_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.provider_model_name is empty".to_string(),
));
}
if name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"public model name is empty".to_string(),
));
}
if display_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"public model display_name is empty".to_string(),
));
}
Ok(Self {
id,
provider_id,
provider_name,
provider_model_name,
name,
display_name,
description,
icon_url,
input_price_per_1m,
output_price_per_1m,
cache_creation_price_per_1m,
cache_read_price_per_1m,
supports_vision,
supports_function_calling,
supports_streaming,
is_active,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct PublicCatalogModelListQuery {
pub provider_id: Option<String>,
pub offset: usize,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PublicCatalogModelSearchQuery {
pub search: String,
pub provider_id: Option<String>,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AdminProviderModelListQuery {
pub provider_id: String,
pub is_active: Option<bool>,
pub offset: usize,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminGlobalModel {
pub id: String,
pub name: String,
pub display_name: String,
pub is_active: bool,
pub default_price_per_request: Option<f64>,
pub default_tiered_pricing: Option<Value>,
pub supported_capabilities: Option<Value>,
pub config: Option<Value>,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
}
impl StoredAdminGlobalModel {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
name: String,
display_name: String,
is_active: bool,
default_price_per_request: Option<f64>,
default_tiered_pricing: Option<Value>,
supported_capabilities: Option<Value>,
config: Option<Value>,
created_at_unix_secs: Option<u64>,
updated_at_unix_secs: Option<u64>,
) -> Result<Self, crate::DataLayerError> {
if id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.id is empty".to_string(),
));
}
if name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.name is empty".to_string(),
));
}
if display_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.display_name is empty".to_string(),
));
}
Ok(Self {
id,
name,
display_name,
is_active,
default_price_per_request,
default_tiered_pricing,
supported_capabilities,
config,
created_at_unix_secs,
updated_at_unix_secs,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct AdminGlobalModelListQuery {
pub offset: usize,
pub limit: usize,
pub is_active: Option<bool>,
pub search: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminProviderModel {
pub id: String,
pub provider_id: String,
pub global_model_id: String,
pub provider_model_name: String,
pub provider_model_mappings: Option<Value>,
pub price_per_request: Option<f64>,
pub tiered_pricing: Option<Value>,
pub supports_vision: Option<bool>,
pub supports_function_calling: Option<bool>,
pub supports_streaming: Option<bool>,
pub supports_extended_thinking: Option<bool>,
pub supports_image_generation: Option<bool>,
pub is_active: bool,
pub is_available: bool,
pub config: Option<Value>,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
pub global_model_name: Option<String>,
pub global_model_display_name: Option<String>,
pub global_model_default_price_per_request: Option<f64>,
pub global_model_default_tiered_pricing: Option<Value>,
pub global_model_config: Option<Value>,
}
impl StoredAdminProviderModel {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
provider_id: String,
global_model_id: String,
provider_model_name: String,
provider_model_mappings: Option<Value>,
price_per_request: Option<f64>,
tiered_pricing: Option<Value>,
supports_vision: Option<bool>,
supports_function_calling: Option<bool>,
supports_streaming: Option<bool>,
supports_extended_thinking: Option<bool>,
supports_image_generation: Option<bool>,
is_active: bool,
is_available: bool,
config: Option<Value>,
created_at_unix_secs: Option<u64>,
updated_at_unix_secs: Option<u64>,
global_model_name: Option<String>,
global_model_display_name: Option<String>,
global_model_default_price_per_request: Option<f64>,
global_model_default_tiered_pricing: Option<Value>,
global_model_config: Option<Value>,
) -> Result<Self, crate::DataLayerError> {
if id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.id is empty".to_string(),
));
}
if provider_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.provider_id is empty".to_string(),
));
}
if global_model_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.global_model_id is empty".to_string(),
));
}
if provider_model_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.provider_model_name is empty".to_string(),
));
}
Ok(Self {
id,
provider_id,
global_model_id,
provider_model_name,
provider_model_mappings,
price_per_request,
tiered_pricing,
supports_vision,
supports_function_calling,
supports_streaming,
supports_extended_thinking,
supports_image_generation,
is_active,
is_available,
config,
created_at_unix_secs,
updated_at_unix_secs,
global_model_name,
global_model_display_name,
global_model_default_price_per_request,
global_model_default_tiered_pricing,
global_model_config,
})
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UpsertAdminProviderModelRecord {
pub id: String,
pub provider_id: String,
pub global_model_id: String,
pub provider_model_name: String,
pub provider_model_mappings: Option<Value>,
pub price_per_request: Option<f64>,
pub tiered_pricing: Option<Value>,
pub supports_vision: Option<bool>,
pub supports_function_calling: Option<bool>,
pub supports_streaming: Option<bool>,
pub supports_extended_thinking: Option<bool>,
pub supports_image_generation: Option<bool>,
pub is_active: bool,
pub is_available: bool,
pub config: Option<Value>,
}
impl UpsertAdminProviderModelRecord {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
provider_id: String,
global_model_id: String,
provider_model_name: String,
provider_model_mappings: Option<Value>,
price_per_request: Option<f64>,
tiered_pricing: Option<Value>,
supports_vision: Option<bool>,
supports_function_calling: Option<bool>,
supports_streaming: Option<bool>,
supports_extended_thinking: Option<bool>,
supports_image_generation: Option<bool>,
is_active: bool,
is_available: bool,
config: Option<Value>,
) -> Result<Self, crate::DataLayerError> {
if id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.id is empty".to_string(),
));
}
if provider_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.provider_id is empty".to_string(),
));
}
if global_model_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.global_model_id is empty".to_string(),
));
}
if provider_model_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"models.provider_model_name is empty".to_string(),
));
}
Ok(Self {
id,
provider_id,
global_model_id,
provider_model_name,
provider_model_mappings,
price_per_request,
tiered_pricing,
supports_vision,
supports_function_calling,
supports_streaming,
supports_extended_thinking,
supports_image_generation,
is_active,
is_available,
config,
})
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CreateAdminGlobalModelRecord {
pub id: String,
pub name: String,
pub display_name: String,
pub is_active: bool,
pub default_price_per_request: Option<f64>,
pub default_tiered_pricing: Option<Value>,
pub supported_capabilities: Option<Value>,
pub config: Option<Value>,
}
impl CreateAdminGlobalModelRecord {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
name: String,
display_name: String,
is_active: bool,
default_price_per_request: Option<f64>,
default_tiered_pricing: Option<Value>,
supported_capabilities: Option<Value>,
config: Option<Value>,
) -> Result<Self, crate::DataLayerError> {
if id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.id is empty".to_string(),
));
}
if name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.name is empty".to_string(),
));
}
if display_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.display_name is empty".to_string(),
));
}
Ok(Self {
id,
name,
display_name,
is_active,
default_price_per_request,
default_tiered_pricing,
supported_capabilities,
config,
})
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UpdateAdminGlobalModelRecord {
pub id: String,
pub display_name: String,
pub is_active: bool,
pub default_price_per_request: Option<f64>,
pub default_tiered_pricing: Option<Value>,
pub supported_capabilities: Option<Value>,
pub config: Option<Value>,
}
impl UpdateAdminGlobalModelRecord {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
display_name: String,
is_active: bool,
default_price_per_request: Option<f64>,
default_tiered_pricing: Option<Value>,
supported_capabilities: Option<Value>,
config: Option<Value>,
) -> Result<Self, crate::DataLayerError> {
if id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.id is empty".to_string(),
));
}
if display_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"global_models.display_name is empty".to_string(),
));
}
Ok(Self {
id,
display_name,
is_active,
default_price_per_request,
default_tiered_pricing,
supported_capabilities,
config,
})
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredPublicGlobalModelPage {
pub items: Vec<StoredPublicGlobalModel>,
pub total: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminGlobalModelPage {
pub items: Vec<StoredAdminGlobalModel>,
pub total: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderModelStats {
pub provider_id: String,
pub total_models: u64,
pub active_models: u64,
}
impl StoredProviderModelStats {
pub fn new(
provider_id: String,
total_models: i64,
active_models: i64,
) -> Result<Self, crate::DataLayerError> {
if provider_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider model stats provider_id is empty".to_string(),
));
}
if total_models < 0 || active_models < 0 {
return Err(crate::DataLayerError::UnexpectedValue(
"provider model stats count is negative".to_string(),
));
}
Ok(Self {
provider_id,
total_models: total_models as u64,
active_models: active_models as u64,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderActiveGlobalModel {
pub provider_id: String,
pub global_model_id: String,
}
impl StoredProviderActiveGlobalModel {
pub fn new(
provider_id: String,
global_model_id: String,
) -> Result<Self, crate::DataLayerError> {
if provider_id.trim().is_empty() || global_model_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider active global model identity is empty".to_string(),
));
}
Ok(Self {
provider_id,
global_model_id,
})
}
}
#[async_trait]
pub trait GlobalModelReadRepository: Send + Sync {
async fn list_public_models(
&self,
query: &PublicGlobalModelQuery,
) -> Result<StoredPublicGlobalModelPage, crate::DataLayerError>;
async fn get_public_model_by_name(
&self,
model_name: &str,
) -> Result<Option<StoredPublicGlobalModel>, crate::DataLayerError>;
async fn list_public_catalog_models(
&self,
query: &PublicCatalogModelListQuery,
) -> Result<Vec<StoredPublicCatalogModel>, crate::DataLayerError>;
async fn search_public_catalog_models(
&self,
query: &PublicCatalogModelSearchQuery,
) -> Result<Vec<StoredPublicCatalogModel>, crate::DataLayerError>;
async fn list_admin_global_models(
&self,
query: &AdminGlobalModelListQuery,
) -> Result<StoredAdminGlobalModelPage, crate::DataLayerError>;
async fn list_admin_provider_models(
&self,
query: &AdminProviderModelListQuery,
) -> Result<Vec<StoredAdminProviderModel>, crate::DataLayerError>;
async fn list_admin_provider_available_source_models(
&self,
provider_id: &str,
) -> Result<Vec<StoredAdminProviderModel>, crate::DataLayerError>;
async fn get_admin_provider_model(
&self,
provider_id: &str,
model_id: &str,
) -> Result<Option<StoredAdminProviderModel>, crate::DataLayerError>;
async fn get_admin_global_model_by_id(
&self,
global_model_id: &str,
) -> Result<Option<StoredAdminGlobalModel>, crate::DataLayerError>;
async fn get_admin_global_model_by_name(
&self,
model_name: &str,
) -> Result<Option<StoredAdminGlobalModel>, crate::DataLayerError>;
async fn list_admin_provider_models_by_global_model_id(
&self,
global_model_id: &str,
) -> Result<Vec<StoredAdminProviderModel>, crate::DataLayerError>;
async fn list_provider_model_stats(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderModelStats>, crate::DataLayerError>;
async fn list_active_global_model_ids_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderActiveGlobalModel>, crate::DataLayerError>;
}
#[async_trait]
pub trait GlobalModelWriteRepository: Send + Sync {
async fn create_admin_provider_model(
&self,
record: &UpsertAdminProviderModelRecord,
) -> Result<Option<StoredAdminProviderModel>, crate::DataLayerError>;
async fn update_admin_provider_model(
&self,
record: &UpsertAdminProviderModelRecord,
) -> Result<Option<StoredAdminProviderModel>, crate::DataLayerError>;
async fn delete_admin_provider_model(
&self,
provider_id: &str,
model_id: &str,
) -> Result<bool, crate::DataLayerError>;
async fn create_admin_global_model(
&self,
record: &CreateAdminGlobalModelRecord,
) -> Result<Option<StoredAdminGlobalModel>, crate::DataLayerError>;
async fn update_admin_global_model(
&self,
record: &UpdateAdminGlobalModelRecord,
) -> Result<Option<StoredAdminGlobalModel>, crate::DataLayerError>;
async fn delete_admin_global_model(
&self,
global_model_id: &str,
) -> Result<bool, crate::DataLayerError>;
}

View File

@@ -0,0 +1,9 @@
pub mod billing;
pub mod candidate_selection;
pub mod candidates;
pub mod global_models;
pub mod provider_catalog;
pub mod quota;
pub mod settlement;
pub mod usage;
pub mod video_tasks;

View File

@@ -0,0 +1,7 @@
mod types;
pub use types::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};

View File

@@ -0,0 +1,595 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderCatalogProvider {
pub id: String,
pub name: String,
pub description: Option<String>,
pub website: Option<String>,
pub provider_type: String,
pub billing_type: Option<String>,
pub monthly_quota_usd: Option<f64>,
pub monthly_used_usd: Option<f64>,
pub quota_reset_day: Option<u64>,
pub quota_last_reset_at_unix_secs: Option<u64>,
pub quota_expires_at_unix_secs: Option<u64>,
pub provider_priority: i32,
pub is_active: bool,
pub keep_priority_on_conversion: bool,
pub enable_format_conversion: bool,
pub concurrent_limit: Option<i32>,
pub max_retries: Option<i32>,
pub proxy: Option<serde_json::Value>,
pub request_timeout_secs: Option<f64>,
pub stream_first_byte_timeout_secs: Option<f64>,
pub config: Option<serde_json::Value>,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
}
impl StoredProviderCatalogProvider {
pub fn new(
id: String,
name: String,
website: Option<String>,
provider_type: String,
) -> Result<Self, crate::DataLayerError> {
if name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"providers.name is empty".to_string(),
));
}
if provider_type.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"providers.provider_type is empty".to_string(),
));
}
Ok(Self {
id,
name,
description: None,
website,
provider_type,
billing_type: None,
monthly_quota_usd: None,
monthly_used_usd: None,
quota_reset_day: None,
quota_last_reset_at_unix_secs: None,
quota_expires_at_unix_secs: None,
provider_priority: 0,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
created_at_unix_secs: None,
updated_at_unix_secs: None,
})
}
#[allow(clippy::too_many_arguments)]
pub fn with_transport_fields(
mut self,
is_active: bool,
keep_priority_on_conversion: bool,
enable_format_conversion: bool,
concurrent_limit: Option<i32>,
max_retries: Option<i32>,
proxy: Option<serde_json::Value>,
request_timeout_secs: Option<f64>,
stream_first_byte_timeout_secs: Option<f64>,
config: Option<serde_json::Value>,
) -> Self {
self.is_active = is_active;
self.keep_priority_on_conversion = keep_priority_on_conversion;
self.enable_format_conversion = enable_format_conversion;
self.concurrent_limit = concurrent_limit;
self.max_retries = max_retries;
self.proxy = proxy;
self.request_timeout_secs = request_timeout_secs;
self.stream_first_byte_timeout_secs = stream_first_byte_timeout_secs;
self.config = config;
self
}
pub fn with_description(mut self, description: Option<String>) -> Self {
self.description = description;
self
}
#[allow(clippy::too_many_arguments)]
pub fn with_billing_fields(
mut self,
billing_type: Option<String>,
monthly_quota_usd: Option<f64>,
monthly_used_usd: Option<f64>,
quota_reset_day: Option<u64>,
quota_last_reset_at_unix_secs: Option<u64>,
quota_expires_at_unix_secs: Option<u64>,
) -> Self {
self.billing_type = billing_type;
self.monthly_quota_usd = monthly_quota_usd;
self.monthly_used_usd = monthly_used_usd;
self.quota_reset_day = quota_reset_day;
self.quota_last_reset_at_unix_secs = quota_last_reset_at_unix_secs;
self.quota_expires_at_unix_secs = quota_expires_at_unix_secs;
self
}
pub fn with_routing_fields(mut self, provider_priority: i32) -> Self {
self.provider_priority = provider_priority;
self
}
pub fn with_timestamps(
mut self,
created_at_unix_secs: Option<u64>,
updated_at_unix_secs: Option<u64>,
) -> Self {
self.created_at_unix_secs = created_at_unix_secs;
self.updated_at_unix_secs = updated_at_unix_secs;
self
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderCatalogEndpoint {
pub id: String,
pub provider_id: String,
pub api_format: String,
pub api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub is_active: bool,
pub health_score: f64,
pub base_url: String,
pub header_rules: Option<serde_json::Value>,
pub body_rules: Option<serde_json::Value>,
pub max_retries: Option<i32>,
pub custom_path: Option<String>,
pub config: Option<serde_json::Value>,
pub format_acceptance_config: Option<serde_json::Value>,
pub proxy: Option<serde_json::Value>,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
}
impl StoredProviderCatalogEndpoint {
pub fn new(
id: String,
provider_id: String,
api_format: String,
api_family: Option<String>,
endpoint_kind: Option<String>,
is_active: bool,
) -> Result<Self, crate::DataLayerError> {
if api_format.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider_endpoints.api_format is empty".to_string(),
));
}
Ok(Self {
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active,
health_score: 1.0,
base_url: String::new(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
created_at_unix_secs: None,
updated_at_unix_secs: None,
})
}
#[allow(clippy::too_many_arguments)]
pub fn with_transport_fields(
mut self,
base_url: String,
header_rules: Option<serde_json::Value>,
body_rules: Option<serde_json::Value>,
max_retries: Option<i32>,
custom_path: Option<String>,
config: Option<serde_json::Value>,
format_acceptance_config: Option<serde_json::Value>,
proxy: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
if base_url.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider_endpoints.base_url is empty".to_string(),
));
}
self.base_url = base_url;
self.header_rules = header_rules;
self.body_rules = body_rules;
self.max_retries = max_retries;
self.custom_path = custom_path;
self.config = config;
self.format_acceptance_config = format_acceptance_config;
self.proxy = proxy;
Ok(self)
}
pub fn with_health_score(mut self, health_score: f64) -> Self {
self.health_score = health_score;
self
}
pub fn with_timestamps(
mut self,
created_at_unix_secs: Option<u64>,
updated_at_unix_secs: Option<u64>,
) -> Self {
self.created_at_unix_secs = created_at_unix_secs;
self.updated_at_unix_secs = updated_at_unix_secs;
self
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderCatalogKey {
pub id: String,
pub provider_id: String,
pub name: String,
pub auth_type: String,
pub capabilities: Option<serde_json::Value>,
pub is_active: bool,
pub api_formats: Option<serde_json::Value>,
pub encrypted_api_key: String,
pub encrypted_auth_config: Option<String>,
pub note: Option<String>,
pub internal_priority: i32,
pub rate_multipliers: Option<serde_json::Value>,
pub global_priority_by_format: Option<serde_json::Value>,
pub allowed_models: Option<serde_json::Value>,
pub expires_at_unix_secs: Option<u64>,
pub cache_ttl_minutes: i32,
pub max_probe_interval_minutes: i32,
pub proxy: Option<serde_json::Value>,
pub fingerprint: Option<serde_json::Value>,
pub rpm_limit: Option<u32>,
pub learned_rpm_limit: Option<u32>,
pub concurrent_429_count: Option<u32>,
pub rpm_429_count: Option<u32>,
pub last_429_at_unix_secs: Option<u64>,
pub last_429_type: Option<String>,
pub adjustment_history: Option<serde_json::Value>,
pub utilization_samples: Option<serde_json::Value>,
pub last_probe_increase_at_unix_secs: Option<u64>,
pub request_count: Option<u32>,
pub success_count: Option<u32>,
pub error_count: Option<u32>,
pub total_response_time_ms: Option<u32>,
pub last_used_at_unix_secs: Option<u64>,
pub auto_fetch_models: bool,
pub last_models_fetch_at_unix_secs: Option<u64>,
pub last_models_fetch_error: Option<String>,
pub locked_models: Option<serde_json::Value>,
pub model_include_patterns: Option<serde_json::Value>,
pub model_exclude_patterns: Option<serde_json::Value>,
pub upstream_metadata: Option<serde_json::Value>,
pub oauth_invalid_at_unix_secs: Option<u64>,
pub oauth_invalid_reason: Option<String>,
pub status_snapshot: Option<serde_json::Value>,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
pub health_by_format: Option<serde_json::Value>,
pub circuit_breaker_by_format: Option<serde_json::Value>,
}
impl StoredProviderCatalogKey {
pub fn new(
id: String,
provider_id: String,
name: String,
auth_type: String,
capabilities: Option<serde_json::Value>,
is_active: bool,
) -> Result<Self, crate::DataLayerError> {
if name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider_api_keys.name is empty".to_string(),
));
}
if auth_type.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider_api_keys.auth_type is empty".to_string(),
));
}
Ok(Self {
id,
provider_id,
name,
auth_type,
capabilities,
is_active,
api_formats: None,
encrypted_api_key: String::new(),
encrypted_auth_config: None,
note: None,
internal_priority: 50,
rate_multipliers: None,
global_priority_by_format: None,
allowed_models: None,
expires_at_unix_secs: None,
cache_ttl_minutes: 5,
max_probe_interval_minutes: 32,
proxy: None,
fingerprint: None,
rpm_limit: None,
learned_rpm_limit: None,
concurrent_429_count: None,
rpm_429_count: None,
last_429_at_unix_secs: None,
last_429_type: None,
adjustment_history: None,
utilization_samples: None,
last_probe_increase_at_unix_secs: None,
request_count: None,
success_count: None,
error_count: None,
total_response_time_ms: None,
last_used_at_unix_secs: None,
auto_fetch_models: false,
last_models_fetch_at_unix_secs: None,
last_models_fetch_error: None,
locked_models: None,
model_include_patterns: None,
model_exclude_patterns: None,
upstream_metadata: None,
oauth_invalid_at_unix_secs: None,
oauth_invalid_reason: None,
status_snapshot: None,
created_at_unix_secs: None,
updated_at_unix_secs: None,
health_by_format: None,
circuit_breaker_by_format: None,
})
}
#[allow(clippy::too_many_arguments)]
pub fn with_transport_fields(
mut self,
api_formats: Option<serde_json::Value>,
encrypted_api_key: String,
encrypted_auth_config: Option<String>,
rate_multipliers: Option<serde_json::Value>,
global_priority_by_format: Option<serde_json::Value>,
allowed_models: Option<serde_json::Value>,
expires_at_unix_secs: Option<u64>,
proxy: Option<serde_json::Value>,
fingerprint: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
if encrypted_api_key.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider_api_keys.api_key is empty".to_string(),
));
}
self.api_formats = api_formats;
self.encrypted_api_key = encrypted_api_key;
self.encrypted_auth_config = encrypted_auth_config;
self.rate_multipliers = rate_multipliers;
self.global_priority_by_format = global_priority_by_format;
self.allowed_models = allowed_models;
self.expires_at_unix_secs = expires_at_unix_secs;
self.proxy = proxy;
self.fingerprint = fingerprint;
Ok(self)
}
#[allow(clippy::too_many_arguments)]
pub fn with_rate_limit_fields(
mut self,
rpm_limit: Option<u32>,
learned_rpm_limit: Option<u32>,
concurrent_429_count: Option<u32>,
rpm_429_count: Option<u32>,
last_429_at_unix_secs: Option<u64>,
adjustment_history: Option<serde_json::Value>,
request_count: Option<u32>,
success_count: Option<u32>,
) -> Self {
self.rpm_limit = rpm_limit;
self.learned_rpm_limit = learned_rpm_limit;
self.concurrent_429_count = concurrent_429_count;
self.rpm_429_count = rpm_429_count;
self.last_429_at_unix_secs = last_429_at_unix_secs;
self.adjustment_history = adjustment_history;
self.request_count = request_count;
self.success_count = success_count;
self
}
pub fn with_usage_fields(
mut self,
error_count: Option<u32>,
total_response_time_ms: Option<u32>,
) -> Self {
self.error_count = error_count;
self.total_response_time_ms = total_response_time_ms;
self
}
pub fn with_health_fields(
mut self,
health_by_format: Option<serde_json::Value>,
circuit_breaker_by_format: Option<serde_json::Value>,
) -> Self {
self.health_by_format = health_by_format;
self.circuit_breaker_by_format = circuit_breaker_by_format;
self
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct ProviderCatalogKeyListQuery {
pub provider_id: String,
pub search: Option<String>,
pub is_active: Option<bool>,
pub offset: usize,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderCatalogKeyPage {
pub items: Vec<StoredProviderCatalogKey>,
pub total: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderCatalogKeyStats {
pub provider_id: String,
pub total_keys: u64,
pub active_keys: u64,
}
impl StoredProviderCatalogKeyStats {
pub fn new(
provider_id: String,
total_keys: i64,
active_keys: i64,
) -> Result<Self, crate::DataLayerError> {
if provider_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider key stats provider_id is empty".to_string(),
));
}
if total_keys < 0 || active_keys < 0 {
return Err(crate::DataLayerError::UnexpectedValue(
"provider key stats count is negative".to_string(),
));
}
Ok(Self {
provider_id,
total_keys: total_keys as u64,
active_keys: active_keys as u64,
})
}
}
#[async_trait]
pub trait ProviderCatalogReadRepository: Send + Sync {
async fn list_providers(
&self,
active_only: bool,
) -> Result<Vec<StoredProviderCatalogProvider>, crate::DataLayerError>;
async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, crate::DataLayerError>;
async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, crate::DataLayerError>;
async fn list_endpoints_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, crate::DataLayerError>;
async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
async fn list_keys_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
) -> Result<StoredProviderCatalogKeyPage, crate::DataLayerError>;
async fn list_key_stats_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyStats>, crate::DataLayerError>;
}
#[async_trait]
pub trait ProviderCatalogWriteRepository: Send + Sync {
async fn create_provider(
&self,
provider: &StoredProviderCatalogProvider,
shift_existing_priorities_from: Option<i32>,
) -> Result<StoredProviderCatalogProvider, crate::DataLayerError>;
async fn update_provider(
&self,
provider: &StoredProviderCatalogProvider,
) -> Result<StoredProviderCatalogProvider, crate::DataLayerError>;
async fn delete_provider(&self, provider_id: &str) -> Result<bool, crate::DataLayerError>;
async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), crate::DataLayerError>;
async fn create_endpoint(
&self,
endpoint: &StoredProviderCatalogEndpoint,
) -> Result<StoredProviderCatalogEndpoint, crate::DataLayerError>;
async fn update_endpoint(
&self,
endpoint: &StoredProviderCatalogEndpoint,
) -> Result<StoredProviderCatalogEndpoint, crate::DataLayerError>;
async fn delete_endpoint(&self, endpoint_id: &str) -> Result<bool, crate::DataLayerError>;
async fn create_key(
&self,
key: &StoredProviderCatalogKey,
) -> Result<StoredProviderCatalogKey, crate::DataLayerError>;
async fn update_key(
&self,
key: &StoredProviderCatalogKey,
) -> Result<StoredProviderCatalogKey, crate::DataLayerError>;
async fn delete_key(&self, key_id: &str) -> Result<bool, crate::DataLayerError>;
async fn clear_key_oauth_invalid_marker(
&self,
key_id: &str,
) -> Result<bool, crate::DataLayerError>;
async fn update_key_oauth_credentials(
&self,
key_id: &str,
encrypted_api_key: &str,
encrypted_auth_config: Option<&str>,
expires_at_unix_secs: Option<u64>,
) -> Result<bool, crate::DataLayerError>;
async fn update_key_health_state(
&self,
key_id: &str,
is_active: bool,
health_by_format: Option<&serde_json::Value>,
circuit_breaker_by_format: Option<&serde_json::Value>,
) -> Result<bool, crate::DataLayerError>;
}

View File

@@ -0,0 +1,6 @@
mod types;
pub use types::{
ProviderQuotaReadRepository, ProviderQuotaRepository, ProviderQuotaWriteRepository,
StoredProviderQuotaSnapshot,
};

View File

@@ -0,0 +1,71 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderQuotaSnapshot {
pub provider_id: String,
pub billing_type: String,
pub monthly_quota_usd: Option<f64>,
pub monthly_used_usd: f64,
pub quota_reset_day: Option<u64>,
pub quota_last_reset_at_unix_secs: Option<u64>,
pub quota_expires_at_unix_secs: Option<u64>,
pub is_active: bool,
}
impl StoredProviderQuotaSnapshot {
#[allow(clippy::too_many_arguments)]
pub fn new(
provider_id: String,
billing_type: String,
monthly_quota_usd: Option<f64>,
monthly_used_usd: f64,
quota_reset_day: Option<i32>,
quota_last_reset_at_unix_secs: Option<i64>,
quota_expires_at_unix_secs: Option<i64>,
is_active: bool,
) -> Result<Self, crate::DataLayerError> {
if provider_id.trim().is_empty() || billing_type.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider quota identity is empty".to_string(),
));
}
if !monthly_used_usd.is_finite() || monthly_quota_usd.is_some_and(|v| !v.is_finite()) {
return Err(crate::DataLayerError::UnexpectedValue(
"provider quota value is not finite".to_string(),
));
}
Ok(Self {
provider_id,
billing_type,
monthly_quota_usd,
monthly_used_usd,
quota_reset_day: quota_reset_day.map(|value| value as u64),
quota_last_reset_at_unix_secs: quota_last_reset_at_unix_secs.map(|value| value as u64),
quota_expires_at_unix_secs: quota_expires_at_unix_secs.map(|value| value as u64),
is_active,
})
}
}
#[async_trait]
pub trait ProviderQuotaReadRepository: Send + Sync {
async fn find_by_provider_id(
&self,
provider_id: &str,
) -> Result<Option<StoredProviderQuotaSnapshot>, crate::DataLayerError>;
}
#[async_trait]
pub trait ProviderQuotaWriteRepository: Send + Sync {
async fn reset_due(&self, now_unix_secs: u64) -> Result<usize, crate::DataLayerError>;
}
pub trait ProviderQuotaRepository:
ProviderQuotaReadRepository + ProviderQuotaWriteRepository + Send + Sync
{
}
impl<T> ProviderQuotaRepository for T where
T: ProviderQuotaReadRepository + ProviderQuotaWriteRepository + Send + Sync
{
}

View File

@@ -0,0 +1,5 @@
mod types;
pub use types::{
SettlementRepository, SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
};

View File

@@ -0,0 +1,83 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UsageSettlementInput {
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub provider_id: Option<String>,
pub status: String,
pub billing_status: String,
pub total_cost_usd: f64,
pub actual_total_cost_usd: f64,
pub finalized_at_unix_secs: Option<u64>,
}
impl UsageSettlementInput {
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
if self.request_id.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"settlement request_id cannot be empty".to_string(),
));
}
if self.status.trim().is_empty() || self.billing_status.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"settlement status cannot be empty".to_string(),
));
}
if !self.total_cost_usd.is_finite() || !self.actual_total_cost_usd.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"settlement cost must be finite".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredUsageSettlement {
pub request_id: String,
pub wallet_id: Option<String>,
pub billing_status: String,
pub wallet_balance_before: Option<f64>,
pub wallet_balance_after: Option<f64>,
pub wallet_recharge_balance_before: Option<f64>,
pub wallet_recharge_balance_after: Option<f64>,
pub wallet_gift_balance_before: Option<f64>,
pub wallet_gift_balance_after: Option<f64>,
pub provider_monthly_used_usd: Option<f64>,
pub finalized_at_unix_secs: Option<u64>,
}
#[async_trait]
pub trait SettlementWriteRepository: Send + Sync {
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, crate::DataLayerError>;
}
pub trait SettlementRepository: SettlementWriteRepository + Send + Sync {}
impl<T> SettlementRepository for T where T: SettlementWriteRepository + Send + Sync {}
#[cfg(test)]
mod tests {
use super::UsageSettlementInput;
#[test]
fn rejects_invalid_settlement_input() {
let input = UsageSettlementInput {
request_id: "".to_string(),
user_id: None,
api_key_id: None,
provider_id: None,
status: "completed".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 0.1,
actual_total_cost_usd: 0.1,
finalized_at_unix_secs: None,
};
assert!(input.validate().is_err());
}
}

View File

@@ -0,0 +1,7 @@
mod types;
pub use types::{
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageRepository,
UsageWriteRepository,
};

View File

@@ -0,0 +1,639 @@
use async_trait::async_trait;
use serde_json::Value;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredRequestUsageAudit {
pub id: String,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub provider_name: String,
pub model: String,
pub target_model: Option<String>,
pub provider_id: Option<String>,
pub provider_endpoint_id: Option<String>,
pub provider_api_key_id: Option<String>,
pub request_type: Option<String>,
pub api_format: Option<String>,
pub api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub endpoint_api_format: Option<String>,
pub provider_api_family: Option<String>,
pub provider_endpoint_kind: Option<String>,
pub has_format_conversion: bool,
pub is_stream: bool,
pub input_tokens: u64,
pub output_tokens: u64,
pub total_tokens: u64,
pub cache_creation_input_tokens: u64,
pub cache_read_input_tokens: u64,
pub cache_creation_cost_usd: f64,
pub cache_read_cost_usd: f64,
pub output_price_per_1m: Option<f64>,
pub total_cost_usd: f64,
pub actual_total_cost_usd: f64,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub error_category: Option<String>,
pub response_time_ms: Option<u64>,
pub first_byte_time_ms: Option<u64>,
pub status: String,
pub billing_status: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_metadata: Option<Value>,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
pub finalized_at_unix_secs: Option<u64>,
}
impl StoredRequestUsageAudit {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
request_id: String,
user_id: Option<String>,
api_key_id: Option<String>,
username: Option<String>,
api_key_name: Option<String>,
provider_name: String,
model: String,
target_model: Option<String>,
provider_id: Option<String>,
provider_endpoint_id: Option<String>,
provider_api_key_id: Option<String>,
request_type: Option<String>,
api_format: Option<String>,
api_family: Option<String>,
endpoint_kind: Option<String>,
endpoint_api_format: Option<String>,
provider_api_family: Option<String>,
provider_endpoint_kind: Option<String>,
has_format_conversion: bool,
is_stream: bool,
input_tokens: i32,
output_tokens: i32,
total_tokens: i32,
total_cost_usd: f64,
actual_total_cost_usd: f64,
status_code: Option<i32>,
error_message: Option<String>,
error_category: Option<String>,
response_time_ms: Option<i32>,
first_byte_time_ms: Option<i32>,
status: String,
billing_status: String,
created_at_unix_secs: i64,
updated_at_unix_secs: i64,
finalized_at_unix_secs: Option<i64>,
) -> Result<Self, crate::DataLayerError> {
if request_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.request_id is empty".to_string(),
));
}
if provider_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.provider_name is empty".to_string(),
));
}
if model.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.model is empty".to_string(),
));
}
if status.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.status is empty".to_string(),
));
}
if billing_status.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.billing_status is empty".to_string(),
));
}
if !total_cost_usd.is_finite() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.total_cost_usd is not finite".to_string(),
));
}
if !actual_total_cost_usd.is_finite() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.actual_total_cost_usd is not finite".to_string(),
));
}
Ok(Self {
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
provider_name,
model,
target_model,
provider_id,
provider_endpoint_id,
provider_api_key_id,
request_type,
api_format,
api_family,
endpoint_kind,
endpoint_api_format,
provider_api_family,
provider_endpoint_kind,
has_format_conversion,
is_stream,
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")?,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
cache_creation_cost_usd: 0.0,
cache_read_cost_usd: 0.0,
output_price_per_1m: None,
total_cost_usd,
actual_total_cost_usd,
status_code: parse_u16(status_code, "usage.status_code")?,
error_message,
error_category,
response_time_ms: parse_optional_u64(response_time_ms, "usage.response_time_ms")?,
first_byte_time_ms: parse_optional_u64(first_byte_time_ms, "usage.first_byte_time_ms")?,
status,
billing_status,
request_headers: None,
request_body: None,
provider_request_headers: None,
provider_request_body: None,
response_headers: None,
response_body: None,
client_response_headers: None,
client_response_body: None,
request_metadata: None,
created_at_unix_secs: parse_timestamp(
created_at_unix_secs,
"usage.created_at_unix_secs",
)?,
updated_at_unix_secs: parse_timestamp(
updated_at_unix_secs,
"usage.updated_at_unix_secs",
)?,
finalized_at_unix_secs: finalized_at_unix_secs
.map(|value| parse_timestamp(value, "usage.finalized_at_unix_secs"))
.transpose()?,
})
}
pub fn with_cache_input_tokens(
mut self,
cache_creation_input_tokens: u64,
cache_read_input_tokens: u64,
) -> Self {
self.cache_creation_input_tokens = cache_creation_input_tokens;
self.cache_read_input_tokens = cache_read_input_tokens;
self
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderUsageWindow {
pub provider_id: String,
pub window_start_unix_secs: u64,
pub total_requests: u64,
pub successful_requests: u64,
pub failed_requests: u64,
pub avg_response_time_ms: f64,
pub total_cost_usd: f64,
}
impl StoredProviderUsageWindow {
#[allow(clippy::too_many_arguments)]
pub fn new(
provider_id: String,
window_start_unix_secs: i64,
total_requests: i64,
successful_requests: i64,
failed_requests: i64,
avg_response_time_ms: f64,
total_cost_usd: f64,
) -> Result<Self, crate::DataLayerError> {
if provider_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider usage window provider_id is empty".to_string(),
));
}
if !avg_response_time_ms.is_finite() || !total_cost_usd.is_finite() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider usage window value is not finite".to_string(),
));
}
Ok(Self {
provider_id,
window_start_unix_secs: parse_timestamp(
window_start_unix_secs,
"provider_usage_tracking.window_start_unix_secs",
)?,
total_requests: parse_timestamp(
total_requests,
"provider_usage_tracking.total_requests",
)?,
successful_requests: parse_timestamp(
successful_requests,
"provider_usage_tracking.successful_requests",
)?,
failed_requests: parse_timestamp(
failed_requests,
"provider_usage_tracking.failed_requests",
)?,
avg_response_time_ms,
total_cost_usd,
})
}
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderUsageSummary {
pub total_requests: u64,
pub successful_requests: u64,
pub failed_requests: u64,
pub avg_response_time_ms: f64,
pub total_cost_usd: f64,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct UsageAuditListQuery {
pub created_from_unix_secs: Option<u64>,
pub created_until_unix_secs: Option<u64>,
pub user_id: Option<String>,
pub provider_name: Option<String>,
pub model: Option<String>,
}
#[async_trait]
pub trait UsageReadRepository: Send + Sync {
async fn find_by_id(
&self,
id: &str,
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
async fn find_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
async fn list_usage_audits(
&self,
query: &UsageAuditListQuery,
) -> Result<Vec<StoredRequestUsageAudit>, crate::DataLayerError>;
async fn list_recent_usage_audits(
&self,
user_id: Option<&str>,
limit: usize,
) -> Result<Vec<StoredRequestUsageAudit>, crate::DataLayerError>;
async fn summarize_total_tokens_by_api_key_ids(
&self,
api_key_ids: &[String],
) -> Result<std::collections::BTreeMap<String, u64>, crate::DataLayerError>;
async fn summarize_provider_usage_since(
&self,
provider_id: &str,
since_unix_secs: u64,
) -> Result<StoredProviderUsageSummary, crate::DataLayerError>;
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UpsertUsageRecord {
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub provider_name: String,
pub model: String,
pub target_model: Option<String>,
pub provider_id: Option<String>,
pub provider_endpoint_id: Option<String>,
pub provider_api_key_id: Option<String>,
pub request_type: Option<String>,
pub api_format: Option<String>,
pub api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub endpoint_api_format: Option<String>,
pub provider_api_family: Option<String>,
pub provider_endpoint_kind: Option<String>,
pub has_format_conversion: Option<bool>,
pub is_stream: Option<bool>,
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub total_tokens: Option<u64>,
pub cache_creation_input_tokens: Option<u64>,
pub cache_read_input_tokens: Option<u64>,
pub cache_creation_cost_usd: Option<f64>,
pub cache_read_cost_usd: Option<f64>,
pub output_price_per_1m: Option<f64>,
pub total_cost_usd: Option<f64>,
pub actual_total_cost_usd: Option<f64>,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub error_category: Option<String>,
pub response_time_ms: Option<u64>,
pub first_byte_time_ms: Option<u64>,
pub status: String,
pub billing_status: String,
pub request_headers: Option<Value>,
pub request_body: Option<Value>,
pub provider_request_headers: Option<Value>,
pub provider_request_body: Option<Value>,
pub response_headers: Option<Value>,
pub response_body: Option<Value>,
pub client_response_headers: Option<Value>,
pub client_response_body: Option<Value>,
pub request_metadata: Option<Value>,
pub finalized_at_unix_secs: Option<u64>,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: u64,
}
impl UpsertUsageRecord {
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
if self.request_id.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert request_id cannot be empty".to_string(),
));
}
if self.provider_name.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert provider_name cannot be empty".to_string(),
));
}
if self.model.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert model cannot be empty".to_string(),
));
}
if self.status.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert status cannot be empty".to_string(),
));
}
if self.billing_status.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert billing_status cannot be empty".to_string(),
));
}
if let Some(value) = self.total_cost_usd {
if !value.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert total_cost_usd must be finite".to_string(),
));
}
}
if let Some(value) = self.cache_creation_cost_usd {
if !value.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert cache_creation_cost_usd must be finite".to_string(),
));
}
}
if let Some(value) = self.cache_read_cost_usd {
if !value.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert cache_read_cost_usd must be finite".to_string(),
));
}
}
if let Some(value) = self.output_price_per_1m {
if !value.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert output_price_per_1m must be finite".to_string(),
));
}
}
if let Some(value) = self.actual_total_cost_usd {
if !value.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"usage upsert actual_total_cost_usd must be finite".to_string(),
));
}
}
Ok(())
}
}
#[async_trait]
pub trait UsageWriteRepository: Send + Sync {
async fn upsert(
&self,
usage: UpsertUsageRecord,
) -> Result<StoredRequestUsageAudit, crate::DataLayerError>;
}
pub trait UsageRepository: UsageReadRepository + UsageWriteRepository + Send + Sync {}
impl<T> UsageRepository for T where T: UsageReadRepository + UsageWriteRepository + Send + Sync {}
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}"))
})
}
fn parse_optional_u64(
value: Option<i32>,
field_name: &str,
) -> Result<Option<u64>, crate::DataLayerError> {
value
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
})
})
.transpose()
}
fn parse_u16(value: Option<i32>, field_name: &str) -> Result<Option<u16>, crate::DataLayerError> {
value
.map(|value| {
u16::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
})
})
.transpose()
}
fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayerError> {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
})
}
#[cfg(test)]
mod tests {
use super::{StoredRequestUsageAudit, UpsertUsageRecord};
use serde_json::json;
#[test]
fn rejects_empty_request_id() {
assert!(StoredRequestUsageAudit::new(
"usage-1".to_string(),
"".to_string(),
None,
None,
None,
None,
"OpenAI".to_string(),
"gpt-4.1".to_string(),
None,
None,
None,
None,
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
false,
false,
10,
20,
30,
0.1,
0.1,
Some(200),
None,
None,
Some(120),
Some(80),
"completed".to_string(),
"settled".to_string(),
100,
101,
Some(102),
)
.is_err());
}
#[test]
fn rejects_negative_token_count() {
assert!(StoredRequestUsageAudit::new(
"usage-1".to_string(),
"req-1".to_string(),
None,
None,
None,
None,
"OpenAI".to_string(),
"gpt-4.1".to_string(),
None,
None,
None,
None,
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
false,
false,
-1,
20,
30,
0.1,
0.1,
Some(200),
None,
None,
Some(120),
Some(80),
"completed".to_string(),
"settled".to_string(),
100,
101,
Some(102),
)
.is_err());
}
#[test]
fn rejects_invalid_upsert_payload() {
let record = UpsertUsageRecord {
request_id: "".to_string(),
user_id: None,
api_key_id: None,
username: None,
api_key_name: None,
provider_name: "openai".to_string(),
model: "gpt-5".to_string(),
target_model: None,
provider_id: None,
provider_endpoint_id: None,
provider_api_key_id: None,
request_type: Some("chat".to_string()),
api_format: Some("openai:chat".to_string()),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
provider_api_family: Some("openai".to_string()),
provider_endpoint_kind: Some("chat".to_string()),
has_format_conversion: Some(false),
is_stream: Some(false),
input_tokens: Some(10),
output_tokens: Some(20),
total_tokens: Some(30),
cache_creation_input_tokens: None,
cache_read_input_tokens: None,
cache_creation_cost_usd: None,
cache_read_cost_usd: None,
output_price_per_1m: None,
total_cost_usd: None,
actual_total_cost_usd: None,
status_code: Some(200),
error_message: None,
error_category: None,
response_time_ms: Some(120),
first_byte_time_ms: None,
status: "completed".to_string(),
billing_status: "pending".to_string(),
request_headers: Some(json!({"authorization": "Bearer test"})),
request_body: Some(json!({"model": "gpt-5"})),
provider_request_headers: None,
provider_request_body: None,
response_headers: None,
response_body: None,
client_response_headers: None,
client_response_body: None,
request_metadata: None,
finalized_at_unix_secs: None,
created_at_unix_secs: Some(100),
updated_at_unix_secs: 101,
};
assert!(record.validate().is_err());
}
}

View File

@@ -0,0 +1,7 @@
mod types;
pub use types::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskRepository, VideoTaskStatus,
VideoTaskStatusCount, VideoTaskWriteRepository,
};

View File

@@ -0,0 +1,611 @@
use async_trait::async_trait;
use serde_json::Value;
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
pub enum VideoTaskStatus {
Pending,
Submitted,
Queued,
Processing,
Completed,
Failed,
Cancelled,
Expired,
Deleted,
}
impl VideoTaskStatus {
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"pending" => Ok(Self::Pending),
"submitted" => Ok(Self::Submitted),
"queued" => Ok(Self::Queued),
"processing" => Ok(Self::Processing),
"completed" => Ok(Self::Completed),
"failed" => Ok(Self::Failed),
"cancelled" => Ok(Self::Cancelled),
"expired" => Ok(Self::Expired),
"deleted" => Ok(Self::Deleted),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unsupported video_tasks.status: {other}"
))),
}
}
pub fn is_active(self) -> bool {
matches!(
self,
Self::Pending | Self::Submitted | Self::Queued | Self::Processing
)
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredVideoTask {
pub id: String,
pub short_id: Option<String>,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub external_task_id: Option<String>,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub client_api_format: Option<String>,
pub provider_api_format: Option<String>,
pub format_converted: bool,
pub model: Option<String>,
pub prompt: Option<String>,
pub original_request_body: Option<Value>,
pub duration_seconds: Option<u32>,
pub resolution: Option<String>,
pub aspect_ratio: Option<String>,
pub size: Option<String>,
pub status: VideoTaskStatus,
pub progress_percent: u16,
pub progress_message: Option<String>,
pub retry_count: u32,
pub poll_interval_seconds: u32,
pub next_poll_at_unix_secs: Option<u64>,
pub poll_count: u32,
pub max_poll_count: u32,
pub created_at_unix_secs: u64,
pub submitted_at_unix_secs: Option<u64>,
pub completed_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: u64,
pub error_code: Option<String>,
pub error_message: Option<String>,
pub video_url: Option<String>,
pub request_metadata: Option<Value>,
}
impl StoredVideoTask {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
short_id: Option<String>,
request_id: String,
user_id: Option<String>,
api_key_id: Option<String>,
username: Option<String>,
api_key_name: Option<String>,
external_task_id: Option<String>,
provider_id: Option<String>,
endpoint_id: Option<String>,
key_id: Option<String>,
client_api_format: Option<String>,
provider_api_format: Option<String>,
format_converted: bool,
model: Option<String>,
prompt: Option<String>,
original_request_body: Option<Value>,
duration_seconds: Option<i32>,
resolution: Option<String>,
aspect_ratio: Option<String>,
size: Option<String>,
status: VideoTaskStatus,
progress_percent: i32,
progress_message: Option<String>,
retry_count: i32,
poll_interval_seconds: i32,
next_poll_at_unix_secs: Option<i64>,
poll_count: i32,
max_poll_count: i32,
created_at_unix_secs: i64,
submitted_at_unix_secs: Option<i64>,
completed_at_unix_secs: Option<i64>,
updated_at_unix_secs: i64,
error_code: Option<String>,
error_message: Option<String>,
video_url: Option<String>,
request_metadata: Option<Value>,
) -> Result<Self, crate::DataLayerError> {
let progress_percent = u16::try_from(progress_percent).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid progress_percent: {progress_percent}"
))
})?;
let retry_count = u32::try_from(retry_count).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid retry_count: {retry_count}"))
})?;
let poll_interval_seconds = u32::try_from(poll_interval_seconds).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid poll_interval_seconds: {poll_interval_seconds}"
))
})?;
let next_poll_at_unix_secs =
coerce_optional_unix_secs(next_poll_at_unix_secs, "next_poll_at_unix_secs")?;
let poll_count = u32::try_from(poll_count).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid poll_count: {poll_count}"))
})?;
let max_poll_count = u32::try_from(max_poll_count).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid max_poll_count: {max_poll_count}"
))
})?;
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid created_at_unix_secs: {created_at_unix_secs}"
))
})?;
let submitted_at_unix_secs =
coerce_optional_unix_secs(submitted_at_unix_secs, "submitted_at_unix_secs")?;
let completed_at_unix_secs =
coerce_optional_unix_secs(completed_at_unix_secs, "completed_at_unix_secs")?;
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
))
})?;
let duration_seconds = match duration_seconds {
Some(value) => Some(u32::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid duration_seconds: {value}"))
})?),
None => None,
};
Ok(Self {
id,
short_id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
external_task_id,
provider_id,
endpoint_id,
key_id,
client_api_format,
provider_api_format,
format_converted,
model,
prompt,
original_request_body,
duration_seconds,
resolution,
aspect_ratio,
size,
status,
progress_percent,
progress_message,
retry_count,
poll_interval_seconds,
next_poll_at_unix_secs,
poll_count,
max_poll_count,
created_at_unix_secs,
submitted_at_unix_secs,
completed_at_unix_secs,
updated_at_unix_secs,
error_code,
error_message,
video_url,
request_metadata,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpsertVideoTask {
pub id: String,
pub short_id: Option<String>,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub external_task_id: Option<String>,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub client_api_format: Option<String>,
pub provider_api_format: Option<String>,
pub format_converted: bool,
pub model: Option<String>,
pub prompt: Option<String>,
pub original_request_body: Option<Value>,
pub duration_seconds: Option<u32>,
pub resolution: Option<String>,
pub aspect_ratio: Option<String>,
pub size: Option<String>,
pub status: VideoTaskStatus,
pub progress_percent: u16,
pub progress_message: Option<String>,
pub retry_count: u32,
pub poll_interval_seconds: u32,
pub next_poll_at_unix_secs: Option<u64>,
pub poll_count: u32,
pub max_poll_count: u32,
pub created_at_unix_secs: u64,
pub submitted_at_unix_secs: Option<u64>,
pub completed_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: u64,
pub error_code: Option<String>,
pub error_message: Option<String>,
pub video_url: Option<String>,
pub request_metadata: Option<Value>,
}
impl UpsertVideoTask {
pub fn into_stored(self) -> StoredVideoTask {
StoredVideoTask {
id: self.id,
short_id: self.short_id,
request_id: self.request_id,
user_id: self.user_id,
api_key_id: self.api_key_id,
username: self.username,
api_key_name: self.api_key_name,
external_task_id: self.external_task_id,
provider_id: self.provider_id,
endpoint_id: self.endpoint_id,
key_id: self.key_id,
client_api_format: self.client_api_format,
provider_api_format: self.provider_api_format,
format_converted: self.format_converted,
model: self.model,
prompt: self.prompt,
original_request_body: self.original_request_body,
duration_seconds: self.duration_seconds,
resolution: self.resolution,
aspect_ratio: self.aspect_ratio,
size: self.size,
status: self.status,
progress_percent: self.progress_percent,
progress_message: self.progress_message,
retry_count: self.retry_count,
poll_interval_seconds: self.poll_interval_seconds,
next_poll_at_unix_secs: self.next_poll_at_unix_secs,
poll_count: self.poll_count,
max_poll_count: self.max_poll_count,
created_at_unix_secs: self.created_at_unix_secs,
submitted_at_unix_secs: self.submitted_at_unix_secs,
completed_at_unix_secs: self.completed_at_unix_secs,
updated_at_unix_secs: self.updated_at_unix_secs,
error_code: self.error_code,
error_message: self.error_message,
video_url: self.video_url,
request_metadata: self.request_metadata,
}
}
}
impl From<StoredVideoTask> for UpsertVideoTask {
fn from(task: StoredVideoTask) -> Self {
Self {
id: task.id,
short_id: task.short_id,
request_id: task.request_id,
user_id: task.user_id,
api_key_id: task.api_key_id,
username: task.username,
api_key_name: task.api_key_name,
external_task_id: task.external_task_id,
provider_id: task.provider_id,
endpoint_id: task.endpoint_id,
key_id: task.key_id,
client_api_format: task.client_api_format,
provider_api_format: task.provider_api_format,
format_converted: task.format_converted,
model: task.model,
prompt: task.prompt,
original_request_body: task.original_request_body,
duration_seconds: task.duration_seconds,
resolution: task.resolution,
aspect_ratio: task.aspect_ratio,
size: task.size,
status: task.status,
progress_percent: task.progress_percent,
progress_message: task.progress_message,
retry_count: task.retry_count,
poll_interval_seconds: task.poll_interval_seconds,
next_poll_at_unix_secs: task.next_poll_at_unix_secs,
poll_count: task.poll_count,
max_poll_count: task.max_poll_count,
created_at_unix_secs: task.created_at_unix_secs,
submitted_at_unix_secs: task.submitted_at_unix_secs,
completed_at_unix_secs: task.completed_at_unix_secs,
updated_at_unix_secs: task.updated_at_unix_secs,
error_code: task.error_code,
error_message: task.error_message,
video_url: task.video_url,
request_metadata: task.request_metadata,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VideoTaskLookupKey<'a> {
Id(&'a str),
ShortId(&'a str),
UserExternal {
user_id: &'a str,
external_task_id: &'a str,
},
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct VideoTaskQueryFilter {
pub user_id: Option<String>,
pub status: Option<VideoTaskStatus>,
pub model_substring: Option<String>,
pub client_api_format: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct VideoTaskStatusCount {
pub status: VideoTaskStatus,
pub count: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct VideoTaskModelCount {
pub model: String,
pub count: u64,
}
#[async_trait]
pub trait VideoTaskReadRepository: Send + Sync {
async fn find(
&self,
key: VideoTaskLookupKey<'_>,
) -> Result<Option<StoredVideoTask>, crate::DataLayerError>;
async fn list_active(
&self,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
async fn list_due(
&self,
now_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
async fn list_page(
&self,
filter: &VideoTaskQueryFilter,
offset: usize,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, crate::DataLayerError>;
async fn count_by_status(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<Vec<VideoTaskStatusCount>, crate::DataLayerError>;
async fn count_distinct_users(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<u64, crate::DataLayerError>;
async fn top_models(
&self,
filter: &VideoTaskQueryFilter,
limit: usize,
) -> Result<Vec<VideoTaskModelCount>, crate::DataLayerError>;
async fn count_created_since(
&self,
filter: &VideoTaskQueryFilter,
created_since_unix_secs: u64,
) -> Result<u64, crate::DataLayerError>;
}
#[async_trait]
pub trait VideoTaskWriteRepository: Send + Sync {
async fn upsert(&self, task: UpsertVideoTask)
-> Result<StoredVideoTask, crate::DataLayerError>;
async fn update_if_active(
&self,
task: UpsertVideoTask,
) -> Result<Option<StoredVideoTask>, crate::DataLayerError>;
async fn claim_due(
&self,
now_unix_secs: u64,
claim_until_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
}
pub trait VideoTaskRepository:
VideoTaskReadRepository + VideoTaskWriteRepository + Send + Sync
{
}
impl<T> VideoTaskRepository for T where
T: VideoTaskReadRepository + VideoTaskWriteRepository + Send + Sync
{
}
fn coerce_optional_unix_secs(
value: Option<i64>,
field: &str,
) -> Result<Option<u64>, crate::DataLayerError> {
match value {
Some(value) => Ok(Some(u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field}: {value}"))
})?)),
None => Ok(None),
}
}
#[cfg(test)]
mod tests {
use super::{StoredVideoTask, VideoTaskStatus};
#[allow(clippy::type_complexity)]
fn base_new_args() -> (
String,
Option<String>,
String,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
bool,
Option<String>,
Option<String>,
Option<serde_json::Value>,
Option<i32>,
Option<String>,
Option<String>,
Option<String>,
VideoTaskStatus,
i32,
Option<String>,
i32,
i32,
Option<i64>,
i32,
i32,
i64,
Option<i64>,
Option<i64>,
i64,
Option<String>,
Option<String>,
Option<String>,
Option<serde_json::Value>,
) {
(
"task-1".to_string(),
None,
"request-1".to_string(),
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
false,
None,
None,
None,
None,
None,
None,
None,
VideoTaskStatus::Submitted,
10,
None,
0,
10,
Some(1),
0,
360,
1,
None,
None,
1,
None,
None,
None,
None,
)
}
#[test]
fn parses_status_from_database_text() {
assert_eq!(
VideoTaskStatus::from_database("processing").expect("status should parse"),
VideoTaskStatus::Processing
);
}
#[test]
fn rejects_invalid_database_status() {
assert!(VideoTaskStatus::from_database("mystery").is_err());
}
#[test]
fn rejects_invalid_numeric_fields() {
let mut args = base_new_args();
args.22 = -1;
assert!(StoredVideoTask::new(
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
)
.is_err());
}
#[test]
fn rejects_negative_updated_at_values() {
let mut args = base_new_args();
args.32 = -1;
assert!(StoredVideoTask::new(
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
)
.is_err());
}
#[test]
fn rejects_negative_created_at_values() {
let mut args = base_new_args();
args.29 = -1;
assert!(StoredVideoTask::new(
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
)
.is_err());
}
#[test]
fn rejects_negative_optional_completed_at_values() {
let mut args = base_new_args();
args.31 = Some(-1);
assert!(StoredVideoTask::new(
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
)
.is_err());
}
}