feat: Monthly token limits per user

Database: - Add monthly_token_limit and token_limit_reset_day to users table - Create ai_usage_monthly view for current month usage - Add check_token_limit() function to verify available tokens - Add get_user_usage_stats() function for historical usage API Endpoints: - PUT /admin/users/{user_id}/token-limit - Set monthly limit - GET /admin/users/{user_id}/token-usage - Get user's current usage - GET /admin/users/{user_id}/token-limit/check - Check if user has tokens Features: - Default limit: 100,000 tokens/month - Reset day: 1st of month (configurable 1-28) - Limit of 0 = unlimited tokens - Enforce limits at API level (check before AI requests) Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
2026-03-23 16:10:27 -03:00
parent ef213a61a0
commit 9e57756528
3 changed files with 265 additions and 1 deletions
@@ -0,0 +1,138 @@
+-- Monthly Token Limits per User
+-- Allows setting and enforcing monthly AI token usage limits
+
+-- Add monthly token limit column to users table
+ALTER TABLE users 
+ADD COLUMN IF NOT EXISTS monthly_token_limit INTEGER DEFAULT 100000,
+ADD COLUMN IF NOT EXISTS token_limit_reset_day INTEGER DEFAULT 1; -- Day of month to reset (1-28)
+
+COMMENT ON COLUMN users.monthly_token_limit IS 'Maximum AI tokens user can consume per month (0 = unlimited)';
+COMMENT ON COLUMN users.token_limit_reset_day IS 'Day of month when token counter resets (1-28, to avoid month-end issues)';
+
+-- Create view for current month usage per user
+CREATE OR REPLACE VIEW ai_usage_monthly AS
+SELECT 
+    au.user_id,
+    au.organization_id,
+    DATE_TRUNC('month', au.created_at + (u.token_limit_reset_day - 1) * INTERVAL '1 day') AS usage_month,
+    COUNT(*) AS total_requests,
+    SUM(au.tokens_used) AS total_tokens,
+    SUM(au.input_tokens) AS input_tokens,
+    SUM(au.output_tokens) AS output_tokens,
+    SUM(au.estimated_cost_usd) AS total_cost_usd
+FROM ai_usage_logs au
+JOIN users u ON au.user_id = u.id
+WHERE au.created_at >= DATE_TRUNC('month', NOW() + (u.token_limit_reset_day - 1) * INTERVAL '1 day') - (u.token_limit_reset_day - 1) * INTERVAL '1 day'
+GROUP BY au.user_id, au.organization_id, usage_month;
+
+COMMENT ON VIEW ai_usage_monthly IS 'Current month AI usage per user with custom reset day';
+
+-- Function to check if user has exceeded token limit
+CREATE OR REPLACE FUNCTION check_token_limit(
+    p_user_id UUID,
+    p_additional_tokens INTEGER DEFAULT 0
+)
+RETURNS TABLE (
+    has_available_tokens BOOLEAN,
+    monthly_limit INTEGER,
+    used_tokens BIGINT,
+    remaining_tokens BIGINT,
+    reset_date TIMESTAMPTZ
+) AS $$
+DECLARE
+    v_limit INTEGER;
+    v_reset_day INTEGER;
+    v_used BIGINT;
+    v_month_start TIMESTAMPTZ;
+    v_next_reset TIMESTAMPTZ;
+BEGIN
+    -- Get user's limit settings
+    SELECT monthly_token_limit, token_limit_reset_day
+    INTO v_limit, v_reset_day
+    FROM users
+    WHERE id = p_user_id;
+    
+    -- Default values if not set
+    v_limit := COALESCE(v_limit, 100000);
+    v_reset_day := COALESCE(v_reset_day, 1);
+    
+    -- If limit is 0, unlimited
+    IF v_limit = 0 THEN
+        RETURN QUERY SELECT 
+            TRUE,
+            v_limit,
+            0::BIGINT,
+            0::BIGINT,
+            NOW() + INTERVAL '1 month';
+        RETURN;
+    END IF;
+    
+    -- Calculate current month start based on reset day
+    v_month_start := DATE_TRUNC('month', NOW() + (v_reset_day - 1) * INTERVAL '1 day') 
+                     - (v_reset_day - 1) * INTERVAL '1 day';
+    
+    -- Calculate next reset date
+    v_next_reset := v_month_start + INTERVAL '1 month';
+    
+    -- Get current usage
+    SELECT COALESCE(SUM(tokens_used), 0)
+    INTO v_used
+    FROM ai_usage_logs
+    WHERE user_id = p_user_id
+      AND created_at >= v_month_start;
+    
+    -- Return results
+    RETURN QUERY SELECT
+        (v_used + p_additional_tokens) < v_limit,
+        v_limit,
+        v_used,
+        (v_limit - v_used)::BIGINT,
+        v_next_reset;
+END;
+$$ LANGUAGE plpgsql;
+
+COMMENT ON FUNCTION check_token_limit IS 'Check if user has available tokens for current period';
+
+-- Function to get usage statistics for a user
+CREATE OR REPLACE FUNCTION get_user_usage_stats(
+    p_user_id UUID,
+    p_months INTEGER DEFAULT 3
+)
+RETURNS TABLE (
+    month_start DATE,
+    total_requests BIGINT,
+    total_tokens BIGINT,
+    input_tokens BIGINT,
+    output_tokens BIGINT,
+    total_cost_usd NUMERIC,
+    monthly_limit INTEGER,
+    percentage_used NUMERIC
+) AS $$
+BEGIN
+    RETURN QUERY
+    SELECT 
+        DATE_TRUNC('month', au.created_at)::DATE AS month_start,
+        COUNT(*) AS total_requests,
+        SUM(au.tokens_used) AS total_tokens,
+        SUM(au.input_tokens) AS input_tokens,
+        SUM(au.output_tokens) AS output_tokens,
+        SUM(au.estimated_cost_usd) AS total_cost_usd,
+        u.monthly_token_limit,
+        CASE 
+            WHEN u.monthly_token_limit = 0 THEN 0
+            ELSE (SUM(au.tokens_used)::NUMERIC / u.monthly_token_limit * 100)::NUMERIC(5,2)
+        END AS percentage_used
+    FROM ai_usage_logs au
+    JOIN users u ON au.user_id = u.id
+    WHERE au.user_id = p_user_id
+      AND au.created_at >= NOW() - (p_months || ' months')::INTERVAL
+    GROUP BY DATE_TRUNC('month', au.created_at), u.monthly_token_limit
+    ORDER BY month_start DESC;
+END;
+$$ LANGUAGE plpgsql;
+
+COMMENT ON FUNCTION get_user_usage_stats IS 'Get AI usage statistics for user over last N months';
+
+-- Add index for faster monthly queries
+CREATE INDEX IF NOT EXISTS idx_ai_usage_user_month 
+ON ai_usage_logs(user_id, created_at);