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

//! LiteLLM client for the Ask Ash AI credits system.
//!
//! Thin HTTP wrapper around LiteLLM's `/chat/completions` endpoint.
//! This is the preferred way to call LLMs - it provides model routing,
//! fallbacks, and unified interface while billing remains in nxtgauge-backend-rust.

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
struct LiteLLMChatMessage {
    role: String,
    content: String,
}

#[derive(Debug, Clone, Serialize)]
struct LiteLLMRequest {
    model: String,
    messages: Vec<LiteLLMChatMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<i32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
}

#[derive(Debug, Clone, Deserialize)]
struct LiteLLMChoice {
    message: LiteLLMChatMessage,
    finish_reason: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct LiteLLMUsage {
    pub prompt_tokens: Option<i32>,
    pub completion_tokens: Option<i32>,
    pub total_tokens: Option<i32>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct LiteLLMResponse {
    pub id: String,
    pub model: String,
    pub choices: Vec<LiteLLMChoice>,
    pub usage: Option<LiteLLMUsage>,
}

impl LiteLLMResponse {
    /// Extract the generated text from the response
    pub fn generated_text(&self) -> String {
        self.choices
            .first()
            .map(|c| c.message.content.clone())
            .unwrap_or_default()
    }

    /// Get token counts if available
    pub fn token_counts(&self) -> Option<(i32, i32, i32)> {
        self.usage.as_ref().map(|u| (
            u.prompt_tokens.unwrap_or(0),
            u.completion_tokens.unwrap_or(0),
            u.total_tokens.unwrap_or(0),
        ))
    }
}

/// Call LiteLLM with a single prompt (simplified interface)
/// 
/// # Arguments
/// * `base_url` - LiteLLM base URL (e.g., "http://litellm.nxtgauge-ai.svc.cluster.local:4000")
/// * `model_alias` - LiteLLM model alias (e.g., "askash-fast", "askash-main")
/// * `prompt` - The user prompt
/// * `api_key` - Optional API key for authentication
/// * `max_tokens` - Optional max output tokens
/// 
/// # Returns
/// * `Ok(LiteLLMResponse)` - The successful response with generated text and metadata
/// * `Err(String)` - Error message if the request failed
pub async fn call_litellm(
    base_url: &str,
    model_alias: &str,
    prompt: &str,
    api_key: Option<&str>,
    max_tokens: Option<i32>,
) -> Result<LiteLLMResponse, String> {
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(60))
        .build()
        .map_err(|e| format!("Failed to build HTTP client: {}", e))?;

    let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/'));

    let request = LiteLLMRequest {
        model: model_alias.to_string(),
        messages: vec![
            LiteLLMChatMessage {
                role: "user".to_string(),
                content: prompt.to_string(),
            },
        ],
        max_tokens,
        temperature: Some(0.7),
    };

    let mut req_builder = client.post(&url)
        .header("Content-Type", "application/json");

    // Add authorization if API key is provided
    if let Some(key) = api_key {
        req_builder = req_builder.header("Authorization", format!("Bearer {}", key));
    }

    let response = req_builder
        .json(&request)
        .send()
        .await
        .map_err(|e| format!("LiteLLM request failed: {}", e))?;

    let status = response.status();
    if !status.is_success() {
        let error_body = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
        return Err(format!("LiteLLM error ({}): {}", status, error_body));
    }

    let litellm_response: LiteLLMResponse = response
        .json()
        .await
        .map_err(|e| format!("Failed to parse LiteLLM response: {}", e))?;

    Ok(litellm_response)
}

/// Call LiteLLM with a system prompt and user message
pub async fn call_litellm_with_system(
    base_url: &str,
    model_alias: &str,
    system_prompt: &str,
    user_message: &str,
    api_key: Option<&str>,
    max_tokens: Option<i32>,
) -> Result<LiteLLMResponse, String> {
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(60))
        .build()
        .map_err(|e| format!("Failed to build HTTP client: {}", e))?;

    let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/'));

    let request = LiteLLMRequest {
        model: model_alias.to_string(),
        messages: vec![
            LiteLLMChatMessage {
                role: "system".to_string(),
                content: system_prompt.to_string(),
            },
            LiteLLMChatMessage {
                role: "user".to_string(),
                content: user_message.to_string(),
            },
        ],
        max_tokens,
        temperature: Some(0.7),
    };

    let mut req_builder = client.post(&url)
        .header("Content-Type", "application/json");

    if let Some(key) = api_key {
        req_builder = req_builder.header("Authorization", format!("Bearer {}", key));
    }

    let response = req_builder
        .json(&request)
        .send()
        .await
        .map_err(|e| format!("LiteLLM request failed: {}", e))?;

    let status = response.status();
    if !status.is_success() {
        let error_body = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
        return Err(format!("LiteLLM error ({}): {}", status, error_body));
    }

    let litellm_response: LiteLLMResponse = response
        .json()
        .await
        .map_err(|e| format!("Failed to parse LiteLLM response: {}", e))?;

    Ok(litellm_response)
}

/// Get LiteLLM configuration from environment
pub fn get_litellm_config() -> (String, String, Option<String>) {
    let base_url = std::env::var("LITELLM_BASE_URL")
        .unwrap_or_else(|_| "http://litellm.nxtgauge-ai.svc.cluster.local:4000".to_string());
    
    let model_alias = std::env::var("LITELLM_MODEL")
        .unwrap_or_else(|_| "askash-fast".to_string());
    
    let api_key = std::env::var("LITELLM_API_KEY").ok();

    (base_url, model_alias, api_key)
}

/// Health check for LiteLLM
pub async fn check_litellm_health(base_url: &str, api_key: Option<&str>) -> Result<bool, String> {
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(5))
        .build()
        .map_err(|e| format!("Failed to build HTTP client: {}", e))?;

    let url = format!("{}/health", base_url.trim_end_matches('/'));

    let mut req_builder = client.get(&url);

    if let Some(key) = api_key {
        req_builder = req_builder.header("Authorization", format!("Bearer {}", key));
    }

    match req_builder.send().await {
        Ok(response) => Ok(response.status().is_success()),
        Err(e) => Err(format!("Health check failed: {}", e)),
    }
}