nxtgauge-gitops/docs/AI_PLANS_IMPLEMENTATION_PLAN.md

AI Plans Implementation Plan

Executive Summary

Implement per-user API keys with usage tracking and plan tiers for Ask Ash AI assistant.

Phase 1: Database Design (Week 1)

1.1 New Tables

-- Migration: 001_add_ai_plans.sql

-- API Keys table (supports rotation)
CREATE TABLE api_keys (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    key_hash VARCHAR(255) UNIQUE NOT NULL,  -- bcrypt hash of key
    key_prefix VARCHAR(20) NOT NULL,         -- sk-nxtgauge-abc...
    is_active BOOLEAN DEFAULT true,
    created_at TIMESTAMP DEFAULT NOW(),
    expires_at TIMESTAMP NULL,
    revoked_at TIMESTAMP NULL,
    revoked_reason TEXT NULL
);

-- AI Plans table (plan definitions)
CREATE TABLE ai_plans (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    name VARCHAR(50) UNIQUE NOT NULL,        -- 'free', 'pro', 'business', 'enterprise'
    display_name VARCHAR(100) NOT NULL,
    monthly_credits INTEGER NOT NULL,
    rate_limit_rpm INTEGER NOT NULL,           -- requests per minute
    rate_limit_rph INTEGER NOT NULL,          -- requests per hour
    max_tokens_per_request INTEGER NOT NULL,
    price_monthly DECIMAL(10,2) NOT NULL,
    features JSONB NOT NULL,                  -- allowed models, etc.
    is_active BOOLEAN DEFAULT true,
    created_at TIMESTAMP DEFAULT NOW()
);

-- User AI Subscriptions
CREATE TABLE user_ai_subscriptions (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    user_id UUID UNIQUE NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    plan_id UUID NOT NULL REFERENCES ai_plans(id),
    credits_remaining INTEGER NOT NULL,
    credits_used_this_month INTEGER DEFAULT 0,
    subscription_status VARCHAR(50) DEFAULT 'active', -- 'active', 'paused', 'cancelled'
    current_period_start TIMESTAMP NOT NULL,
    current_period_end TIMESTAMP NOT NULL,
    cancelled_at TIMESTAMP NULL,
    created_at TIMESTAMP DEFAULT NOW(),
    updated_at TIMESTAMP DEFAULT NOW()
);

-- AI Usage Logs (for tracking & billing)
CREATE TABLE ai_usage_logs (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    user_id UUID NOT NULL REFERENCES users(id),
    api_key_id UUID REFERENCES api_keys(id),
    model VARCHAR(100) NOT NULL,              -- 'askash-fast', 'askash-main', etc.
    request_type VARCHAR(100) NOT NULL,       -- 'help', 'resume', 'jd', 'cover_letter'
    tokens_input INTEGER NOT NULL,
    tokens_output INTEGER NOT NULL,
    tokens_total INTEGER NOT NULL,
    cost_estimate DECIMAL(10,6),              -- calculated cost
    request_duration_ms INTEGER,              -- response time
    was_successful BOOLEAN DEFAULT true,
    error_message TEXT NULL,
    ip_address INET,
    user_agent TEXT,
    created_at TIMESTAMP DEFAULT NOW()
);

-- Rate Limit Tracking (Redis alternative)
CREATE TABLE rate_limit_windows (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    user_id UUID NOT NULL REFERENCES users(id),
    window_start TIMESTAMP NOT NULL,
    window_end TIMESTAMP NOT NULL,
    requests_count INTEGER DEFAULT 0,
    UNIQUE(user_id, window_start)
);

-- Insert default plans
INSERT INTO ai_plans (name, display_name, monthly_credits, rate_limit_rpm, rate_limit_rph, max_tokens_per_request, price_monthly, features) VALUES
('free', 'Free', 100, 10, 100, 1000, 0.00, '{"models": ["askash-fast", "help-assistant", "messenger"]}'),
('pro', 'Pro', 1000, 60, 1000, 4000, 9.00, '{"models": ["askash-fast", "askash-main", "help-assistant", "messenger", "recommender", "safety-check"]}'),
('business', 'Business', 5000, 120, 5000, 8000, 29.00, '{"models": ["askash-fast", "askash-main", "jd-generator", "profile-writer", "service-writer", "requirement-writer"]}'),
('enterprise', 'Enterprise', 50000, 0, 0, 32000, 99.00, '{"models": ["all"], "priority": true}');

-- Indexes for performance
CREATE INDEX idx_api_keys_user_id ON api_keys(user_id);
CREATE INDEX idx_api_keys_key_hash ON api_keys(key_hash);
CREATE INDEX idx_user_ai_subscriptions_user_id ON user_ai_subscriptions(user_id);
CREATE INDEX idx_ai_usage_logs_user_id ON ai_usage_logs(user_id);
CREATE INDEX idx_ai_usage_logs_created_at ON ai_usage_logs(created_at);
CREATE INDEX idx_ai_usage_logs_model ON ai_usage_logs(model);

1.2 Migration Strategy

# Run migrations
sqlx migrate run --source ./migrations

# Or manual SQL execution
psql $DATABASE_URL < migrations/001_add_ai_plans.sql

Phase 2: Core Implementation (Week 1-2)

2.1 Models/DTOs

// src/models/ai_plan.rs

use serde::{Deserialize, Serialize};
use sqlx::FromRow;
use uuid::Uuid;

#[derive(Debug, Serialize, Deserialize, FromRow)]
pub struct AiPlan {
    pub id: Uuid,
    pub name: String,
    pub display_name: String,
    pub monthly_credits: i32,
    pub rate_limit_rpm: i32,
    pub rate_limit_rph: i32,
    pub max_tokens_per_request: i32,
    pub price_monthly: f64,
    pub features: serde_json::Value,
}

#[derive(Debug, Serialize, Deserialize, FromRow)]
pub struct ApiKey {
    pub id: Uuid,
    pub user_id: Uuid,
    pub key_hash: String,
    pub key_prefix: String,
    pub is_active: bool,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
}

#[derive(Debug, Serialize, Deserialize, FromRow)]
pub struct UserAiSubscription {
    pub id: Uuid,
    pub user_id: Uuid,
    pub plan_id: Uuid,
    pub credits_remaining: i32,
    pub credits_used_this_month: i32,
    pub subscription_status: String,
    pub current_period_start: chrono::DateTime<chrono::Utc>,
    pub current_period_end: chrono::DateTime<chrono::Utc>,
}

#[derive(Debug, Serialize, Deserialize, FromRow)]
pub struct AiUsageLog {
    pub id: Uuid,
    pub user_id: Uuid,
    pub api_key_id: Option<Uuid>,
    pub model: String,
    pub request_type: String,
    pub tokens_input: i32,
    pub tokens_output: i32,
    pub tokens_total: i32,
    pub cost_estimate: Option<f64>,
    pub request_duration_ms: Option<i32>,
    pub was_successful: bool,
    pub created_at: chrono::DateTime<chrono::Utc>,
}

2.2 API Key Generation Service

// src/services/api_key_service.rs

use bcrypt::{hash, verify, DEFAULT_COST};
use rand::{distributions::Alphanumeric, Rng};
use uuid::Uuid;

pub struct ApiKeyService;

impl ApiKeyService {
    /// Generate a new API key
    /// Returns: (full_key, key_hash, key_prefix)
    pub fn generate_key(user_id: Uuid) -> (String, String, String) {
        let prefix = "sk-nxtgauge";
        let user_part = user_id.to_string().split('-').next().unwrap_or("");
        let random_suffix: String = rand::thread_rng()
            .sample_iter(&Alphanumeric)
            .take(24)
            .map(char::from)
            .collect();
        
        let full_key = format!("{}-{}-{}", prefix, user_part, random_suffix);
        let key_prefix = format!("{}-{}-", prefix, user_part);
        
        // Hash for storage (use first 8 chars as salt identifier)
        let key_hash = hash(&full_key, DEFAULT_COST).unwrap();
        
        (full_key, key_hash, key_prefix)
    }
    
    /// Validate an API key against hash
    pub fn validate_key(provided_key: &str, stored_hash: &str) -> bool {
        verify(provided_key, stored_hash).unwrap_or(false)
    }
    
    /// Extract user ID from API key (without validation)
    pub fn extract_user_id_from_key(key: &str) -> Option<String> {
        let parts: Vec<&str> = key.split('-').collect();
        if parts.len() >= 3 && parts[0] == "sk" && parts[1] == "nxtgauge" {
            Some(parts[2].to_string())
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_generate_key() {
        let user_id = Uuid::new_v4();
        let (full_key, hash, prefix) = ApiKeyService::generate_key(user_id);
        
        assert!(full_key.starts_with("sk-nxtgauge-"));
        assert_eq!(full_key.len(), 45); // sk-nxtgauge- + 8 + - + 24
        assert!(ApiKeyService::validate_key(&full_key, &hash));
    }
}

2.3 Usage Tracking Service

// src/services/ai_usage_service.rs

use sqlx::PgPool;
use uuid::Uuid;
use chrono::{DateTime, Utc, Duration};

pub struct AiUsageService {
    db: PgPool,
}

impl AiUsageService {
    pub fn new(db: PgPool) -> Self {
        Self { db }
    }
    
    /// Check if user has credits and rate limit allows request
    pub async fn can_make_request(
        &self,
        user_id: Uuid,
        model: &str,
    ) -> Result<RequestAllowance, AiError> {
        // Get user's subscription
        let subscription = sqlx::query_as::<_, UserAiSubscription>(
            "SELECT * FROM user_ai_subscriptions WHERE user_id = $1"
        )
        .bind(user_id)
        .fetch_optional(&self.db)
        .await?;
        
        let subscription = subscription.ok_or(AiError::NoSubscription)?;
        
        // Check subscription status
        if subscription.subscription_status != "active" {
            return Err(AiError::SubscriptionInactive);
        }
        
        // Check credits
        if subscription.credits_remaining <= 0 {
            return Err(AiError::InsufficientCredits);
        }
        
        // Get plan details
        let plan = sqlx::query_as::<_, AiPlan>(
            "SELECT * FROM ai_plans WHERE id = $1"
        )
        .bind(subscription.plan_id)
        .fetch_one(&self.db)
        .await?;
        
        // Check rate limit (current window)
        let window_start = Utc::now() - Duration::minutes(1);
        let recent_requests: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM ai_usage_logs 
             WHERE user_id = $1 AND created_at > $2"
        )
        .bind(user_id)
        .bind(window_start)
        .fetch_one(&self.db)
        .await?;
        
        if recent_requests >= plan.rate_limit_rpm as i64 {
            return Err(AiError::RateLimitExceeded);
        }
        
        Ok(RequestAllowance {
            user_id,
            plan,
            subscription,
            remaining_requests: plan.rate_limit_rpm - recent_requests as i32,
        })
    }
    
    /// Log AI usage and deduct credits
    pub async fn log_usage(
        &self,
        user_id: Uuid,
        api_key_id: Option<Uuid>,
        model: &str,
        request_type: &str,
        tokens_input: i32,
        tokens_output: i32,
        duration_ms: i32,
        success: bool,
    ) -> Result<(), AiError> {
        let tokens_total = tokens_input + tokens_output;
        
        // Calculate cost (example pricing)
        let cost = match model {
            "askash-fast" | "help-assistant" | "messenger" => tokens_total as f64 * 0.000001, // $0.001 per 1K tokens
            "askash-main" | "jd-generator" | "profile-writer" => tokens_total as f64 * 0.000005, // $0.005 per 1K tokens
            _ => tokens_total as f64 * 0.000005,
        };
        
        // Insert usage log
        sqlx::query(
            "INSERT INTO ai_usage_logs 
             (user_id, api_key_id, model, request_type, tokens_input, tokens_output, 
              tokens_total, cost_estimate, request_duration_ms, was_successful)
             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)"
        )
        .bind(user_id)
        .bind(api_key_id)
        .bind(model)
        .bind(request_type)
        .bind(tokens_input)
        .bind(tokens_output)
        .bind(tokens_total)
        .bind(cost)
        .bind(duration_ms)
        .bind(success)
        .execute(&self.db)
        .await?;
        
        // Deduct credits (1 credit per 1000 tokens, minimum 1)
        let credits_to_deduct = ((tokens_total as f64 / 1000.0).ceil() as i32).max(1);
        
        sqlx::query(
            "UPDATE user_ai_subscriptions 
             SET credits_remaining = credits_remaining - $1,
                 credits_used_this_month = credits_used_this_month + $1,
                 updated_at = NOW()
             WHERE user_id = $2"
        )
        .bind(credits_to_deduct)
        .bind(user_id)
        .execute(&self.db)
        .await?;
        
        Ok(())
    }
    
    /// Get usage statistics for user
    pub async fn get_user_usage(
        &self,
        user_id: Uuid,
        start_date: DateTime<Utc>,
        end_date: DateTime<Utc>,
    ) -> Result<UsageStats, AiError> {
        let stats = sqlx::query_as::<_, UsageStats>(
            "SELECT 
                COUNT(*) as total_requests,
                SUM(tokens_input) as total_tokens_input,
                SUM(tokens_output) as total_tokens_output,
                SUM(tokens_total) as total_tokens,
                SUM(cost_estimate) as total_cost,
                COUNT(CASE WHEN was_successful = false THEN 1 END) as failed_requests
             FROM ai_usage_logs 
             WHERE user_id = $1 AND created_at BETWEEN $2 AND $3"
        )
        .bind(user_id)
        .bind(start_date)
        .bind(end_date)
        .fetch_one(&self.db)
        .await?;
        
        Ok(stats)
    }
}

#[derive(Debug)]
pub struct RequestAllowance {
    pub user_id: Uuid,
    pub plan: AiPlan,
    pub subscription: UserAiSubscription,
    pub remaining_requests: i32,
}

#[derive(Debug)]
pub struct UsageStats {
    pub total_requests: i64,
    pub total_tokens_input: i64,
    pub total_tokens_output: i64,
    pub total_tokens: i64,
    pub total_cost: Option<f64>,
    pub failed_requests: i64,
}

#[derive(Debug, thiserror::Error)]
pub enum AiError {
    #[error("No subscription found")]
    NoSubscription,
    #[error("Subscription inactive")]
    SubscriptionInactive,
    #[error("Insufficient credits")]
    InsufficientCredits,
    #[error("Rate limit exceeded")]
    RateLimitExceeded,
    #[error("Database error: {0}")]
    Database(#[from] sqlx::Error),
}

Phase 3: API Endpoints (Week 2)

3.1 AI Controller

// src/controllers/ai_controller.rs

use actix_web::{web, HttpResponse, HttpRequest};
use crate::services::AiUsageService;
use crate::middleware::AiAuth;

pub fn ai_routes(cfg: &mut web::ServiceConfig) {
    cfg.service(
        web::scope("/api/v1/ai")
            // These require API key auth
            .route("/chat", web::post().to(chat_completion))
            .route("/usage", web::get().to(get_usage))
            // These require user JWT auth
            .route("/keys", web::post().to(generate_api_key))
            .route("/keys", web::get().to(list_api_keys))
            .route("/keys/{key_id}", web::delete().to(revoke_api_key))
            .route("/upgrade", web::post().to(upgrade_plan))
    );
}

/// Main chat endpoint (requires API key)
async fn chat_completion(
    req: HttpRequest,
    body: web::Json<ChatRequest>,
    usage_service: web::Data<AiUsageService>,
    litellm_client: web::Data<LitellmClient>,
) -> HttpResponse {
    let start_time = std::time::Instant::now();
    
    // Extract API key from header
    let api_key = match extract_api_key(&req) {
        Some(key) => key,
        None => return HttpResponse::Unauthorized().json(ErrorResponse {
            error: "Missing API key".to_string(),
        }),
    };
    
    // Validate and get user info
    let (user_id, api_key_id, plan) = match validate_api_key(&api_key).await {
        Ok(info) => info,
        Err(e) => return HttpResponse::Unauthorized().json(ErrorResponse {
            error: e.to_string(),
        }),
    };
    
    // Check if user can make request
    let allowance = match usage_service.can_make_request(user_id, &body.model).await {
        Ok(a) => a,
        Err(AiError::InsufficientCredits) => {
            return HttpResponse::PaymentRequired().json(ErrorResponse {
                error: "Insufficient credits. Please upgrade your plan.".to_string(),
            });
        }
        Err(AiError::RateLimitExceeded) => {
            return HttpResponse::TooManyRequests().json(ErrorResponse {
                error: "Rate limit exceeded. Please slow down.".to_string(),
            });
        }
        Err(e) => return HttpResponse::InternalServerError().json(ErrorResponse {
            error: e.to_string(),
        }),
    };
    
    // Forward to LiteLLM
    let litellm_response = match litellm_client.chat_completion(&body).await {
        Ok(resp) => resp,
        Err(e) => {
            // Log failed request
            let _ = usage_service.log_usage(
                user_id, api_key_id, &body.model, &body.request_type,
                0, 0, start_time.elapsed().as_millis() as i32, false,
            ).await;
            
            return HttpResponse::InternalServerError().json(ErrorResponse {
                error: format!("LiteLLM error: {}", e),
            });
        }
    };
    
    // Log successful usage
    let duration_ms = start_time.elapsed().as_millis() as i32;
    let tokens_input = litellm_response.usage.prompt_tokens;
    let tokens_output = litellm_response.usage.completion_tokens;
    
    let _ = usage_service.log_usage(
        user_id, api_key_id, &body.model, &body.request_type,
        tokens_input, tokens_output, duration_ms, true,
    ).await;
    
    // Return response with headers
    HttpResponse::Ok()
        .insert_header(("X-RateLimit-Remaining", allowance.remaining_requests.to_string()))
        .insert_header(("X-Credits-Remaining", allowance.subscription.credits_remaining.to_string()))
        .json(litellm_response)
}

/// Get user's usage statistics
async fn get_usage(
    auth: web::ReqData<UserClaims>, // JWT auth
    usage_service: web::Data<AiUsageService>,
    query: web::Query<UsageQuery>,
) -> HttpResponse {
    let start_date = query.start_date.unwrap_or_else(|| {
        Utc::now() - Duration::days(30)
    });
    let end_date = query.end_date.unwrap_or(Utc::now());
    
    match usage_service.get_user_usage(auth.user_id, start_date, end_date).await {
        Ok(stats) => HttpResponse::Ok().json(stats),
        Err(e) => HttpResponse::InternalServerError().json(ErrorResponse {
            error: e.to_string(),
        }),
    }
}

/// Generate new API key
async fn generate_api_key(
    auth: web::ReqData<UserClaims>,
    body: web::Json<GenerateKeyRequest>,
    api_key_service: web::Data<ApiKeyService>,
) -> HttpResponse {
    let (full_key, key_hash, key_prefix) = ApiKeyService::generate_key(auth.user_id);
    
    // Save to database
    match api_key_service.save_key(auth.user_id, &key_hash, &key_prefix).await {
        Ok(_) => HttpResponse::Ok().json(GenerateKeyResponse {
            api_key: full_key,
            prefix: key_prefix,
            created_at: Utc::now(),
        }),
        Err(e) => HttpResponse::InternalServerError().json(ErrorResponse {
            error: e.to_string(),
        }),
    }
}

Phase 4: Middleware (Week 2)

4.1 API Key Extraction

// src/middleware/ai_auth.rs

use actix_web::{dev::ServiceRequest, Error, HttpMessage};
use actix_web::dev::{Transform, Service};
use futures::future::{LocalBoxFuture, ok, Ready};
use std::task::{Context, Poll};

pub struct AiAuth;

impl<S, B> Transform<S, ServiceRequest> for AiAuth
where
    S: Service<ServiceRequest, Response = actix_web::dev::ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
    B: 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Transform = AiAuthMiddleware<S>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ok(AiAuthMiddleware { service })
    }
}

pub struct AiAuthMiddleware<S> {
    service: S,
}

impl<S, B> Service<ServiceRequest> for AiAuthMiddleware<S>
where
    S: Service<ServiceRequest, Response = actix_web::dev::ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
    B: 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&self, req: ServiceRequest) -> Self::Future {
        // Extract API key from Authorization header
        if let Some(auth_header) = req.headers().get("Authorization") {
            if let Ok(auth_str) = auth_header.to_str() {
                if auth_str.starts_with("Bearer ") {
                    let api_key = &auth_str[7..];
                    // Store in request extensions for later use
                    req.extensions_mut().insert(api_key.to_string());
                }
            }
        }

        let fut = self.service.call(req);
        Box::pin(async move {
            fut.await
        })
    }
}

Phase 5: Integration (Week 3)

5.1 Frontend Changes

// Frontend API client

class AiClient {
  private apiKey: string;
  
  constructor(apiKey: string) {
    this.apiKey = apiKey;
  }
  
  async chatCompletion(model: string, message: string): Promise<AIResponse> {
    const response = await fetch('/api/v1/ai/chat', {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${this.apiKey}`,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        model,
        messages: [{ role: 'user', content: message }],
        request_type: 'help',
      }),
    });
    
    // Check for credit/rate limit headers
    const creditsRemaining = response.headers.get('X-Credits-Remaining');
    const rateLimitRemaining = response.headers.get('X-RateLimit-Remaining');
    
    if (response.status === 402) {
      throw new Error('Insufficient credits. Please upgrade.');
    }
    
    if (response.status === 429) {
      throw new Error('Rate limit exceeded. Please slow down.');
    }
    
    return response.json();
  }
  
  async getUsageStats(): Promise<UsageStats> {
    const response = await fetch('/api/v1/ai/usage', {
      headers: {
        'Authorization': `Bearer ${this.jwtToken}`, // JWT for authenticated endpoints
      },
    });
    return response.json();
  }
}

Phase 6: Testing & Deployment (Week 3-4)

6.1 Testing Checklist

• Unit tests for API key generation
• Unit tests for usage tracking
• Integration tests for rate limiting
• Load tests for concurrent requests
• Security tests (key validation, SQL injection)

6.2 Migration Steps

# 1. Backup database
pg_dump $DATABASE_URL > backup_pre_ai_plans.sql

# 2. Run migrations
sqlx migrate run

# 3. Seed default plans
psql $DATABASE_URL < seed_plans.sql

# 4. Deploy new backend version
cargo build --release

# 5. Generate API keys for existing users
./scripts/migrate_existing_users.sh

# 6. Verify deployment
./scripts/verify_ai_plans.sh

Timeline

Week	Phase	Deliverables
Week 1	Database + Core Services	Tables, models, services
Week 2	API + Middleware	Endpoints, auth, rate limiting
Week 3	Integration + Frontend	React components, testing
Week 4	Testing + Deployment	Load tests, monitoring, docs

Cost Estimation

Component	Cost
Database storage	~$5/mo (10GB)
Compute (tracking)	Minimal (async)
Total additional	~$5/mo

Next Steps

1. Review this plan - Any changes needed?
2. Approve database schema - Are all fields needed?
3. Set priority - Which features are must-have for MVP?

Ready to start Phase 1?