chore: remove broken company_ai_credits stub
- Delete apps/companies/src/handlers/ai.rs (broken placeholder code)
- Remove ai module export from handlers/mod.rs
- Remove /api/companies/ai route from main.rs
The broken stub had:
- Uuid::parse_str("placeholder") that always errored
- Uuid::new_v4() generating random IDs instead of using auth
- Queries to non-existent company_ai_credits/ai_usage_log tables
AI credits are now properly handled by the users service with
the new ai_credits module (wallet, ledger, LiteLLM integration).
This commit is contained in:
Ashwin Kumar Sivakumar
parent
0dd5045676
commit
c795242040
3 changed files with 1 additions and 362 deletions
|
|
@ -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<AppState> {
|
||||
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<chrono::Utc>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
pub created_at: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct UsageHistoryResponse {
|
||||
pub total_entries: i64,
|
||||
pub entries: Vec<UsageEntry>,
|
||||
pub total_credits_used: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct UsageQueryParams {
|
||||
pub page: Option<i64>,
|
||||
pub per_page: Option<i64>,
|
||||
pub request_type: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct CompanyAICredits {
|
||||
company_id: Uuid,
|
||||
credits_balance: i32,
|
||||
updated_at: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
// ============== Route Handlers ==============
|
||||
|
||||
/// GET /api/companies/ai/credits
|
||||
/// Get current AI credit balance
|
||||
async fn get_ai_credits(
|
||||
_auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||
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<AppState>,
|
||||
Json(request): Json<GenerateAiRequest>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||
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::<String>();
|
||||
let result_preview = "AI generated response".chars().take(100).collect::<String>();
|
||||
|
||||
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<AppState>,
|
||||
Query(query): Query<UsageQueryParams>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||
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<String, String> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,4 @@
|
|||
pub mod admin;
|
||||
pub mod ai;
|
||||
|
||||
use axum::{
|
||||
extract::{Multipart, Path, Query, State},
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue