gemini_crate 0.1.0

//! The main client for interacting with the Gemini API.
use crate::config::Config;
use crate::errors::Error;
use crate::types::{Content, GenerateContentRequest, GenerateContentResponse, Part, Role};
use backon::{ExponentialBuilder, Retryable};
use reqwest::Client;
use serde::{de::DeserializeOwned, Deserialize, Serialize};

const API_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";

/// The `GeminiClient` is the main entry point for interacting with the Gemini API.
///
/// It handles the configuration, HTTP requests, and error handling for the API calls.
///
/// # Examples
///
/// ```no_run
/// use gemini_crate::client::GeminiClient;
///
/// #[tokio::main]
/// async fn main() {
///     let client = GeminiClient::new().unwrap();
///     let models = client.list_models().await.unwrap();
///     println!("Available models: {:?}", models);
/// }
/// ```
pub struct GeminiClient {
    client: Client,
    config: Config,
    base_url: String,
}

impl GeminiClient {
    /// Creates a new `GeminiClient`.
    ///
    /// This function will attempt to load the `GEMINI_API_KEY` from the environment.
    ///
    /// # Errors
    ///
    /// This function will return an error if the `GEMINI_API_KEY` is not set.
    pub fn new() -> Result<Self, Error> {
        let config = Config::new()?;
        let client = Client::new();
        Ok(Self {
            client,
            config,
            base_url: API_BASE_URL.to_string(),
        })
    }

    /// Creates a new `GeminiClient` with a specific config for testing.
    #[cfg(test)]
    pub(crate) fn with_config(config: Config) -> Self {
        Self {
            client: Client::new(),
            config,
            base_url: API_BASE_URL.to_string(),
        }
    }

    /// Creates a new `GeminiClient` with a specific base URL for testing.
    #[cfg(test)]
    pub(crate) fn with_base_url(mut self, base_url: &str) -> Self {
        self.base_url = base_url.to_string();
        self
    }

    /// Lists the available models.
    ///
    /// This function sends a request to the Gemini API to retrieve a list of the
    /// available models.
    ///
    /// # Errors
    ///
    /// This function will return an error if the request fails or if the API
    /// returns an error.
    pub async fn list_models(&self) -> Result<ListModelsResponse, Error> {
        let url = format!("{}/models", self.base_url);
        self.get_with_retry(&url).await
    }

    /// Generates text from a given prompt using the specified model.
    ///
    /// # Arguments
    ///
    /// * `model` - The name of the model to use for generation (e.g., "gemini-pro").
    /// * `prompt` - The text prompt to send to the model.
    ///
    /// # Errors
    ///
    /// This function will return an error if the request fails or if the API
    /// returns an error.
    pub async fn generate_text(&self, model: &str, prompt: &str) -> Result<String, Error> {
        let url = format!("{}/models/{}:generateContent", self.base_url, model);

        let request = GenerateContentRequest {
            contents: vec![Content {
                role: Role::User,
                parts: vec![Part {
                    text: prompt.to_string(),
                }],
            }],
        };

        let response: GenerateContentResponse = self.post_with_retry(&url, &request).await?;

        if let Some(candidate) = response.candidates.first() {
            if let Some(part) = candidate.content.parts.first() {
                return Ok(part.text.clone());
            }
        }

        Err(Error::Api("No content found in response".to_string()))
    }

    async fn get_with_retry<R>(&self, url: &str) -> Result<R, Error>
    where
        R: DeserializeOwned,
    {
        let url = url.to_string();
        let api_key = self.config.api_key().to_string();
        let client = self.client.clone();

        (|| async {
            let response = client
                .get(&url)
                .header("x-goog-api-key", &api_key)
                .send()
                .await
                .map_err(Error::Network)?;

            if response.status().is_server_error() {
                let status_code = response.status().as_u16();
                let error_text = response.text().await.unwrap_or_default();
                return Err(Error::Api(format!("HTTP {}: {}", status_code, error_text)));
            }
            if !response.status().is_success() {
                return Err(Error::Api(response.text().await.unwrap_or_default()));
            }

            response.json::<R>().await.map_err(Error::Json)
        })
        .retry(ExponentialBuilder::default())
        .when(|e| match e {
            Error::Network(_) => true,
            Error::Api(msg) => {
                // Check if it's a server error response that we should retry
                msg.contains("503")
                    || msg.contains("502")
                    || msg.contains("500")
                    || msg.contains("504")
            }
            _ => false,
        })
        .await
    }
    async fn post_with_retry<T, R>(&self, url: &str, body: &T) -> Result<R, Error>
    where
        T: Serialize + Clone,
        R: DeserializeOwned,
    {
        let url = url.to_string();
        let api_key = self.config.api_key().to_string();
        let client = self.client.clone();
        let body = body.clone();

        (|| async {
            let response = client
                .post(&url)
                .header("x-goog-api-key", &api_key)
                .json(&body)
                .send()
                .await
                .map_err(Error::Network)?;

            if response.status().is_server_error() {
                let status_code = response.status().as_u16();
                let error_text = response.text().await.unwrap_or_default();
                return Err(Error::Api(format!("HTTP {}: {}", status_code, error_text)));
            }
            if !response.status().is_success() {
                return Err(Error::Api(response.text().await.unwrap_or_default()));
            }

            response.json::<R>().await.map_err(Error::Json)
        })
        .retry(ExponentialBuilder::default())
        .when(|e| match e {
            Error::Network(_) => true,
            Error::Api(msg) => {
                // Check if it's a server error response that we should retry
                msg.contains("503")
                    || msg.contains("502")
                    || msg.contains("500")
                    || msg.contains("504")
            }
            _ => false,
        })
        .await
    }
}

/// The response from the `list_models` endpoint.
#[derive(Deserialize, Debug)]
pub struct ListModelsResponse {
    /// The list of available models.
    pub models: Vec<Model>,
}

/// Represents a model available through the Gemini API.
#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Model {
    /// The name of the model.
    pub name: String,
    /// The version of the model.
    #[serde(default)]
    pub version: String,
    /// The display name of the model.
    #[serde(default)]
    pub display_name: String,
    /// A description of the model.
    #[serde(default)]
    pub description: String,
    /// The input token limit for the model.
    #[serde(default)]
    pub input_token_limit: u32,
    /// The output token limit for the model.
    #[serde(default)]
    pub output_token_limit: u32,
    /// The supported generation methods for the model.
    #[serde(default)]
    pub supported_generation_methods: Vec<String>,
}