feat: Enforce Automatic Token Limits in AI Handlers
CMS Service: - generate_quiz: Check 2000 tokens limit - generate_course: Check 5000 tokens limit - generate_hotspots: Check 2000 tokens limit - generate_role_play: Check 2500 tokens limit - summarize_lesson: Check 1500 tokens limit LMS Service: - chat_with_tutor: Check 1000 tokens limit - evaluate_audio_response: Check 1500 tokens limit Behavior: - Returns HTTP 429 Too Many Requests when limit exceeded - Error message indicates monthly limit exceeded - Users must contact admin for limit increase - Works with automatic alert system (Phase 3) Total Functions Protected: 7 Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
This commit is contained in:
@@ -1074,6 +1074,11 @@ pub async fn summarize_lesson(
|
||||
Path(id): Path<Uuid>,
|
||||
) -> Result<Json<Lesson>, StatusCode> {
|
||||
tracing::info!("Received summarization request for lesson: {}", id);
|
||||
|
||||
// Check token limit before proceeding (estimate 1500 tokens for summary)
|
||||
if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 1500).await {
|
||||
return Err(StatusCode::TOO_MANY_REQUESTS);
|
||||
}
|
||||
// 1. Fetch lesson
|
||||
let lesson =
|
||||
sqlx::query_as::<_, Lesson>("SELECT * FROM lessons WHERE id = $1 AND organization_id = $2")
|
||||
@@ -1213,6 +1218,11 @@ pub async fn generate_quiz(
|
||||
Json(quiz_req): Json<QuizAIRequest>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
tracing::info!("Received quiz generation request for lesson: {}", id);
|
||||
|
||||
// Check token limit before proceeding (estimate 2000 tokens for quiz)
|
||||
if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 2000).await {
|
||||
return Err(StatusCode::TOO_MANY_REQUESTS);
|
||||
}
|
||||
// 1. Fetch lesson
|
||||
let lesson =
|
||||
sqlx::query_as::<_, Lesson>("SELECT * FROM lessons WHERE id = $1 AND organization_id = $2")
|
||||
@@ -2086,11 +2096,16 @@ pub struct GenerateHotspotsPayload {
|
||||
|
||||
pub async fn generate_hotspots(
|
||||
Org(org_ctx): Org,
|
||||
_claims: common::auth::Claims,
|
||||
claims: common::auth::Claims,
|
||||
State(pool): State<PgPool>,
|
||||
Path(lesson_id): Path<Uuid>,
|
||||
Json(payload): Json<GenerateHotspotsPayload>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
// Check token limit before proceeding (estimate 2000 tokens for hotspots)
|
||||
if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 2000).await {
|
||||
return Err(StatusCode::TOO_MANY_REQUESTS);
|
||||
}
|
||||
|
||||
// 1. Resolve image path
|
||||
// imageUrl in frontend is like "/assets/filename.ext"
|
||||
// We need to map it to "uploads/filename.ext"
|
||||
@@ -2267,7 +2282,7 @@ pub async fn generate_hotspots(
|
||||
let total_tokens = input_tokens + output_tokens;
|
||||
|
||||
let _ = sqlx::query("SELECT log_ai_usage($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)")
|
||||
.bind(_claims.sub)
|
||||
.bind(claims.sub)
|
||||
.bind(org_ctx.id)
|
||||
.bind(total_tokens)
|
||||
.bind(input_tokens)
|
||||
@@ -2294,11 +2309,16 @@ pub struct GenerateRolePlayPayload {
|
||||
|
||||
pub async fn generate_role_play(
|
||||
Org(org_ctx): Org,
|
||||
_claims: common::auth::Claims,
|
||||
claims: common::auth::Claims,
|
||||
State(pool): State<PgPool>,
|
||||
Path(lesson_id): Path<Uuid>,
|
||||
Json(payload): Json<GenerateRolePlayPayload>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
// Check token limit before proceeding (estimate 2500 tokens for role-play)
|
||||
if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 2500).await {
|
||||
return Err(StatusCode::TOO_MANY_REQUESTS);
|
||||
}
|
||||
|
||||
// 1. Fetch lesson context
|
||||
let lesson = sqlx::query_as::<_, Lesson>("SELECT * FROM lessons WHERE id = $1 AND organization_id = $2")
|
||||
.bind(lesson_id)
|
||||
@@ -2396,7 +2416,7 @@ pub async fn generate_role_play(
|
||||
let total_tokens = input_tokens + output_tokens;
|
||||
|
||||
let _ = sqlx::query("SELECT log_ai_usage($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)")
|
||||
.bind(_claims.sub)
|
||||
.bind(claims.sub)
|
||||
.bind(org_ctx.id)
|
||||
.bind(total_tokens)
|
||||
.bind(input_tokens)
|
||||
@@ -4032,6 +4052,11 @@ pub async fn generate_course(
|
||||
"Starting AI course generation for prompt: {}",
|
||||
payload.prompt
|
||||
);
|
||||
|
||||
// Check token limit before proceeding (estimate 5000 tokens for course generation)
|
||||
if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 5000).await {
|
||||
return Err(StatusCode::TOO_MANY_REQUESTS);
|
||||
}
|
||||
|
||||
// 1. Determine target org
|
||||
let target_org_id = payload.target_organization_id.unwrap_or(org_ctx.id);
|
||||
|
||||
@@ -2102,9 +2102,14 @@ pub async fn get_recommendations(
|
||||
|
||||
pub async fn evaluate_audio_response(
|
||||
Org(_org_ctx): Org,
|
||||
_claims: Claims,
|
||||
claims: Claims,
|
||||
Json(payload): Json<AudioGradingPayload>,
|
||||
) -> Result<Json<AudioGradingResponse>, (StatusCode, String)> {
|
||||
// Check token limit before proceeding (estimate 1500 tokens for audio evaluation)
|
||||
if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 1500).await {
|
||||
return Err((StatusCode::TOO_MANY_REQUESTS, "Token limit exceeded".to_string()));
|
||||
}
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let provider = env::var("AI_PROVIDER").unwrap_or_else(|_| "openai".to_string());
|
||||
@@ -2496,6 +2501,11 @@ pub async fn chat_with_tutor(
|
||||
Path(lesson_id): Path<Uuid>,
|
||||
Json(payload): Json<ChatPayload>,
|
||||
) -> Result<Json<ChatResponse>, (StatusCode, String)> {
|
||||
// Check token limit before proceeding (estimate 1000 tokens for chat)
|
||||
if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 1000).await {
|
||||
return Err((StatusCode::TOO_MANY_REQUESTS, "Monthly AI token limit exceeded. Please contact your administrator.".to_string()));
|
||||
}
|
||||
|
||||
// 1. Fetch lesson context with access check (matches get_lesson_content)
|
||||
let is_preview = claims.token_type.as_deref() == Some("preview");
|
||||
|
||||
|
||||
Reference in New Issue
Block a user