7902b265a9
- Add comprehensive AI plans implementation documentation - Add LiteLLM gateway Kubernetes manifests - Update PostgreSQL and Forgejo deployment configs - Add build-from-binaries script
754 lines
23 KiB
Markdown
754 lines
23 KiB
Markdown
# 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<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
|
|
|
|
```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<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
|
|
|
|
```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<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
|
|
|
|
```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<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
|
|
|
|
```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<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
|
|
|
|
```typescript
|
|
// 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
|
|
|
|
```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?
|