use crate::ai::gemini::{GeminiClient, GeminiError, GenerateContentRequest};
use crate::ai::quota::QuotaManager;
use axum::{
extract::{Json, State},
http::StatusCode,
response::IntoResponse,
};
use std::sync::Arc;
use tracing::error;
pub struct ProxyState {
pub client: Arc<GeminiClient>,
pub quota_manager: Arc<QuotaManager>,
}
pub async fn handle_generate(
State(state): State<Arc<ProxyState>>,
Json(request): Json<GenerateContentRequest>,
) -> impl IntoResponse {
loop {
let _slept = state.quota_manager.wait_for_access().await;
match state.client.generate_content_single(&request).await {
Ok(response) => {
state.quota_manager.report_success().await;
return (StatusCode::OK, Json(response)).into_response();
}
Err(e) => {
if let Some(err) = e.downcast_ref::<GeminiError>() {
match err {
GeminiError::QuotaExceeded(duration) => {
state.quota_manager.report_quota_error(*duration).await;
continue;
}
GeminiError::TransientError(_, _) => {
state.quota_manager.report_transient_error().await;
continue;
}
_ => {}
}
}
error!("Gemini Proxy Error: {}", e);
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(serde_json::json!({"error": format!("{:#}", e)})),
)
.into_response();
}
}
}
}