easy_llm 0.2.0

use async_trait::async_trait;
use llm_api_rs::{
    Anthropic, ChatCompletionRequest, ChatMessage, DeepSeek, Gemini, LlmProvider, OpenAI,
};
use ollama_rs::{Ollama, generation::completion::request::GenerationRequest};

/// Represents the possible errors that can occur.
#[derive(Debug)]
pub enum Error {
    /// Represents a provider that is not supported.
    NotSupport(String),
    /// Represents that an API key is required.
    RequireApiKey,
    /// Represents an error from the Ollama service.
    Ollama(ollama_rs::error::OllamaError),
    /// Represents an error from the LLM service.
    LlmService(String),
    Io(std::io::Error),
}

impl From<ollama_rs::error::OllamaError> for Error {
    fn from(value: ollama_rs::error::OllamaError) -> Self {
        Self::Ollama(value)
    }
}
impl From<std::io::Error> for Error {
    fn from(value: std::io::Error) -> Self {
        Self::Io(value)
    }
}

/// supported LLM providers.
#[derive(Debug, PartialEq, Eq)]
pub enum Provider {
    Ollama,
    OpenAi,
    Gemini,
    Anthropic,
    Deepseek,
}

impl Provider {
    pub fn parse_str<T: AsRef<str>>(s: T) -> Result<Self, Error> {
        match s.as_ref().trim().to_lowercase().as_str() {
            "ollama" => Ok(Self::Ollama),
            "openai" => Ok(Self::OpenAi),
            "gemini" => Ok(Self::Gemini),
            "anthropic" => Ok(Self::Anthropic),
            "deepseek" => Ok(Self::Deepseek),
            _ => Err(Error::NotSupport(s.as_ref().to_string())),
        }
    }
    fn is_ollama(&self) -> bool {
        self == &Self::Ollama
    }
}

/// # Examples
///
/// ```rust
/// use easy_llm::{Llm, LlmControl, Provider};
///
/// // Create an OpenAI client
/// let openai_llm = Llm::new(
///     Provider::OpenAi,
///     "gpt-4".to_string(),
///     Some("your-api-key".to_string()),
///     Some(0.7),
///     Some(1000),
///     None,
/// );
///
/// // Create an Ollama client (no API key required)
/// let ollama_llm = Llm::new(
///     Provider::Ollama,
///     "llama2".to_string(),
///     None,
///     Some(0.8),
///     None,
///     None, // if it's None, use localhost and 11434
/// );
/// ```
pub struct Llm {
    provider: Provider,
    model: String,
    api_key: Option<String>,
    temperature: Option<f32>,
    max_tokens: Option<u32>,
    /// if it's None, use localhost and 11434 as default value
    base_url: Option<(String, u16)>,
}

#[async_trait]
pub trait LlmControl {
    fn check_api_key(&self) -> bool;
    fn new(
        provider: Provider,
        model: String,
        api_key: Option<String>,
        temperature: Option<f32>,
        max_tokens: Option<u32>,
        base_url: Option<(String, u16)>,
    ) -> Self;
    async fn call<T>(&self, prompt: T) -> Result<String, Error>
    where
        T: AsRef<str> + Send + Sync + 'static;
    fn sync_call<T>(&self, prompt: T) -> Result<String, Error>
    where
        T: AsRef<str> + Send + Sync + 'static;
}

#[async_trait]
impl LlmControl for Llm {
    fn check_api_key(&self) -> bool {
        match self.api_key.clone() {
            Some(_) => true,
            None => !self.provider.is_ollama(),
        }
    }

    fn new(
        provider: Provider,
        model: String,
        api_key: Option<String>,
        temperature: Option<f32>,
        max_tokens: Option<u32>,
        base_url: Option<(String, u16)>,
    ) -> Self {
        Self {
            provider,
            model,
            api_key,
            temperature,
            max_tokens,
            base_url,
        }
    }

    async fn call<T>(&self, prompt: T) -> Result<String, Error>
    where
        T: AsRef<str> + Send + Sync + 'static,
    {
        async {
            if self.provider == Provider::Ollama {
                let base_url = self
                    .base_url
                    .clone()
                    .unwrap_or(("http://localhost:11434".to_string(), 11434));
                let ollama_result = Ollama::new(base_url.0, base_url.1);
                let res = ollama_result
                    .generate(GenerationRequest::new(self.model.clone(), prompt.as_ref()))
                    .await;
                match res {
                    Ok(v) => Ok(v.response.to_string()),
                    Err(e) => Err(Error::Ollama(e)),
                }
            } else {
                let api = self.api_key.clone().unwrap();
                let client: Box<dyn LlmProvider + Send> = match self.provider {
                    Provider::OpenAi => Box::new(OpenAI::new(api)),
                    Provider::Gemini => Box::new(Gemini::new(api)),
                    Provider::Anthropic => Box::new(Anthropic::new(api)),
                    Provider::Deepseek => Box::new(DeepSeek::new(api)),
                    _ => unreachable!(),
                };

                let req = ChatCompletionRequest {
                    model: self.model.clone(),
                    messages: vec![ChatMessage {
                        role: "user".to_string(),
                        content: prompt.as_ref().to_string(),
                    }],
                    temperature: self.temperature,
                    max_tokens: self.max_tokens,
                };

                match client.chat_completion(req).await {
                    Ok(res) => {
                        if let Some(v) = res.choices.first() {
                            Ok(v.message.content.clone())
                        } else {
                            Ok(String::new())
                        }
                    }
                    Err(e) => Err(Error::LlmService(e.to_string())),
                }
            }
        }
        .await
    }
    fn sync_call<T>(&self, prompt: T) -> Result<String, Error>
    where
        T: AsRef<str> + Send + Sync + 'static,
    {
        let rt = tokio::runtime::Runtime::new()?;
        rt.block_on(self.call(prompt))
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        // Llm, LlmControl,
        Provider,
    };

    #[test]
    fn parse_test() {
        let sources = [
            "ollama",
            "Ollama",
            "openai",
            "OpenAI",
            "gemini",
            "GEmini",
            "anthropic",
            "deepseek",
        ];
        let right_cl = [
            Provider::Ollama,
            Provider::Ollama,
            Provider::OpenAi,
            Provider::OpenAi,
            Provider::Gemini,
            Provider::Gemini,
            Provider::Anthropic,
            Provider::Deepseek,
        ];

        let res = sources
            .iter()
            .zip(right_cl)
            .map(|f| (Provider::parse_str(f.0).unwrap(), f.1))
            .all(|f| f.0 == f.1);
        assert!(res);
    }
}