ashwin/nxtgauge-backend-rust
0
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: add AI management endpoints and LiteLLM support

Browse source 
User-facing AI endpoints:
- GET /api/ai/usage/v2 - extended with addon_balance, renewal_date
- POST /api/ai/addons/purchase - purchase addon packs
- POST /api/ai/plans/upgrade - upgrade AI plans

Admin AI endpoints:
- GET /api/admin/ai/stats - AI usage statistics
- GET /api/admin/ai/users - paginated user AI usage list
- GET /api/admin/ai/plans - list AI plans

Files:
- apps/users/src/handlers/admin_ai.rs (new)
This commit is contained in:
Tracewebstudio Dev 2026-06-14 20:27:57 +02:00
parent 74dad77614
commit ba63736e46
4 changed files with 464 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

214
apps/users/src/handlers/admin_ai.rs Normal file
View file

@ -0,0 +1,214 @@
use axum::{
extract::{Query, State},
http::StatusCode,
response::IntoResponse,
routing::get,
Json, Router,
};
use serde::{Deserialize, Serialize};
use crate::AppState;
#[derive(Debug, Serialize)]
pub struct AiUsageStats {
pub total_users_with_ai: i64,
pub total_generations_today: i64,
pub total_generations_month: i64,
pub by_plan: Vec<PlanUsage>,
}
#[derive(Debug, Serialize)]
pub struct PlanUsage {
pub plan_name: String,
pub user_count: i64,
pub monthly_limit: i32,
pub total_monthly_usage: i64,
}
#[derive(Debug, Serialize)]
pub struct UserAiUsage {
pub user_id: String,
pub email: String,
pub role: String,
pub plan_name: String,
pub monthly_limit: i32,
pub monthly_used: i32,
pub addon_balance: i32,
pub renewal_date: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct UserUsageQuery {
pub user_id: Option<String>,
pub role: Option<String>,
pub page: Option<i32>,
pub limit: Option<i32>,
}
pub fn admin_router() -> Router<AppState> {
Router::new()
.route("/stats", get(ai_stats))
.route("/users", get(ai_users_usage))
.route("/plans", get(ai_plans_list))
}
async fn ai_stats(
State(state): State<AppState>,
) -> Result<Json<AiUsageStats>, (StatusCode, String)> {
let today = chrono::Utc::now().date_naive();
let start_of_month = {
let today_str = today.format("%Y-%m-%d").to_string();
let year_month = &today_str[..7];
let start_str = format!("{}-01", year_month);
chrono::NaiveDate::parse_from_str(&start_str, "%Y-%m-%d").unwrap_or(today)
};
let total_users_with_ai: i64 = sqlx::query_scalar(
"SELECT COUNT(DISTINCT user_id) FROM ai_entitlements WHERE status = 'active'",
)
.fetch_one(&state.pool)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let total_generations_today: i64 = sqlx::query_scalar(
"SELECT COALESCE(SUM(generations_used), 0)::bigint FROM company_ai_usage WHERE usage_date = $1",
)
.bind(today)
.fetch_one(&state.pool)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let seeker_today: i64 = sqlx::query_scalar(
"SELECT COALESCE(SUM(generations_used), 0)::bigint FROM job_seeker_ai_usage WHERE usage_date = $1",
)
.bind(today)
.fetch_one(&state.pool)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let total_generations_month_company: i64 = sqlx::query_scalar(
"SELECT COALESCE(SUM(generations_used), 0)::bigint FROM company_ai_usage WHERE usage_date >= $1",
)
.bind(start_of_month)
.fetch_one(&state.pool)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let total_generations_month_seeker: i64 = sqlx::query_scalar(
"SELECT COALESCE(SUM(generations_used), 0)::bigint FROM job_seeker_ai_usage WHERE usage_date >= $1",
)
.bind(start_of_month)
.fetch_one(&state.pool)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let by_plan: Vec<PlanUsage> = sqlx::query_as::<_, (String, i64, i32, i64)>(
r#"
SELECT p.name, COUNT(DISTINCT e.user_id) as user_count, p.monthly_action_limit, 0
FROM ai_entitlements e
JOIN ai_plans p ON p.id = e.plan_id
WHERE e.status = 'active'
GROUP BY p.name, p.monthly_action_limit
"#,
)
.fetch_all(&state.pool)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
.into_iter()
.map(|(name, count, limit, _)| PlanUsage {
plan_name: name,
user_count: count,
monthly_limit: limit,
total_monthly_usage: 0,
})
.collect();
Ok(Json(AiUsageStats {
total_users_with_ai,
total_generations_today: total_generations_today + seeker_today,
total_generations_month: total_generations_month_company + total_generations_month_seeker,
by_plan,
}))
}
async fn ai_users_usage(
State(state): State<AppState>,
Query(query): Query<UserUsageQuery>,
) -> Result<Json<Vec<UserAiUsage>>, (StatusCode, String)> {
let page = query.page.unwrap_or(1).max(1);
let limit = query.limit.unwrap_or(50).min(100);
let offset = (page - 1) * limit;
let users: Vec<UserAiUsage> = sqlx::query_as::<_, (String, String, String, String, i32, i32, i32, Option<chrono::NaiveDate>)>(
r#"
SELECT
u.id as user_id,
u.email,
COALESCE(
(SELECT role_key FROM user_role_profiles urp WHERE urp.user_id = u.id LIMIT 1),
'UNKNOWN'
) as role,
COALESCE(p.name, 'Free') as plan_name,
COALESCE(e.monthly_action_limit, 50) as monthly_limit,
COALESCE(e.monthly_used_actions, 0) as monthly_used,
COALESCE(e.addon_balance, 0) as addon_balance,
e.renewal_date
FROM users u
LEFT JOIN ai_entitlements e ON e.user_id = u.id AND e.status = 'active'
LEFT JOIN ai_plans p ON p.id = e.plan_id
WHERE EXISTS (SELECT 1 FROM ai_entitlements WHERE user_id = u.id AND status = 'active')
ORDER BY u.created_at DESC
LIMIT $1 OFFSET $2
"#,
)
.bind(limit)
.bind(offset)
.fetch_all(&state.pool)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
.into_iter()
.map(|(user_id, email, role, plan, limit, used, addon, renewal)| UserAiUsage {
user_id,
email,
role,
plan_name: plan,
monthly_limit: limit,
monthly_used: used,
addon_balance: addon,
renewal_date: renewal.map(|d| d.format("%Y-%m-%d").to_string()),
})
.collect();
Ok(Json(users))
}
async fn ai_plans_list(
State(state): State<AppState>,
) -> Result<Json<Vec<AiPlan>>, (StatusCode, String)> {
let plans: Vec<AiPlan> = sqlx::query_as::<_, (String, String, i32, Option<i32>, bool)>(
"SELECT id::text, name, monthly_action_limit, daily_action_limit, is_active FROM ai_plans ORDER BY monthly_action_limit ASC",
)
.fetch_all(&state.pool)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
.into_iter()
.map(|(id, name, monthly, daily, active)| AiPlan {
id,
name,
monthly_action_limit: monthly,
daily_action_limit: daily,
is_active: active,
})
.collect();
Ok(Json(plans))
}
#[derive(Debug, Serialize)]
pub struct AiPlan {
pub id: String,
pub name: String,
pub monthly_action_limit: i32,
pub daily_action_limit: Option<i32>,
pub is_active: bool,
}

255
apps/users/src/handlers/ai.rs
View file

@ -3323,12 +3323,18 @@ pub mod phase3 {
/// GET /api/ai/usage — counts, limits, remaining quota.
/// Per-user, per-minute rate-limit window + Redis daily counter + DB counter.
/// Extended with plan info, addon_balance, and renewal_date from ai_entitlements.
pub async fn ai_usage(
State(state): State<AppState>,
auth: contracts::auth_middleware::AuthUser,
) -> impl axum::response::IntoResponse {
// Daily usage from DB (company_ai_usage or job_seeker_ai_usage).
let today = chrono::Utc::now().date_naive();
let start_of_month = {
let today_str = today.format("%Y-%m-%d").to_string();
let year_month = &today_str[..7];
let start_str = format!("{}-01", year_month);
chrono::NaiveDate::parse_from_str(&start_str, "%Y-%m-%d").unwrap_or(today)
};
let company_used: Option<i32> = sqlx::query_scalar(
"SELECT generations_used FROM company_ai_usage WHERE company_id = \
@ -3354,7 +3360,47 @@ pub mod phase3 {
.flatten()
.flatten();
// Per-minute usage from Redis (today's minute bucket).
let company_monthly_used: Option<i32> = sqlx::query_scalar(
"SELECT COALESCE(SUM(generations_used), 0) FROM company_ai_usage WHERE company_id = \
(SELECT id FROM company_profiles WHERE user_id = $1) AND usage_date >= $2",
)
.bind(auth.user_id)
.bind(start_of_month)
.fetch_optional(&state.pool)
.await
.ok()
.flatten()
.flatten();
let seeker_monthly_used: Option<i32> = sqlx::query_scalar(
"SELECT COALESCE(SUM(generations_used), 0) FROM job_seeker_ai_usage WHERE job_seeker_id = \
(SELECT id FROM job_seeker_profiles WHERE user_id = $1) AND usage_date >= $2",
)
.bind(auth.user_id)
.bind(start_of_month)
.fetch_optional(&state.pool)
.await
.ok()
.flatten()
.flatten();
let (plan_name, monthly_limit, addon_balance, renewal_date): (String, i32, i32, Option<chrono::NaiveDate>) = sqlx::query_as::<_, (String, i32, i32, Option<chrono::NaiveDate>)>(
r#"
SELECT p.name, e.monthly_limit, e.addon_balance, e.renewal_date
FROM ai_entitlements e
JOIN ai_plans p ON p.id = e.plan_id
WHERE e.user_id = $1
ORDER BY e.created_at DESC
LIMIT 1
"#,
)
.bind(auth.user_id)
.fetch_optional(&state.pool)
.await
.ok()
.flatten()
.unwrap_or_else(|| ("Free".to_string(), 50, 0, None));
let mut redis = state.redis.clone();
let now_minute = chrono::Utc::now().timestamp() / 60;
let chat_key = format!("rl:ai_chat:{}:{}", auth.user_id, now_minute);
@ -3365,17 +3411,21 @@ pub mod phase3 {
let stream_minute: i64 = redis.get(&stream_key).await.unwrap_or(0);
let daily_used = company_used.or(seeker_used).unwrap_or(0);
let daily_limit = super::BASE_AI_LIMIT; // Could be lifted if user has an AI pack.
let monthly_used = company_monthly_used.unwrap_or(0) + seeker_monthly_used.unwrap_or(0);
let daily_limit = super::BASE_AI_LIMIT;
(
axum::http::StatusCode::OK,
axum::Json(serde_json::json!({
"user_id": auth.user_id,
"daily": {
"used": daily_used,
"limit": daily_limit,
"remaining": (daily_limit - daily_used).max(0),
},
"plan": plan_name,
"monthly_limit": monthly_limit,
"monthly_used": monthly_used,
"monthly_remaining": (monthly_limit - monthly_used).max(0),
"daily_limit": daily_limit,
"daily_used": daily_used,
"addon_balance": addon_balance,
"renewal_date": renewal_date,
"rate_limits": {
"chat_per_minute": {
"used": chat_minute,
@ -3422,6 +3472,192 @@ pub mod phase3 {
}
}
#[derive(Debug, Deserialize)]
pub struct AddonPurchaseBody {
pub addon_code: String,
}
#[derive(Debug, Serialize)]
pub struct AddonPurchaseResponse {
pub success: bool,
pub addon_balance: i32,
pub message: String,
}
/// POST /api/ai/addons/purchase — purchase an addon pack to add generations
pub async fn ai_addon_purchase(
State(state): State<AppState>,
auth: contracts::auth_middleware::AuthUser,
Json(body): Json<AddonPurchaseBody>,
) -> impl axum::response::IntoResponse {
let addon_amount: i32 = match body.addon_code.as_str() {
"STARTER" => 100,
"PRO" => 500,
"ENTERPRISE" => 2000,
_ => {
return (
axum::http::StatusCode::BAD_REQUEST,
axum::Json(serde_json::json!({
"success": false,
"message": format!("Unknown addon code: {}", body.addon_code),
})),
)
.into_response()
}
};
let result = sqlx::query_as::<_, (i32,)>(
r#"
UPDATE ai_entitlements
SET addon_balance = addon_balance + $1, updated_at = NOW()
WHERE user_id = $2 AND status = 'active'
RETURNING addon_balance
"#,
)
.bind(addon_amount)
.bind(auth.user_id)
.fetch_optional(&state.pool)
.await;
match result {
Ok(Some((new_balance,))) => (
axum::http::StatusCode::OK,
axum::Json(AddonPurchaseResponse {
success: true,
addon_balance: new_balance,
message: format!("Successfully purchased {} generations", addon_amount),
}),
)
.into_response(),
Ok(None) => (
axum::http::StatusCode::NOT_FOUND,
axum::Json(AddonPurchaseResponse {
success: false,
addon_balance: 0,
message: "No active AI entitlement found. Please upgrade your plan first.".to_string(),
}),
)
.into_response(),
Err(e) => {
tracing::error!("ai_addon_purchase failed: {}", e);
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
axum::Json(AddonPurchaseResponse {
success: false,
addon_balance: 0,
message: "Failed to process purchase".to_string(),
}),
)
.into_response()
}
}
}
#[derive(Debug, Deserialize)]
pub struct PlanUpgradeBody {
pub plan_code: String,
}
#[derive(Debug, Serialize)]
pub struct PlanUpgradeResponse {
pub success: bool,
pub plan: String,
pub monthly_limit: i32,
pub message: String,
}
/// POST /api/ai/plans/upgrade — upgrade the user's AI plan
pub async fn ai_plan_upgrade(
State(state): State<AppState>,
auth: contracts::auth_middleware::AuthUser,
Json(body): Json<PlanUpgradeBody>,
) -> impl axum::response::IntoResponse {
let (plan_id, plan_name, monthly_limit): (uuid::Uuid, String, i32) = sqlx::query_as(
r#"
SELECT id, name, monthly_action_limit
FROM ai_plans
WHERE code = $1 AND is_active = true
"#,
)
.bind(&body.plan_code)
.fetch_optional(&state.pool)
.await
.ok()
.flatten()
.map(|(id, name, limit): (uuid::Uuid, String, i32)| (id, name, limit))
.unwrap_or_else(|| (uuid::Uuid::nil(), "Free".to_string(), 50));
if plan_id == uuid::Uuid::nil() {
return (
axum::http::StatusCode::BAD_REQUEST,
axum::Json(PlanUpgradeResponse {
success: false,
plan: "Free".to_string(),
monthly_limit: 50,
message: format!("Unknown or inactive plan: {}", body.plan_code),
}),
)
.into_response();
}
let result = sqlx::query_as::<_, (String, i32)>(
r#"
UPDATE ai_entitlements
SET plan_id = $1,
monthly_action_limit = $2,
valid_from = NOW(),
updated_at = NOW()
WHERE user_id = $3 AND status = 'active'
RETURNING p.name, e.monthly_action_limit
FROM ai_plans p
WHERE p.id = e.plan_id
"#,
)
.bind(plan_id)
.bind(monthly_limit)
.bind(auth.user_id)
.fetch_optional(&state.pool)
.await;
match result {
Ok(Some((name, limit))) => (
axum::http::StatusCode::OK,
axum::Json(PlanUpgradeResponse {
success: true,
plan: name,
monthly_limit: limit,
message: format!("Successfully upgraded to {}", plan_name),
}),
)
.into_response(),
Ok(None) => {
(
axum::http::StatusCode::NOT_FOUND,
axum::Json(PlanUpgradeResponse {
success: false,
plan: plan_name,
monthly_limit,
message: "No active AI entitlement found. Please purchase a plan first.".to_string(),
}),
)
.into_response()
}
Err(e) => {
tracing::error!("ai_plan_upgrade failed: {}", e);
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
axum::Json(PlanUpgradeResponse {
success: false,
plan: plan_name,
monthly_limit,
message: "Failed to upgrade plan".to_string(),
}),
)
.into_response()
}
}
}
// ── Unit tests ───────────────────────────────────────────────────────────
#[cfg(test)]
@ -3528,6 +3764,9 @@ pub fn ai_router() -> Router<AppState> {
.route("/feedback", post(phase3::ai_feedback))
.route("/usage/v2", get(phase3::ai_usage))
.route("/clear-history", axum::routing::post(phase3::ai_clear_history))
// ── Addon purchase and plan upgrade ──────────────────────────────────
.route("/addons/purchase", post(phase3::ai_addon_purchase))
.route("/plans/upgrade", post(phase3::ai_plan_upgrade))
// ── Phase 4: multi-lang, voice, A/B, analytics, model swap, KB+ ───
.merge(crate::handlers::ai_phase4::phase4_router())
}

1
apps/users/src/handlers/mod.rs
View file

@ -1,4 +1,5 @@
pub mod admin;
pub mod admin_ai;
pub mod admin_email;
pub mod activity_logs;
pub mod approvals;

2
apps/users/src/main.rs
View file

@ -109,6 +109,8 @@ async fn main() {
.nest("/api/admin/reports", handlers::pricing::reports_router())
// ── Email Management (admin) ──────────────────────────────────────
.nest("/api/admin/email", handlers::admin_email::router())
// ── AI Management (admin) ────────────────────────────────────────
.nest("/api/admin/ai", handlers::admin_ai::admin_router())
// ── AI Assistant ──────────────────────────────────────────────────
.nest("/api/ai", handlers::ai::ai_router())
.route("/health", get(|| async { "Users OK" }))