nxtgauge-backend-rust/apps/users/src/ai_credits.rs

//! Reusable "charge an AI feature" primitive for handler code.
//!
//! Wraps `db::models::ai_credits::AiCreditsRepository`'s reserve/capture/
//! release calls around an arbitrary unit of work (in practice, a LiteLLM
//! call): reserve credits up front, run the work, capture on success or
//! release on failure/panic-unwind-safe-error. No handler should call
//! `try_reserve_credits`/`try_capture_reservation`/`try_release_reservation`
//! directly -- always go through `charge_ai_feature` so a hold can never be
//! left unresolved by a forgotten capture/release call at a call site.

use axum::http::StatusCode;
use db::models::ai_credits::{AiCreditsError, AiCreditsRepository, AiFeatureCost};
use sqlx::PgPool;
use uuid::Uuid;

/// Maximum prompt length in characters (approx 4 chars per token as per architecture doc)
/// Used for input validation before reserving credits
pub const APPROX_CHARS_PER_TOKEN: i32 = 4;

/// Validate prompt/response lengths against feature limits before making LLM calls
/// Returns None if valid, or Some((status, error_json)) if validation fails
pub fn validate_prompt_length(prompt: &str, feature: &AiFeatureCost) -> Option<(StatusCode, serde_json::Value)> {
    // Check input length if max_input_tokens is configured
    if let Some(max_tokens) = feature.max_input_tokens {
        let max_chars = max_tokens * APPROX_CHARS_PER_TOKEN;
        if prompt.len() as i32 > max_chars {
            return Some((
                StatusCode::BAD_REQUEST,
                serde_json::json!({
                    "error": format!(
                        "Prompt too long: {} characters exceeds maximum of {} (approx {} tokens)",
                        prompt.len(),
                        max_chars,
                        max_tokens
                    ),
                    "code": "PROMPT_TOO_LONG"
                }),
            ));
        }
    }
    None
}

/// Truncate text for preview/logging (first N chars)
pub fn truncate_for_preview(text: &str, max_chars: usize) -> String {
    if text.len() <= max_chars {
        text.to_string()
    } else {
        format!("{}...", &text[..max_chars])
    }
}

#[derive(Debug)]
pub enum ChargeError<E> {
    Credits(AiCreditsError),
    UnknownFeature(String),
    Work(E),
}

impl<E: std::fmt::Display> std::fmt::Display for ChargeError<E> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ChargeError::Credits(e) => write!(f, "AI credits error: {e}"),
            ChargeError::UnknownFeature(f2) => write!(f, "Unknown AI feature '{f2}'"),
            ChargeError::Work(e) => write!(f, "{e}"),
        }
    }
}

pub struct ChargeOutcome<T> {
    pub result: T,
    pub feature: AiFeatureCost,
    pub credits_charged: i32,
}

/// Reserve credits for `feature_code`, run `work`, and settle the
/// reservation based on whether `work` succeeded. `work` receives the
/// resolved feature cost (so it knows which model/timeout to use) and
/// returns `Ok(T)` to capture the charge or `Err(E)` to release it
/// uncharged.
pub async fn charge_ai_feature<T, E, F, Fut>(
    pool: &PgPool,
    user_id: Uuid,
    feature_code: &str,
    request_id: Option<&str>,
    idempotency_key: Option<&str>,
    work: F,
) -> Result<ChargeOutcome<T>, ChargeError<E>>
where
    F: FnOnce(AiFeatureCost) -> Fut,
    Fut: std::future::Future<Output = Result<T, E>>,
{
    let feature = AiCreditsRepository::get_feature_cost(pool, feature_code)
        .await
        .map_err(|e| ChargeError::Credits(AiCreditsError::Db(e)))?
        .ok_or_else(|| ChargeError::UnknownFeature(feature_code.to_string()))?;

    let hold = AiCreditsRepository::try_reserve_credits(
        pool,
        user_id,
        feature_code,
        feature.credit_cost,
        request_id,
        idempotency_key,
    )
    .await
    .map_err(ChargeError::Credits)?;

    match work(feature.clone()).await {
        Ok(result) => {
            AiCreditsRepository::try_capture_reservation(pool, hold.id, None)
                .await
                .map_err(ChargeError::Credits)?;
            log_usage(pool, user_id, &feature, hold.credits_held, "success", request_id, None, None, None, None, None).await;
            Ok(ChargeOutcome {
                result,
                feature,
                credits_charged: hold.credits_held,
            })
        }
        Err(e) => {
            // Best-effort release -- if this fails, the hold's expires_at
            // (Section 4.2 of the architecture doc) is the backstop until
            // a reaper job exists to sweep abandoned holds.
            let _ = AiCreditsRepository::try_release_reservation(pool, hold.id).await;
            log_usage(pool, user_id, &feature, 0, "error", request_id, None, None, None, None, None).await;
            Err(ChargeError::Work(e))
        }
    }
}

/// Best-effort usage log write -- failures here must never fail the
/// caller's request (the credit charge/release already happened; a
/// missing analytics row is not worth surfacing as an error to the user).
#[allow(clippy::too_many_arguments)]
async fn log_usage(
    pool: &PgPool,
    user_id: Uuid,
    feature: &AiFeatureCost,
    credits_charged: i32,
    status: &str,
    request_id: Option<&str>,
    prompt_preview: Option<&str>,
    response_preview: Option<&str>,
    input_tokens: Option<i32>,
    output_tokens: Option<i32>,
    total_tokens: Option<i32>,
) {
    let result = sqlx::query(
        r#"
        INSERT INTO ai_usage_logs 
            (user_id, feature_code, model_alias, credits_charged, status, request_id, 
             prompt_preview, response_preview, input_tokens, output_tokens, total_tokens)
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
        "#,
    )
    .bind(user_id)
    .bind(&feature.feature_code)
    .bind(&feature.default_model)
    .bind(credits_charged)
    .bind(status)
    .bind(request_id)
    .bind(prompt_preview)
    .bind(response_preview)
    .bind(input_tokens)
    .bind(output_tokens)
    .bind(total_tokens)
    .execute(pool)
    .await;

    if let Err(e) = result {
        tracing::error!("Failed to write ai_usage_logs row: {e}");
    }
}

/// Charge an AI feature using LiteLLM with full audit logging and token capture
/// 
/// This is the preferred way to charge AI features - it uses LiteLLM for model routing
/// and captures token counts for cost tracking.
/// 
/// # Arguments
/// * `pool` - Database connection pool
/// * `user_id` - The user being charged
/// * `feature_code` - The AI feature being used (e.g., "jd_generate", "cover_letter_generate")
/// * `request_id` - Optional request tracking ID
/// * `idempotency_key` - Optional idempotency key
/// * `prompt` - The prompt being sent to the LLM
/// * `litellm_response` - The LiteLLM response containing generated text and token counts
/// 
/// # Returns
/// * `Ok(ChargeOutcome<String>)` - The generated text from the LLM
/// * `Err(ChargeError)` - Error if credits insufficient or DB failure
pub async fn charge_ai_feature_litellm(
    pool: &PgPool,
    user_id: Uuid,
    feature_code: &str,
    request_id: Option<&str>,
    idempotency_key: Option<&str>,
    prompt: &str,
    litellm_response: &crate::litellm::LiteLLMResponse,
) -> Result<ChargeOutcome<String>, ChargeError<String>> {
    let feature = AiCreditsRepository::get_feature_cost(pool, feature_code)
        .await
        .map_err(|e| ChargeError::Credits(AiCreditsError::Db(e)))?
        .ok_or_else(|| ChargeError::UnknownFeature(feature_code.to_string()))?;

    // Validate prompt length before reserving
    if let Some((_, error_json)) = validate_prompt_length(prompt, &feature) {
        return Err(ChargeError::Work(format!(
            "Validation failed: {}",
            error_json["error"]
        )));
    }

    let prompt_preview = truncate_for_preview(prompt, 200);
    let response_text = litellm_response.generated_text();
    let response_preview = truncate_for_preview(&response_text, 200);

    let hold = AiCreditsRepository::try_reserve_credits(
        pool,
        user_id,
        feature_code,
        feature.credit_cost,
        request_id,
        idempotency_key,
    )
    .await
    .map_err(ChargeError::Credits)?;

    // Capture the charge since we have a successful LiteLLM response
    AiCreditsRepository::try_capture_reservation(pool, hold.id, None)
        .await
        .map_err(ChargeError::Credits)?;

    // Log with token counts if available
    let input_tokens = litellm_response.usage.as_ref().and_then(|u| u.prompt_tokens);
    let output_tokens = litellm_response.usage.as_ref().and_then(|u| u.completion_tokens);
    let total_tokens = litellm_response.usage.as_ref().and_then(|u| u.total_tokens);

    log_usage(
        pool,
        user_id,
        &feature,
        hold.credits_held,
        "success",
        request_id,
        Some(&prompt_preview),
        Some(&response_preview),
        input_tokens,
        output_tokens,
        total_tokens,
    )
    .await;

    Ok(ChargeOutcome {
        result: response_text,
        feature,
        credits_charged: hold.credits_held,
    })
}

/// Convenience function to call LiteLLM and charge credits in one operation
/// 
/// This is the recommended pattern for handlers:
/// ```rust
/// let response = call_litellm_and_charge(
///     &state.pool,
///     auth.user_id,
///     "jd_generate",
///     None,
///     None,
///     &prompt,
/// ).await?;
/// ```
pub async fn call_litellm_and_charge(
    pool: &PgPool,
    user_id: Uuid,
    feature_code: &str,
    request_id: Option<&str>,
    idempotency_key: Option<&str>,
    prompt: &str,
) -> Result<ChargeOutcome<String>, ChargeError<String>> {
    // First reserve credits
    let feature = AiCreditsRepository::get_feature_cost(pool, feature_code)
        .await
        .map_err(|e| ChargeError::Credits(AiCreditsError::Db(e)))?
        .ok_or_else(|| ChargeError::UnknownFeature(feature_code.to_string()))?;

    // Validate prompt length before reserving
    if let Some((_, error_json)) = validate_prompt_length(prompt, &feature) {
        return Err(ChargeError::Work(format!(
            "Validation failed: {}",
            error_json["error"]
        )));
    }

    let prompt_preview = truncate_for_preview(prompt, 200);

    let hold = AiCreditsRepository::try_reserve_credits(
        pool,
        user_id,
        feature_code,
        feature.credit_cost,
        request_id,
        idempotency_key,
    )
    .await
    .map_err(ChargeError::Credits)?;

    // Get LiteLLM config and call
    let (base_url, _default_model, api_key) = crate::litellm::get_litellm_config();
    let model_alias = feature.default_model.clone();
    
    // Use feature's max_output_tokens if configured
    let max_tokens = feature.max_output_tokens;

    match crate::litellm::call_litellm(&base_url, &model_alias, prompt, api_key.as_deref(), max_tokens).await {
        Ok(litellm_response) => {
            // Capture the charge
            AiCreditsRepository::try_capture_reservation(pool, hold.id, None)
                .await
                .map_err(ChargeError::Credits)?;

            let response_text = litellm_response.generated_text();
            let response_preview = truncate_for_preview(&response_text, 200);

            // Log with token counts
            let input_tokens = litellm_response.usage.as_ref().and_then(|u| u.prompt_tokens);
            let output_tokens = litellm_response.usage.as_ref().and_then(|u| u.completion_tokens);
            let total_tokens = litellm_response.usage.as_ref().and_then(|u| u.total_tokens);

            log_usage(
                pool,
                user_id,
                &feature,
                hold.credits_held,
                "success",
                request_id,
                Some(&prompt_preview),
                Some(&response_preview),
                input_tokens,
                output_tokens,
                total_tokens,
            )
            .await;

            Ok(ChargeOutcome {
                result: response_text,
                feature,
                credits_charged: hold.credits_held,
            })
        }
        Err(e) => {
            // Release the hold on error
            let _ = AiCreditsRepository::try_release_reservation(pool, hold.id).await;
            log_usage(
                pool,
                user_id,
                &feature,
                0,
                "error",
                request_id,
                Some(&prompt_preview),
                None,
                None,
                None,
                None,
            )
            .await;
            Err(ChargeError::Work(e))
        }
    }
}