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

use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::Value;

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiteLlmChatRequest {
    pub model: String,
    pub messages: Vec<LiteLlmChatMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<i32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
}

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiteLlmChoice {
    pub index: Option<i32>,
    pub message: Option<LiteLlmChatMessage>,
    pub finish_reason: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiteLlmChatResponse {
    pub id: Option<String>,
    pub model: Option<String>,
    pub choices: Vec<LiteLlmChoice>,
    pub usage: Option<LiteLlmUsage>,
}

#[derive(Debug, thiserror::Error)]
pub enum LiteLlmError {
    #[error("HTTP error: {0}")]
    Http(#[from] reqwest::Error),
    #[error("LiteLLM returned error: {status} - {body}")]
    Api { status: u16, body: String },
    #[error("No completion returned")]
    NoCompletion,
    #[error("Configuration error: missing LiteLLM base URL")]
    MissingBaseUrl,
}

impl LiteLlmError {
    pub fn status_code(&self) -> axum::http::StatusCode {
        use axum::http::StatusCode;
        match self {
            LiteLlmError::Http(_) => StatusCode::BAD_GATEWAY,
            LiteLlmError::Api { status, .. } => StatusCode::from_u16(*status)
                .unwrap_or(StatusCode::BAD_GATEWAY),
            LiteLlmError::NoCompletion => StatusCode::BAD_GATEWAY,
            LiteLlmError::MissingBaseUrl => StatusCode::INTERNAL_SERVER_ERROR,
        }
    }

    /// Client-facing error body. Deliberately does NOT include the raw
    /// upstream response body (`LiteLlmError::Api.body`) — that can contain
    /// internal routing details or diagnostics from the LiteLLM proxy /
    /// underlying model provider that shouldn't be exposed to end users.
    /// Log the full error server-side at the call site instead.
    pub fn error_body(&self) -> serde_json::Value {
        let message = match self {
            LiteLlmError::Api { status, .. } => {
                format!("The AI service returned an error (status {status}).")
            }
            LiteLlmError::Http(_) => "Could not reach the AI service.".to_string(),
            LiteLlmError::NoCompletion => "The AI service returned no result.".to_string(),
            LiteLlmError::MissingBaseUrl => "The AI service is not configured.".to_string(),
        };
        serde_json::json!({
            "error": message,
            "code": "LITELLM_ERROR"
        })
    }
}

pub struct LiteLlmClient {
    client: Client,
    base_url: String,
    api_key: Option<String>,
}

impl LiteLlmClient {
    pub fn new() -> Result<Self, LiteLlmError> {
        let base_url = std::env::var("LITELLM_BASE_URL")
            .unwrap_or_else(|_| "http://litellm.nxtgauge-ai.svc.cluster.local:4000".to_string());

        if base_url.is_empty() {
            return Err(LiteLlmError::MissingBaseUrl);
        }

        let api_key = std::env::var("LITELLM_API_KEY").ok();

        Ok(Self {
            client: Client::new(),
            base_url,
            api_key,
        })
    }

    pub fn from_url(base_url: String, api_key: Option<String>) -> Self {
        Self {
            client: Client::new(),
            base_url,
            api_key,
        }
    }

    pub async fn chat_completion(
        &self,
        request: LiteLlmChatRequest,
    ) -> Result<LiteLlmChatResponse, LiteLlmError> {
        let url = format!("{}/v1/chat/completions", self.base_url);

        let mut req = self.client.post(&url).json(&request);
        if let Some(key) = &self.api_key {
            req = req.header("Authorization", format!("Bearer {}", key));
        }

        let response = req.send().await?;
        let status = response.status();

        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            tracing::error!("LiteLLM API error: status={} body={}", status.as_u16(), body);
            return Err(LiteLlmError::Api {
                status: status.as_u16(),
                body,
            });
        }

        let result = response.json::<LiteLlmChatResponse>().await?;
        Ok(result)
    }

    pub async fn chat_completion_text(
        &self,
        model: &str,
        system_prompt: Option<&str>,
        user_message: &str,
        max_tokens: Option<i32>,
    ) -> Result<(String, LiteLlmUsage, Option<String>), LiteLlmError> {
        let mut messages = Vec::new();
        if let Some(system) = system_prompt {
            messages.push(LiteLlmChatMessage {
                role: "system".to_string(),
                content: system.to_string(),
            });
        }
        messages.push(LiteLlmChatMessage {
            role: "user".to_string(),
            content: user_message.to_string(),
        });

        let request = LiteLlmChatRequest {
            model: model.to_string(),
            messages,
            temperature: Some(0.7),
            max_tokens,
            stream: Some(false),
            user: None,
        };

        let response = self.chat_completion(request).await?;
        let request_id = response.id.clone();
        let usage = response.usage.unwrap_or(LiteLlmUsage {
            prompt_tokens: None,
            completion_tokens: None,
            total_tokens: None,
        });

        let content = response
            .choices
            .into_iter()
            .next()
            .and_then(|c| c.message)
            .map(|m| m.content)
            .ok_or(LiteLlmError::NoCompletion)?;

        Ok((content, usage, request_id))
    }

    /// Issues a per-account LiteLLM virtual key via `/key/generate`, scoped to `user_id`.
    /// Requires `self.api_key` to be the LiteLLM master key.
    pub async fn generate_key(&self, user_id: uuid::Uuid) -> Result<String, LiteLlmError> {
        let url = format!("{}/key/generate", self.base_url);

        let mut req = self.client.post(&url).json(&serde_json::json!({
            "user_id": user_id.to_string(),
            "key_alias": format!("user-{}", user_id),
        }));
        if let Some(key) = &self.api_key {
            req = req.header("Authorization", format!("Bearer {}", key));
        }

        let response = req.send().await?;
        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            tracing::error!("LiteLLM API error: status={} body={}", status.as_u16(), body);
            return Err(LiteLlmError::Api {
                status: status.as_u16(),
                body,
            });
        }

        let body: Value = response.json().await?;
        body.get("key")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string())
            .ok_or(LiteLlmError::NoCompletion)
    }

    /// Direct pass-through for callers that want the raw JSON response.
    #[allow(dead_code)]
    pub async fn raw_chat_completion(&self, body: Value) -> Result<Value, LiteLlmError> {
        let url = format!("{}/v1/chat/completions", self.base_url);

        let mut req = self.client.post(&url).json(&body);
        if let Some(key) = &self.api_key {
            req = req.header("Authorization", format!("Bearer {}", key));
        }

        let response = req.send().await?;
        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            tracing::error!("LiteLLM API error: status={} body={}", status.as_u16(), body);
            return Err(LiteLlmError::Api {
                status: status.as_u16(),
                body,
            });
        }

        let result = response.json::<Value>().await?;
        Ok(result)
    }
}

impl Default for LiteLlmClient {
    fn default() -> Self {
        Self::new().unwrap_or_else(|e| {
            tracing::error!("Failed to create default LiteLlmClient: {}", e);
            Self::from_url("http://localhost:4000".to_string(), None)
        })
    }
}