GitLite/src-tauri/crates/commit_ai/src/cloud.rs

use std::time::Duration;

use serde::{Deserialize, Serialize};

use crate::{build_messages, sanitize_message};

// Generous sizing so a detailed body with bullet points isn't cut off.
const DEFAULT_MAX_TOKENS: u32 = 1500;
// Without a timeout, a hanging endpoint would permanently block the "AI" button.
const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);

fn http_client() -> Result<reqwest::Client, String> {
    reqwest::Client::builder()
        .timeout(REQUEST_TIMEOUT)
        .build()
        .map_err(|err| format!("Could not create HTTP client: {err}"))
}

#[derive(Serialize)]
struct OpenAiMessage {
    role: &'static str,
    content: String,
}

#[derive(Serialize)]
struct OpenAiRequest {
    model: String,
    messages: Vec<OpenAiMessage>,
    temperature: f32,
}

#[derive(Deserialize)]
struct OpenAiResponseMessage {
    content: Option<String>,
}

#[derive(Deserialize)]
struct OpenAiChoice {
    message: OpenAiResponseMessage,
}

#[derive(Deserialize)]
struct OpenAiResponse {
    #[serde(default)]
    choices: Vec<OpenAiChoice>,
}

async fn openai_compatible_request(
    url: String,
    bearer: Option<&str>,
    model: &str,
    diff: &str,
    notes: Option<&str>,
) -> Result<String, String> {
    let (system, user) = build_messages(diff, notes)?;
    let body = OpenAiRequest {
        model: model.to_string(),
        messages: vec![
            OpenAiMessage { role: "system", content: system },
            OpenAiMessage { role: "user", content: user },
        ],
        temperature: 0.3,
    };

    let client = http_client()?;
    let mut request = client.post(url).json(&body);
    if let Some(key) = bearer.filter(|k| !k.trim().is_empty()) {
        request = request.bearer_auth(key);
    }

    let response = request
        .send()
        .await
        .map_err(|err| format!("Request to the AI model failed: {err}"))?;
    let status = response.status();
    let text = response
        .text()
        .await
        .map_err(|err| format!("Could not read response: {err}"))?;

    if !status.is_success() {
        return Err(format!("API error ({status}): {text}"));
    }

    let parsed: OpenAiResponse = serde_json::from_str(&text)
        .map_err(|err| format!("Could not process response: {err}"))?;

    parsed
        .choices
        .into_iter()
        .next()
        .and_then(|choice| choice.message.content)
        .map(|content| sanitize_message(&content))
        .filter(|content| !content.is_empty())
        .ok_or_else(|| "The model did not return a response.".to_string())
}

pub async fn generate_openai(
    api_key: &str,
    model: &str,
    diff: &str,
    notes: Option<&str>,
) -> Result<String, String> {
    if api_key.trim().is_empty() {
        return Err("OpenAI API key is missing.".to_string());
    }
    openai_compatible_request(
        "https://api.openai.com/v1/chat/completions".to_string(),
        Some(api_key),
        model,
        diff,
        notes,
    )
    .await
}

pub async fn generate_custom(
    base_url: &str,
    api_key: Option<&str>,
    model: &str,
    diff: &str,
    notes: Option<&str>,
) -> Result<String, String> {
    if base_url.trim().is_empty() {
        return Err("Endpoint URL is missing.".to_string());
    }
    let url = format!("{}/chat/completions", base_url.trim_end_matches('/'));
    openai_compatible_request(url, api_key, model, diff, notes).await
}

#[derive(Serialize)]
struct AnthropicMessage {
    role: &'static str,
    content: String,
}

#[derive(Serialize)]
struct AnthropicRequest {
    model: String,
    max_tokens: u32,
    system: String,
    messages: Vec<AnthropicMessage>,
}

#[derive(Deserialize)]
struct AnthropicContentBlock {
    #[serde(default)]
    text: Option<String>,
}

#[derive(Deserialize)]
struct AnthropicResponse {
    #[serde(default)]
    content: Vec<AnthropicContentBlock>,
}

pub async fn generate_anthropic(
    api_key: &str,
    model: &str,
    diff: &str,
    notes: Option<&str>,
) -> Result<String, String> {
    if api_key.trim().is_empty() {
        return Err("Anthropic API key is missing.".to_string());
    }
    let (system, user) = build_messages(diff, notes)?;
    let body = AnthropicRequest {
        model: model.to_string(),
        max_tokens: DEFAULT_MAX_TOKENS,
        system,
        messages: vec![AnthropicMessage { role: "user", content: user }],
    };

    let client = http_client()?;
    let response = client
        .post("https://api.anthropic.com/v1/messages")
        .header("x-api-key", api_key)
        .header("anthropic-version", "2023-06-01")
        .json(&body)
        .send()
        .await
        .map_err(|err| format!("Request to Anthropic failed: {err}"))?;
    let status = response.status();
    let text = response
        .text()
        .await
        .map_err(|err| format!("Could not read response: {err}"))?;

    if !status.is_success() {
        return Err(format!("API error ({status}): {text}"));
    }

    let parsed: AnthropicResponse = serde_json::from_str(&text)
        .map_err(|err| format!("Could not process response: {err}"))?;

    parsed
        .content
        .into_iter()
        .find_map(|block| block.text)
        .map(|text| sanitize_message(&text))
        .filter(|text| !text.is_empty())
        .ok_or_else(|| "The model did not return a response.".to_string())
}