diff --git a/apps/companies/src/handlers/ai.rs b/apps/companies/src/handlers/ai.rs deleted file mode 100644 index ce0ac12..0000000 --- a/apps/companies/src/handlers/ai.rs +++ /dev/null @@ -1,360 +0,0 @@ -use crate::AppState; -use axum::{ - extract::{Query, State}, - http::StatusCode, - response::IntoResponse, - routing::{get, post}, - Json, Router, -}; -use contracts::auth_middleware::AuthUser; -use serde::{Deserialize, Serialize}; -use sqlx::FromRow; -use uuid::Uuid; - -/// AI credit and generation endpoints for companies -pub fn ai_router() -> Router { - Router::new() - .route("/credits", get(get_ai_credits)) - .route("/usage-history", get(get_usage_history)) - .route("/generate", post(generate_ai)) -} - -// ============== Request/Response Types ============== - -#[derive(Debug, Deserialize)] -pub struct GenerateAiRequest { - pub prompt: String, - pub request_type: String, -} - -#[derive(Debug, Serialize)] -pub struct GenerateAiResponse { - pub success: bool, - pub content: String, - pub credits_remaining: i32, - pub request_id: Uuid, -} - -#[derive(Debug, Serialize)] -pub struct CreditsResponse { - pub company_id: Uuid, - pub credits_balance: i32, - pub updated_at: chrono::DateTime, -} - -#[derive(Debug, Serialize, FromRow)] -pub struct UsageEntry { - pub id: Uuid, - pub request_type: String, - pub credits_used: i32, - pub prompt_preview: String, - pub result_preview: String, - pub model_used: String, - pub status: String, - pub error_message: Option, - pub created_at: chrono::DateTime, -} - -#[derive(Debug, Serialize)] -pub struct UsageHistoryResponse { - pub total_entries: i64, - pub entries: Vec, - pub total_credits_used: i64, -} - -#[derive(Debug, Deserialize)] -pub struct UsageQueryParams { - pub page: Option, - pub per_page: Option, - pub request_type: Option, -} - -#[derive(Debug, FromRow)] -struct CompanyAICredits { - company_id: Uuid, - credits_balance: i32, - updated_at: chrono::DateTime, -} - -// ============== Route Handlers ============== - -/// GET /api/companies/ai/credits -/// Get current AI credit balance -async fn get_ai_credits( - _auth: AuthUser, - State(state): State, -) -> Result { - let company_id = Uuid::parse_str("placeholder").map_err(|_| { - (StatusCode::BAD_REQUEST, "Invalid company ID".to_string()) - })?; - - let credits = sqlx::query_as::<_, CompanyAICredits>( - r#" - SELECT company_id, credits_balance, updated_at - FROM company_ai_credits - WHERE company_id = $1 - "#, - ) - .bind(company_id) - .fetch_optional(&state.pool) - .await - .map_err(|e| { - tracing::error!("Failed to fetch AI credits: {}", e); - (StatusCode::INTERNAL_SERVER_ERROR, "Database error".to_string()) - })?; - - let balance = credits.map(|c| c.credits_balance).unwrap_or(0); - - let response = CreditsResponse { - company_id, - credits_balance: balance, - updated_at: chrono::Utc::now(), - }; - - Ok((StatusCode::OK, Json(response))) -} - -/// POST /api/companies/ai/generate -/// Generate AI content with credit deduction -async fn generate_ai( - _auth: AuthUser, - State(state): State, - Json(request): Json, -) -> Result { - let company_id = Uuid::new_v4(); // Placeholder - should extract from auth - - // Validate request - if request.prompt.is_empty() { - return Err((StatusCode::BAD_REQUEST, "Prompt cannot be empty".to_string())); - } - - tracing::info!( - company_id = %company_id, - request_type = %request.request_type, - "AI generate request received" - ); - - // Check credits - let credits = sqlx::query_scalar::<_, i32>( - r#" - SELECT credits_balance - FROM company_ai_credits - WHERE company_id = $1 - FOR UPDATE - "#, - ) - .bind(company_id) - .fetch_optional(&state.pool) - .await - .map_err(|e| { - tracing::error!("Failed to check credits: {}", e); - (StatusCode::INTERNAL_SERVER_ERROR, "Database error".to_string()) - })?; - - let credits_before = credits.unwrap_or(0); - - if credits_before < 1 { - return Err((StatusCode::PAYMENT_REQUIRED, "Insufficient AI credits".to_string())); - } - - // Deduct credit - sqlx::query( - r#" - UPDATE company_ai_credits - SET credits_balance = credits_balance - 1, - updated_at = NOW() - WHERE company_id = $1 - RETURNING credits_balance - "#, - ) - .bind(company_id) - .fetch_one(&state.pool) - .await - .map_err(|e| { - tracing::error!("Failed to deduct credits: {}", e); - (StatusCode::INTERNAL_SERVER_ERROR, "Database error".to_string()) - })?; - - // Log usage - let request_id = Uuid::new_v4(); - let prompt_preview = request.prompt.chars().take(100).collect::(); - let result_preview = "AI generated response".chars().take(100).collect::(); - - sqlx::query( - r#" - INSERT INTO ai_usage_log (id, company_id, request_type, credits_used, prompt_preview, result_preview, model_used, status, created_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, NOW()) - "#, - ) - .bind(request_id) - .bind(company_id) - .bind(&request.request_type) - .bind(1_i32) - .bind(prompt_preview) - .bind(result_preview) - .bind("gemma3:270m") - .bind("success") - .execute(&state.pool) - .await - .map_err(|e| { - tracing::error!("Failed to log usage: {}", e); - }).ok(); - - tracing::info!( - company_id = %company_id, - request_id = %request_id, - credits_before = credits_before, - credits_after = credits_before - 1, - "AI generation completed" - ); - - // Call Ollama service - let ollama_base = std::env::var("OLLAMA_BASE_URL") - .unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string()); - - let generated_content = call_ollama_generate(&ollama_base, &request.prompt).await - .map_err(|e| { - tracing::error!("Ollama call failed: {}", e); - (StatusCode::SERVICE_UNAVAILABLE, "AI service unavailable".to_string()) - })?; - - let response = GenerateAiResponse { - success: true, - content: generated_content, - credits_remaining: credits_before - 1, - request_id, - }; - - Ok((StatusCode::OK, Json(response))) -} - -/// GET /api/companies/ai/usage-history -/// Get AI usage history for a company -async fn get_usage_history( - _auth: AuthUser, - State(state): State, - Query(query): Query, -) -> Result { - let company_id = Uuid::new_v4(); // Placeholder - let page = query.page.unwrap_or(1).max(1); - let per_page = query.per_page.unwrap_or(20).clamp(1, 100); - let offset = (page - 1) * per_page; - - // Get total count - let total = sqlx::query_scalar::<_, i64>( - r#" - SELECT COUNT(*) FROM ai_usage_log WHERE company_id = $1 - "#, - ) - .bind(company_id) - .fetch_one(&state.pool) - .await - .map_err(|e| { - tracing::error!("Failed to count usage entries: {}", e); - 0_i64 - }).unwrap_or(0); - - // Get entries - let entries = sqlx::query_as::<_, UsageEntry>( - r#" - SELECT id, request_type, credits_used, prompt_preview, result_preview, - model_used, status, error_message, created_at - FROM ai_usage_log - WHERE company_id = $1 - ORDER BY created_at DESC - LIMIT $2 OFFSET $3 - "#, - ) - .bind(company_id) - .bind(per_page) - .bind(offset) - .fetch_all(&state.pool) - .await - .map_err(|e| { - tracing::error!("Failed to fetch usage history: {}", e); - (StatusCode::INTERNAL_SERVER_ERROR, "Database error".to_string()) - })?; - - let total_credits = entries.iter().map(|e| e.credits_used as i64).sum(); - - let response = UsageHistoryResponse { - total_entries: total, - entries, - total_credits_used: total_credits, - }; - - Ok((StatusCode::OK, Json(response))) -} - -// ============== Helper Functions ============== - -#[derive(Serialize)] -struct OllamaGenerateRequest { - model: String, - prompt: String, - stream: bool, -} - -#[derive(Deserialize)] -struct OllamaGenerateResponse { - response: String, -} - -async fn call_ollama_generate(base_url: &str, prompt: &str) -> Result { - let url = format!("{}/api/generate", base_url); - - let req = OllamaGenerateRequest { - model: "gemma3:270m".to_string(), - prompt: prompt.to_string(), - stream: false, - }; - - let client = reqwest::Client::new(); - let response = client - .post(&url) - .json(&req) - .send() - .await - .map_err(|e| format!("Ollama request failed: {}", e))?; - - if !response.status().is_success() { - return Err(format!("Ollama returned status: {}", response.status())); - } - - let result: OllamaGenerateResponse = response - .json() - .await - .map_err(|e| format!("Failed to parse Ollama response: {}", e))?; - - Ok(result.response) -} - -// ============== Tests ============== -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_request_deserialization() { - let json = serde_json::json!({ - "prompt": "Generate a job description", - "request_type": "job_description" - }); - let req: GenerateAiRequest = serde_json::from_value(json).unwrap(); - assert_eq!(req.prompt, "Generate a job description"); - assert_eq!(req.request_type, "job_description"); - } - - #[test] - fn test_response_serialization() { - let resp = GenerateAiResponse { - success: true, - content: "Generated content".to_string(), - credits_remaining: 5, - request_id: Uuid::new_v4(), - }; - let json = serde_json::to_value(&resp).unwrap(); - assert_eq!(json["success"], true); - assert_eq!(json["credits_remaining"], 5); - } -} diff --git a/apps/companies/src/handlers/mod.rs b/apps/companies/src/handlers/mod.rs index 878f7ca..5c2722e 100644 --- a/apps/companies/src/handlers/mod.rs +++ b/apps/companies/src/handlers/mod.rs @@ -1,5 +1,4 @@ pub mod admin; -pub mod ai; use axum::{ extract::{Multipart, Path, Query, State}, diff --git a/apps/companies/src/main.rs b/apps/companies/src/main.rs index a7ff9d9..f5b27c7 100644 --- a/apps/companies/src/main.rs +++ b/apps/companies/src/main.rs @@ -45,7 +45,7 @@ async fn main() { let app = Router::new() .nest("/api/companies", handlers::router()) .nest("/api/admin/companies", handlers::admin::router()) - .nest("/api/companies/ai", handlers::ai::ai_router()) + .route("/health", get(|| async { "Companies OK" })) .with_state(state);