use chrono::{DateTime, Datelike, Timelike, Utc};
use reqwest::{
header::{HeaderMap, HeaderValue, AUTHORIZATION, USER_AGENT},
Client,
};
use std::sync::Arc;
use std::time::Duration;
use crate::{
cache::FileCache,
config::{ApiConfig, CacheConfig},
Error, Result,
};
use super::{
AnthropicConfig, AnthropicUsageStats, CostBreakdown, HttpClient, MetricData, RateLimitInfo,
StatisticsData, UsageSummary,
};
#[derive(Debug, Clone)]
pub struct ApiClient {
http: HttpClient,
anthropic: Option<AnthropicApiClient>,
}
impl ApiClient {
pub fn new(config: ApiConfig) -> Result<Self> {
let http = HttpClient::new(config.clone())?;
let anthropic = if let Some(anthropic_config) = config.anthropic {
Some(AnthropicApiClient::new(anthropic_config)?.with_usage_tracking())
} else {
None
};
Ok(Self { http, anthropic })
}
pub fn with_cache(config: ApiConfig, cache_config: CacheConfig) -> Result<Self> {
let http = HttpClient::new(config.clone())?;
let cache = Arc::new(FileCache::new(cache_config));
let anthropic = if let Some(anthropic_config) = config.anthropic {
if !anthropic_config.api_key.is_empty() {
Some(
AnthropicApiClient::new(anthropic_config)?
.with_cache(cache.clone())
.with_usage_tracking(),
)
} else {
None
}
} else {
None
};
Ok(Self { http, anthropic })
}
pub async fn from_config_with_cache(config: crate::config::Config) -> Result<Self> {
let api_config = config.api;
let cache_config = config.cache;
let client = Self::with_cache(api_config, cache_config)?;
if let Some(anthropic_client) = &client.anthropic {
if let Some(cache) = &anthropic_client.cache {
cache.init().await?;
}
}
Ok(client)
}
pub async fn from_env_with_cache() -> Result<Self> {
let config = crate::config::Config::load().await?;
Self::from_config_with_cache(config).await
}
pub fn anthropic(&self) -> Option<&AnthropicApiClient> {
self.anthropic.as_ref()
}
pub async fn fetch_claude_usage_stats(
&self,
start_time: DateTime<Utc>,
end_time: DateTime<Utc>,
) -> Result<AnthropicUsageStats> {
let anthropic_client = self
.anthropic
.as_ref()
.ok_or_else(|| Error::config("Anthropic API not configured"))?;
anthropic_client
.fetch_usage_stats(start_time, end_time)
.await
}
pub async fn fetch_rate_limit_info(&self) -> Result<RateLimitInfo> {
let anthropic_client = self
.anthropic
.as_ref()
.ok_or_else(|| Error::config("Anthropic API not configured"))?;
anthropic_client.fetch_rate_limit_info().await
}
pub async fn fetch_billing_info(
&self,
start_time: DateTime<Utc>,
end_time: DateTime<Utc>,
) -> Result<CostBreakdown> {
let anthropic_client = self
.anthropic
.as_ref()
.ok_or_else(|| Error::config("Anthropic API not configured"))?;
anthropic_client
.fetch_billing_info(start_time, end_time)
.await
}
pub async fn fetch_daily_usage_stats(&self) -> Result<AnthropicUsageStats> {
let end_time = Utc::now();
let start_time = end_time - chrono::Duration::days(1);
self.fetch_claude_usage_stats(start_time, end_time).await
}
pub async fn fetch_weekly_usage_stats(&self) -> Result<AnthropicUsageStats> {
let end_time = Utc::now();
let start_time = end_time - chrono::Duration::days(7);
self.fetch_claude_usage_stats(start_time, end_time).await
}
pub async fn fetch_monthly_usage_stats(&self) -> Result<AnthropicUsageStats> {
let end_time = Utc::now();
let start_time = end_time - chrono::Duration::days(30);
self.fetch_claude_usage_stats(start_time, end_time).await
}
pub async fn fetch_current_month_usage_stats(&self) -> Result<AnthropicUsageStats> {
let now = Utc::now();
let start_time = now
.with_day(1)
.unwrap()
.with_hour(0)
.unwrap()
.with_minute(0)
.unwrap()
.with_second(0)
.unwrap();
self.fetch_claude_usage_stats(start_time, now).await
}
pub async fn fetch_current_month_billing(&self) -> Result<CostBreakdown> {
let now = Utc::now();
let start_time = now
.with_day(1)
.unwrap()
.with_hour(0)
.unwrap()
.with_minute(0)
.unwrap()
.with_second(0)
.unwrap();
self.fetch_billing_info(start_time, now).await
}
pub async fn get_usage_summary(&self) -> Result<UsageSummary> {
let _anthropic_client = self
.anthropic
.as_ref()
.ok_or_else(|| Error::config("Anthropic API not configured"))?;
let (daily_stats, weekly_stats, monthly_stats, rate_limit) = tokio::try_join!(
self.fetch_daily_usage_stats(),
self.fetch_weekly_usage_stats(),
self.fetch_monthly_usage_stats(),
self.fetch_rate_limit_info()
)?;
Ok(UsageSummary {
daily: daily_stats,
weekly: weekly_stats,
monthly: monthly_stats,
rate_limit,
timestamp: Utc::now(),
})
}
pub async fn submit_statistics(&self, data: &StatisticsData) -> Result<()> {
if let Some(base_url) = &self.http.config().base_url {
let url = format!("{}/api/v1/statistics", base_url);
let response = self.http.client().post(&url).json(data).send().await?;
if !response.status().is_success() {
return Err(Error::api(format!(
"Failed to submit statistics: {}",
response.status()
)));
}
}
Ok(())
}
pub async fn get_metrics(&self, query: &str) -> Result<Vec<MetricData>> {
if let Some(base_url) = &self.http.config().base_url {
let url = format!("{}/api/v1/metrics?q={}", base_url, query);
let response = self.http.client().get(&url).send().await?;
if !response.status().is_success() {
return Err(Error::api(format!(
"Failed to get metrics: {}",
response.status()
)));
}
let metrics: Vec<MetricData> = response.json().await?;
return Ok(metrics);
}
Ok(vec![])
}
pub async fn health_check(&self) -> Result<bool> {
if let Some(base_url) = &self.http.config().base_url {
let url = format!("{}/health", base_url);
let response = self.http.client().get(&url).send().await?;
Ok(response.status().is_success())
} else {
Ok(true) }
}
}
#[derive(Debug, Clone)]
pub struct AnthropicApiClient {
client: Client,
config: AnthropicConfig,
cache: Option<Arc<FileCache>>,
usage_tracker: Option<Arc<super::LocalUsageTracker>>,
}
impl AnthropicApiClient {
pub fn new(config: AnthropicConfig) -> Result<Self> {
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {}", config.api_key))
.map_err(|e| Error::config(format!("Invalid API key format: {}", e)))?,
);
headers.insert(USER_AGENT, HeaderValue::from_static("cstats-client/1.0.0"));
headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
let client = Client::builder()
.timeout(Duration::from_secs(config.timeout_seconds))
.default_headers(headers)
.build()
.map_err(Error::Http)?;
Ok(Self {
client,
config,
cache: None,
usage_tracker: None,
})
}
pub fn from_env() -> Result<Self> {
let api_key = std::env::var("ANTHROPIC_API_KEY")
.map_err(|_| Error::config("ANTHROPIC_API_KEY environment variable not set"))?;
let mut config = AnthropicConfig {
api_key,
..Default::default()
};
if let Ok(base_url) = std::env::var("ANTHROPIC_BASE_URL") {
config.base_url = base_url;
}
if let Ok(timeout) = std::env::var("ANTHROPIC_TIMEOUT_SECONDS") {
if let Ok(timeout_val) = timeout.parse() {
config.timeout_seconds = timeout_val;
}
}
Self::new(config)
}
pub fn with_cache(mut self, cache: Arc<FileCache>) -> Self {
self.cache = Some(cache);
self
}
pub fn with_usage_tracking(mut self) -> Self {
self.usage_tracker = Some(Arc::new(super::LocalUsageTracker::new()));
self
}
pub fn usage_tracker(&self) -> Option<&Arc<super::LocalUsageTracker>> {
self.usage_tracker.as_ref()
}
pub async fn fetch_usage_stats(
&self,
start_time: DateTime<Utc>,
end_time: DateTime<Utc>,
) -> Result<AnthropicUsageStats> {
if let Some(tracker) = &self.usage_tracker {
tracing::info!(
"Returning locally tracked usage statistics for period {} to {}",
start_time,
end_time
);
return tracker.get_usage_stats(start_time, end_time).await;
}
tracing::warn!(
"Anthropic API does not provide public usage statistics endpoints. \
Returning mock data for period {} to {}. \
Enable local usage tracking for actual statistics.",
start_time,
end_time
);
Ok(self.create_mock_usage_stats(start_time, end_time))
}
pub async fn fetch_rate_limit_info(&self) -> Result<RateLimitInfo> {
if let Some(tracker) = &self.usage_tracker {
return tracker.get_rate_limit_info().await;
}
tracing::warn!(
"Anthropic API does not provide public rate limit endpoints. \
Returning estimated rate limits based on service tier. \
Actual rate limits are enforced through response headers."
);
Ok(RateLimitInfo {
requests_per_minute: 1000, requests_remaining: 1000, reset_time: Utc::now() + chrono::Duration::seconds(60),
tokens_per_minute: Some(50_000), tokens_remaining: Some(50_000), })
}
pub async fn fetch_billing_info(
&self,
start_time: DateTime<Utc>,
end_time: DateTime<Utc>,
) -> Result<CostBreakdown> {
tracing::warn!(
"Anthropic API does not provide public billing endpoints. \
Returning mock data for period {} to {}. \
For actual billing information, visit the Anthropic Console at console.anthropic.com.",
start_time,
end_time
);
Ok(self.create_mock_cost_breakdown(start_time, end_time))
}
pub async fn health_check(&self) -> Result<bool> {
let test_payload = serde_json::json!({
"model": "claude-3-haiku-20240307",
"messages": [{
"role": "user",
"content": "Hello"
}]
});
match self
.client
.post(format!("{}/v1/messages/count-tokens", self.config.base_url))
.json(&test_payload)
.send()
.await
{
Ok(response) => {
let status = response.status();
Ok(status.is_success() || status == reqwest::StatusCode::UNAUTHORIZED)
}
Err(_) => Ok(false),
}
}
fn create_mock_usage_stats(
&self,
start_time: DateTime<Utc>,
end_time: DateTime<Utc>,
) -> AnthropicUsageStats {
use super::{ApiCallStats, TokenUsage, UsagePeriod};
use std::collections::HashMap;
AnthropicUsageStats {
token_usage: TokenUsage {
input_tokens: 0,
output_tokens: 0,
total_tokens: 0,
by_model: HashMap::new(),
},
api_calls: ApiCallStats {
total_calls: 0,
successful_calls: 0,
failed_calls: 0,
avg_response_time_ms: 0.0,
by_model: HashMap::new(),
hourly_breakdown: vec![],
},
costs: self.create_mock_cost_breakdown(start_time, end_time),
model_usage: vec![],
period: UsagePeriod {
start: start_time,
end: end_time,
period_type: "mock".to_string(),
},
}
}
fn create_mock_cost_breakdown(
&self,
_start_time: DateTime<Utc>,
_end_time: DateTime<Utc>,
) -> CostBreakdown {
use super::TokenCostBreakdown;
use std::collections::HashMap;
CostBreakdown {
total_cost_usd: 0.0,
by_model: HashMap::new(),
by_token_type: TokenCostBreakdown {
input_cost_usd: 0.0,
output_cost_usd: 0.0,
},
estimated_monthly_cost_usd: 0.0,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::{ApiConfig, Config};
#[tokio::test]
#[ignore = "Integration test - requires ANTHROPIC_API_KEY environment variable"]
async fn example_anthropic_api_usage() -> Result<()> {
let client = ApiClient::from_env_with_cache().await?;
if let Some(anthropic_client) = client.anthropic() {
let is_healthy = anthropic_client.health_check().await?;
println!("API health check: {}", is_healthy);
match client.fetch_rate_limit_info().await {
Ok(rate_limit) => {
println!(
"Rate limit - Remaining: {}, Reset: {:?}",
rate_limit.requests_remaining, rate_limit.reset_time
);
}
Err(e) => println!("Failed to fetch rate limit info: {}", e),
}
match client.fetch_daily_usage_stats().await {
Ok(daily_stats) => {
println!(
"Daily usage - Total tokens: {}",
daily_stats.token_usage.total_tokens
);
println!(
"Daily usage - API calls: {}",
daily_stats.api_calls.total_calls
);
println!(
"Daily usage - Total cost: ${:.4}",
daily_stats.costs.total_cost_usd
);
}
Err(e) => println!("Failed to fetch daily usage stats: {}", e),
}
match client.get_usage_summary().await {
Ok(summary) => {
println!("Usage Summary:");
println!(" Daily tokens: {}", summary.daily.token_usage.total_tokens);
println!(
" Weekly tokens: {}",
summary.weekly.token_usage.total_tokens
);
println!(
" Monthly tokens: {}",
summary.monthly.token_usage.total_tokens
);
println!(
" Rate limit remaining: {}",
summary.rate_limit.requests_remaining
);
}
Err(e) => println!("Failed to fetch usage summary: {}", e),
}
match client.fetch_current_month_billing().await {
Ok(billing) => {
println!("Current month billing: ${:.4}", billing.total_cost_usd);
for (model, cost) in billing.by_model.iter() {
println!(" {}: ${:.4}", model, cost.cost_usd);
}
}
Err(e) => println!("Failed to fetch billing info: {}", e),
}
} else {
println!("Anthropic API not configured");
}
Ok(())
}
#[tokio::test]
async fn test_config_from_env() -> Result<()> {
let config = Config::from_env()?;
assert_eq!(config.api.timeout_seconds, 30);
assert_eq!(config.api.retry_attempts, 3);
Ok(())
}
#[tokio::test]
async fn test_client_without_anthropic() -> Result<()> {
let config = ApiConfig::default();
let client = ApiClient::new(config)?;
assert!(client.anthropic().is_none());
let result = client.fetch_daily_usage_stats().await;
assert!(result.is_err());
Ok(())
}
}