oxios_kernel/token_maxing/
live_quota.rs1use async_trait::async_trait;
12use chrono::{DateTime, Utc};
13use serde::Serialize;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Default)]
17#[serde(rename_all = "lowercase")]
18pub enum PlanType {
19 Subscription,
21 Metered,
23 #[default]
25 Unknown,
26}
27
28#[derive(Debug, Clone, Serialize)]
30pub struct RateWindow {
31 pub name: String,
32 pub used: Option<f64>,
33 pub limit: Option<f64>,
34 pub remaining_percent: Option<f64>,
35 pub resets_at: Option<DateTime<Utc>>,
36}
37
38#[derive(Debug, Clone, Serialize)]
40pub struct QuotaSnapshot {
41 pub provider: String,
42 pub plan: Option<String>,
43 #[serde(default)]
44 pub plan_type: PlanType,
45 #[serde(default)]
46 pub token_limit: Option<f64>,
47 pub rate_windows: Vec<RateWindow>,
48 pub fetched_at: DateTime<Utc>,
49 pub error: Option<String>,
50}
51
52impl QuotaSnapshot {
53 pub fn blank(provider: impl Into<String>) -> Self {
54 Self {
55 provider: provider.into(),
56 plan: None,
57 plan_type: PlanType::Unknown,
58 token_limit: None,
59 rate_windows: Vec::new(),
60 fetched_at: Utc::now(),
61 error: None,
62 }
63 }
64
65 pub fn is_subscription_signal(&self) -> bool {
66 if self.error.is_some() || self.plan_type != PlanType::Subscription {
67 return false;
68 }
69 self.rate_windows
70 .iter()
71 .any(|w| w.remaining_percent.is_some() || w.resets_at.is_some())
72 }
73
74 pub fn best_remaining_percent(&self) -> Option<f64> {
75 self.rate_windows
76 .iter()
77 .filter_map(|w| w.remaining_percent)
78 .next()
79 }
80
81 pub fn best_resets_at(&self) -> Option<DateTime<Utc>> {
82 self.rate_windows.iter().filter_map(|w| w.resets_at).next()
83 }
84}
85
86#[async_trait]
88pub trait QuotaFetcher: Send + Sync {
89 fn provider(&self) -> &str;
90 fn has_credentials(&self, api_key: Option<&str>) -> bool {
91 api_key.is_some_and(|k| !k.is_empty())
92 }
93 async fn fetch(&self, api_key: Option<&str>) -> anyhow::Result<QuotaSnapshot>;
94}