use std::time::Duration; use serde::{Deserialize, Serialize}; use crate::{build_messages, looks_like_diff_echo, 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::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, temperature: f32, } #[derive(Deserialize)] struct OpenAiResponseMessage { content: Option, } #[derive(Deserialize)] struct OpenAiChoice { message: OpenAiResponseMessage, } #[derive(Deserialize)] struct OpenAiResponse { #[serde(default)] choices: Vec, } async fn openai_compatible_request( url: String, bearer: Option<&str>, model: &str, diff: &str, notes: Option<&str>, ) -> Result { 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}"))?; let message = 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())?; if looks_like_diff_echo(&message) { return Err("The model returned the diff instead of a commit message.".to_string()); } Ok(message) } pub async fn generate_openai( api_key: &str, model: &str, diff: &str, notes: Option<&str>, ) -> Result { 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 { 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, } #[derive(Deserialize)] struct AnthropicContentBlock { #[serde(default)] text: Option, } #[derive(Deserialize)] struct AnthropicResponse { #[serde(default)] content: Vec, } pub async fn generate_anthropic( api_key: &str, model: &str, diff: &str, notes: Option<&str>, ) -> Result { 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}"))?; let message = 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())?; if looks_like_diff_echo(&message) { return Err("The model returned the diff instead of a commit message.".to_string()); } Ok(message) }