# 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 ```sql -- 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 ```bash # 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 ```rust // 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, pub expires_at: Option>, } #[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, pub current_period_end: chrono::DateTime, } #[derive(Debug, Serialize, Deserialize, FromRow)] pub struct AiUsageLog { pub id: Uuid, pub user_id: Uuid, pub api_key_id: Option, pub model: String, pub request_type: String, pub tokens_input: i32, pub tokens_output: i32, pub tokens_total: i32, pub cost_estimate: Option, pub request_duration_ms: Option, pub was_successful: bool, pub created_at: chrono::DateTime, } ``` ### 2.2 API Key Generation Service ```rust // 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 { 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 ```rust // 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 { // 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, 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, end_date: DateTime, ) -> Result { 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, 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 ```rust // 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, usage_service: web::Data, litellm_client: web::Data, ) -> 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, // JWT auth usage_service: web::Data, query: web::Query, ) -> 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, body: web::Json, api_key_service: web::Data, ) -> 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 ```rust // 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 Transform for AiAuth where S: Service, Error = Error>, S::Future: 'static, B: 'static, { type Response = S::Response; type Error = S::Error; type Transform = AiAuthMiddleware; type InitError = (); type Future = Ready>; fn new_transform(&self, service: S) -> Self::Future { ok(AiAuthMiddleware { service }) } } pub struct AiAuthMiddleware { service: S, } impl Service for AiAuthMiddleware where S: Service, Error = Error>, S::Future: 'static, B: 'static, { type Response = S::Response; type Error = S::Error; type Future = LocalBoxFuture<'static, Result>; fn poll_ready(&self, cx: &mut Context<'_>) -> Poll> { 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 ```typescript // Frontend API client class AiClient { private apiKey: string; constructor(apiKey: string) { this.apiKey = apiKey; } async chatCompletion(model: string, message: string): Promise { 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 { 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 ```bash # 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?