1use async_trait::async_trait;
2use serde::Deserialize;
3use serde::de::{self, Deserializer};
4
5use crate::error::SpendPanelError;
6use crate::model::{CreditsSnapshot, PlanInfo, RateWindow, RateWindowStatus, UsageSnapshot};
7use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
8
9#[derive(Debug, serde::Deserialize)]
10struct BalanceResponse {
11 #[serde(rename = "canConsume")]
12 can_consume: bool,
13 #[serde(rename = "consumptionCurrency")]
14 consumption_currency: Option<String>,
15 balances: Balances,
16 #[serde(
17 rename = "diemEpochAllocation",
18 default,
19 deserialize_with = "deserialize_opt_f64"
20 )]
21 diem_epoch_allocation: Option<f64>,
22}
23
24#[derive(Debug, serde::Deserialize)]
25struct Balances {
26 #[serde(default, deserialize_with = "deserialize_opt_f64")]
27 diem: Option<f64>,
28 #[serde(default, deserialize_with = "deserialize_opt_f64")]
29 usd: Option<f64>,
30}
31
32#[derive(Debug, Clone, PartialEq)]
33struct VeniceUsage {
34 can_consume: bool,
35 consumption_currency: Option<String>,
36 diem_balance: Option<f64>,
37 usd_balance: Option<f64>,
38 diem_epoch_allocation: Option<f64>,
39}
40
41pub struct VeniceProvider {
42 metadata: ProviderMetadata,
43 balance_url: Option<String>,
44}
45
46impl VeniceProvider {
47 pub fn new() -> Self {
48 Self {
49 metadata: ProviderMetadata {
50 id: "venice",
51 name: "Venice",
52 description: "Venice DIEM/USD API balance monitor",
53 auth_methods: &["api_key", "env"],
54 website: Some("https://venice.ai"),
55 },
56 balance_url: None,
57 }
58 }
59
60 pub fn with_balance_url(url: &str) -> Self {
61 let mut p = Self::new();
62 p.balance_url = Some(url.to_string());
63 p
64 }
65
66 fn clean(raw: &str) -> String {
67 let mut value = raw.trim();
68 if value.len() >= 2
69 && ((value.starts_with('"') && value.ends_with('"'))
70 || (value.starts_with('\'') && value.ends_with('\'')))
71 {
72 value = &value[1..value.len() - 1];
73 }
74 value.trim().to_string()
75 }
76
77 fn resolve_api_key(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
78 for key in ["api_key", "token"] {
79 if let Some(value) = ctx.config.get(key) {
80 let cleaned = Self::clean(value);
81 if !cleaned.is_empty() {
82 return Ok(cleaned);
83 }
84 }
85 }
86 for env in ["VENICE_API_KEY", "VENICE_KEY"] {
87 if let Ok(value) = std::env::var(env) {
88 let cleaned = Self::clean(&value);
89 if !cleaned.is_empty() {
90 return Ok(cleaned);
91 }
92 }
93 }
94 Err(SpendPanelError::AuthFailed(
95 "venice".into(),
96 "no API key found in config, token, VENICE_API_KEY, or VENICE_KEY".into(),
97 ))
98 }
99
100 fn balance_url(&self, ctx: &ProviderContext) -> String {
101 ctx.config
102 .get("balance_url")
103 .or_else(|| ctx.config.get("api_url"))
104 .or_else(|| ctx.config.get("base_url"))
105 .map(String::as_str)
106 .filter(|v| !v.is_empty())
107 .map(Self::clean)
108 .or_else(|| self.balance_url.clone())
109 .unwrap_or_else(|| "https://api.venice.ai/api/v1/billing/balance".into())
110 }
111
112 fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
113 reqwest::Client::builder()
114 .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
115 .build()
116 .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
117 }
118
119 fn parse_response(body: &str) -> Result<VeniceUsage, SpendPanelError> {
120 let decoded: BalanceResponse = serde_json::from_str(body)
121 .map_err(|e| SpendPanelError::ParseError("venice".into(), e.to_string()))?;
122 Ok(VeniceUsage {
123 can_consume: decoded.can_consume,
124 consumption_currency: decoded.consumption_currency,
125 diem_balance: decoded.balances.diem,
126 usd_balance: decoded.balances.usd,
127 diem_epoch_allocation: decoded.diem_epoch_allocation,
128 })
129 }
130
131 async fn fetch_balance(
132 client: &reqwest::Client,
133 url: String,
134 api_key: &str,
135 ) -> Result<VeniceUsage, SpendPanelError> {
136 let resp = client
137 .get(url)
138 .header("Authorization", format!("Bearer {}", api_key))
139 .header("Accept", "application/json")
140 .send()
141 .await
142 .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
143 let status = resp.status();
144 let body = resp
145 .text()
146 .await
147 .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
148 if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
149 return Err(SpendPanelError::AuthFailed(
150 "venice".into(),
151 format!("invalid API key (HTTP {})", status.as_u16()),
152 ));
153 }
154 if !status.is_success() {
155 return Err(SpendPanelError::ProviderError(
156 "venice".into(),
157 format!("HTTP {}", status),
158 ));
159 }
160 Self::parse_response(&body)
161 }
162
163 fn balance_window(usage: &VeniceUsage) -> RateWindow {
164 let active = usage
165 .consumption_currency
166 .as_deref()
167 .map(str::to_ascii_uppercase);
168 let (label, ratio) = if !usage.can_consume {
169 ("Balance unavailable for API calls".into(), 1.0)
170 } else if active.as_deref() == Some("USD") && usage.usd_balance.unwrap_or(0.0) > 0.0 {
171 (
172 format!("${:.2} USD remaining", usage.usd_balance.unwrap()),
173 0.0,
174 )
175 } else if active.as_deref() != Some("USD") {
176 if let (Some(diem), Some(allocation)) =
177 (usage.diem_balance, usage.diem_epoch_allocation)
178 && allocation > 0.0
179 {
180 let used = ((allocation - diem) / allocation).clamp(0.0, 1.0);
181 (
182 format!("DIEM {:.2} / {:.2} epoch allocation", diem, allocation),
183 used,
184 )
185 } else if usage.diem_balance.unwrap_or(0.0) > 0.0 {
186 (
187 format!("DIEM {:.2} remaining", usage.diem_balance.unwrap()),
188 0.0,
189 )
190 } else if usage.usd_balance.unwrap_or(0.0) > 0.0 {
191 (
192 format!("${:.2} USD remaining", usage.usd_balance.unwrap()),
193 0.0,
194 )
195 } else {
196 ("No Venice API balance available".into(), 1.0)
197 }
198 } else if usage.usd_balance.unwrap_or(0.0) > 0.0 {
199 (
200 format!("${:.2} USD remaining", usage.usd_balance.unwrap()),
201 0.0,
202 )
203 } else {
204 ("No Venice API balance available".into(), 1.0)
205 };
206
207 RateWindow {
208 label,
209 window_minutes: 0,
210 usage_ratio: ratio,
211 limit: None,
212 used: None,
213 remaining: None,
214 resets_at: None,
215 status: RateWindowStatus::from_ratio(ratio),
216 }
217 }
218
219 fn snapshot_from_usage(usage: VeniceUsage) -> UsageSnapshot {
220 let mut snapshot = UsageSnapshot::new("venice");
221 snapshot.primary_rate_window = Some(Self::balance_window(&usage));
222 if let Some(usd) = usage.usd_balance {
223 snapshot.credits = Some(CreditsSnapshot::new(usd, "USD"));
224 } else if let Some(diem) = usage.diem_balance {
225 let mut credits = CreditsSnapshot::new(diem, "DIEM");
226 credits.total = usage.diem_epoch_allocation;
227 snapshot.credits = Some(credits);
228 }
229 let mut features = Vec::new();
230 if let Some(currency) = &usage.consumption_currency {
231 features.push(format!("consumption currency: {}", currency));
232 }
233 features.push(format!("can consume: {}", usage.can_consume));
234 snapshot.plan = Some(PlanInfo {
235 name: "Venice API".into(),
236 tier: None,
237 features,
238 price: None,
239 currency: usage.consumption_currency.clone(),
240 billing_period: None,
241 });
242 snapshot
243 }
244}
245
246impl Default for VeniceProvider {
247 fn default() -> Self {
248 Self::new()
249 }
250}
251
252#[async_trait]
253impl UsageProvider for VeniceProvider {
254 fn metadata(&self) -> &ProviderMetadata {
255 &self.metadata
256 }
257
258 fn detect_credentials(&self) -> bool {
259 ["VENICE_API_KEY", "VENICE_KEY"]
260 .iter()
261 .any(|env| std::env::var(env).is_ok_and(|v| !Self::clean(&v).is_empty()))
262 }
263
264 async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
265 let api_key = Self::resolve_api_key(ctx)?;
266 let client = Self::build_client(ctx)?;
267 let usage = Self::fetch_balance(&client, self.balance_url(ctx), &api_key).await?;
268 Ok(Self::snapshot_from_usage(usage))
269 }
270}
271
272fn deserialize_opt_f64<'de, D>(deserializer: D) -> Result<Option<f64>, D::Error>
273where
274 D: Deserializer<'de>,
275{
276 let value = Option::<serde_json::Value>::deserialize(deserializer)?;
277 match value {
278 None | Some(serde_json::Value::Null) => Ok(None),
279 Some(serde_json::Value::Number(n)) => n
280 .as_f64()
281 .ok_or_else(|| de::Error::custom("number cannot be represented as f64"))
282 .map(Some),
283 Some(serde_json::Value::String(s)) => {
284 let trimmed = s.trim();
285 if trimmed.is_empty() {
286 Ok(None)
287 } else {
288 trimmed.parse::<f64>().map(Some).map_err(de::Error::custom)
289 }
290 }
291 Some(other) => Err(de::Error::custom(format!(
292 "expected number/string/null, got {other}"
293 ))),
294 }
295}
296
297#[cfg(test)]
298mod tests {
299 use super::*;
300 use pretty_assertions::assert_eq;
301 use wiremock::matchers::{header, method, path};
302 use wiremock::{Mock, MockServer, ResponseTemplate};
303
304 #[test]
305 fn test_provider_metadata() {
306 let meta = VeniceProvider::new().metadata().clone();
307 assert_eq!(meta.id, "venice");
308 assert_eq!(meta.name, "Venice");
309 }
310
311 #[test]
312 fn test_parses_string_encoded_balances_and_allocation() {
313 let json = r#"{"canConsume":true,"consumptionCurrency":"DIEM","balances":{"diem":"90.50","usd":"25.75"},"diemEpochAllocation":"100.0"}"#;
314 let usage = VeniceProvider::parse_response(json).unwrap();
315 assert_eq!(usage.diem_balance, Some(90.50));
316 assert_eq!(usage.usd_balance, Some(25.75));
317 assert_eq!(usage.diem_epoch_allocation, Some(100.0));
318 }
319
320 #[test]
321 fn test_diem_allocation_progress() {
322 let usage = VeniceUsage {
323 can_consume: true,
324 consumption_currency: Some("DIEM".into()),
325 diem_balance: Some(75.0),
326 usd_balance: None,
327 diem_epoch_allocation: Some(100.0),
328 };
329 let window = VeniceProvider::balance_window(&usage);
330 assert_eq!(window.label, "DIEM 75.00 / 100.00 epoch allocation");
331 assert_eq!(window.usage_ratio, 0.25);
332 }
333
334 #[test]
335 fn test_usd_display_when_active_currency_usd() {
336 let usage = VeniceUsage {
337 can_consume: true,
338 consumption_currency: Some("USD".into()),
339 diem_balance: Some(50.0),
340 usd_balance: Some(12.34),
341 diem_epoch_allocation: Some(100.0),
342 };
343 let window = VeniceProvider::balance_window(&usage);
344 assert_eq!(window.label, "$12.34 USD remaining");
345 assert_eq!(window.usage_ratio, 0.0);
346 }
347
348 #[test]
349 fn test_can_consume_false_exhausts_window() {
350 let usage = VeniceUsage {
351 can_consume: false,
352 consumption_currency: Some("USD".into()),
353 diem_balance: None,
354 usd_balance: Some(100.0),
355 diem_epoch_allocation: None,
356 };
357 let window = VeniceProvider::balance_window(&usage);
358 assert_eq!(window.label, "Balance unavailable for API calls");
359 assert_eq!(window.usage_ratio, 1.0);
360 }
361
362 #[test]
363 fn test_zero_balances() {
364 let usage = VeniceProvider::parse_response(
365 r#"{"canConsume":true,"consumptionCurrency":"USD","balances":{"diem":0,"usd":0},"diemEpochAllocation":null}"#,
366 )
367 .unwrap();
368 let window = VeniceProvider::balance_window(&usage);
369 assert_eq!(window.label, "No Venice API balance available");
370 assert_eq!(window.usage_ratio, 1.0);
371 }
372
373 #[tokio::test]
374 async fn test_fetch_usage_sends_bearer_token() {
375 let server = MockServer::start().await;
376 Mock::given(method("GET"))
377 .and(path("/balance"))
378 .and(header("authorization", "Bearer ven-test"))
379 .and(header("accept", "application/json"))
380 .respond_with(ResponseTemplate::new(200).set_body_raw(
381 r#"{"canConsume":true,"consumptionCurrency":"USD","balances":{"diem":null,"usd":15.5},"diemEpochAllocation":null}"#,
382 "application/json",
383 ))
384 .mount(&server)
385 .await;
386 let provider = VeniceProvider::with_balance_url(&format!("{}/balance", server.uri()));
387 let snapshot = provider
388 .fetch_usage(&ProviderContext::with_api_key("ven-test"))
389 .await
390 .unwrap();
391 assert_eq!(
392 snapshot.primary_rate_window.unwrap().label,
393 "$15.50 USD remaining"
394 );
395 assert_eq!(snapshot.credits.unwrap().balance, 15.5);
396 }
397
398 #[tokio::test]
399 async fn test_fetch_usage_401_is_auth_failed() {
400 let server = MockServer::start().await;
401 Mock::given(method("GET"))
402 .and(path("/balance"))
403 .respond_with(ResponseTemplate::new(401))
404 .mount(&server)
405 .await;
406 let provider = VeniceProvider::with_balance_url(&format!("{}/balance", server.uri()));
407 let err = provider
408 .fetch_usage(&ProviderContext::with_api_key("bad"))
409 .await
410 .unwrap_err();
411 assert!(matches!(err, SpendPanelError::AuthFailed(_, _)));
412 }
413}