flomo-ai/flomo-ai-desktop/src/api/llm_client.rs

use crate::config::AppSettings;
use serde::{Deserialize, Serialize};

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionRequest {
    pub model: String,
    pub messages: Vec<ChatMessage>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatChoice {
    pub message: ChatMessage,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionResponse {
    pub choices: Vec<ChatChoice>,
}

pub fn call_llm(settings: &AppSettings, user_input: String, selected_prompt: Option<String>) -> Result<String, String> {
    let full_content = if let Some(prompt) = selected_prompt {
        if !prompt.is_empty() {
            format!("{}{}", prompt, user_input)
        } else {
            user_input
        }
    } else {
        user_input
    };

    let messages = vec![ChatMessage {
        role: "user".to_string(),
        content: full_content,
    }];

    let request = ChatCompletionRequest {
        model: settings.llm_config.model.clone(),
        messages,
    };

    let client = reqwest::blocking::Client::builder()
        .timeout(std::time::Duration::from_secs(120))
        .build()
        .map_err(|e| format!("Failed to create HTTP client: {}", e))?;

    let mut req_builder = client
        .post(format!("{}/chat/completions", settings.llm_config.base_url))
        .header("Content-Type", "application/json");

    if !settings.llm_config.api_key.is_empty() {
        req_builder = req_builder.header("Authorization", format!("Bearer {}", settings.llm_config.api_key));
    }

    for header in &settings.header_configs {
        if !header.key.is_empty() {
            req_builder = req_builder.header(&header.key, &header.value);
        }
    }

    let body = serde_json::to_string(&request)
        .map_err(|e| format!("Failed to serialize request: {}", e))?;

    let response = req_builder
        .body(body)
        .send()
        .map_err(|e| format!("Request failed: {}", e))?;

    if !response.status().is_success() {
        let status = response.status();
        let error_body = response.text().unwrap_or_default();
        return Err(format!("API error {}: {}", status, error_body));
    }

    let response_text = response
        .text()
        .map_err(|e| format!("Failed to read response: {}", e))?;

    let completion: ChatCompletionResponse = serde_json::from_str(&response_text)
        .map_err(|e| format!("Failed to parse response: {}", e))?;

    if completion.choices.is_empty() {
        return Err("No response from API".to_string());
    }

    Ok(completion.choices[0].message.content.clone())
}